Feat: Use internal state instead of Airflow Variables to store Plan DAG specs (#1533)

Author: izeigerman

sqlmesh/core/plan/evaluator.py

@@ -350,7 +350,7 @@ def client(self) -> BaseAirflowClient:
         return self._mwaa_client
 
     def _apply_plan(self, plan: Plan, plan_request_id: str) -> None:
-        from sqlmesh.schedulers.airflow.plan import create_plan_dag_spec
+        from sqlmesh.schedulers.airflow.plan import PlanDagState, create_plan_dag_spec
 
         plan_application_request = airflow_common.PlanApplicationRequest(
             new_snapshots=list(plan.new_snapshots),
@@ -367,13 +367,7 @@ def _apply_plan(self, plan: Plan, plan_request_id: str) -> None:
             forward_only=plan.forward_only,
         )
         plan_dag_spec = create_plan_dag_spec(plan_application_request, self.state_sync)
-
-        _, stderr = self._mwaa_client.set_variable(
-            airflow_common.plan_dag_spec_key(plan_request_id), plan_dag_spec.json()
-        )
-
-        if stderr:
-            logger.warning("MWAA CLI stderr:\n%s", stderr)
+        PlanDagState.from_state_sync(self.state_sync).add_dag_spec(plan_dag_spec)
 
 
 def can_evaluate_before_promote(

sqlmesh/core/state_sync/engine_adapter.py

@@ -85,6 +85,7 @@ def __init__(
         self.environments_table = nullsafe_join(".", self.schema, "_environments")
         self.seeds_table = nullsafe_join(".", self.schema, "_seeds")
         self.intervals_table = nullsafe_join(".", self.schema, "_intervals")
+        self.plan_dags_table = nullsafe_join(".", self.schema, "_plan_dags")
         self.versions_table = nullsafe_join(".", self.schema, "_versions")
 
         self._snapshot_columns_to_types = {
@@ -306,7 +307,7 @@ def _update_snapshot(self, snapshot: Snapshot) -> None:
         snapshot.updated_ts = now_timestamp()
         self.engine_adapter.update_table(
             self.snapshots_table,
-            {"snapshot": snapshot.json()},
+            {"snapshot": _snapshot_to_json(snapshot)},
             where=self._snapshot_id_filter([snapshot.snapshot_id]),
             contains_json=True,
         )
@@ -726,22 +727,22 @@ def rollback(self) -> None:
         """Rollback to the previous migration."""
         logger.info("Starting migration rollback.")
         tables = (self.snapshots_table, self.environments_table, self.versions_table)
+        optional_tables = (self.seeds_table, self.intervals_table, self.plan_dags_table)
         versions = self.get_versions(validate=False)
         if versions.schema_version == 0:
             # Clean up state tables
-            for table in tables + (self.seeds_table, self.intervals_table):
+            for table in tables + optional_tables:
                 self.engine_adapter.drop_table(table)
         else:
             if not all(self.engine_adapter.table_exists(f"{table}_backup") for table in tables):
                 raise SQLMeshError("There are no prior migrations to roll back to.")
             for table in tables:
                 self._restore_table(table, _backup_table_name(table))
 
-            if self.engine_adapter.table_exists(_backup_table_name(self.seeds_table)):
-                self._restore_table(self.seeds_table, _backup_table_name(self.seeds_table))
+            for optional_table in optional_tables:
+                if self.engine_adapter.table_exists(_backup_table_name(optional_table)):
+                    self._restore_table(optional_table, _backup_table_name(optional_table))
 
-            if self.engine_adapter.table_exists(_backup_table_name(self.intervals_table)):
-                self._restore_table(self.intervals_table, _backup_table_name(self.intervals_table))
         logger.info("Migration rollback successful.")
 
     def _backup_state(self) -> None:
@@ -751,6 +752,7 @@ def _backup_state(self) -> None:
             self.versions_table,
             self.seeds_table,
             self.intervals_table,
+            self.plan_dags_table,
         ):
             if self.engine_adapter.table_exists(table):
                 with self.engine_adapter.transaction(TransactionType.DDL):
@@ -1000,7 +1002,7 @@ def _snapshots_to_df(snapshots: t.Iterable[Snapshot]) -> pd.DataFrame:
                 "name": snapshot.name,
                 "identifier": snapshot.identifier,
                 "version": snapshot.version,
-                "snapshot": snapshot.json(exclude={"intervals", "dev_intervals"}),
+                "snapshot": _snapshot_to_json(snapshot),
                 "kind_name": snapshot.model_kind_name.value if snapshot.model_kind_name else None,
             }
             for snapshot in snapshots
@@ -1033,3 +1035,7 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
 
 def _backup_table_name(table_name: str) -> str:
     return f"{table_name}_backup"
+
+
+def _snapshot_to_json(snapshot: Snapshot) -> str:
+    return snapshot.json(exclude={"intervals", "dev_intervals"})

sqlmesh/migrations/v0028_add_plan_dags_table.py

@@ -0,0 +1,28 @@
+"""Creates the '_plan_dags' table if Airflow is used."""
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type
+
+
+def migrate(state_sync):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    plan_dags_table = "_plan_dags"
+
+    if schema:
+        engine_adapter.create_schema(schema)
+        plan_dags_table = f"{schema}.{plan_dags_table}"
+
+    text_type = index_text_type(engine_adapter.dialect)
+
+    engine_adapter.create_state_table(
+        plan_dags_table,
+        {
+            "request_id": exp.DataType.build(text_type),
+            "dag_id": exp.DataType.build(text_type),
+            "dag_spec": exp.DataType.build("text"),
+        },
+        primary_key=("request_id",),
+    )
+
+    engine_adapter.create_index(plan_dags_table, "dag_id_idx", ("dag_id",))

sqlmesh/schedulers/airflow/api.py

@@ -5,14 +5,13 @@
 from functools import wraps
 
 from airflow.api_connexion import security
-from airflow.models import Variable
 from airflow.www.app import csrf
 from flask import Blueprint, Response, jsonify, make_response, request
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.snapshot import SnapshotId, SnapshotNameVersion
 from sqlmesh.schedulers.airflow import common, util
-from sqlmesh.schedulers.airflow.plan import create_plan_dag_spec
+from sqlmesh.schedulers.airflow.plan import PlanDagState, create_plan_dag_spec
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
 
@@ -40,13 +39,11 @@ def apply_plan() -> Response:
         plan = common.PlanApplicationRequest.parse_obj(request.json or {})
         with util.scoped_state_sync() as state_sync:
             spec = create_plan_dag_spec(plan, state_sync)
+            PlanDagState.from_state_sync(state_sync).add_dag_spec(spec)
+            return make_response(jsonify(request_id=spec.request_id), 201)
     except Exception as ex:
         return _error(str(ex))
 
-    Variable.set(common.plan_dag_spec_key(spec.request_id), spec.json())
-
-    return make_response(jsonify(request_id=spec.request_id), 201)
-
 
 @sqlmesh_api_v1.route("/environments/<name>")
 @csrf.exempt

sqlmesh/schedulers/airflow/common.py

@@ -68,17 +68,17 @@ class PlanDagSpec(PydanticModel):
     promoted_snapshots: t.List[SnapshotTableInfo]
     demoted_snapshots: t.List[SnapshotTableInfo]
     start: TimeLike
-    end: t.Optional[TimeLike]
-    unpaused_dt: t.Optional[TimeLike]
+    end: t.Optional[TimeLike] = None
+    unpaused_dt: t.Optional[TimeLike] = None
     no_gaps: bool
     plan_id: str
-    previous_plan_id: t.Optional[str]
+    previous_plan_id: t.Optional[str] = None
     notification_targets: t.List[NotificationTarget]
     backfill_concurrent_tasks: int
     ddl_concurrent_tasks: int
     users: t.List[User]
     is_dev: bool
-    forward_only: t.Optional[bool]
+    forward_only: t.Optional[bool] = None
     environment_expiration_ts: t.Optional[int] = None
     dag_start_ts: t.Optional[int] = None
 
@@ -143,10 +143,6 @@ def plan_application_dag_id(environment: str, request_id: str) -> str:
     return f"sqlmesh_plan_application__{environment}__{request_id}"
 
 
-def environment_key(env: str) -> str:
-    return f"{ENV_KEY_PREFIX}__{env}"
-
-
 def plan_dag_spec_key(request_id: str) -> str:
     return f"{PLAN_DAG_SPEC_KEY_PREFIX}__{request_id}"
 

sqlmesh/schedulers/airflow/integration.py

@@ -15,6 +15,7 @@
 from sqlmesh.schedulers.airflow import common, util
 from sqlmesh.schedulers.airflow.dag_generator import SnapshotDagGenerator
 from sqlmesh.schedulers.airflow.operators import targets
+from sqlmesh.schedulers.airflow.plan import PlanDagState
 
 logger = logging.getLogger(__name__)
 
@@ -99,6 +100,9 @@ def dags(self) -> t.List[DAG]:
         """
         with util.scoped_state_sync() as state_sync:
             stored_snapshots = state_sync.get_snapshots(None)
+            plan_dag_specs = PlanDagState.from_state_sync(state_sync).get_dag_specs()
+            # TODO: Remove this once all DAG specs are moved into the internal state (after about 1 week)
+            plan_dag_specs += _get_plan_dag_specs_from_variables()
 
         dag_generator = SnapshotDagGenerator(
             self._engine_operator,
@@ -112,7 +116,7 @@ def dags(self) -> t.List[DAG]:
         cadence_dags = dag_generator.generate_cadence_dags(prod_env.snapshots) if prod_env else []
 
         plan_application_dags = [
-            dag_generator.generate_plan_application_dag(s) for s in _get_plan_dag_specs()
+            dag_generator.generate_plan_application_dag(s) for s in plan_dag_specs
         ]
 
         system_dags = [
@@ -187,6 +191,8 @@ def _janitor_task(
             ttl=plan_application_dag_ttl, session=session
         )
         logger.info("Deleting expired Plan Application DAGs: %s", plan_application_dag_ids)
+        PlanDagState.from_state_sync(state_sync).delete_dag_specs(plan_application_dag_ids)
+        # TODO: Remove this once all DAG specs are moved into the internal state (after about 1 week)
         util.delete_variables(
             {common.plan_dag_spec_key_from_dag_id(dag_id) for dag_id in plan_application_dag_ids},
             session=session,
@@ -197,7 +203,7 @@ def _janitor_task(
 
 
 @provide_session
-def _get_plan_dag_specs(
+def _get_plan_dag_specs_from_variables(
     session: Session = util.PROVIDED_SESSION,
 ) -> t.List[common.PlanDagSpec]:
     records = (

sqlmesh/schedulers/airflow/mwaa_client.py

@@ -21,19 +21,6 @@ def __init__(self, airflow_url: str, auth_token: str, console: t.Optional[Consol
             {"Authorization": f"Bearer {auth_token}", "Content-Type": "text/plain"}
         )
 
-    def set_variable(self, key: str, value: str) -> t.Tuple[str, str]:
-        """Sets the Airflow variable.
-
-        Args:
-            key: The name of the variable.
-            value: The value of the variable.
-
-        Returns:
-            A tuple of stdout and stderr from the MWAA CLI.
-        """
-        value = value.replace("\\", "\\\\").replace('"', '\\"')
-        return self._post(f'variables set {key} "{value}"')
-
     def get_first_dag_run_id(self, dag_id: str) -> t.Optional[str]:
         dag_runs = self._list_dag_runs(dag_id)
         if dag_runs:

sqlmesh/schedulers/airflow/operators/hwm_sensor.py

@@ -60,6 +60,6 @@ def _compute_target_high_water_mark(
         self, dag_run: DagRun, target_snapshot: Snapshot
     ) -> datetime:
         target_date = to_datetime(dag_run.data_interval_end)
-        target_prev = to_datetime(target_snapshot.model.cron_floor(target_date))
-        this_prev = to_datetime(self.this_snapshot.model.cron_floor(target_date))
+        target_prev = to_datetime(target_snapshot.node.cron_floor(target_date))
+        this_prev = to_datetime(self.this_snapshot.node.cron_floor(target_date))
         return min(target_prev, this_prev)

sqlmesh/schedulers/airflow/plan.py

@@ -2,16 +2,84 @@
 
 import typing as t
 
+import pandas as pd
+from sqlglot import exp
+
 from sqlmesh.core import scheduler
+from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.plan import can_evaluate_before_promote
 from sqlmesh.core.snapshot import SnapshotTableInfo
-from sqlmesh.core.state_sync import StateSync
+from sqlmesh.core.state_sync import EngineAdapterStateSync, StateSync
+from sqlmesh.core.state_sync.base import DelegatingStateSync
 from sqlmesh.schedulers.airflow import common
 from sqlmesh.utils.date import now, now_timestamp
 from sqlmesh.utils.errors import SQLMeshError
 
 
+class PlanDagState:
+    def __init__(self, engine_adapter: EngineAdapter, plan_dags_table: str):
+        self.engine_adapter = engine_adapter
+
+        self._plan_dags_table = plan_dags_table
+        self._plan_dag_columns_to_types = {
+            "request_id": exp.DataType.build("text"),
+            "dag_id": exp.DataType.build("text"),
+            "dag_spec": exp.DataType.build("text"),
+        }
+
+    @classmethod
+    def from_state_sync(cls, state_sync: StateSync) -> PlanDagState:
+        while isinstance(state_sync, DelegatingStateSync):
+            state_sync = state_sync.state_sync
+        if not isinstance(state_sync, EngineAdapterStateSync):
+            raise SQLMeshError(f"Unsupported state sync {state_sync.__class__.__name__}")
+        return cls(state_sync.engine_adapter, state_sync.plan_dags_table)
+
+    def add_dag_spec(self, spec: common.PlanDagSpec) -> None:
+        """Adds a new DAG spec to the state.
+
+        Args:
+            spec: the plan DAG spec to add.
+        """
+        df = pd.DataFrame(
+            [
+                {
+                    "request_id": spec.request_id,
+                    "dag_id": common.plan_application_dag_id(
+                        spec.environment_naming_info.name, spec.request_id
+                    ),
+                    "dag_spec": spec.json(),
+                }
+            ]
+        )
+        self.engine_adapter.insert_append(
+            self._plan_dags_table,
+            df,
+            columns_to_types=self._plan_dag_columns_to_types,
+            contains_json=True,
+        )
+
+    def get_dag_specs(self) -> t.List[common.PlanDagSpec]:
+        """Returns all DAG specs in the state."""
+        query = exp.select("dag_spec").from_(self._plan_dags_table)
+        return [
+            common.PlanDagSpec.parse_raw(row[0])
+            for row in self.engine_adapter.fetchall(
+                query, ignore_unsupported_errors=True, quote_identifiers=True
+            )
+        ]
+
+    def delete_dag_specs(self, dag_ids: t.Collection[str]) -> None:
+        """Deletes the DAG specs with the given DAG IDs."""
+        if not dag_ids:
+            return
+        self.engine_adapter.delete_from(
+            self._plan_dags_table,
+            where=exp.column("dag_id").isin(*dag_ids),
+        )
+
+
 def create_plan_dag_spec(
     request: common.PlanApplicationRequest, state_sync: StateSync
 ) -> common.PlanDagSpec:

tests/core/test_plan_evaluator.py

@@ -125,18 +125,21 @@ def test_mwaa_evaluator(sushi_plan: Plan, mocker: MockerFixture):
 
     state_sync_mock = mocker.Mock()
 
-    plan_dag_spec_json = """{"request_id": "test_request_id"}"""
-
     plan_dag_spec_mock = mocker.Mock()
-    plan_dag_spec_mock.json.return_value = plan_dag_spec_json
 
     create_plan_dag_spec_mock = mocker.patch("sqlmesh.schedulers.airflow.plan.create_plan_dag_spec")
     create_plan_dag_spec_mock.return_value = plan_dag_spec_mock
 
+    plan_dag_state_mock = mocker.Mock()
+    mocker.patch(
+        "sqlmesh.schedulers.airflow.plan.PlanDagState.from_state_sync",
+        return_value=plan_dag_state_mock,
+    )
+
     evaluator = MWAAPlanEvaluator(mwaa_client_mock, state_sync_mock)
     evaluator.evaluate(sushi_plan)
 
-    mwaa_client_mock.set_variable.assert_called_once_with(mocker.ANY, plan_dag_spec_json)
+    plan_dag_state_mock.add_dag_spec.assert_called_once_with(plan_dag_spec_mock)
 
     mwaa_client_mock.wait_for_dag_run_completion.assert_called_once()
     mwaa_client_mock.wait_for_first_dag_run.assert_called_once()