Fix: Evaluate a snapshot after promotion if it has paused forward-only parents (#1522)

Author: izeigerman

sqlmesh/core/plan/__init__.py

@@ -9,4 +9,5 @@
     BuiltInPlanEvaluator,
     MWAAPlanEvaluator,
     PlanEvaluator,
+    can_evaluate_before_promote,
 )

sqlmesh/core/plan/evaluator.py

@@ -72,13 +72,16 @@ def __init__(
         self.__all_snapshots: t.Dict[str, t.Dict[SnapshotId, Snapshot]] = {}
 
     def evaluate(self, plan: Plan) -> None:
-        snapshots = plan.snapshots
+        snapshots = {s.snapshot_id: s for s in plan.snapshots}
+        all_names = {s.name for s in plan.snapshots}
         if plan.is_dev:
-            before_promote_snapshots = {s.name for s in snapshots}
+            before_promote_snapshots = all_names
             after_promote_snapshots = set()
         else:
-            before_promote_snapshots = {s.name for s in snapshots if not s.is_paused_forward_only}
-            after_promote_snapshots = {s.name for s in snapshots if s.is_paused_forward_only}
+            before_promote_snapshots = {
+                s.name for s in snapshots.values() if can_evaluate_before_promote(s, snapshots)
+            }
+            after_promote_snapshots = all_names - before_promote_snapshots
 
         self._push(plan)
         self._restate(plan)
@@ -368,3 +371,11 @@ def _apply_plan(self, plan: Plan, plan_request_id: str) -> None:
 
         if stderr:
             raise SQLMeshError(f"Failed to submit a plan application request:\n{stderr}")
+
+
+def can_evaluate_before_promote(
+    snapshot: Snapshot, snapshots: t.Dict[SnapshotId, Snapshot]
+) -> bool:
+    return not snapshot.is_paused_forward_only and not any(
+        snapshots[p_id].is_paused_forward_only for p_id in snapshot.parents
+    )

sqlmesh/schedulers/airflow/plan.py

@@ -4,6 +4,7 @@
 
 from sqlmesh.core import scheduler
 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.schedulers.airflow import common
@@ -59,11 +60,11 @@ def create_plan_dag_spec(
 
     backfill_intervals_per_snapshot = [
         common.BackfillIntervalsPerSnapshot(
-            snapshot_id=snapshot.snapshot_id,
+            snapshot_id=s.snapshot_id,
             intervals=intervals,
-            before_promote=request.is_dev or not snapshot.is_paused_forward_only,
+            before_promote=request.is_dev or can_evaluate_before_promote(s, all_snapshots),
         )
-        for snapshot, intervals in backfill_batches.items()
+        for s, intervals in backfill_batches.items()
     ]
 
     return common.PlanDagSpec(

tests/core/test_plan_evaluator.py

@@ -3,14 +3,16 @@
 from sqlglot import parse_one
 
 from sqlmesh.core.context import Context
-from sqlmesh.core.model import FullKind, SqlModel, ViewKind
+from sqlmesh.core.model import FullKind, IncrementalByTimeRangeKind, SqlModel, ViewKind
 from sqlmesh.core.plan import (
     AirflowPlanEvaluator,
     BuiltInPlanEvaluator,
     MWAAPlanEvaluator,
     Plan,
+    can_evaluate_before_promote,
 )
 from sqlmesh.core.snapshot import SnapshotChangeCategory
+from sqlmesh.utils.date import now_timestamp
 from sqlmesh.utils.errors import SQLMeshError
 
 
@@ -162,3 +164,61 @@ def test_mwaa_evaluator_error_from_cli(sushi_plan: Plan, mocker: MockerFixture):
         evaluator.evaluate(sushi_plan)
 
     mwaa_client_mock.set_variable.assert_called_once_with(mocker.ANY, plan_dag_spec_json)
+
+
+def test_can_evaluate_before_promote(sushi_context: Context):
+    parent_model_a = SqlModel(
+        name="sushi.new_test_model_a",
+        kind=IncrementalByTimeRangeKind(time_column="ds"),
+        cron="@daily",
+        start="2020-01-01",
+        query=parse_one("SELECT 1::INT AS one, '2023-01-01' as ds"),
+    )
+    parent_model_b = SqlModel(
+        name="sushi.new_test_model_b",
+        kind=IncrementalByTimeRangeKind(time_column="ds"),
+        cron="@daily",
+        start="2020-01-01",
+        query=parse_one("SELECT 2::INT AS two, '2023-01-01' as ds"),
+    )
+    child_model = SqlModel(
+        name="sushi.new_test_model_child",
+        kind=FullKind(),
+        start="2020-01-01",
+        query=parse_one("SELECT one, two FROM sushi.new_test_model_a, sushi.new_test_model_b"),
+    )
+
+    sushi_context.upsert_model(parent_model_a)
+    sushi_context.upsert_model(parent_model_b)
+    sushi_context.upsert_model(child_model)
+
+    snapshots = sushi_context.snapshots
+
+    parent_snapshot_a = snapshots[parent_model_a.name]
+    parent_snapshot_b = snapshots[parent_model_b.name]
+    child_snapshot = snapshots[child_model.name]
+
+    all_snapshots = {
+        s.snapshot_id: s for s in [parent_snapshot_a, parent_snapshot_b, child_snapshot]
+    }
+
+    parent_snapshot_a.change_category = SnapshotChangeCategory.BREAKING
+    parent_snapshot_b.change_category = SnapshotChangeCategory.BREAKING
+    child_snapshot.change_category = SnapshotChangeCategory.BREAKING
+    assert can_evaluate_before_promote(child_snapshot, all_snapshots)
+
+    parent_snapshot_a.change_category = SnapshotChangeCategory.FORWARD_ONLY
+    parent_snapshot_b.change_category = SnapshotChangeCategory.FORWARD_ONLY
+    assert not can_evaluate_before_promote(child_snapshot, all_snapshots)
+
+    parent_snapshot_a.unpaused_ts = now_timestamp()
+    assert not can_evaluate_before_promote(child_snapshot, all_snapshots)
+
+    parent_snapshot_b.unpaused_ts = now_timestamp()
+    assert can_evaluate_before_promote(child_snapshot, all_snapshots)
+
+    child_snapshot.change_category = SnapshotChangeCategory.FORWARD_ONLY
+    assert not can_evaluate_before_promote(child_snapshot, all_snapshots)
+
+    child_snapshot.unpaused_ts = now_timestamp()
+    assert can_evaluate_before_promote(child_snapshot, all_snapshots)