Pass snapshots and deployability_index to session properties renderer

Author: treysp

sqlmesh/core/plan/evaluator.py

@@ -336,6 +336,8 @@ def visit_virtual_layer_update_stage(
                 self._demote_snapshots(
                     [stage.all_snapshots[s.snapshot_id] for s in stage.demoted_snapshots],
                     stage.demoted_environment_naming_info,
+                    snapshots=stage.all_snapshots,
+                    deployability_index=stage.deployability_index,
                     on_complete=lambda s: self.console.update_promotion_progress(s, False),
                 )
 
@@ -378,10 +380,16 @@ def _demote_snapshots(
         self,
         target_snapshots: t.Iterable[Snapshot],
         environment_naming_info: EnvironmentNamingInfo,
+        snapshots: t.Dict[SnapshotId, Snapshot],
+        deployability_index: t.Optional[DeployabilityIndex] = None,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
         self.snapshot_evaluator.demote(
-            target_snapshots, environment_naming_info, on_complete=on_complete
+            target_snapshots,
+            environment_naming_info,
+            snapshots=snapshots,
+            deployability_index=deployability_index,
+            on_complete=on_complete,
         )
 
     def _restatement_intervals_across_all_environments(

sqlmesh/core/snapshot/evaluator.py

@@ -277,6 +277,8 @@ def demote(
         self,
         target_snapshots: t.Iterable[Snapshot],
         environment_naming_info: EnvironmentNamingInfo,
+        snapshots: t.Dict[SnapshotId, Snapshot],
+        deployability_index: t.Optional[DeployabilityIndex] = None,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
         """Demotes the given collection of snapshots in the target environment by removing its view.
@@ -289,7 +291,9 @@ def demote(
         with self.concurrent_context():
             concurrent_apply_to_snapshots(
                 target_snapshots,
-                lambda s: self._demote_snapshot(s, environment_naming_info, on_complete),
+                lambda s: self._demote_snapshot(
+                    s, environment_naming_info, snapshots, deployability_index, on_complete
+                ),
                 self.ddl_concurrent_tasks,
             )
 
@@ -1013,6 +1017,8 @@ def _demote_snapshot(
         self,
         snapshot: Snapshot,
         environment_naming_info: EnvironmentNamingInfo,
+        snapshots: t.Dict[SnapshotId, Snapshot],
+        deployability_index: t.Optional[DeployabilityIndex],
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]],
     ) -> None:
         if snapshot.is_model:
@@ -1028,7 +1034,10 @@ def _demote_snapshot(
                 adapter.transaction(),
                 adapter.session(
                     snapshot.model.render_session_properties(
-                        engine_adapter=adapter, runtime_stage=RuntimeStage.PROMOTING
+                        engine_adapter=adapter,
+                        snapshots=snapshots,
+                        deployability_index=deployability_index,
+                        runtime_stage=RuntimeStage.PROMOTING,
                     )
                 ),
             ):

sqlmesh/engines/commands.py

@@ -48,8 +48,10 @@ class PromoteCommandPayload(PydanticModel):
 
 
 class DemoteCommandPayload(PydanticModel):
-    snapshots: t.List[Snapshot]
+    target_snapshots: t.List[Snapshot]
     environment_naming_info: EnvironmentNamingInfo
+    snapshots: t.Dict[SnapshotId, Snapshot]
+    deployability_index: DeployabilityIndex
 
 
 class CleanupCommandPayload(PydanticModel):
@@ -128,8 +130,10 @@ def demote(
     if isinstance(command_payload, str):
         command_payload = DemoteCommandPayload.parse_raw(command_payload)
     evaluator.demote(
-        command_payload.snapshots,
+        command_payload.target_snapshots,
         command_payload.environment_naming_info,
+        snapshots=command_payload.snapshots,
+        deployability_index=command_payload.deployability_index,
     )
 
 

tests/core/test_snapshot_evaluator.py

@@ -320,7 +320,7 @@ def test_demote(mocker: MockerFixture, adapter_mock, make_snapshot):
     snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
-    evaluator.demote([snapshot], EnvironmentNamingInfo(name="test_env"))
+    evaluator.demote([snapshot], EnvironmentNamingInfo(name="test_env"), snapshots={})
 
     adapter_mock.transaction.assert_called()
     adapter_mock.session.assert_called()
@@ -2711,7 +2711,7 @@ def test_standalone_audit(mocker: MockerFixture, adapter_mock, make_snapshot):
     adapter_mock.session.assert_not_called()
 
     # Demote
-    evaluator.demote([snapshot], EnvironmentNamingInfo(name="test_env"))
+    evaluator.demote([snapshot], EnvironmentNamingInfo(name="test_env"), snapshots={})
 
     adapter_mock.assert_not_called()
     adapter_mock.transaction.assert_not_called()
@@ -3854,7 +3854,7 @@ def test_multiple_engine_creation(snapshot: Snapshot, adapters, make_snapshot):
     assert len(post_calls) == 1
     assert post_calls[0].sql(dialect="postgres") == expected_call
 
-    evaluator.demote([snapshot_2], EnvironmentNamingInfo(name="test_env"))
+    evaluator.demote([snapshot_2], EnvironmentNamingInfo(name="test_env"), snapshots={})
     engine_adapters["secondary"].drop_view.assert_not_called()
     engine_adapters["default"].drop_view.assert_called_once_with(
         "test_schema__test_env.test_model",
@@ -4100,7 +4100,7 @@ def model_with_statements(context, **kwargs):
     engine_adapters["default"].get_catalog_type.assert_not_called()
     assert len(engine_adapters["secondary"].get_catalog_type.call_args_list) == 2
 
-    evaluator.demote([snapshot], environment_naming_info)
+    evaluator.demote([snapshot], environment_naming_info, snapshots={})
     engine_adapters["default"].drop_view.assert_called_once_with(
         "db__test_env.multi_engine_test_model",
         cascade=False,
@@ -4114,7 +4114,7 @@ def model_with_statements(context, **kwargs):
     assert view_args[0][0][0] == "db__test_env.multi_engine_test_model"
 
     # Similarly for demotion
-    evaluator.demote([snapshot], environment_naming_info_gw)
+    evaluator.demote([snapshot], environment_naming_info_gw, snapshots={})
     engine_adapters["secondary"].drop_view.assert_called_once_with(
         "db__test_env.multi_engine_test_model",
         cascade=False,
@@ -4172,7 +4172,7 @@ def test_multiple_engine_virtual_layer(snapshot: Snapshot, adapters, make_snapsh
     assert view_args_secondary[0][0][0] == "test_schema__test_env.test_model"
 
     # Demotion will follow with the same pattern
-    evaluator.demote([snapshot_2, snapshot], environment_naming_info)
+    evaluator.demote([snapshot_2, snapshot], environment_naming_info, snapshots={})
     engine_adapters["default"].drop_view.assert_called_once_with(
         "db__test_env.model",
         cascade=False,