Fix: Selectively widen restatement intervals if they affect a downstream model that cannot be partially restated (#3902)

Author: erindru

sqlmesh/core/plan/evaluator.py

@@ -438,7 +438,9 @@ def _restate(self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapsho
         # Without this rule, its possible that promoting a dev table to prod will introduce old data to prod
         snapshot_intervals_to_restate.update(
             self._restatement_intervals_across_all_environments(
-                plan.restatements, plan.disabled_restatement_models
+                prod_restatements=plan.restatements,
+                disable_restatement_models=plan.disabled_restatement_models,
+                loaded_snapshots={s.snapshot_id: s for s in snapshots_by_name.values()},
             )
         )
 
@@ -448,7 +450,10 @@ def _restate(self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapsho
         )
 
     def _restatement_intervals_across_all_environments(
-        self, prod_restatements: t.Dict[str, Interval], disable_restatement_models: t.Set[str]
+        self,
+        prod_restatements: t.Dict[str, Interval],
+        disable_restatement_models: t.Set[str],
+        loaded_snapshots: t.Dict[SnapshotId, Snapshot],
     ) -> t.Set[t.Tuple[SnapshotTableInfo, Interval]]:
         """
         Given a map of snapshot names + intervals to restate in prod:
@@ -462,7 +467,7 @@ def _restatement_intervals_across_all_environments(
         if not prod_restatements:
             return set()
 
-        snapshots_to_restate: t.Set[t.Tuple[SnapshotTableInfo, Interval]] = set()
+        snapshots_to_restate: t.Dict[SnapshotId, t.Tuple[SnapshotTableInfo, Interval]] = {}
 
         for env in self.state_sync.get_environments():
             keyed_snapshots = {s.name: s.table_info for s in env.snapshots}
@@ -480,10 +485,47 @@ def _restatement_intervals_across_all_environments(
                     if x not in disable_restatement_models
                 ]
                 snapshots_to_restate.update(
-                    {(keyed_snapshots[a], intervals) for a in affected_snapshot_names}
+                    {
+                        keyed_snapshots[a].snapshot_id: (keyed_snapshots[a], intervals)
+                        for a in affected_snapshot_names
+                    }
+                )
+
+        # for any affected full_history_restatement_only snapshots, we need to widen the intervals being restated to
+        # include the whole time range for that snapshot. This requires a call to state to load the full snapshot record,
+        # so we only do it if necessary
+        if full_history_restatement_snapshot_ids := [
+            # FIXME: full_history_restatement_only is just one indicator that the snapshot can only be fully refreshed, the other one is Model.depends_on_self
+            # however, to figure out depends_on_self, we have to render all the model queries which, alongside having to fetch full snapshots from state,
+            # is problematic in secure environments that are deliberately isolated from arbitrary user code (since rendering a query may require user macros to be present)
+            # So for now, these are not considered
+            s_id
+            for s_id, s in snapshots_to_restate.items()
+            if s[0].full_history_restatement_only
+        ]:
+            # only load full snapshot records that we havent already loaded
+            additional_snapshots = self.state_sync.get_snapshots(
+                [
+                    s.snapshot_id
+                    for s in full_history_restatement_snapshot_ids
+                    if s.snapshot_id not in loaded_snapshots
+                ]
+            )
+
+            all_snapshots = loaded_snapshots | additional_snapshots
+
+            for full_snapshot_id in full_history_restatement_snapshot_ids:
+                full_snapshot = all_snapshots[full_snapshot_id]
+                _, original_intervals = snapshots_to_restate[full_snapshot_id]
+                original_start, original_end = original_intervals
+
+                # get_removal_interval() widens intervals if necessary
+                new_intervals = full_snapshot.get_removal_interval(
+                    start=original_start, end=original_end
                 )
+                snapshots_to_restate[full_snapshot_id] = (full_snapshot.table_info, new_intervals)
 
-        return snapshots_to_restate
+        return set(snapshots_to_restate.values())
 
 
 class BaseAirflowPlanEvaluator(PlanEvaluator):

tests/core/test_integration.py

@@ -3027,6 +3027,143 @@ def _dates_in_table(table_name: str) -> t.List[str]:
         ]
 
 
+def test_prod_restatement_plan_causes_dev_intervals_to_be_widened_on_full_restatement_only_model(
+    tmp_path,
+):
+    """
+    Scenario:
+        I have am INCREMENTAL_BY_TIME_RANGE model A[daily] in prod
+        I create dev and add a INCREMENTAL_BY_UNIQUE_KEY model B (which supports full restatement only)
+        I prod, I restate one day of A which should cause intervals in dev to be cleared (but not processed)
+        In dev, I run a plan
+
+    Outcome:
+        In the dev plan, the entire model for B should be rebuilt because it does not support partial restatement
+    """
+
+    model_a = """
+    MODEL (
+        name test.a,
+        kind INCREMENTAL_BY_TIME_RANGE (
+            time_column "ts"
+        ),
+        start '2024-01-01 00:00:00',
+        cron '@daily'
+    );
+
+    select account_id, ts from test.external_table where ts between @start_ts and @end_ts;
+    """
+
+    model_b = """
+        MODEL (
+            name test.b,
+            kind INCREMENTAL_BY_UNIQUE_KEY (
+                unique_key (account_id, ts)
+            ),
+            cron '@daily'
+        );
+
+        select account_id, ts from test.a where ts between @start_ts and @end_ts;
+        """
+
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+
+    with open(models_dir / "a.sql", "w") as f:
+        f.write(model_a)
+
+    config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
+    ctx = Context(paths=[tmp_path], config=config)
+
+    engine_adapter = ctx.engine_adapter
+    engine_adapter.create_schema("test")
+
+    # source data
+    df = pd.DataFrame(
+        {
+            "account_id": [1001, 1002, 1003, 1004],
+            "ts": [
+                "2024-01-01 00:30:00",
+                "2024-01-02 01:30:00",
+                "2024-01-03 02:30:00",
+                "2024-01-04 00:30:00",
+            ],
+        }
+    )
+    columns_to_types = {
+        "account_id": exp.DataType.build("int"),
+        "ts": exp.DataType.build("timestamp"),
+    }
+    external_table = exp.table_(table="external_table", db="test", quoted=True)
+    engine_adapter.create_table(table_name=external_table, columns_to_types=columns_to_types)
+    engine_adapter.insert_append(
+        table_name=external_table, query_or_df=df, columns_to_types=columns_to_types
+    )
+
+    # plan + apply A[daily] in prod
+    ctx.plan(auto_apply=True)
+
+    # add B[daily] in dev
+    with open(models_dir / "b.sql", "w") as f:
+        f.write(model_b)
+
+    # plan + apply dev
+    ctx.load()
+    ctx.plan(environment="dev", auto_apply=True)
+
+    def _dates_in_table(table_name: str) -> t.List[str]:
+        return [
+            str(r[0]) for r in engine_adapter.fetchall(f"select ts from {table_name} order by ts")
+        ]
+
+    # verify initial state
+    for tbl in ["test.a", "test__dev.b"]:
+        assert _dates_in_table(tbl) == [
+            "2024-01-01 00:30:00",
+            "2024-01-02 01:30:00",
+            "2024-01-03 02:30:00",
+            "2024-01-04 00:30:00",
+        ]
+
+    # restate A in prod
+    engine_adapter.execute("delete from test.external_table where ts = '2024-01-02 01:30:00'")
+    ctx.plan(
+        restate_models=["test.a"],
+        start="2024-01-02 00:00:00",
+        end="2024-01-03 00:00:00",
+        auto_apply=True,
+        no_prompts=True,
+    )
+
+    # verify result
+    assert _dates_in_table("test.a") == [
+        "2024-01-01 00:30:00",
+        "2024-01-03 02:30:00",
+        "2024-01-04 00:30:00",
+    ]
+
+    # dev shouldnt have been affected yet
+    assert _dates_in_table("test__dev.b") == [
+        "2024-01-01 00:30:00",
+        "2024-01-02 01:30:00",
+        "2024-01-03 02:30:00",
+        "2024-01-04 00:30:00",
+    ]
+
+    # plan dev which should trigger the missing intervals to get repopulated
+    ctx.plan(environment="dev", auto_apply=True)
+
+    # dev should have fully refreshed
+    # this is proven by the fact that INCREMENTAL_BY_UNIQUE_KEY cant propagate deletes, so if the
+    # model was not fully rebuilt, the deleted record would still be present
+    for tbl in ["test.a", "test__dev.b"]:
+        assert _dates_in_table(tbl) == [
+            "2024-01-01 00:30:00",
+            "2024-01-03 02:30:00",
+            "2024-01-04 00:30:00",
+        ]
+
+
 def test_prod_restatement_plan_missing_model_in_dev(
     tmp_path: Path,
 ):