Fix: Introduce environment finalization (#764)

Author: izeigerman

sqlmesh/core/context.py

@@ -892,7 +892,7 @@ def _model_tables(self) -> t.Dict[str, str]:
 
     def _context_diff(
         self,
-        environment: str | Environment,
+        environment: str,
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         create_from: t.Optional[str] = None,
     ) -> ContextDiff:

sqlmesh/core/context_diff.py

@@ -13,7 +13,6 @@
 
 import typing as t
 
-from sqlmesh.core.environment import Environment
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotChangeCategory,
@@ -38,6 +37,8 @@ class ContextDiff(PydanticModel):
     """The environment to diff."""
     is_new_environment: bool
     """Whether the target environment is new."""
+    is_unfinalized_environment: bool
+    """Whether the currently stored environment record is in unfinalized state."""
     create_from: str
     """The name of the environment the target environment will be created from if new."""
     added: t.Set[str]
@@ -56,7 +57,7 @@ class ContextDiff(PydanticModel):
     @classmethod
     def create(
         cls,
-        environment: str | Environment,
+        environment: str,
         snapshots: t.Dict[str, Snapshot],
         create_from: str,
         state_reader: StateReader,
@@ -73,12 +74,8 @@ def create(
         Returns:
             The ContextDiff object.
         """
-        if isinstance(environment, str):
-            environment = environment.lower()
-            env = state_reader.get_environment(environment)
-        else:
-            env = environment
-            environment = env.name.lower()
+        environment = environment.lower()
+        env = state_reader.get_environment(environment)
 
         if env is None:
             env = state_reader.get_environment(create_from.lower())
@@ -151,6 +148,7 @@ def create(
         return ContextDiff(
             environment=environment,
             is_new_environment=is_new_environment,
+            is_unfinalized_environment=bool(env and not env.finalized_ts),
             create_from=create_from,
             added=added,
             removed=removed,
@@ -162,7 +160,9 @@ def create(
 
     @property
     def has_changes(self) -> bool:
-        return self.has_snapshot_changes or self.is_new_environment
+        return (
+            self.has_snapshot_changes or self.is_new_environment or self.is_unfinalized_environment
+        )
 
     @property
     def has_snapshot_changes(self) -> bool:

sqlmesh/core/environment.py

@@ -24,6 +24,7 @@ class Environment(PydanticModel):
     plan_id: str
     previous_plan_id: t.Optional[str]
     expiration_ts: t.Optional[int]
+    finalized_ts: t.Optional[int]
 
     @validator("snapshots", pre=True)
     @classmethod

sqlmesh/core/plan/evaluator.py

@@ -137,6 +137,7 @@ def on_complete(snapshot: SnapshotInfoLike) -> None:
                 environment=environment.name,
                 on_complete=on_complete,
             )
+            self.state_sync.finalize(environment)
             completed = True
         finally:
             self.console.stop_promotion_progress(success=completed)

sqlmesh/core/state_sync/base.py

@@ -337,6 +337,15 @@ def promote(
            A tuple of (added snapshot table infos, removed snapshot table infos)
         """
 
+    @abc.abstractmethod
+    def finalize(self, environment: Environment) -> None:
+        """Finalize the target environment, indicating that this environment has been
+        fully promoted and is ready for use.
+
+        Args:
+            environment: The target environment to finalize.
+        """
+
     @abc.abstractmethod
     def delete_expired_environments(self) -> t.List[Environment]:
         """Removes expired environments.

sqlmesh/core/state_sync/common.py

@@ -17,7 +17,7 @@
     SnapshotTableInfo,
 )
 from sqlmesh.core.state_sync.base import StateSync
-from sqlmesh.utils.date import TimeLike, now, to_datetime
+from sqlmesh.utils.date import TimeLike, now, now_timestamp, to_datetime
 from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -129,6 +129,26 @@ def promote(
         self._update_environment(environment)
         return table_infos, [existing_table_infos[name] for name in missing_models]
 
+    @transactional()
+    def finalize(self, environment: Environment) -> None:
+        """Finalize the target environment, indicating that this environment has been
+        fully promoted and is ready for use.
+
+        Args:
+            environment: The target environment to finalize.
+        """
+        logger.info("Finalizing environment '%s'", environment)
+
+        stored_environment = self._get_environment(environment.name, lock_for_update=True)
+        if stored_environment and stored_environment.plan_id != environment.plan_id:
+            raise SQLMeshError(
+                f"Plan '{environment.plan_id}' is no longer valid for the target environment '{environment.name}'. "
+                f"Stored plan ID: '{stored_environment.plan_id}'. Please recreate the plan and try again"
+            )
+
+        environment.finalized_ts = now_timestamp()
+        self._update_environment(environment)
+
     @transactional()
     def delete_expired_snapshots(self) -> t.List[Snapshot]:
         current_time = now()

sqlmesh/core/state_sync/engine_adapter.py

@@ -90,6 +90,7 @@ def environment_columns_to_types(self) -> t.Dict[str, exp.DataType]:
             "plan_id": exp.DataType.build("text"),
             "previous_plan_id": exp.DataType.build("text"),
             "expiration_ts": exp.DataType.build("bigint"),
+            "finalized_ts": exp.DataType.build("bigint"),
         }
 
     @property
@@ -236,6 +237,7 @@ def _update_environment(self, environment: Environment) -> None:
                             environment.plan_id,
                             environment.previous_plan_id,
                             environment.expiration_ts,
+                            environment.finalized_ts,
                         )
                     ],
                     columns_to_types=self.environment_columns_to_types,

sqlmesh/migrations/v0004_environmnent_add_finalized_at.py

@@ -0,0 +1,19 @@
+"""Add support for environment finalization."""
+from sqlglot import exp
+
+
+def migrate(state_sync):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    environments_table = state_sync.environments_table
+
+    alter_table_exp = exp.AlterTable(
+        this=exp.to_table(environments_table),
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("finalized_ts"),
+                kind=exp.DataType.build("bigint"),
+            )
+        ],
+    )
+
+    engine_adapter.execute(alter_table_exp)

sqlmesh/schedulers/airflow/dag_generator.py

@@ -242,23 +242,34 @@ def _create_promotion_demotion_tasks(
         start_task = EmptyOperator(task_id="snapshot_promotion_start")
         end_task = EmptyOperator(task_id="snapshot_promotion_end")
 
+        environment = Environment(
+            name=request.environment_name,
+            snapshots=request.promoted_snapshots,
+            start_at=request.start,
+            end_at=request.end,
+            plan_id=request.plan_id,
+            previous_plan_id=request.previous_plan_id,
+            expiration_ts=request.environment_expiration_ts,
+        )
+
         update_state_task = PythonOperator(
             task_id="snapshot_promotion__update_state",
             python_callable=promotion_update_state_task,
             op_kwargs={
-                "snapshots": request.promoted_snapshots,
-                "environment_name": request.environment_name,
-                "start": request.start,
-                "end": request.end,
+                "environment": environment,
                 "unpaused_dt": request.unpaused_dt,
                 "no_gaps": request.no_gaps,
-                "plan_id": request.plan_id,
-                "previous_plan_id": request.previous_plan_id,
-                "environment_expiration_ts": request.environment_expiration_ts,
             },
         )
 
+        finalize_task = PythonOperator(
+            task_id="snapshot_promotion__finalize",
+            python_callable=promotion_finalize_task,
+            op_kwargs={"environment": environment},
+        )
+
         start_task >> update_state_task
+        finalize_task >> end_task
 
         if request.promoted_snapshots:
             create_views_task = self._create_snapshot_promotion_operator(
@@ -268,7 +279,7 @@ def _create_promotion_demotion_tasks(
                 request.is_dev,
                 "snapshot_promotion__create_views",
             )
-            create_views_task >> end_task
+            create_views_task >> finalize_task
 
             if not request.is_dev and request.unpaused_dt:
                 migrate_tables_task = self._create_snapshot_migrate_tables_operator(
@@ -289,10 +300,10 @@ def _create_promotion_demotion_tasks(
                 "snapshot_promotion__delete_views",
             )
             update_state_task >> delete_views_task
-            delete_views_task >> end_task
+            delete_views_task >> finalize_task
 
         if not request.promoted_snapshots and not request.demoted_snapshots:
-            update_state_task >> end_task
+            update_state_task >> finalize_task
 
         return (start_task, end_task)
 
@@ -479,26 +490,16 @@ def creation_update_state_task(new_snapshots: t.Iterable[Snapshot]) -> None:
 
 
 def promotion_update_state_task(
-    snapshots: t.List[SnapshotTableInfo],
-    environment_name: str,
-    start: TimeLike,
-    end: t.Optional[TimeLike],
+    environment: Environment,
     unpaused_dt: t.Optional[TimeLike],
     no_gaps: bool,
-    plan_id: str,
-    previous_plan_id: t.Optional[str],
-    environment_expiration_ts: t.Optional[int],
 ) -> None:
-    environment = Environment(
-        name=environment_name,
-        snapshots=snapshots,
-        start_at=start,
-        end_at=end,
-        plan_id=plan_id,
-        previous_plan_id=previous_plan_id,
-        expiration_ts=environment_expiration_ts,
-    )
     with util.scoped_state_sync() as state_sync:
         state_sync.promote(environment, no_gaps=no_gaps)
-        if snapshots and not end and unpaused_dt:
-            state_sync.unpause_snapshots(snapshots, unpaused_dt)
+        if environment.snapshots and not environment.end_at and unpaused_dt:
+            state_sync.unpause_snapshots(environment.snapshots, unpaused_dt)
+
+
+def promotion_finalize_task(environment: Environment) -> None:
+    with util.scoped_state_sync() as state_sync:
+        state_sync.finalize(environment)

tests/core/test_state_sync.py

@@ -1,5 +1,6 @@
 import typing as t
 
+import duckdb
 import pandas as pd
 import pytest
 from pytest_mock.plugin import MockerFixture
@@ -445,6 +446,34 @@ def test_promote_snapshots_no_gaps(state_sync: EngineAdapterStateSync, make_snap
     promote_snapshots(state_sync, [new_snapshot_same_interval], "prod", no_gaps=True)
 
 
+def test_finalize(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            query=parse_one("select 1, ds"),
+        ),
+    )
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    state_sync.push_snapshots([snapshot_a])
+    promote_snapshots(state_sync, [snapshot_a], "prod")
+
+    env = state_sync.get_environment("prod")
+    assert env
+    state_sync.finalize(env)
+
+    env = state_sync.get_environment("prod")
+    assert env
+    assert env.finalized_ts is not None
+
+    env.plan_id = "different_plan_id"
+    with pytest.raises(
+        SQLMeshError,
+        match=r"Plan 'different_plan_id' is no longer valid for the target environment 'prod'.*",
+    ):
+        state_sync.finalize(env)
+
+
 def test_start_date_gap(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
     model = SqlModel(
         name="a",
@@ -569,8 +598,8 @@ def test_get_version(state_sync: EngineAdapterStateSync) -> None:
         schema_version=SCHEMA_VERSION, sqlglot_version=SQLGLOT_VERSION
     )
 
-    # old install does not have this table / row
-    delete_versions(state_sync)
+    # Start with a clean slate.
+    state_sync = EngineAdapterStateSync(create_engine_adapter(duckdb.connect, "duckdb"))
 
     with pytest.raises(
         SQLMeshError,
@@ -626,7 +655,9 @@ def test_migrate(state_sync: EngineAdapterStateSync, mocker: MockerFixture) -> N
     state_sync.migrate()
     mock.assert_not_called()
 
-    delete_versions(state_sync)
+    # Start with a clean slate.
+    state_sync = EngineAdapterStateSync(create_engine_adapter(duckdb.connect, "duckdb"))
+
     state_sync.migrate()
     mock.assert_called_once()
     assert state_sync.get_versions() == Versions(