Fix!: In Airflow only generate cadence DAGs for snapshots promoted to prod (#1514)

Author: izeigerman

sqlmesh/schedulers/airflow/dag_generator.py

@@ -13,7 +13,12 @@
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
 from sqlmesh.core.notification_target import NotificationTarget
 from sqlmesh.core.plan import PlanStatus
-from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotTableInfo
+from sqlmesh.core.snapshot import (
+    Snapshot,
+    SnapshotId,
+    SnapshotIdLike,
+    SnapshotTableInfo,
+)
 from sqlmesh.schedulers.airflow import common, util
 from sqlmesh.schedulers.airflow.operators import targets
 from sqlmesh.schedulers.airflow.operators.hwm_sensor import HighWaterMarkSensor
@@ -61,12 +66,13 @@ def __init__(
         self._ddl_engine_operator_args = ddl_engine_operator_args or {}
         self._snapshots = snapshots
 
-    def generate_cadence_dags(self) -> t.List[DAG]:
-        return [
-            self._create_cadence_dag_for_snapshot(s)
-            for s in self._snapshots.values()
-            if s.unpaused_ts and not s.is_symbolic and not s.is_seed
-        ]
+    def generate_cadence_dags(self, snapshots: t.Iterable[SnapshotIdLike]) -> t.List[DAG]:
+        dags = []
+        for s in snapshots:
+            snapshot = self._snapshots[s.snapshot_id]
+            if snapshot.unpaused_ts and not snapshot.is_symbolic and not snapshot.is_seed:
+                dags.append(self._create_cadence_dag_for_snapshot(snapshot))
+        return dags
 
     def generate_plan_application_dag(self, spec: common.PlanDagSpec) -> DAG:
         return self._create_plan_application_dag(spec)

sqlmesh/schedulers/airflow/integration.py

@@ -10,6 +10,7 @@
 from airflow.utils.session import provide_session
 from sqlalchemy.orm import Session
 
+from sqlmesh.core import constants as c
 from sqlmesh.engines import commands
 from sqlmesh.schedulers.airflow import common, util
 from sqlmesh.schedulers.airflow.dag_generator import SnapshotDagGenerator
@@ -107,7 +108,8 @@ def dags(self) -> t.List[DAG]:
             stored_snapshots,
         )
 
-        cadence_dags = dag_generator.generate_cadence_dags()
+        prod_env = state_sync.get_environment(c.PROD)
+        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()