Fix: Evaluate snapshots after pomoting the environment if the plan contains paused forward-only snapshots (#912)

Author: izeigerman

sqlmesh/core/plan/definition.py

@@ -541,6 +541,7 @@ def _ensure_no_forward_only_revert(self) -> None:
                 candidate.snapshot_id not in self.context_diff.new_snapshots
                 and promoted.is_forward_only
                 and not candidate.is_forward_only
+                and not candidate.is_indirect_non_breaking
                 and (
                     promoted.version == candidate.version
                     or candidate.data_version in promoted.previous_versions

sqlmesh/core/plan/evaluator.py

@@ -20,7 +20,11 @@
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.plan.definition import Plan
 from sqlmesh.core.scheduler import Scheduler
-from sqlmesh.core.snapshot import SnapshotEvaluator, SnapshotInfoLike
+from sqlmesh.core.snapshot import (
+    SnapshotEvaluator,
+    SnapshotInfoLike,
+    has_paused_forward_only,
+)
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.core.user import User
 from sqlmesh.schedulers.airflow import common as airflow_common
@@ -61,7 +65,7 @@ def __init__(
     def evaluate(self, plan: Plan) -> None:
         tasks = (
             [self._push, self._restate, self._backfill, self._promote]
-            if not plan.forward_only or plan.is_dev
+            if not has_paused_forward_only(plan.snapshots, plan.snapshots) or plan.is_dev
             else [self._push, self._restate, self._promote, self._backfill]
         )
 

sqlmesh/core/snapshot/__init__.py

@@ -14,6 +14,7 @@
     SnapshotNameVersionLike,
     SnapshotTableInfo,
     fingerprint_from_model,
+    has_paused_forward_only,
     merge_intervals,
     table_name,
     to_table_mapping,

sqlmesh/core/snapshot/definition.py

@@ -1070,3 +1070,16 @@ def to_table_mapping(snapshots: t.Iterable[Snapshot], is_dev: bool) -> t.Dict[st
         for snapshot in snapshots
         if snapshot.version and not snapshot.is_symbolic
     }
+
+
+def has_paused_forward_only(
+    targets: t.Iterable[SnapshotIdLike],
+    snapshots: t.Union[t.List[Snapshot], t.Dict[SnapshotId, Snapshot]],
+) -> bool:
+    if not isinstance(snapshots, dict):
+        snapshots = {s.snapshot_id: s for s in snapshots}
+    for target in targets:
+        target_snapshot = snapshots[target.snapshot_id]
+        if target_snapshot.is_paused and target_snapshot.is_forward_only:
+            return True
+    return False

sqlmesh/schedulers/airflow/dag_generator.py

@@ -12,7 +12,12 @@
 from sqlmesh.core._typing import NotificationTarget
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.plan import PlanStatus
-from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotTableInfo
+from sqlmesh.core.snapshot import (
+    Snapshot,
+    SnapshotId,
+    SnapshotTableInfo,
+    has_paused_forward_only,
+)
 from sqlmesh.schedulers.airflow import common, util
 from sqlmesh.schedulers.airflow.operators import targets
 from sqlmesh.schedulers.airflow.operators.hwm_sensor import HighWaterMarkSensor
@@ -160,7 +165,10 @@ def _create_plan_application_dag(self, plan_dag_spec: common.PlanDagSpec) -> DAG
             ) = self._create_promotion_demotion_tasks(plan_dag_spec)
 
             start_task >> create_start_task
-            if not plan_dag_spec.forward_only or plan_dag_spec.is_dev:
+            if (
+                not has_paused_forward_only(plan_dag_spec.promoted_snapshots, all_snapshots)
+                or plan_dag_spec.is_dev
+            ):
                 create_end_task >> backfill_start_task
                 backfill_end_task >> promote_start_task
                 latest_end_task = promote_end_task

tests/core/test_snapshot.py

@@ -24,6 +24,7 @@
     SnapshotFingerprint,
     categorize_change,
     fingerprint_from_model,
+    has_paused_forward_only,
 )
 from sqlmesh.utils.date import to_datetime, to_timestamp
 from sqlmesh.utils.errors import SQLMeshError
@@ -915,3 +916,13 @@ def test_physical_schema(snapshot: Snapshot):
     assert new_snapshot.physical_schema == "custom_schema"
     assert new_snapshot.data_version.physical_schema == "custom_schema"
     assert new_snapshot.table_info.physical_schema == "custom_schema"
+
+
+def test_has_paused_forward_only(snapshot: Snapshot):
+    assert not has_paused_forward_only([snapshot], [snapshot])
+
+    snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+    assert has_paused_forward_only([snapshot], [snapshot])
+
+    snapshot.set_unpaused_ts("2023-01-01")
+    assert not has_paused_forward_only([snapshot], [snapshot])