Feat: prevent other processes seeing missing intervals during restatement

Author: erindru

sqlmesh/core/plan/evaluator.py

@@ -284,31 +284,34 @@ def visit_audit_only_run_stage(
     def visit_restatement_stage(
         self, stage: stages.RestatementStage, plan: EvaluatablePlan
     ) -> None:
-        snapshot_intervals_to_restate = {(s, i) for s, i in stage.snapshot_intervals.items()}
-
-        # Restating intervals on prod plans should mean that the intervals are cleared across
-        # all environments, not just the version currently in prod
+        # Restating intervals on prod plans means that once the data for the intervals being restated has been refreshed
+        # (which happens in the backfill stage) then we need to clear those intervals *from state* across all other environments.
+        #
         # This ensures that work done in dev environments can still be promoted to prod
-        # by forcing dev environments to re-run intervals that changed in prod
+        # by forcing dev environments to re-run intervals that changed in prod (because after this stage runs they show as missing)
+        #
+        # It also means that any new dev environments created while this restatement plan was running also get the
+        # correct intervals cleared because we look up matching snapshots as at right now and not as at the time the plan
+        # was created, which could have been several hours ago if there was a lot of data to restate.
         #
         # Without this rule, its possible that promoting a dev table to prod will introduce old data to prod
-        snapshot_intervals_to_restate.update(
-            {
-                (s.table_info, s.interval)
-                for s in identify_restatement_intervals_across_snapshot_versions(
-                    state_reader=self.state_sync,
-                    prod_restatements=plan.restatements,
-                    disable_restatement_models=plan.disabled_restatement_models,
-                    loaded_snapshots={s.snapshot_id: s for s in stage.all_snapshots.values()},
-                    current_ts=to_timestamp(plan.execution_time or now()),
-                ).values()
-            }
-        )
-
-        self.state_sync.remove_intervals(
-            snapshot_intervals=list(snapshot_intervals_to_restate),
-            remove_shared_versions=plan.is_prod,
-        )
+        snapshot_intervals_to_restate = {
+            (s.table_info, s.interval)
+            for s in identify_restatement_intervals_across_snapshot_versions(
+                state_reader=self.state_sync,
+                prod_restatements=plan.restatements,
+                # TODO: we need to ensure that only dev environments are affected
+                disable_restatement_models=plan.disabled_restatement_models,
+                loaded_snapshots={s.snapshot_id: s for s in stage.all_snapshots.values()},
+                current_ts=to_timestamp(plan.execution_time or now()),
+            ).values()
+        }
+
+        if snapshot_intervals_to_restate:
+            self.state_sync.remove_intervals(
+                snapshot_intervals=list(snapshot_intervals_to_restate),
+                remove_shared_versions=plan.is_prod,
+            )
 
     def visit_environment_record_update_stage(
         self, stage: stages.EnvironmentRecordUpdateStage, plan: EvaluatablePlan

sqlmesh/core/plan/explainer.py

@@ -41,7 +41,9 @@ def evaluate(
         plan: EvaluatablePlan,
         circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
     ) -> None:
-        plan_stages = stages.build_plan_stages(plan, self.state_reader, self.default_catalog)
+        plan_stages = stages.build_plan_stages(
+            plan, self.state_reader, self.default_catalog, explain=True
+        )
         explainer_console = _get_explainer_console(
             self.console, plan.environment, self.default_catalog
         )
@@ -148,9 +150,11 @@ def visit_audit_only_run_stage(self, stage: stages.AuditOnlyRunStage) -> Tree:
 
     def visit_restatement_stage(self, stage: stages.RestatementStage) -> Tree:
         tree = Tree("[bold]Invalidate data intervals as part of restatement[/bold]")
-        for snapshot_table_info, interval in stage.snapshot_intervals.items():
-            display_name = self._display_name(snapshot_table_info)
-            tree.add(f"{display_name} [{to_ts(interval[0])} - {to_ts(interval[1])}]")
+        if snapshot_intervals := stage.snapshot_intervals_to_clear:
+            for clear_request in snapshot_intervals.values():
+                display_name = self._display_name(clear_request.table_info)
+                interval = clear_request.interval
+                tree.add(f"{display_name} [{to_ts(interval[0])} - {to_ts(interval[1])}]")
         return tree
 
     def visit_backfill_stage(self, stage: stages.BackfillStage) -> Tree:

sqlmesh/core/plan/stages.py

@@ -3,7 +3,11 @@
 from dataclasses import dataclass
 from sqlmesh.core import constants as c
 from sqlmesh.core.environment import EnvironmentStatements, EnvironmentNamingInfo, Environment
-from sqlmesh.core.plan.common import should_force_rebuild
+from sqlmesh.core.plan.common import (
+    should_force_rebuild,
+    identify_restatement_intervals_across_snapshot_versions,
+    SnapshotIntervalClearRequest,
+)
 from sqlmesh.core.plan.definition import EvaluatablePlan
 from sqlmesh.core.state_sync import StateReader
 from sqlmesh.core.scheduler import merged_missing_intervals, SnapshotToIntervals
@@ -12,7 +16,6 @@
     Snapshot,
     SnapshotTableInfo,
     SnapshotId,
-    Interval,
 )
 
 
@@ -98,16 +101,22 @@ class AuditOnlyRunStage:
 
 @dataclass
 class RestatementStage:
-    """Restate intervals for given snapshots.
+    """Clear intervals from state for given snapshots.
 
     Args:
-        snapshot_intervals: Intervals to restate.
+        snapshot_intervals_to_clear: Intervals to clear from state for snapshots in dev environments, keyed by snapshot name
+        deployability_index: Deployability of the snapshots in :snapshot_intervals_to_clear, used for
+            expaining exactly which physical tables are affected. Note that these snapshots are from
+            other environments that are not present in this plan, so we cant use the plan DeployabilityIndex
         all_snapshots: All snapshots in the plan by name.
     """
 
-    snapshot_intervals: t.Dict[SnapshotTableInfo, Interval]
     all_snapshots: t.Dict[str, Snapshot]
 
+    # Only used for --explain so may not be populated
+    snapshot_intervals_to_clear: t.Optional[t.Dict[str, SnapshotIntervalClearRequest]]
+    deployability_index: t.Optional[DeployabilityIndex]
+
 
 @dataclass
 class BackfillStage:
@@ -219,15 +228,17 @@ class PlanStagesBuilder:
     Args:
         state_reader: The state reader to use to read the snapshots and environment.
         default_catalog: The default catalog to use for the snapshots.
+        explain: Whether the stages are bing built for PlanExplainer or for the PlanExecutor.
+            This allows the ability to conditionally attach additional metadata to each stage that can be helpful for a user,
+            but not necessarily required in order to execute the plan.
     """
 
     def __init__(
-        self,
-        state_reader: StateReader,
-        default_catalog: t.Optional[str],
+        self, state_reader: StateReader, default_catalog: t.Optional[str], explain: bool = False
     ):
         self.state_reader = state_reader
         self.default_catalog = default_catalog
+        self.explain = explain
 
     def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
         """Builds the plan stages for the given plan.
@@ -321,10 +332,6 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
         if audit_only_snapshots:
             stages.append(AuditOnlyRunStage(snapshots=list(audit_only_snapshots.values())))
 
-        restatement_stage = self._get_restatement_stage(plan, snapshots_by_name)
-        if restatement_stage:
-            stages.append(restatement_stage)
-
         if missing_intervals_before_promote:
             stages.append(
                 BackfillStage(
@@ -349,6 +356,15 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
                 )
             )
 
+        # note: "restatement stage" (which is clearing intervals in state - not actually performing the restatements, that's the backfill stage)
+        # needs to come *after* the backfill stage so that at no time do other plans / runs see empty prod intervals and compete with this plan to try to fill them.
+        # in addition, when we update intervals in state, we only clear intervals from dev snapshots to force dev models to be backfilled based on the new prod data.
+        # we can leave prod intervals alone because by the time this plan finishes, the intervals in state have not actually changed, since restatement replaces
+        # data for existing intervals and does not produce new ones
+        restatement_stage = self._get_restatement_stage(plan, snapshots_by_name)
+        if restatement_stage:
+            stages.append(restatement_stage)
+
         stages.append(
             EnvironmentRecordUpdateStage(
                 no_gaps_snapshot_names={s.name for s in before_promote_snapshots}
@@ -443,15 +459,50 @@ def _get_after_all_stage(
     def _get_restatement_stage(
         self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapshot]
     ) -> t.Optional[RestatementStage]:
-        snapshot_intervals_to_restate = {}
-        for name, interval in plan.restatements.items():
-            restated_snapshot = snapshots_by_name[name]
-            restated_snapshot.remove_interval(interval)
-            snapshot_intervals_to_restate[restated_snapshot.table_info] = interval
-        if not snapshot_intervals_to_restate or plan.is_dev:
+        if not plan.restatements or plan.is_dev:
+            # The RestatementStage to clear intervals from state across all environments is not needed for plans against dev, only prod
             return None
+
+        snapshot_intervals_to_clear: t.Optional[t.Dict[str, SnapshotIntervalClearRequest]] = None
+        deployability_index: t.Optional[DeployabilityIndex] = None
+
+        if self.explain:
+            # This is re-calculated in the PlanEvaluator because there can be a time lag between:
+            #  - a plan being generated and a plan being executed
+            #  - restatement backfill starting and restatement backfill completing
+            # During this time, someone may create a new dev environment based on partially restated data.
+            # So we look up the intervals to clear in other environments at the last minute after restatement has occurred.
+            #
+            # However, if the user wants to know what's going on (via --explain), we look them up here as well so we
+            # can output them to the console as part of the `--explain` output
+            intervals_to_clear = identify_restatement_intervals_across_snapshot_versions(
+                state_reader=self.state_reader,
+                prod_restatements=plan.restatements,
+                disable_restatement_models=plan.disabled_restatement_models,
+                loaded_snapshots={s.snapshot_id: s for s in snapshots_by_name.values()},
+            )
+
+            if intervals_to_clear:
+                snapshot_intervals_to_clear = {
+                    s_id.name: r for s_id, r in intervals_to_clear.items()
+                }
+
+                # creating a deployability index over the "snapshot intervals to clear"
+                # allows us to print the physical names of the tables affected in the console output
+                # note that we can't use the DeployabilityIndex on the plan because it only includes
+                # snapshots for the current environment, not across all environments
+                deployability_index = DeployabilityIndex.create(
+                    snapshots=self.state_reader.get_snapshots(
+                        [s.snapshot_id for s in snapshot_intervals_to_clear.values()]
+                    ),
+                    start=plan.start,
+                    start_override_per_model=plan.start_override_per_model,
+                )
+
         return RestatementStage(
-            snapshot_intervals=snapshot_intervals_to_restate, all_snapshots=snapshots_by_name
+            snapshot_intervals_to_clear=snapshot_intervals_to_clear,
+            deployability_index=deployability_index,
+            all_snapshots=snapshots_by_name,
         )
 
     def _get_physical_layer_update_stage(
@@ -678,5 +729,6 @@ def build_plan_stages(
     plan: EvaluatablePlan,
     state_reader: StateReader,
     default_catalog: t.Optional[str],
+    explain: bool = False,
 ) -> t.List[PlanStage]:
-    return PlanStagesBuilder(state_reader, default_catalog).build(plan)
+    return PlanStagesBuilder(state_reader, default_catalog, explain=explain).build(plan)