Feat: Support restatements for the incremental unmanaged model kind (#2423)

Author: izeigerman

docs/concepts/models/model_kinds.md

@@ -166,7 +166,7 @@ WHERE
   event_date BETWEEN @start_date AND @end_date;
 ```
 
-**Note:** Models of the `INCREMENTAL_BY_UNIQUE_KEY` kind are inherently [non-idempotent](../glossary.md#idempotency), which should be taken into consideration during data [restatement](../plans.md#restatement-plans).
+**Note:** Models of the `INCREMENTAL_BY_UNIQUE_KEY` kind are inherently [non-idempotent](../glossary.md#idempotency), which should be taken into consideration during data [restatement](../plans.md#restatement-plans). As a result, partial data restatement is not supported for this model kind, which means that the entire table will be recreated from scratch if restated.
 
 ### Unique Key Expressions
 
@@ -320,7 +320,9 @@ SQLMesh achieves this by adding a `valid_from` and `valid_to` column to your mod
 
 Therefore you can use these models to not only tell you what the latest value is for a given record but also what the values were anytime in the past. Note that maintaining this history does come at a cost of increased storage and compute and this may not be a good fit for sources that change frequently since the history could get very large.
 
-There are two ways to tracking changes: By Time (Recommended) or By Column. 
+**Note**: Partial data [restatement](../plans.md#restatement-plans) is not supported for this model kind, which means that the entire table will be recreated from scratch if restated. This may lead to data loss, which is why data restatement is disabled for models of this kind by default.
+
+There are two ways to tracking changes: By Time (Recommended) or By Column.
 
 ### SCD Type 2 By Time (Recommended)
 

docs/concepts/plans.md

@@ -133,6 +133,10 @@ For this reason, the `plan` command supports the `--restate-model`, which allows
 
 Application of a plan will trigger a cascading backfill for all specified models (other than external tables), as well as all models downstream from them. The plan's date range determines the data intervals that will be affected.
 
+Please note that models of kinds [INCREMENTAL_BY_UNIQUE_KEY](models/model_kinds.md#INCREMENTAL_BY_UNIQUE_KEY), [SCD_TYPE_2_BY_TIME](models/model_kinds.md#scd-type-2), and [SCD_TYPE_2_BY_COLUMN](models/model_kinds.md#scd-type-2) cannot be partially restated. Therefore, such models will be fully refreshed regardless of the start/end dates provided by a user in the plan.
+
+To prevent models from ever being restated, set the [disable_restatement](models/overview.md#disable_restatement) attribute to `true`.
+
 See examples below for how to restate both based on model names and model tags.
 
 

sqlmesh/core/engine_adapter/base.py

@@ -1159,6 +1159,7 @@ def scd_type_2_by_time(
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        truncate: bool = False,
         **kwargs: t.Any,
     ) -> None:
         self._scd_type_2(
@@ -1174,6 +1175,7 @@ def scd_type_2_by_time(
             columns_to_types=columns_to_types,
             table_description=table_description,
             column_descriptions=column_descriptions,
+            truncate=truncate,
         )
 
     def scd_type_2_by_column(
@@ -1190,6 +1192,7 @@ def scd_type_2_by_column(
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        truncate: bool = False,
         **kwargs: t.Any,
     ) -> None:
         self._scd_type_2(
@@ -1205,6 +1208,7 @@ def scd_type_2_by_column(
             execution_time_as_valid_from=execution_time_as_valid_from,
             table_description=table_description,
             column_descriptions=column_descriptions,
+            truncate=truncate,
         )
 
     def _scd_type_2(
@@ -1223,6 +1227,7 @@ def _scd_type_2(
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        truncate: bool = False,
     ) -> None:
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             source_table, columns_to_types, target_table=target_table, batch_size=0
@@ -1366,6 +1371,11 @@ def _scd_type_2(
                 .when(exp.column(f"t_{valid_from_name}").is_(exp.Null()), update_valid_from_start)
                 .else_(exp.column(f"t_{valid_from_name}"))
             ).as_(valid_from_name)
+
+        existing_rows_query = exp.select(*table_columns).from_(target_table)
+        if truncate:
+            existing_rows_query = existing_rows_query.limit(0)
+
         with source_queries[0] as source_query:
             query = (
                 exp.Select()  # type: ignore
@@ -1378,16 +1388,12 @@ def _scd_type_2(
                 # Historical Records that Do Not Change
                 .with_(
                     "static",
-                    exp.select(*table_columns)
-                    .from_(target_table)
-                    .where(f"{valid_to_name} IS NOT NULL"),
+                    existing_rows_query.where(f"{valid_to_name} IS NOT NULL"),
                 )
                 # Latest Records that can be updated
                 .with_(
                     "latest",
-                    exp.select(*table_columns)
-                    .from_(target_table)
-                    .where(f"{valid_to_name} IS NULL"),
+                    existing_rows_query.where(f"{valid_to_name} IS NULL"),
                 )
                 # Deleted records which can be used to determine `valid_from` for undeleted source records
                 .with_(

sqlmesh/core/engine_adapter/trino.py

@@ -216,6 +216,7 @@ def _scd_type_2(
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        truncate: bool = False,
     ) -> None:
         if columns_to_types and self.current_catalog_type == "delta_lake":
             columns_to_types = self._to_delta_ts(columns_to_types)
@@ -235,6 +236,7 @@ def _scd_type_2(
             columns_to_types,
             table_description,
             column_descriptions,
+            truncate,
         )
 
     # delta_lake only supports two timestamp data types. This method converts other

sqlmesh/core/model/kind.py

@@ -115,6 +115,13 @@ def only_execution_time(self) -> bool:
         """Whether or not this model only cares about execution time to render."""
         return self.is_view or self.is_full or self.is_scd_type_2
 
+    @property
+    def full_history_restatement_only(self) -> bool:
+        """Whether or not this model only supports restatement of full history."""
+        return (
+            self.is_incremental_unmanaged or self.is_incremental_by_unique_key or self.is_scd_type_2
+        )
+
 
 class ModelKindName(str, ModelKindMixin, Enum):
     """The kind of model, determining how this data is computed and stored in the warehouse."""
@@ -342,7 +349,7 @@ class IncrementalUnmanagedKind(_ModelKind):
     name: Literal[ModelKindName.INCREMENTAL_UNMANAGED] = ModelKindName.INCREMENTAL_UNMANAGED
     insert_overwrite: SQLGlotBool = False
     forward_only: SQLGlotBool = True
-    disable_restatement: Literal[True] = True
+    disable_restatement: SQLGlotBool = True
 
     @property
     def data_hash_values(self) -> t.List[t.Optional[str]]:

sqlmesh/core/snapshot/definition.py

@@ -675,7 +675,13 @@ def get_removal_interval(
         *,
         strict: bool = True,
     ) -> Interval:
-        end = execution_time or now() if self.depends_on_past else end
+        end = (
+            execution_time or now()
+            if self.depends_on_past or self.full_history_restatement_only
+            else end
+        )
+        if self.full_history_restatement_only and self.intervals:
+            start = self.intervals[0][0]
         return self.inclusive_exclusive(start, end, strict)
 
     def inclusive_exclusive(

sqlmesh/core/snapshot/evaluator.py

@@ -49,6 +49,7 @@
 )
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
+    Intervals,
     QualifiedViewName,
     Snapshot,
     SnapshotChangeCategory,
@@ -845,7 +846,7 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         """Inserts the given query or a DataFrame into the target table or replaces a view.
@@ -865,7 +866,7 @@ def append(
         table_name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         """Appends the given query or a DataFrame to the existing table.
@@ -952,6 +953,27 @@ def demote(
             environment_naming_info: Naming information for the target environment.
         """
 
+    def _replace_query_for_model(self, model: Model, name: str, query_or_df: QueryOrDF) -> None:
+        """Replaces the table for the given model.
+
+        Args:
+            model: The target model.
+            name: The name of the target table.
+            query_or_df: The query or DataFrame to replace the target table with.
+        """
+        self.adapter.replace_query(
+            name,
+            query_or_df,
+            columns_to_types=model.columns_to_types if model.annotated else None,
+            storage_format=model.storage_format,
+            partitioned_by=model.partitioned_by,
+            partition_interval_unit=model.interval_unit,
+            clustered_by=model.clustered_by,
+            table_properties=model.table_properties,
+            table_description=model.description,
+            column_descriptions=model.column_descriptions,
+        )
+
 
 class SymbolicStrategy(EvaluationStrategy):
     def insert(
@@ -960,7 +982,7 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         pass
@@ -971,7 +993,7 @@ def append(
         table_name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         pass
@@ -1063,7 +1085,7 @@ def append(
         table_name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
@@ -1138,7 +1160,7 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
@@ -1160,17 +1182,20 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
-        self.adapter.merge(
-            name,
-            query_or_df,
-            columns_to_types=model.columns_to_types,
-            unique_key=model.unique_key,
-            when_matched=model.when_matched,
-        )
+        if not _intervals(snapshot, deployability_index):
+            self._replace_query_for_model(model, name, query_or_df)
+        else:
+            self.adapter.merge(
+                name,
+                query_or_df,
+                columns_to_types=model.columns_to_types,
+                unique_key=model.unique_key,
+                when_matched=model.when_matched,
+            )
 
     def append(
         self,
@@ -1198,11 +1223,13 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
-        if isinstance(model.kind, IncrementalUnmanagedKind) and model.kind.insert_overwrite:
+        if not _intervals(snapshot, deployability_index):
+            self._replace_query_for_model(model, name, query_or_df)
+        elif isinstance(model.kind, IncrementalUnmanagedKind) and model.kind.insert_overwrite:
             self.adapter.insert_overwrite_by_partition(
                 name,
                 query_or_df,
@@ -1227,22 +1254,11 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
-        self.adapter.replace_query(
-            name,
-            query_or_df,
-            columns_to_types=model.columns_to_types if model.annotated else None,
-            storage_format=model.storage_format,
-            partitioned_by=model.partitioned_by,
-            partition_interval_unit=model.interval_unit,
-            clustered_by=model.clustered_by,
-            table_properties=model.table_properties,
-            table_description=model.description,
-            column_descriptions=model.column_descriptions,
-        )
+        self._replace_query_for_model(model, name, query_or_df)
 
 
 class SCDType2Strategy(MaterializableStrategy):
@@ -1290,10 +1306,11 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
+        truncate = not _intervals(snapshot, deployability_index)
         if isinstance(model.kind, SCDType2ByTimeKind):
             self.adapter.scd_type_2_by_time(
                 target_table=name,
@@ -1307,6 +1324,7 @@ def insert(
                 columns_to_types=model.columns_to_types,
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
+                truncate=truncate,
                 **kwargs,
             )
         elif isinstance(model.kind, SCDType2ByColumnKind):
@@ -1322,6 +1340,7 @@ def insert(
                 execution_time_as_valid_from=model.kind.execution_time_as_valid_from,
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
+                truncate=truncate,
                 **kwargs,
             )
         else:
@@ -1335,7 +1354,7 @@ def append(
         table_name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
@@ -1382,7 +1401,7 @@ def insert(
         name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         model = snapshot.model
@@ -1425,7 +1444,7 @@ def append(
         table_name: str,
         query_or_df: QueryOrDF,
         snapshots: t.Dict[str, Snapshot],
-        deployability_index: t.Optional[DeployabilityIndex],
+        deployability_index: DeployabilityIndex,
         **kwargs: t.Any,
     ) -> None:
         raise ConfigError(f"Cannot append to a view '{table_name}'.")
@@ -1485,3 +1504,11 @@ def delete(self, name: str) -> None:
 
     def _is_materialized_view(self, model: Model) -> bool:
         return isinstance(model.kind, ViewKind) and model.kind.materialized
+
+
+def _intervals(snapshot: Snapshot, deployability_index: DeployabilityIndex) -> Intervals:
+    return (
+        snapshot.intervals
+        if deployability_index.is_deployable(snapshot)
+        else snapshot.dev_intervals
+    )