Establish session before snapshot promotion/demotion

Author: treysp

sqlmesh/core/plan/evaluator.py

@@ -334,7 +334,7 @@ def visit_virtual_layer_update_stage(
             )
             if stage.demoted_environment_naming_info:
                 self._demote_snapshots(
-                    stage.demoted_snapshots,
+                    [stage.all_snapshots[s.snapshot_id] for s in stage.demoted_snapshots],
                     stage.demoted_environment_naming_info,
                     on_complete=lambda s: self.console.update_promotion_progress(s, False),
                 )
@@ -376,7 +376,7 @@ def _promote_snapshots(
 
     def _demote_snapshots(
         self,
-        target_snapshots: t.Iterable[SnapshotTableInfo],
+        target_snapshots: t.Iterable[Snapshot],
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:

sqlmesh/core/snapshot/evaluator.py

@@ -275,7 +275,7 @@ def promote(
 
     def demote(
         self,
-        target_snapshots: t.Iterable[SnapshotInfoLike],
+        target_snapshots: t.Iterable[Snapshot],
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
@@ -988,25 +988,30 @@ def _promote_snapshot(
                 table_mapping=table_mapping,
                 runtime_stage=RuntimeStage.PROMOTING,
             )
-            _evaluation_strategy(snapshot, adapter).promote(
-                table_name=table_name,
-                view_name=view_name,
-                model=snapshot.model,
-                environment=environment_naming_info.name,
-                snapshots=snapshots,
-                **render_kwargs,
-            )
 
-            snapshot_by_name = {s.name: s for s in (snapshots or {}).values()}
-            render_kwargs["snapshots"] = snapshot_by_name
-            adapter.execute(snapshot.model.render_on_virtual_update(**render_kwargs))
+            with (
+                adapter.transaction(),
+                adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
+            ):
+                _evaluation_strategy(snapshot, adapter).promote(
+                    table_name=table_name,
+                    view_name=view_name,
+                    model=snapshot.model,
+                    environment=environment_naming_info.name,
+                    snapshots=snapshots,
+                    **render_kwargs,
+                )
+
+                snapshot_by_name = {s.name: s for s in (snapshots or {}).values()}
+                render_kwargs["snapshots"] = snapshot_by_name
+                adapter.execute(snapshot.model.render_on_virtual_update(**render_kwargs))
 
         if on_complete is not None:
             on_complete(snapshot)
 
     def _demote_snapshot(
         self,
-        snapshot: SnapshotInfoLike,
+        snapshot: Snapshot,
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]],
     ) -> None:
@@ -1018,7 +1023,18 @@ def _demote_snapshot(
         view_name = snapshot.qualified_view_name.for_environment(
             environment_naming_info, dialect=adapter.dialect
         )
-        _evaluation_strategy(snapshot, adapter).demote(view_name)
+        session_properties = (
+            snapshot.model.render_session_properties(
+                engine_adapter=adapter, runtime_stage=RuntimeStage.PROMOTING
+            )
+            if snapshot.is_model
+            else {}
+        )
+        with (
+            adapter.transaction(),
+            adapter.session(session_properties),
+        ):
+            _evaluation_strategy(snapshot, adapter).demote(view_name)
 
         if on_complete is not None:
             on_complete(snapshot)

sqlmesh/engines/commands.py

@@ -9,7 +9,6 @@
     SnapshotEvaluator,
     SnapshotId,
     SnapshotTableCleanupTask,
-    SnapshotTableInfo,
 )
 from sqlmesh.core.state_sync import cleanup_expired_views
 from sqlmesh.utils.date import TimeLike
@@ -49,7 +48,7 @@ class PromoteCommandPayload(PydanticModel):
 
 
 class DemoteCommandPayload(PydanticModel):
-    snapshots: t.List[SnapshotTableInfo]
+    snapshots: t.List[Snapshot]
     environment_naming_info: EnvironmentNamingInfo