Move steps into a separate module

Author: izeigerman

sqlmesh/core/plan/explainer.py

@@ -1,33 +1,30 @@
 import abc
 import typing as t
+import logging
 
-from dataclasses import dataclass
 from rich.console import Console as RichConsole
 from rich.tree import Tree
 from sqlglot.dialects.dialect import DialectType
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, TerminalConsole, get_console
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.plan.definition import EvaluatablePlan, SnapshotIntervals
+from sqlmesh.core.plan import steps
 from sqlmesh.core.plan.evaluator import (
     PlanEvaluator,
-    get_audit_only_snapshots,
-    get_snapshots_to_create,
 )
 from sqlmesh.core.state_sync import StateReader
-from sqlmesh.core.scheduler import merged_missing_intervals, SnapshotToIntervals
 from sqlmesh.core.snapshot.definition import (
-    DeployabilityIndex,
-    Snapshot,
     SnapshotInfoMixin,
-    SnapshotTableInfo,
-    Interval,
 )
 from sqlmesh.utils import Verbosity, rich as srich
 from sqlmesh.utils.date import to_ts
 from sqlmesh.utils.errors import SQLMeshError
 
 
+logger = logging.getLogger(__name__)
+
+
 class PlanExplainer(PlanEvaluator):
     def __init__(
         self,
@@ -42,244 +39,16 @@ def __init__(
     def evaluate(
         self, plan: EvaluatablePlan, circuit_breaker: t.Optional[t.Callable[[], bool]] = None
     ) -> None:
-        new_snapshots = {s.snapshot_id: s for s in plan.new_snapshots}
-        stored_snapshots = self.state_reader.get_snapshots(plan.environment.snapshots)
-        snapshots = {**new_snapshots, **stored_snapshots}
-        snapshots_by_name = {s.name: s for s in snapshots.values()}
-
-        all_selected_for_backfill_snapshots = {
-            s.snapshot_id for s in snapshots.values() if plan.is_selected_for_backfill(s.name)
-        }
-
-        deployability_index = DeployabilityIndex.create(snapshots, start=plan.start)
-        deployability_index_for_creation = deployability_index
-        if plan.is_dev:
-            before_promote_snapshots = all_selected_for_backfill_snapshots
-            after_promote_snapshots = set()
-            snapshots_with_schema_migration = []
-        else:
-            before_promote_snapshots = {
-                s.snapshot_id
-                for s in snapshots.values()
-                if deployability_index.is_representative(s)
-                and plan.is_selected_for_backfill(s.name)
-            }
-            after_promote_snapshots = all_selected_for_backfill_snapshots - before_promote_snapshots
-            deployability_index = DeployabilityIndex.all_deployable()
-
-            snapshots_with_schema_migration = [
-                s
-                for s in snapshots.values()
-                if s.is_paused
-                and s.is_materialized
-                and not deployability_index_for_creation.is_representative(s)
-            ]
-
-        snapshots_to_intervals = self._missing_intervals(
-            plan, snapshots_by_name, deployability_index
-        )
-
-        steps: t.List[PlanStep] = []
-
-        before_all = [
-            statement
-            for environment_statements in plan.environment_statements or []
-            for statement in environment_statements.before_all
-        ]
-        if before_all:
-            steps.append(BeforeAllStep(statements=before_all))
-
-        snapshots_to_create = [
-            s
-            for s in get_snapshots_to_create(plan, snapshots)
-            if s in snapshots_to_intervals and s.is_model and not s.is_symbolic
-        ]
-        if snapshots_to_create:
-            steps.append(
-                PhysicalLayerUpdateStep(
-                    snapshots=snapshots_to_create,
-                    deployability_index=deployability_index_for_creation,
-                )
-            )
-
-        audit_only_snapshots = get_audit_only_snapshots(new_snapshots, self.state_reader)
-        if audit_only_snapshots:
-            steps.append(AuditOnlyRunStep(snapshots=list(audit_only_snapshots.values())))
-
-        if plan.restatements and not plan.is_dev:
-            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
-            steps.append(RestatementStep(snapshot_intervals=snapshot_intervals_to_restate))
-
-        if before_promote_snapshots and not plan.empty_backfill and not plan.skip_backfill:
-            missing_intervals_before_promote = {
-                s: i
-                for s, i in snapshots_to_intervals.items()
-                if s.snapshot_id in before_promote_snapshots
-            }
-            if missing_intervals_before_promote:
-                steps.append(
-                    BackfillStep(
-                        snapshot_to_intervals=missing_intervals_before_promote,
-                        deployability_index=deployability_index,
-                    )
-                )
-
-        steps.append(UpdateEnvironmentRecordStep())
-
-        if snapshots_with_schema_migration:
-            steps.append(MigrateSchemasStep(snapshots=snapshots_with_schema_migration))
-
-        if after_promote_snapshots and not plan.empty_backfill and not plan.skip_backfill:
-            missing_intervals_after_promote = {
-                s: i
-                for s, i in snapshots_to_intervals.items()
-                if s.snapshot_id in after_promote_snapshots
-            }
-            if missing_intervals_after_promote:
-                steps.append(
-                    BackfillStep(
-                        snapshot_to_intervals=missing_intervals_after_promote,
-                        deployability_index=deployability_index,
-                    )
-                )
-
-        promoted_snapshots, demoted_snapshots = self._get_promoted_demoted_snapshots(plan)
-        if promoted_snapshots or demoted_snapshots:
-            steps.append(
-                UpdateVirtualLayerStep(
-                    promoted_snapshots=promoted_snapshots,
-                    demoted_snapshots=demoted_snapshots,
-                    deployability_index=deployability_index,
-                )
-            )
-
-        after_all = [
-            statement
-            for environment_statements in plan.environment_statements or []
-            for statement in environment_statements.after_all
-        ]
-        if after_all:
-            steps.append(AfterAllStep(statements=after_all))
-
+        plan_steps = steps.build_plan_steps(plan, self.state_reader, self.default_catalog)
         explainer_console = _get_explainer_console(
             self.console, plan.environment, self.default_catalog
         )
-        explainer_console.explain(steps)
-
-    def _get_promoted_demoted_snapshots(
-        self, plan: EvaluatablePlan
-    ) -> t.Tuple[t.Set[SnapshotTableInfo], t.Set[SnapshotTableInfo]]:
-        existing_environment = self.state_reader.get_environment(plan.environment.name)
-        if existing_environment:
-            snapshots_by_name = {s.name: s for s in existing_environment.snapshots}
-            demoted_snapshot_names = {s.name for s in existing_environment.promoted_snapshots} - {
-                s.name for s in plan.environment.promoted_snapshots
-            }
-            demoted_snapshots = {snapshots_by_name[name] for name in demoted_snapshot_names}
-        else:
-            demoted_snapshots = set()
-        promoted_snapshots = set(plan.environment.promoted_snapshots)
-        if existing_environment and plan.environment.can_partially_promote(existing_environment):
-            promoted_snapshots -= set(existing_environment.promoted_snapshots)
-
-        def _snapshot_filter(snapshot: SnapshotTableInfo) -> bool:
-            return snapshot.is_model and not snapshot.is_symbolic
-
-        return {s for s in promoted_snapshots if _snapshot_filter(s)}, {
-            s for s in demoted_snapshots if _snapshot_filter(s)
-        }
-
-    def _missing_intervals(
-        self,
-        plan: EvaluatablePlan,
-        snapshots_by_name: t.Dict[str, Snapshot],
-        deployability_index: DeployabilityIndex,
-    ) -> SnapshotToIntervals:
-        return merged_missing_intervals(
-            snapshots=snapshots_by_name.values(),
-            start=plan.start,
-            end=plan.end,
-            execution_time=plan.execution_time,
-            restatements={
-                snapshots_by_name[name].snapshot_id: interval
-                for name, interval in plan.restatements.items()
-            },
-            deployability_index=deployability_index,
-            end_bounded=plan.end_bounded,
-            interval_end_per_model=plan.interval_end_per_model,
-        )
-
-
-@dataclass
-class BeforeAllStep:
-    statements: t.List[str]
-
-
-@dataclass
-class AfterAllStep:
-    statements: t.List[str]
-
-
-@dataclass
-class PhysicalLayerUpdateStep:
-    snapshots: t.List[Snapshot]
-    deployability_index: DeployabilityIndex
-
-
-@dataclass
-class AuditOnlyRunStep:
-    snapshots: t.List[Snapshot]
-
-
-@dataclass
-class RestatementStep:
-    snapshot_intervals: t.Dict[SnapshotTableInfo, Interval]
-
-
-@dataclass
-class BackfillStep:
-    snapshot_to_intervals: SnapshotToIntervals
-    deployability_index: DeployabilityIndex
-    before_promote: bool = True
-
-
-@dataclass
-class MigrateSchemasStep:
-    snapshots: t.List[Snapshot]
-
-
-@dataclass
-class UpdateVirtualLayerStep:
-    promoted_snapshots: t.Set[SnapshotTableInfo]
-    demoted_snapshots: t.Set[SnapshotTableInfo]
-    deployability_index: DeployabilityIndex
-
-
-@dataclass
-class UpdateEnvironmentRecordStep:
-    pass
-
-
-PlanStep = t.Union[
-    BeforeAllStep,
-    AfterAllStep,
-    PhysicalLayerUpdateStep,
-    AuditOnlyRunStep,
-    RestatementStep,
-    BackfillStep,
-    MigrateSchemasStep,
-    UpdateVirtualLayerStep,
-    UpdateEnvironmentRecordStep,
-]
+        explainer_console.explain(plan_steps)
 
 
 class ExplainerConsole(abc.ABC):
     @abc.abstractmethod
-    def explain(self, steps: t.List[PlanStep]) -> None:
+    def explain(self, steps: t.List[steps.PlanStep]) -> None:
         pass
 
 
@@ -301,31 +70,35 @@ def __init__(
         self.verbosity = verbosity
         self.console: RichConsole = console or srich.console
 
-    def explain(self, steps: t.List[PlanStep]) -> None:
+    def explain(self, steps: t.List[steps.PlanStep]) -> None:
         tree = Tree("[bold]Explained plan[/bold]")
         for step in steps:
             handler_name = f"visit_{_to_snake_case(step.__class__.__name__)}"
             if not hasattr(self, handler_name):
-                raise SQLMeshError(f"Unexpected step: {step.__class__.__name__}")
+                logger.error("Unexpected step: %s", step.__class__.__name__)
+                continue
             handler = getattr(self, handler_name)
             result = handler(step)
             if result:
                 tree.add(self._limit_tree(result))
         self.console.print(tree)
 
-    def visit_before_all_step(self, step: BeforeAllStep) -> Tree:
+    def visit_before_all_step(self, step: steps.BeforeAllStep) -> Tree:
         tree = Tree("[bold]Execute before all statements[/bold]")
         for statement in step.statements:
             tree.add(statement)
         return tree
 
-    def visit_after_all_step(self, step: AfterAllStep) -> Tree:
+    def visit_after_all_step(self, step: steps.AfterAllStep) -> Tree:
         tree = Tree("[bold]Execute after all statements[/bold]")
         for statement in step.statements:
             tree.add(statement)
         return tree
 
-    def visit_physical_layer_update_step(self, step: PhysicalLayerUpdateStep) -> Tree:
+    def visit_physical_layer_update_step(self, step: steps.PhysicalLayerUpdateStep) -> Tree:
+        if not step.snapshots:
+            return Tree("[bold]SKIP: No physical layer updates to perform[/bold]")
+
         tree = Tree("[bold]Validate SQL and create physical tables if they do not exist[/bold]")
         for snapshot in step.snapshots:
             is_deployable = (
@@ -363,21 +136,24 @@ def visit_physical_layer_update_step(self, step: PhysicalLayerUpdateStep) -> Tre
             tree.add(model_tree)
         return tree
 
-    def visit_audit_only_run_step(self, step: AuditOnlyRunStep) -> Tree:
+    def visit_audit_only_run_step(self, step: steps.AuditOnlyRunStep) -> Tree:
         tree = Tree("[bold]Audit-only execution[/bold]")
         for snapshot in step.snapshots:
             display_name = self._display_name(snapshot)
             tree.add(display_name)
         return tree
 
-    def visit_restatement_step(self, step: RestatementStep) -> Tree:
+    def visit_restatement_step(self, step: steps.RestatementStep) -> Tree:
         tree = Tree("[bold]Invalidate data intervals as part of restatement[/bold]")
         for snapshot_table_info, interval in step.snapshot_intervals.items():
             display_name = self._display_name(snapshot_table_info)
             tree.add(f"{display_name} [{to_ts(interval[0])} - {to_ts(interval[1])}]")
         return tree
 
-    def visit_backfill_step(self, step: BackfillStep) -> Tree:
+    def visit_backfill_step(self, step: steps.BackfillStep) -> Tree:
+        if not step.snapshot_to_intervals:
+            return Tree("[bold]SKIP: No model batches to execute[/bold]")
+
         tree = Tree(
             "[bold]Backfill models by running their queries and run standalone audits[/bold]"
         )
@@ -426,7 +202,7 @@ def visit_backfill_step(self, step: BackfillStep) -> Tree:
                 tree.add(f"{display_name} \[standalone audit]")
         return tree
 
-    def visit_migrate_schemas_step(self, step: MigrateSchemasStep) -> Tree:
+    def visit_migrate_schemas_step(self, step: steps.MigrateSchemasStep) -> Tree:
         tree = Tree(
             "[bold]Update schemas (add, drop, alter columns) of production physical tables to reflect forward-only changes[/bold]"
         )
@@ -436,7 +212,7 @@ def visit_migrate_schemas_step(self, step: MigrateSchemasStep) -> Tree:
             tree.add(f"{display_name} -> {table_name}")
         return tree
 
-    def visit_update_virtual_layer_step(self, step: UpdateVirtualLayerStep) -> Tree:
+    def visit_virtual_layer_update_step(self, step: steps.VirtualLayerUpdateStep) -> Tree:
         tree = Tree(
             f"[bold]Update the virtual layer for environment '{self.environment_naming_info.name}'[/bold]"
         )
@@ -461,8 +237,8 @@ def visit_update_virtual_layer_step(self, step: UpdateVirtualLayerStep) -> Tree:
             tree.add(self._limit_tree(demote_tree))
         return tree
 
-    def visit_update_environment_record_step(
-        self, step: UpdateEnvironmentRecordStep
+    def visit_environment_record_update_step(
+        self, step: steps.EnvironmentRecordUpdateStep
     ) -> t.Optional[Tree]:
         return None