Feat: Improve CLI and --explain output for restatements

Author: erindru

sqlmesh/core/console.py

@@ -2022,7 +2022,34 @@ def _prompt_categorize(
         plan = plan_builder.build()
 
         if plan.restatements:
-            self._print("\n[bold]Restating models\n")
+            # A plan can have restatements for the following reasons:
+            # - The user specifically called `sqlmesh plan` with --restate-model.
+            #   This creates a "restatement plan" which disallows all other changes and simply force-backfills
+            #   the selected models and their downstream dependencies using the versions of the models stored in state.
+            # - There are no specific restatements (so changes are allowed) AND dev previews need to be computed.
+            #   The "restatements" feature is currently reused for dev previews.
+            if plan.selected_models_to_restate:
+                # There were legitimate restatements, no dev previews
+                tree = Tree(
+                    "[bold]Models selected for restatement[/bold]\n"
+                    "This causes backfill of the model itself as well as affected downstream models"
+                )
+                model_fqn_to_snapshot = {s.name: s for s in plan.snapshots.values()}
+                for model_fqn in plan.selected_models_to_restate:
+                    snapshot = model_fqn_to_snapshot[model_fqn]
+                    display_name = snapshot.display_name(
+                        plan.environment_naming_info,
+                        default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+                        dialect=self.dialect,
+                    )
+                    tree.add(
+                        display_name
+                    )  # note: we deliberately dont show any intervals here; they get shown in the backfill section
+                self._print(tree)
+            else:
+                # We are computing dev previews, do not confuse the user by printing out something to do
+                # with restatements. Dev previews are already highlighted in the backfill step
+                pass
         else:
             self.show_environment_difference_summary(
                 plan.context_diff,

sqlmesh/core/plan/builder.py

@@ -338,6 +338,7 @@ def build(self) -> Plan:
             directly_modified=directly_modified,
             indirectly_modified=indirectly_modified,
             deployability_index=deployability_index,
+            selected_models_to_restate=self._restate_models,
             restatements=restatements,
             start_override_per_model=self._start_override_per_model,
             end_override_per_model=end_override_per_model,

sqlmesh/core/plan/definition.py

@@ -58,7 +58,16 @@ class Plan(PydanticModel, frozen=True):
     indirectly_modified: t.Dict[SnapshotId, t.Set[SnapshotId]]
 
     deployability_index: DeployabilityIndex
+    selected_models_to_restate: t.Optional[t.Set[str]] = None
+    """Models that have been explicitly selected for restatement by a user"""
     restatements: t.Dict[SnapshotId, Interval]
+    """
+    All models being restated, which are typically the explicitly selected ones + their downstream dependencies.
+    
+    Note that dev previews are also considered restatements, so :selected_models_to_restate can be empty
+    while :restatements is still populated with dev previews
+    """
+
     start_override_per_model: t.Optional[t.Dict[str, datetime]]
     end_override_per_model: t.Optional[t.Dict[str, datetime]]
 

sqlmesh/core/plan/explainer.py

@@ -4,6 +4,7 @@
 import typing as t
 import logging
 from dataclasses import dataclass
+from collections import defaultdict
 
 from rich.console import Console as RichConsole
 from rich.tree import Tree
@@ -22,7 +23,11 @@
     PlanEvaluator,
 )
 from sqlmesh.core.state_sync import StateReader
-from sqlmesh.core.snapshot.definition import SnapshotInfoMixin, SnapshotNameVersionLike
+from sqlmesh.core.snapshot.definition import (
+    Snapshot,
+    SnapshotInfoMixin,
+    SnapshotNameVersionLike,
+)
 from sqlmesh.utils import Verbosity, rich as srich, to_snake_case
 from sqlmesh.utils.date import to_ts
 from sqlmesh.utils.errors import SQLMeshError
@@ -76,8 +81,10 @@ class ExplainableRestatementStage(stages.RestatementStage):
     of what might happen when they ask for the plan to be explained
     """
 
-    snapshot_intervals_to_clear: t.Dict[str, SnapshotIntervalClearRequest]
-    """Which snapshots from other environments would have intervals cleared as part of restatement, keyed by name"""
+    snapshot_intervals_to_clear: t.Dict[
+        str, t.List[t.Tuple[Snapshot, SnapshotIntervalClearRequest]]
+    ]
+    """Which snapshots from other environments would have intervals cleared as part of restatement, grouped by name."""
 
     @classmethod
     def from_restatement_stage(
@@ -86,17 +93,30 @@ def from_restatement_stage(
         state_reader: StateReader,
         plan: EvaluatablePlan,
     ) -> ExplainableRestatementStage:
+        loaded_snapshots = {s.snapshot_id: s for s in stage.all_snapshots.values()}
+
         all_restatement_intervals = identify_restatement_intervals_across_snapshot_versions(
             state_reader=state_reader,
             prod_restatements=plan.restatements,
             disable_restatement_models=plan.disabled_restatement_models,
-            loaded_snapshots={s.snapshot_id: s for s in stage.all_snapshots.values()},
+            loaded_snapshots=loaded_snapshots,
         )
 
+        # extend loaded_snapshots with the remaining full Snapshot objects from all_restatement_intervals
+        # so that we can generate physical table names for them while explaining what's going on
+        remaining_snapshot_ids_to_load = set(all_restatement_intervals).difference(loaded_snapshots)
+        loaded_snapshots.update(
+            state_reader.get_snapshots(snapshot_ids=remaining_snapshot_ids_to_load)
+        )
+
+        snapshot_intervals_to_clear = defaultdict(list)
+        for snapshot_id, clear_request in all_restatement_intervals.items():
+            snapshot_intervals_to_clear[clear_request.snapshot.name].append(
+                (loaded_snapshots[snapshot_id], clear_request)
+            )
+
         return cls(
-            snapshot_intervals_to_clear={
-                s.snapshot.name: s for s in all_restatement_intervals.values()
-            },
+            snapshot_intervals_to_clear=snapshot_intervals_to_clear,
             all_snapshots=stage.all_snapshots,
         )
 
@@ -199,15 +219,50 @@ def visit_explainable_restatement_stage(self, stage: ExplainableRestatementStage
     def visit_restatement_stage(
         self, stage: t.Union[ExplainableRestatementStage, stages.RestatementStage]
     ) -> Tree:
-        tree = Tree("[bold]Invalidate data intervals as part of restatement[/bold]")
+        tree = Tree(
+            "[bold]Invalidate data intervals in state for development environments to prevent old data from being promoted[/bold]\n"
+            "This only affects state and will not clear physical data from the tables until the next plan for each environment"
+        )
 
         if isinstance(stage, ExplainableRestatementStage) and (
             snapshot_intervals := stage.snapshot_intervals_to_clear
         ):
-            for clear_request in snapshot_intervals.values():
-                display_name = self._display_name(clear_request.snapshot)
-                interval = clear_request.interval
-                tree.add(f"{display_name} [{to_ts(interval[0])} - {to_ts(interval[1])}]")
+            for name, requests in snapshot_intervals.items():
+                display_name = model_display_name(
+                    name, self.environment_naming_info, self.default_catalog, self.dialect
+                )
+
+                # group by environment for the console output
+                by_environment: t.Dict[t.Optional[str], t.List[Snapshot]] = defaultdict(list)
+
+                interval_start = None
+                interval_end = None
+
+                for snapshot, clear_request in requests:
+                    # used for the top level tree node
+                    interval_start, interval_end = clear_request.interval
+
+                    if clear_request.sorted_environment_names:
+                        # snapshot is promoted in these environments
+                        for env in clear_request.sorted_environment_names:
+                            by_environment[env].append(snapshot)
+                    else:
+                        # snapshot is not currently promoted in any environment
+                        by_environment[None].append(snapshot)
+
+                if not interval_start or not interval_end:
+                    continue
+
+                node = tree.add(f"{display_name} [{to_ts(interval_start)} - {to_ts(interval_end)}]")
+
+                for env_name, snapshots_to_clear in by_environment.items():
+                    env_name = env_name or "(no env)"
+                    for snapshot in snapshots_to_clear:
+                        # note: we dont need a DeployabilityIndex and can just hardcode is_deployable=True.
+                        # The reason is that non-deployable data can never be restated so we only need to
+                        # bother clearing intervals for the deployable version of the table
+                        physical_table_name = snapshot.table_name(True)
+                        node.add(f"{env_name} -> {physical_table_name}")
 
         return tree
 

tests/cli/test_cli.py

@@ -247,7 +247,7 @@ def test_plan_restate_model(runner, tmp_path):
     )
     assert result.exit_code == 0
     assert_duckdb_test(result)
-    assert "Restating models" in result.output
+    assert "Models selected for restatement" in result.output
     assert "sqlmesh_example.full_model   [full refresh" in result.output
     assert_model_batches_executed(result)
     assert "Virtual layer updated" not in result.output

tests/core/test_plan_stages.py

@@ -771,9 +771,12 @@ def _get_snapshots(snapshot_ids: t.Iterable[SnapshotIdLike]):
     # note: we only clear the intervals from state for "a" in dev, we leave prod alone
     assert restatement_stage.snapshot_intervals_to_clear
     assert len(restatement_stage.snapshot_intervals_to_clear) == 1
-    snapshot_name, clear_request = list(restatement_stage.snapshot_intervals_to_clear.items())[0]
+    snapshot_name, clear_requests = list(restatement_stage.snapshot_intervals_to_clear.items())[0]
+    assert len(clear_requests) == 1
+    full_snapshot, clear_request = clear_requests[0]
     assert isinstance(clear_request, SnapshotIntervalClearRequest)
     assert snapshot_name == '"a"'
+    assert full_snapshot.name == snapshot_name
     assert clear_request.snapshot_id == snapshot_a_dev.snapshot_id
     assert clear_request.snapshot == snapshot_a_dev.id_and_version
     assert clear_request.interval == (to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))