Fix: restatements for incremental by partition (#3656)

Author: treysp

sqlmesh/core/snapshot/evaluator.py

@@ -1414,12 +1414,15 @@ def insert(
         is_first_insert: bool,
         **kwargs: t.Any,
     ) -> None:
-        self.adapter.insert_overwrite_by_partition(
-            table_name,
-            query_or_df,
-            partitioned_by=model.partitioned_by,
-            columns_to_types=model.columns_to_types,
-        )
+        if is_first_insert:
+            self._replace_query_for_model(model, table_name, query_or_df)
+        else:
+            self.adapter.insert_overwrite_by_partition(
+                table_name,
+                query_or_df,
+                partitioned_by=model.partitioned_by,
+                columns_to_types=model.columns_to_types,
+            )
 
 
 class IncrementalByTimeRangeStrategy(MaterializableStrategy):

tests/core/test_integration.py

@@ -1622,9 +1622,10 @@ def test_incremental_by_partition(init_and_plan_context: t.Callable):
         f"""
         MODEL (
             name {model_name},
-            kind INCREMENTAL_BY_PARTITION,
+            kind INCREMENTAL_BY_PARTITION (disable_restatement false),
             partitioned_by [key],
             allow_partials true,
+            start '2023-01-07',
         );
 
         SELECT key, value FROM {source_name};
@@ -1665,6 +1666,16 @@ def test_incremental_by_partition(init_and_plan_context: t.Callable):
         ("key_a", 2),
     ]
 
+    # model should fully refresh on restatement
+    context.engine_adapter.replace_query(
+        source_name,
+        d.parse_one("SELECT 'key_c' AS key, 3 AS value"),
+    )
+    context.plan(auto_apply=True, no_prompts=True, restate_models=[model_name])
+    assert context.engine_adapter.fetchall(f"SELECT * FROM {model_name}") == [
+        ("key_c", 3),
+    ]
+
 
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_custom_materialization(init_and_plan_context: t.Callable):

tests/core/test_snapshot_evaluator.py

@@ -3012,6 +3012,8 @@ def test_evaluate_incremental_by_partition(mocker: MockerFixture, make_snapshot,
     snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
+    adapter_mock.columns.return_value = model.columns_to_types
+
     evaluator = SnapshotEvaluator(adapter_mock)
     evaluator.evaluate(
         snapshot,
@@ -3021,6 +3023,36 @@ def test_evaluate_incremental_by_partition(mocker: MockerFixture, make_snapshot,
         snapshots={},
     )
 
+    # uses `replace_query` on first model execution
+    adapter_mock.replace_query.assert_called_once_with(
+        snapshot.table_name(),
+        model.render_query(),
+        partitioned_by=[
+            exp.to_column("ds", quoted=True),
+            exp.to_column("b", quoted=True),
+        ],
+        columns_to_types=model.columns_to_types,
+        clustered_by=[],
+        table_properties={},
+        column_descriptions={},
+        partition_interval_unit=None,
+        storage_format=None,
+        table_description=None,
+        table_format=None,
+    )
+
+    adapter_mock.reset_mock()
+    snapshot.intervals = [(to_timestamp("2020-01-01"), to_timestamp("2020-01-02"))]
+
+    evaluator.evaluate(
+        snapshot,
+        start="2020-01-02",
+        end="2020-01-03",
+        execution_time="2020-01-03",
+        snapshots={},
+    )
+
+    # uses `insert_overwrite_by_partition` on all subsequent model executions
     adapter_mock.insert_overwrite_by_partition.assert_called_once_with(
         snapshot.table_name(),
         model.render_query(),