Fix: Child snapshot targeting a wrong parent forward-only snapshot table when the parent was created in dev and the child in prod (#2019)

Author: izeigerman

sqlmesh/core/plan/evaluator.py

@@ -83,32 +83,39 @@ def evaluate(
             all_names = {
                 s.name for s in snapshots.values() if plan.is_selected_for_backfill(s.name)
             }
-            deployability_index = DeployabilityIndex.create(snapshots)
+            deployability_index_for_evaluation = DeployabilityIndex.create(snapshots)
+            deployability_index_for_creation = deployability_index_for_evaluation
             if plan.is_dev:
                 before_promote_snapshots = all_names
                 after_promote_snapshots = set()
             else:
                 before_promote_snapshots = {
                     s.name
                     for s in snapshots.values()
-                    if deployability_index.is_representative(s)
+                    if deployability_index_for_evaluation.is_representative(s)
                     and plan.is_selected_for_backfill(s.name)
                 }
                 after_promote_snapshots = all_names - before_promote_snapshots
-                deployability_index = DeployabilityIndex.all_deployable()
+                deployability_index_for_evaluation = DeployabilityIndex.all_deployable()
 
             update_intervals_for_new_snapshots(plan.new_snapshots, self.state_sync)
 
-            self._push(plan, deployability_index)
+            self._push(plan, deployability_index_for_creation)
             self._restate(plan)
             self._backfill(
-                plan, before_promote_snapshots, deployability_index, circuit_breaker=circuit_breaker
+                plan,
+                before_promote_snapshots,
+                deployability_index_for_evaluation,
+                circuit_breaker=circuit_breaker,
             )
             promotion_result = self._promote(plan, before_promote_snapshots)
             self._backfill(
-                plan, after_promote_snapshots, deployability_index, circuit_breaker=circuit_breaker
+                plan,
+                after_promote_snapshots,
+                deployability_index_for_evaluation,
+                circuit_breaker=circuit_breaker,
             )
-            self._update_views(plan, promotion_result, deployability_index)
+            self._update_views(plan, promotion_result, deployability_index_for_evaluation)
 
             if not plan.requires_backfill:
                 self.console.log_success("Virtual Update executed successfully")

sqlmesh/core/snapshot/definition.py

@@ -1062,7 +1062,7 @@ class DeployabilityIndex(PydanticModel, frozen=True):
 
     indexed_ids: t.FrozenSet[str]
     is_opposite_index: bool = False
-    representative_shared_version_ids: t.Optional[t.FrozenSet[str]] = None
+    representative_shared_version_ids: t.FrozenSet[str] = frozenset()
 
     @field_validator("indexed_ids", "representative_shared_version_ids", mode="before")
     @classmethod
@@ -1109,10 +1109,7 @@ def is_representative(self, snapshot: SnapshotIdLike) -> bool:
             True if the snapshot is representative, False otherwise.
         """
         snapshot_id = self._snapshot_id_key(snapshot.snapshot_id)
-        representative = (
-            self.representative_shared_version_ids is not None
-            and snapshot_id in self.representative_shared_version_ids
-        )
+        representative = snapshot_id in self.representative_shared_version_ids
         return representative or self.is_deployable(snapshot)
 
     def with_non_deployable(self, snapshot: SnapshotIdLike) -> DeployabilityIndex:

sqlmesh/core/snapshot/evaluator.py

@@ -586,7 +586,7 @@ def _create_snapshot(
                     self.adapter.drop_table(tmp_table_name)
             else:
                 table_deployability_flags = [False]
-                if not snapshot.is_indirect_non_breaking:
+                if not snapshot.is_indirect_non_breaking and not snapshot.is_forward_only:
                     table_deployability_flags.append(True)
                 for is_table_deployable in table_deployability_flags:
                     evaluation_strategy.create(

sqlmesh/schedulers/airflow/common.py

@@ -83,6 +83,7 @@ class PlanDagSpec(PydanticModel):
     environment_expiration_ts: t.Optional[int] = None
     dag_start_ts: t.Optional[int] = None
     deployability_index: DeployabilityIndex = DeployabilityIndex.all_deployable()
+    deployability_index_for_creation: DeployabilityIndex = DeployabilityIndex.all_deployable()
     no_gaps_snapshot_names: t.Optional[t.Set[str]] = None
 
 

sqlmesh/schedulers/airflow/dag_generator.py

@@ -172,7 +172,7 @@ def _create_plan_application_dag(self, plan_dag_spec: common.PlanDagSpec) -> DAG
             (create_start_task, create_end_task) = self._create_creation_tasks(
                 plan_dag_spec.new_snapshots,
                 plan_dag_spec.ddl_concurrent_tasks,
-                plan_dag_spec.deployability_index,
+                plan_dag_spec.deployability_index_for_creation,
             )
 
             (

sqlmesh/schedulers/airflow/plan.py

@@ -114,9 +114,9 @@ def create_plan_dag_spec(
         for s, interval in intervals_to_remove:
             all_snapshots[s.snapshot_id].remove_interval(interval)
 
-    initial_deployability_index = DeployabilityIndex.create(all_snapshots)
-    deployability_index = (
-        initial_deployability_index if request.is_dev else DeployabilityIndex.all_deployable()
+    deployability_index_for_creation = DeployabilityIndex.create(all_snapshots)
+    deployability_index_for_evaluation = (
+        deployability_index_for_creation if request.is_dev else DeployabilityIndex.all_deployable()
     )
 
     if not request.skip_backfill:
@@ -125,7 +125,7 @@ def create_plan_dag_spec(
             start=request.environment.start_at,
             end=end,
             execution_time=now(),
-            deployability_index=deployability_index,
+            deployability_index=deployability_index_for_evaluation,
             restatements=request.restatements,
             ignore_cron=True,
         )
@@ -136,13 +136,17 @@ def create_plan_dag_spec(
         common.BackfillIntervalsPerSnapshot(
             snapshot_id=s.snapshot_id,
             intervals=intervals,
-            before_promote=request.is_dev or initial_deployability_index.is_representative(s),
+            before_promote=request.is_dev or deployability_index_for_creation.is_representative(s),
         )
         for s, intervals in backfill_batches.items()
     ]
 
     no_gaps_snapshot_names = (
-        {s.name for s in all_snapshots.values() if initial_deployability_index.is_representative(s)}
+        {
+            s.name
+            for s in all_snapshots.values()
+            if deployability_index_for_creation.is_representative(s)
+        }
         if request.no_gaps and not request.is_dev
         else None
         if request.no_gaps
@@ -170,7 +174,8 @@ def create_plan_dag_spec(
         forward_only=request.forward_only,
         environment_expiration_ts=request.environment.expiration_ts,
         dag_start_ts=to_timestamp(now_dt),
-        deployability_index=deployability_index,
+        deployability_index=deployability_index_for_evaluation,
+        deployability_index_for_creation=deployability_index_for_creation,
         no_gaps_snapshot_names=no_gaps_snapshot_names,
     )
 

tests/core/test_integration.py

@@ -283,6 +283,60 @@ def test_forward_only_model_regular_plan(init_and_plan_context: t.Callable):
     assert not prod_df["event_date"].tolist()
 
 
+@freeze_time("2023-01-08 15:00:00")
+def test_forward_only_parent_created_in_dev_child_created_in_prod(
+    init_and_plan_context: t.Callable,
+):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    waiter_revenue_by_day_model = context.get_model("sushi.waiter_revenue_by_day")
+    waiter_revenue_by_day_model = add_projection_to_model(
+        t.cast(SqlModel, waiter_revenue_by_day_model)
+    )
+    forward_only_kind = waiter_revenue_by_day_model.kind.copy(update={"forward_only": True})
+    waiter_revenue_by_day_model = waiter_revenue_by_day_model.copy(
+        update={"kind": forward_only_kind}
+    )
+    context.upsert_model(waiter_revenue_by_day_model)
+
+    waiter_revenue_by_day_snapshot = context.get_snapshot(
+        waiter_revenue_by_day_model, raise_if_missing=True
+    )
+    top_waiters_snapshot = context.get_snapshot("sushi.top_waiters", raise_if_missing=True)
+
+    plan = context.plan("dev", no_prompts=True, skip_tests=True)
+    assert len(plan.new_snapshots) == 2
+    assert (
+        plan.context_diff.snapshots[waiter_revenue_by_day_snapshot.snapshot_id].change_category
+        == SnapshotChangeCategory.FORWARD_ONLY
+    )
+    assert (
+        plan.context_diff.snapshots[top_waiters_snapshot.snapshot_id].change_category
+        == SnapshotChangeCategory.FORWARD_ONLY
+    )
+    assert plan.start == to_datetime("2023-01-01")
+    assert not plan.missing_intervals
+
+    context.apply(plan)
+
+    # Update the child to refer to a newly added column.
+    top_waiters_model = context.get_model("sushi.top_waiters")
+    top_waiters_model = add_projection_to_model(t.cast(SqlModel, top_waiters_model), literal=False)
+    context.upsert_model(top_waiters_model)
+
+    top_waiters_snapshot = context.get_snapshot("sushi.top_waiters", raise_if_missing=True)
+
+    plan = context.plan("prod", no_prompts=True, skip_tests=True)
+    assert len(plan.new_snapshots) == 1
+    assert (
+        plan.context_diff.snapshots[top_waiters_snapshot.snapshot_id].change_category
+        == SnapshotChangeCategory.NON_BREAKING
+    )
+
+    context.apply(plan)
+
+
 @freeze_time("2023-01-08 15:00:00")
 def test_plan_set_choice_is_reflected_in_missing_intervals(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")

tests/core/test_snapshot_evaluator.py

@@ -954,10 +954,6 @@ def test_forward_only_snapshot_for_added_model(mocker: MockerFixture, adapter_mo
                 f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp",
                 **common_create_args,
             ),
-            call(
-                f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}",
-                **common_create_args,
-            ),
         ]
     )
 

tests/schedulers/airflow/test_plan.py

@@ -181,6 +181,9 @@ def test_create_plan_dag_spec(
         forward_only=True,
         dag_start_ts=to_timestamp("2023-01-01"),
         no_gaps_snapshot_names=expected_no_gaps_snapshot_names,
+        deployability_index_for_creation=DeployabilityIndex.all_deployable()
+        if not paused_forward_only
+        else DeployabilityIndex.none_deployable(),
     )
 
     state_sync_mock.get_snapshots.assert_called_once()