Fix: No gaps check for non-representative snapshots (#1704)

Author: izeigerman

sqlmesh/core/plan/evaluator.py

@@ -101,7 +101,7 @@ def evaluate(self, plan: Plan) -> None:
             self._push(plan, deployability_index)
             self._restate(plan)
             self._backfill(plan, before_promote_snapshots, deployability_index)
-            promotion_result = self._promote(plan, deployability_index)
+            promotion_result = self._promote(plan, before_promote_snapshots)
             self._backfill(plan, after_promote_snapshots, deployability_index)
             self._update_views(plan, promotion_result, deployability_index)
 
@@ -176,16 +176,18 @@ def on_complete(snapshot: SnapshotInfoLike) -> None:
         self.state_sync.push_snapshots(plan.new_snapshots)
 
     def _promote(
-        self, plan: Plan, deployability_index: t.Optional[DeployabilityIndex] = None
+        self, plan: Plan, no_gaps_snapshot_names: t.Optional[t.Set[str]] = None
     ) -> PromotionResult:
         """Promote a plan.
 
         Args:
             plan: The plan to promote.
-            deployability_index: Indicates which snapshots are deployable in the context of this promotion.
+            no_gaps_snapshot_names: The names of snapshots to check for gaps if the no gaps check is enabled in the plan.
+            If not provided, all snapshots are checked.
         """
         promotion_result = self.state_sync.promote(
-            plan.environment, deployability_index=deployability_index, no_gaps=plan.no_gaps
+            plan.environment,
+            no_gaps_snapshot_names=no_gaps_snapshot_names if plan.no_gaps else set(),
         )
 
         if not plan.is_dev:

sqlmesh/core/state_sync/base.py

@@ -11,7 +11,6 @@
 from sqlmesh import migrations
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
 from sqlmesh.core.snapshot import (
-    DeployabilityIndex,
     Snapshot,
     SnapshotId,
     SnapshotIdLike,
@@ -326,17 +325,16 @@ def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.Lis
     def promote(
         self,
         environment: Environment,
-        deployability_index: t.Optional[DeployabilityIndex] = None,
-        no_gaps: bool = False,
+        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
     ) -> PromotionResult:
         """Update the environment to reflect the current state.
 
         This method verifies that snapshots have been pushed.
 
         Args:
             environment: The environment to promote.
-            deployability_index: Determines snapshots that are deployable in the context of this promotion.
-            no_gaps:  Whether to ensure that new snapshots for models that are already a
+            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
+                all snapshots will be checked. The data gap check ensures that models that are already a
                 part of the target environment have no data gaps when compared against previous
                 snapshots for same models.
 

sqlmesh/core/state_sync/common.py

@@ -10,7 +10,6 @@
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.snapshot import (
-    DeployabilityIndex,
     Snapshot,
     SnapshotId,
     SnapshotIdLike,
@@ -100,17 +99,16 @@ def get_environment(self, environment: str) -> t.Optional[Environment]:
     def promote(
         self,
         environment: Environment,
-        deployability_index: t.Optional[DeployabilityIndex] = None,
-        no_gaps: bool = False,
+        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
     ) -> PromotionResult:
         """Update the environment to reflect the current state.
 
         This method verifies that snapshots have been pushed.
 
         Args:
             environment: The environment to promote.
-            deployability_index: Determines snapshots that are deployable in the context of this promotion.
-            no_gaps:  Whether to ensure that new snapshots for models that are already a
+            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
+                all snapshots will be checked. The data gap check ensures that models that are already a
                 part of the target environment have no data gaps when compared against previous
                 snapshots for same models.
 
@@ -144,12 +142,12 @@ def promote(
                     "Please recreate the plan and try again"
                 )
 
-            if no_gaps:
+            if no_gaps_snapshot_names != set():
                 snapshots = self._get_snapshots(environment.snapshots).values()
                 self._ensure_no_gaps(
                     snapshots,
                     existing_environment,
-                    deployability_index or DeployabilityIndex.all_deployable(),
+                    no_gaps_snapshot_names,
                 )
 
             existing_table_infos = {
@@ -295,7 +293,7 @@ def _ensure_no_gaps(
         self,
         target_snapshots: t.Iterable[Snapshot],
         target_environment: Environment,
-        deployability_index: DeployabilityIndex,
+        snapshot_names: t.Optional[t.Set[str]],
     ) -> None:
         target_snapshots_by_name = {s.name: s for s in target_snapshots}
 
@@ -314,7 +312,7 @@ def _ensure_no_gaps(
         for prev_snapshot in prev_snapshots:
             target_snapshot = target_snapshots_by_name[prev_snapshot.name]
             if (
-                deployability_index.is_representative(target_snapshot)
+                (snapshot_names is None or prev_snapshot.name in snapshot_names)
                 and target_snapshot.is_incremental
                 and prev_snapshot.is_incremental
                 and prev_snapshot.intervals

sqlmesh/schedulers/airflow/common.py

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

sqlmesh/schedulers/airflow/dag_generator.py

@@ -297,8 +297,9 @@ def _create_promotion_demotion_tasks(
             python_callable=promotion_update_state_task,
             op_kwargs={
                 "environment": environment,
-                "deployability_index": request.deployability_index,
-                "no_gaps": request.no_gaps,
+                "no_gaps_snapshot_names": request.no_gaps_snapshot_names
+                if request.no_gaps
+                else set(),
             },
         )
 
@@ -571,11 +572,10 @@ def creation_update_state_task(new_snapshots: t.Iterable[Snapshot]) -> None:
 
 def promotion_update_state_task(
     environment: Environment,
-    deployability_index: DeployabilityIndex,
-    no_gaps: bool,
+    no_gaps_snapshot_names: t.Optional[t.Set[str]],
 ) -> None:
     with util.scoped_state_sync() as state_sync:
-        state_sync.promote(environment, deployability_index=deployability_index, no_gaps=no_gaps)
+        state_sync.promote(environment, no_gaps_snapshot_names=no_gaps_snapshot_names)
 
 
 def promotion_unpause_snapshots_task(

sqlmesh/schedulers/airflow/plan.py

@@ -146,6 +146,14 @@ def create_plan_dag_spec(
         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)}
+        if request.no_gaps and not request.is_dev
+        else None
+        if request.no_gaps
+        else set()
+    )
+
     return common.PlanDagSpec(
         request_id=request.request_id,
         environment_naming_info=request.environment.naming_info,
@@ -168,6 +176,7 @@ def create_plan_dag_spec(
         environment_expiration_ts=request.environment.expiration_ts,
         dag_start_ts=now_timestamp(),
         deployability_index=deployability_index,
+        no_gaps_snapshot_names=no_gaps_snapshot_names,
     )
 
 

sqlmesh/schedulers/airflow/state_sync.py

@@ -6,7 +6,6 @@
 from sqlmesh.core.console import Console
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.snapshot import (
-    DeployabilityIndex,
     Snapshot,
     SnapshotId,
     SnapshotIdLike,
@@ -231,17 +230,16 @@ def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.Lis
     def promote(
         self,
         environment: Environment,
-        deployability_index: t.Optional[DeployabilityIndex] = None,
-        no_gaps: bool = False,
+        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
     ) -> PromotionResult:
         """Update the environment to reflect the current state.
 
         This method verifies that snapshots have been pushed.
 
         Args:
             environment: The environment to promote.
-            deployability_index: Determines snapshots that are deployable in the context of this promotion.
-            no_gaps:  Whether to ensure that new snapshots for models that are already a
+            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
+                all snapshots will be checked. The data gap check ensures that models that are already a
                 part of the target environment have no data gaps when compared against previous
                 snapshots for same models.
 

tests/core/test_state_sync.py

@@ -22,7 +22,6 @@
     SqlModel,
 )
 from sqlmesh.core.snapshot import (
-    DeployabilityIndex,
     Snapshot,
     SnapshotChangeCategory,
     SnapshotId,
@@ -77,8 +76,8 @@ def promote_snapshots(
     state_sync: EngineAdapterStateSync,
     snapshots: t.List[Snapshot],
     environment: str,
-    deployability_index: t.Optional[DeployabilityIndex] = None,
     no_gaps: bool = False,
+    no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
     environment_suffix_target: EnvironmentSuffixTarget = EnvironmentSuffixTarget.SCHEMA,
 ) -> PromotionResult:
     env = Environment(
@@ -90,7 +89,9 @@ def promote_snapshots(
         plan_id="test_plan_id",
         previous_plan_id="test_plan_id",
     )
-    return state_sync.promote(env, deployability_index=deployability_index, no_gaps=no_gaps)
+    return state_sync.promote(
+        env, no_gaps_snapshot_names=no_gaps_snapshot_names if no_gaps else set()
+    )
 
 
 def delete_versions(state_sync: EngineAdapterStateSync) -> None:
@@ -671,8 +672,8 @@ def test_promote_snapshots_no_gaps(state_sync: EngineAdapterStateSync, make_snap
         state_sync,
         [new_snapshot_missing_interval],
         "prod",
-        deployability_index=DeployabilityIndex.none_deployable(),
         no_gaps=True,
+        no_gaps_snapshot_names=set(),
     )
 
 

tests/schedulers/airflow/test_plan.py

@@ -141,6 +141,8 @@ def test_create_plan_dag_spec(
     state_sync_mock.get_snapshot_intervals.return_value = []
     state_sync_mock.refresh_snapshot_intervals.return_value = []
 
+    expected_no_gaps_snapshot_names = {the_snapshot.name} if not paused_forward_only else set()
+
     with mock.patch(
         "sqlmesh.schedulers.airflow.plan.now_timestamp",
         side_effect=lambda: to_timestamp("2023-01-01"),
@@ -174,7 +176,7 @@ def test_create_plan_dag_spec(
         is_dev=False,
         forward_only=True,
         dag_start_ts=to_timestamp("2023-01-01"),
-        deployability_index=DeployabilityIndex.all_deployable(),
+        no_gaps_snapshot_names=expected_no_gaps_snapshot_names,
     )
 
     state_sync_mock.get_snapshots.assert_called_once()
@@ -286,7 +288,7 @@ def test_restatement(
         is_dev=False,
         forward_only=True,
         dag_start_ts=to_timestamp(now_value),
-        deployability_index=DeployabilityIndex.all_deployable(),
+        no_gaps_snapshot_names={the_snapshot.name},
     )
 
     state_sync_mock.get_snapshots.assert_called_once()
@@ -392,6 +394,7 @@ def test_select_models_for_backfill(mocker: MockerFixture, random_name, make_sna
         forward_only=True,
         dag_start_ts=to_timestamp("2023-01-01"),
         deployability_index=DeployabilityIndex.all_deployable(),
+        no_gaps_snapshot_names={"a", "b"},
     )