PR feedback

Author: erindru

sqlmesh/core/console.py

@@ -555,7 +555,6 @@ def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
     def log_models_updated_during_restatement(
         self,
         snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
-        environment: EnvironmentSummary,
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
@@ -791,7 +790,6 @@ def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
     def log_models_updated_during_restatement(
         self,
         snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
-        environment: EnvironmentSummary,
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
@@ -2254,13 +2252,12 @@ def log_failed_models(self, errors: t.List[NodeExecutionFailedError]) -> None:
     def log_models_updated_during_restatement(
         self,
         snapshots: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]],
-        environment: EnvironmentSummary,
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str] = None,
     ) -> None:
         if snapshots:
             tree = Tree(
-                f"[yellow]The following models had new versions deployed in plan '{environment.plan_id}' while data was being restated:[/yellow]"
+                f"[yellow]The following models had new versions deployed while data was being restated:[/yellow]"
             )
 
             for restated_snapshot, updated_snapshot in snapshots:
@@ -2274,13 +2271,7 @@ def log_models_updated_during_restatement(
                 current_branch.add(f"currently active version: '{updated_snapshot.version}'")
 
             self._print(tree)
-
-            self.log_warning(
-                f"\nThe '{environment.name}' environment currently points to [bold]different[/bold] versions of these models, not the versions that just got restated."
-            )
-            self._print(
-                "[yellow]If this is undesirable, please re-run this restatement plan which will apply it to the most recent versions of these models.[/yellow]\n"
-            )
+            self._print("")  # newline spacer
 
     def log_destructive_change(
         self,

sqlmesh/core/plan/evaluator.py

@@ -40,7 +40,7 @@
 from sqlmesh.core.plan.common import identify_restatement_intervals_across_snapshot_versions
 from sqlmesh.utils import CorrelationId
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
-from sqlmesh.utils.errors import PlanError, SQLMeshError
+from sqlmesh.utils.errors import PlanError, ConflictingPlanError, SQLMeshError
 from sqlmesh.utils.date import now, to_timestamp
 
 logger = logging.getLogger(__name__)
@@ -311,37 +311,53 @@ def visit_restatement_stage(
             # Nothing to do
             return
 
-        self.state_sync.remove_intervals(
-            snapshot_intervals=[(s.table_info, s.interval) for s in intervals_to_clear.values()],
-            remove_shared_versions=plan.is_prod,
-        )
-
         # While the restatements were being processed, did any of the snapshots being restated get new versions deployed?
         # If they did, they will not reflect the data that just got restated, so we need to notify the user
+        deployed_during_restatement: t.Dict[
+            str, t.Tuple[SnapshotTableInfo, SnapshotTableInfo]
+        ] = {}  # tuple of (restated_snapshot, current_prod_snapshot)
+
         if deployed_env := self.state_sync.get_environment(plan.environment.name):
             promoted_snapshots_by_name = {s.name: s for s in deployed_env.snapshots}
 
-            deployed_during_restatement: t.List[t.Tuple[SnapshotTableInfo, SnapshotTableInfo]] = []
-
             for name in plan.restatements:
                 snapshot = stage.all_snapshots[name]
                 version = snapshot.table_info.version
                 if (
                     prod_snapshot := promoted_snapshots_by_name.get(name)
                 ) and prod_snapshot.version != version:
-                    deployed_during_restatement.append(
-                        (snapshot.table_info, prod_snapshot.table_info)
+                    deployed_during_restatement[name] = (
+                        snapshot.table_info,
+                        prod_snapshot.table_info,
                     )
 
-            if deployed_during_restatement:
-                self.console.log_models_updated_during_restatement(
-                    deployed_during_restatement,
-                    deployed_env.summary,
-                    plan.environment.naming_info,
-                    self.default_catalog,
-                )
-                # note: the plan will automatically fail at the promotion stage with a ConflictingPlanError because the environment was changed by another plan
-                # so there is no need to explicitly fail the plan here
+        # we need to *not* clear the intervals on the snapshots where new versions were deployed while the restatement was running in order to prevent
+        # subsequent plans from having unexpected intervals to backfill.
+        # we instead list the affected models and abort the plan with an error so the user can decide what to do
+        # (either re-attempt the restatement plan or leave things as they are)
+        filtered_intervals_to_clear = [
+            (s.snapshot, s.interval)
+            for s in intervals_to_clear.values()
+            if s.snapshot.name not in deployed_during_restatement
+        ]
+
+        if filtered_intervals_to_clear:
+            # We still clear intervals in other envs for models that were successfully restated without having new versions promoted during restatement
+            self.state_sync.remove_intervals(
+                snapshot_intervals=filtered_intervals_to_clear,
+                remove_shared_versions=plan.is_prod,
+            )
+
+        if deployed_env and deployed_during_restatement:
+            self.console.log_models_updated_during_restatement(
+                list(deployed_during_restatement.values()),
+                plan.environment.naming_info,
+                self.default_catalog,
+            )
+            raise ConflictingPlanError(
+                f"Another plan ({deployed_env.summary.plan_id}) deployed new versions of {len(deployed_during_restatement)} models in the target environment '{plan.environment.name}' while they were being restated by this plan.\n"
+                "Please re-apply your plan if these new versions should be restated."
+            )
 
     def visit_environment_record_update_stage(
         self, stage: stages.EnvironmentRecordUpdateStage, plan: EvaluatablePlan

sqlmesh/core/plan/explainer.py

@@ -11,7 +11,7 @@
 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.snapshot.definition import DeployabilityIndex
+from sqlmesh.core.snapshot.definition import model_display_name
 from sqlmesh.core.plan.common import (
     SnapshotIntervalClearRequest,
     identify_restatement_intervals_across_snapshot_versions,
@@ -22,9 +22,7 @@
     PlanEvaluator,
 )
 from sqlmesh.core.state_sync import StateReader
-from sqlmesh.core.snapshot.definition import (
-    SnapshotInfoMixin,
-)
+from sqlmesh.core.snapshot.definition import 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
@@ -81,10 +79,6 @@ class ExplainableRestatementStage(stages.RestatementStage):
     snapshot_intervals_to_clear: t.Dict[str, SnapshotIntervalClearRequest]
     """Which snapshots from other environments would have intervals cleared as part of restatement, keyed by name"""
 
-    deployability_index: DeployabilityIndex
-    """Deployability of those snapshots (which arent necessarily present in the current plan so we cant use the
-    plan deployability index), used for outputting physical table names"""
-
     @classmethod
     def from_restatement_stage(
         cls: t.Type[ExplainableRestatementStage],
@@ -99,29 +93,10 @@ def from_restatement_stage(
             loaded_snapshots={s.snapshot_id: s for s in stage.all_snapshots.values()},
         )
 
-        snapshot_intervals_to_clear = {}
-        deployability_index = DeployabilityIndex.all_deployable()
-
-        if all_restatement_intervals:
-            snapshot_intervals_to_clear = {
-                s_id.name: r for s_id, r in all_restatement_intervals.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=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 cls(
-            snapshot_intervals_to_clear=snapshot_intervals_to_clear,
-            deployability_index=deployability_index,
+            snapshot_intervals_to_clear={
+                s.snapshot.name: s for s in all_restatement_intervals.values()
+            },
             all_snapshots=stage.all_snapshots,
         )
 
@@ -230,7 +205,7 @@ def visit_restatement_stage(
             snapshot_intervals := stage.snapshot_intervals_to_clear
         ):
             for clear_request in snapshot_intervals.values():
-                display_name = self._display_name(clear_request.table_info)
+                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])}]")
 
@@ -348,15 +323,22 @@ def visit_finalize_environment_stage(
 
     def _display_name(
         self,
-        snapshot: SnapshotInfoMixin,
+        snapshot: t.Union[SnapshotInfoMixin, SnapshotNameVersionLike],
         environment_naming_info: t.Optional[EnvironmentNamingInfo] = None,
     ) -> str:
-        return snapshot.display_name(
-            environment_naming_info or self.environment_naming_info,
-            self.default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+        naming_kwargs: t.Any = dict(
+            environment_naming_info=environment_naming_info or self.environment_naming_info,
+            default_catalog=self.default_catalog
+            if self.verbosity < Verbosity.VERY_VERBOSE
+            else None,
             dialect=self.dialect,
         )
 
+        if isinstance(snapshot, SnapshotInfoMixin):
+            return snapshot.display_name(**naming_kwargs)
+
+        return model_display_name(node_name=snapshot.name, **naming_kwargs)
+
     def _limit_tree(self, tree: Tree) -> Tree:
         tree_length = len(tree.children)
         if tree_length <= MAX_TREE_LENGTH:

sqlmesh/core/snapshot/definition.py

@@ -1788,7 +1788,19 @@ def display_name(
     """
     if snapshot_info_like.is_audit:
         return snapshot_info_like.name
-    view_name = exp.to_table(snapshot_info_like.name)
+
+    return model_display_name(
+        snapshot_info_like.name, environment_naming_info, default_catalog, dialect
+    )
+
+
+def model_display_name(
+    node_name: str,
+    environment_naming_info: EnvironmentNamingInfo,
+    default_catalog: t.Optional[str],
+    dialect: DialectType = None,
+) -> str:
+    view_name = exp.to_table(node_name)
 
     catalog = (
         None

tests/core/test_integration.py

@@ -10440,8 +10440,9 @@ def test_restatement_plan_detects_prod_deployment_during_restatement(tmp_path: P
         - During restatement, someone else deploys A(dev) to prod, replacing the model that is currently being restated.
 
     Outcome:
-        - The deployment plan for dev -> prod should succeed in deploying the new version
+        - The deployment plan for dev -> prod should succeed in deploying the new version of A
         - The prod restatement plan should fail with a ConflictingPlanError and warn about the model that got updated while undergoing restatement
+        - The new version of A should have no intervals cleared. The user needs to rerun the restatement if the intervals should still be cleared
     """
     orig_console = get_console()
     console = CaptureTerminalConsole()
@@ -10514,6 +10515,7 @@ def entrypoint(evaluator: MacroEvaluator) -> str:
     ctx.load()
     plan = ctx.plan(environment="dev", auto_apply=True)
     assert len(plan.modified_snapshots) == 1
+    new_model_a_snapshot_id = list(plan.modified_snapshots)[0]
 
     # now, trigger a prod restatement plan in a different thread and block it to simulate a long restatement
     def _run_restatement_plan(tmp_path: Path, config: Config, q: queue.Queue):
@@ -10572,12 +10574,20 @@ def _run_restatement_plan(tmp_path: Path, config: Config, q: queue.Queue):
 
     plan_error = restatement_plan_future.result()
     assert isinstance(plan_error, ConflictingPlanError)
+    assert "please re-apply your plan" in repr(plan_error).lower()
 
     output = " ".join(re.split("\s+", console.captured_output, flags=re.UNICODE))
     assert (
-        f"The following models had new versions deployed in plan '{new_prod.plan_id}' while data was being restated: └── test.model_a"
+        f"The following models had new versions deployed while data was being restated: └── test.model_a"
         in output
     )
-    assert "please re-run this restatement plan" in output
+
+    # check that no intervals have been cleared from the model_a currently in prod
+    model_a = ctx.state_sync.get_snapshots(snapshot_ids=[new_model_a_snapshot_id])[
+        new_model_a_snapshot_id
+    ]
+    assert isinstance(model_a.node, SqlModel)
+    assert model_a.node.render_query_or_raise().sql() == 'SELECT 1 AS "changed"'
+    assert len(model_a.intervals)
 
     set_console(orig_console)

tests/core/test_plan_stages.py

@@ -612,9 +612,6 @@ def test_build_plan_stages_restatement_prod_only(
         restatement_stage, state_reader, plan
     )
     assert not restatement_stage.snapshot_intervals_to_clear
-    assert (
-        restatement_stage.deployability_index == DeployabilityIndex.all_deployable()
-    )  # default index
 
     # Verify EnvironmentRecordUpdateStage
     assert isinstance(stages[3], EnvironmentRecordUpdateStage)
@@ -774,12 +771,11 @@ 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
-    assert restatement_stage.deployability_index is not None
     snapshot_name, clear_request = list(restatement_stage.snapshot_intervals_to_clear.items())[0]
     assert isinstance(clear_request, SnapshotIntervalClearRequest)
     assert snapshot_name == '"a"'
     assert clear_request.snapshot_id == snapshot_a_dev.snapshot_id
-    assert clear_request.table_info == snapshot_a_dev.table_info
+    assert clear_request.snapshot == snapshot_a_dev.id_and_version
     assert clear_request.interval == (to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))
 
     # Verify EnvironmentRecordUpdateStage