Repository: TobikoData/sqlmesh
Pull request #4796 (2025-06-23 22:09:47)

Chore: Add metadata_updated_snapshots to EvaluatablePlan


Patch:
diff --git a/sqlmesh/core/plan/definition.py b/sqlmesh/core/plan/definition.py
index f3fb088e..7325dc35 100644
--- a/sqlmesh/core/plan/definition.py
+++ b/sqlmesh/core/plan/definition.py
@@ -261,6 +261,7 @@ class Plan(PydanticModel, frozen=True):
             indirectly_modified_snapshots={
                 s.name: sorted(snapshot_ids) for s, snapshot_ids in self.indirectly_modified.items()
             },
+            metadata_updated_snapshots=sorted(self.metadata_updated),
             removed_snapshots=sorted(self.context_diff.removed_snapshots),
             requires_backfill=self.requires_backfill,
             models_to_backfill=self.models_to_backfill,
@@ -298,6 +299,7 @@ class EvaluatablePlan(PydanticModel):
     ensure_finalized_snapshots: bool
     directly_modified_snapshots: t.List[SnapshotId]
     indirectly_modified_snapshots: t.Dict[str, t.List[SnapshotId]]
+    metadata_updated_snapshots: t.List[SnapshotId]
     removed_snapshots: t.List[SnapshotId]
     requires_backfill: bool
     models_to_backfill: t.Optional[t.Set[str]] = None
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index f2042583..d3374863 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -911,39 +911,40 @@ class SnapshotEvaluator:
         ):
             return
 
+        deployability_index = DeployabilityIndex.all_deployable()
+        render_kwargs: t.Dict[str, t.Any] = dict(
+            engine_adapter=adapter,
+            snapshots=parent_snapshots_by_name(snapshot, snapshots),
+            runtime_stage=RuntimeStage.CREATING,
+            deployability_index=deployability_index,
+        )
         target_table_name = snapshot.table_name()
-        if adapter.table_exists(target_table_name):
-            evaluation_strategy = _evaluation_strategy(snapshot, adapter)
-            tmp_table_name = snapshot.table_name(is_deployable=False)
-            logger.info(
-                "Migrating table schema from '%s' to '%s'",
-                tmp_table_name,
-                target_table_name,
-            )
-            evaluation_strategy.migrate(
-                target_table_name=target_table_name,
-                source_table_name=tmp_table_name,
-                snapshot=snapshot,
-                snapshots=parent_snapshots_by_name(snapshot, snapshots),
-                allow_destructive_snapshots=allow_destructive_snapshots,
-            )
-        else:
-            logger.info(
-                "Creating table '%s' for the snapshot of the forward-only model %s",
-                target_table_name,
-                snapshot.snapshot_id,
-            )
-            deployability_index = DeployabilityIndex.all_deployable()
-            render_kwargs: t.Dict[str, t.Any] = dict(
-                engine_adapter=adapter,
-                snapshots=parent_snapshots_by_name(snapshot, snapshots),
-                runtime_stage=RuntimeStage.CREATING,
-                deployability_index=deployability_index,
-            )
-            with (
-                adapter.transaction(),
-                adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
-            ):
+
+        with (
+            adapter.transaction(),
+            adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
+        ):
+            if adapter.table_exists(target_table_name):
+                evaluation_strategy = _evaluation_strategy(snapshot, adapter)
+                tmp_table_name = snapshot.table_name(is_deployable=False)
+                logger.info(
+                    "Migrating table schema from '%s' to '%s'",
+                    tmp_table_name,
+                    target_table_name,
+                )
+                evaluation_strategy.migrate(
+                    target_table_name=target_table_name,
+                    source_table_name=tmp_table_name,
+                    snapshot=snapshot,
+                    snapshots=parent_snapshots_by_name(snapshot, snapshots),
+                    allow_destructive_snapshots=allow_destructive_snapshots,
+                )
+            else:
+                logger.info(
+                    "Creating table '%s' for the snapshot of the forward-only model %s",
+                    target_table_name,
+                    snapshot.snapshot_id,
+                )
                 self._execute_create(
                     snapshot=snapshot,
                     table_name=target_table_name,


Test patch:
diff --git a/tests/core/test_plan_stages.py b/tests/core/test_plan_stages.py
index 5fa140cb..c8989f9e 100644
--- a/tests/core/test_plan_stages.py
+++ b/tests/core/test_plan_stages.py
@@ -108,6 +108,7 @@ def test_build_plan_stages_basic(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[snapshot_a.snapshot_id, snapshot_b.snapshot_id],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -215,6 +216,7 @@ def test_build_plan_stages_with_before_all_and_after_all(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[snapshot_a.snapshot_id, snapshot_b.snapshot_id],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -323,6 +325,7 @@ def test_build_plan_stages_select_models(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[snapshot_a.snapshot_id, snapshot_b.snapshot_id],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill={snapshot_a.name},
@@ -422,6 +425,7 @@ def test_build_plan_stages_basic_no_backfill(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[snapshot_a.snapshot_id, snapshot_b.snapshot_id],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -531,6 +535,7 @@ def test_build_plan_stages_restatement(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[],  # No changes
         indirectly_modified_snapshots={},  # No changes
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -642,6 +647,7 @@ def test_build_plan_stages_forward_only(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_b.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -770,6 +776,7 @@ def test_build_plan_stages_forward_only_dev(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_b.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -893,6 +900,7 @@ def test_build_plan_stages_audit_only(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_b.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -1016,6 +1024,7 @@ def test_build_plan_stages_forward_only_ensure_finalized_snapshots(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_b.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -1088,6 +1097,7 @@ def test_build_plan_stages_removed_model(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[snapshot_b.snapshot_id],
         requires_backfill=False,
         models_to_backfill=None,
@@ -1169,6 +1179,7 @@ def test_build_plan_stages_environment_suffix_target_changed(
         ensure_finalized_snapshots=False,
         directly_modified_snapshots=[],
         indirectly_modified_snapshots={},
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=False,
         models_to_backfill=None,
@@ -1268,6 +1279,7 @@ def test_build_plan_stages_indirect_non_breaking_no_migration(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_b.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
@@ -1355,6 +1367,7 @@ def test_build_plan_stages_indirect_non_breaking_view_migration(
         indirectly_modified_snapshots={
             new_snapshot_a.name: [new_snapshot_c.snapshot_id],
         },
+        metadata_updated_snapshots=[],
         removed_snapshots=[],
         requires_backfill=True,
         models_to_backfill=None,
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index d572fc7e..d131e6aa 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -54,7 +54,7 @@ from sqlmesh.core.snapshot import (
     SnapshotTableCleanupTask,
 )
 from sqlmesh.core.snapshot.definition import to_view_mapping
-from sqlmesh.core.snapshot.evaluator import CustomMaterialization
+from sqlmesh.core.snapshot.evaluator import CustomMaterialization, SnapshotCreationFailedError
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.date import to_timestamp
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, DestructiveChangeError
@@ -92,13 +92,16 @@ def date_kwargs() -> t.Dict[str, str]:
 
 @pytest.fixture
 def adapter_mock(mocker: MockerFixture):
+    def mock_exit(self, exc_type, exc_value, traceback):
+        pass
+
     transaction_mock = mocker.Mock()
     transaction_mock.__enter__ = mocker.Mock()
-    transaction_mock.__exit__ = mocker.Mock()
+    transaction_mock.__exit__ = mock_exit
 
     session_mock = mocker.Mock()
     session_mock.__enter__ = mocker.Mock()
-    session_mock.__exit__ = mocker.Mock()
+    session_mock.__exit__ = mock_exit
 
     adapter_mock = mocker.Mock()
     adapter_mock.transaction.return_value = transaction_mock
@@ -1160,6 +1163,7 @@ def test_migrate(mocker: MockerFixture, make_snapshot):
     cursor_mock = mocker.Mock()
     connection_mock.cursor.return_value = cursor_mock
     adapter = EngineAdapter(lambda: connection_mock, "")
+    session_spy = mocker.spy(adapter, "session")
 
     current_table = "sqlmesh__test_schema.test_schema__test_model__1"
 
@@ -1201,6 +1205,8 @@ def test_migrate(mocker: MockerFixture, make_snapshot):
         ]
     )
 
+    session_spy.assert_called_once()
+
 
 def test_migrate_missing_table(mocker: MockerFixture, make_snapshot):
     connection_mock = mocker.NonCallableMock()
@@ -1596,7 +1602,8 @@ def test_drop_clone_in_dev_when_migration_fails(mocker: MockerFixture, adapter_m
         ),
     ]
 
-    evaluator.create([snapshot], {})
+    with pytest.raises(SnapshotCreationFailedError):
+        evaluator.create([snapshot], {})
 
     adapter_mock.clone_table.assert_called_once_with(
         f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__dev",
@@ -2537,7 +2544,9 @@ def test_create_seed_on_error(mocker: MockerFixture, adapter_mock, make_snapshot
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     evaluator = SnapshotEvaluator(adapter_mock)
-    evaluator.create([snapshot], {})
+
+    with pytest.raises(SnapshotCreationFailedError):
+        evaluator.create([snapshot], {})
 
     adapter_mock.replace_query.assert_called_once_with(
         f"sqlmesh__db.db__seed__{snapshot.version}",