Establish session before snapshot promotion/demotion

Author: treysp

sqlmesh/core/plan/evaluator.py

@@ -340,7 +340,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),
                 )
@@ -382,7 +382,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

@@ -276,7 +276,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:
@@ -989,25 +989,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:
@@ -1019,7 +1024,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)