PR feedback

Author: erindru

sqlmesh/core/plan/explainer.py

@@ -59,7 +59,12 @@ def evaluate(
 
         # add extra metadata that's only needed at this point for better --explain output
         plan_stages = [
-            ExplainableRestatementStage.from_restatement_stage(stage, self.state_reader, plan)
+            ExplainableRestatementStage.from_restatement_stage(
+                stage,
+                self.state_reader,
+                plan,
+                fetch_full_snapshots=explainer_console.verbosity == Verbosity.VERY_VERBOSE,
+            )
             if isinstance(stage, stages.RestatementStage)
             else stage
             for stage in plan_stages
@@ -73,6 +78,11 @@ class ExplainerConsole(abc.ABC):
     def explain(self, stages: t.List[stages.PlanStage]) -> None:
         pass
 
+    @property
+    @abc.abstractmethod
+    def verbosity(self) -> Verbosity:
+        pass
+
 
 @dataclass
 class ExplainableRestatementStage(stages.RestatementStage):
@@ -82,7 +92,7 @@ class ExplainableRestatementStage(stages.RestatementStage):
     """
 
     snapshot_intervals_to_clear: t.Dict[
-        str, t.List[t.Tuple[Snapshot, SnapshotIntervalClearRequest]]
+        str, t.List[t.Tuple[t.Optional[Snapshot], SnapshotIntervalClearRequest]]
     ]
     """Which snapshots from other environments would have intervals cleared as part of restatement, grouped by name."""
 
@@ -92,6 +102,7 @@ def from_restatement_stage(
         stage: stages.RestatementStage,
         state_reader: StateReader,
         plan: EvaluatablePlan,
+        fetch_full_snapshots: bool = False,
     ) -> ExplainableRestatementStage:
         loaded_snapshots = {s.snapshot_id: s for s in stage.all_snapshots.values()}
 
@@ -102,17 +113,20 @@ def from_restatement_stage(
             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)
-        )
+        if fetch_full_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)
+                (loaded_snapshots.get(snapshot_id), clear_request)
             )
 
         return cls(
@@ -136,9 +150,13 @@ def __init__(
         self.environment_naming_info = environment_naming_info
         self.dialect = dialect
         self.default_catalog = default_catalog
-        self.verbosity = verbosity
+        self._verbosity = verbosity
         self.console: RichConsole = console or srich.console
 
+    @property
+    def verbosity(self) -> Verbosity:
+        return self._verbosity
+
     def explain(self, stages: t.List[stages.PlanStage]) -> None:
         tree = Tree("[bold]Explained plan[/bold]")
         for stage in stages:
@@ -228,41 +246,49 @@ def visit_restatement_stage(
             snapshot_intervals := stage.snapshot_intervals_to_clear
         ):
             for name, requests in snapshot_intervals.items():
+                if not requests:
+                    # ensure that there is at least one SnapshotIntervalClearRequest in the list
+                    continue
+
                 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
-
+                _, clear_request = requests[0]
+                interval_start, interval_end = clear_request.interval
                 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}")
+                if not self.verbosity.is_very_verbose:
+                    # In normal mode we just indicate which environments are affected at a high level
+                    all_environment_names = sorted(
+                        set(env_name for (_, cr) in requests for env_name in cr.environment_names)
+                    )
+                    node.add("in environments: " + ", ".join(all_environment_names))
+                else:
+                    # In "very verbose" mode, we print all the affected physical tables
+
+                    # group by environment for the console output
+                    by_environment: t.Dict[t.Optional[str], t.List[Snapshot]] = defaultdict(list)
+
+                    for snapshot, clear_request in requests:
+                        if not snapshot:
+                            # Snapshots are None (not loaded) unless
+                            continue
+                        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)
+
+                    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/core/test_plan_stages.py

@@ -765,7 +765,7 @@ def _get_snapshots(snapshot_ids: t.Iterable[SnapshotIdLike]):
     restatement_stage = stages[2]
     assert isinstance(restatement_stage, RestatementStage)
     restatement_stage = ExplainableRestatementStage.from_restatement_stage(
-        restatement_stage, state_reader, plan
+        restatement_stage, state_reader, plan, fetch_full_snapshots=True
     )
 
     # note: we only clear the intervals from state for "a" in dev, we leave prod alone
@@ -774,6 +774,7 @@ def _get_snapshots(snapshot_ids: t.Iterable[SnapshotIdLike]):
     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 full_snapshot is not None
     assert isinstance(clear_request, SnapshotIntervalClearRequest)
     assert snapshot_name == '"a"'
     assert full_snapshot.name == snapshot_name