Fix: Initialize dev intervals with prod intervals for new forward-only snapshots (#1561)

Author: izeigerman

sqlmesh/core/console.py

@@ -609,9 +609,6 @@ def _prompt_effective_from(self, plan: Plan, auto_apply: bool) -> None:
             if effective_from:
                 plan.effective_from = effective_from
 
-        if plan.is_dev and plan.effective_from:
-            plan.set_start(plan.effective_from)
-
     def _prompt_backfill(self, plan: Plan, auto_apply: bool) -> None:
         is_forward_only_dev = plan.is_dev and plan.forward_only
         backfill_or_preview = "preview" if is_forward_only_dev else "backfill"

sqlmesh/core/plan/__init__.py

@@ -10,4 +10,5 @@
     MWAAPlanEvaluator,
     PlanEvaluator,
     can_evaluate_before_promote,
+    update_intervals_for_new_snapshots,
 )

sqlmesh/core/plan/definition.py

@@ -112,7 +112,7 @@ def __init__(
         self.__dag: t.Optional[DAG[str]] = None
 
         self._start = start
-        if not self._start and is_dev and (forward_only or self._has_paused_forward_only()):
+        if not self._start and is_dev and forward_only:
             self._start = default_start or yesterday_ds()
 
         self._end = end if end or not is_dev else (default_end or now())
@@ -167,11 +167,8 @@ def start(self) -> TimeLike:
     @start.setter
     def start(self, new_start: TimeLike) -> None:
         self._ensure_valid_date_range(new_start, self._end)
-        self.set_start(new_start)
-        self.override_start = True
-
-    def set_start(self, new_start: TimeLike) -> None:
         self._start = new_start
+        self.override_start = True
         self.__missing_intervals = None
         self._refresh_dag_and_ignored_snapshots()
 
@@ -371,6 +368,7 @@ def set_choice(self, snapshot: Snapshot, choice: SnapshotChangeCategory) -> None
         # Invalidate caches.
         self._categorized = None
         self._uncategorized = None
+        self.__missing_intervals = None
 
     @property
     def effective_from(self) -> t.Optional[TimeLike]:
@@ -393,6 +391,9 @@ def effective_from(self, effective_from: t.Optional[TimeLike]) -> None:
             effective_from: The effective date to set.
         """
         self._set_effective_from(effective_from)
+        if effective_from and self.is_dev and not self.override_start:
+            self._start = effective_from
+            self._refresh_dag_and_ignored_snapshots()
 
     def _set_effective_from(self, effective_from: t.Optional[TimeLike]) -> None:
         if not self.forward_only:
@@ -410,32 +411,30 @@ def _set_effective_from(self, effective_from: t.Optional[TimeLike]) -> None:
     @property
     def _missing_intervals(self) -> t.Dict[t.Tuple[str, str], Intervals]:
         if self.__missing_intervals is None:
-            # we need previous snapshots because this method is cached and users have the option
-            # to choose non-breaking / forward only. this will change the version of the snapshot on the fly
-            # thus changing the missing intervals. additionally we replace any snapshots with the old copies
-            # because they have intervals and the ephemeral ones don't
-            snapshots = {
-                (snapshot.name, snapshot.version_get_or_generate()): snapshot
-                for snapshot in self.snapshots
+            old_snapshots = {
+                (old.name, old.version_get_or_generate()): old
+                for _, old in self.context_diff.modified_snapshots.values()
             }
 
-            for new, old in self.context_diff.modified_snapshots.values():
-                # Never override forward-only snapshots to preserve the effect
-                # of the effective_from setting. Instead re-merge the intervals.
-                if not new.is_forward_only:
-                    snapshots[(old.name, old.version_get_or_generate())] = old
-                else:
-                    new.intervals = []
-                    new.merge_intervals(old)
+            for new in self.context_diff.new_snapshots.values():
+                new.intervals = []
+                new.dev_intervals = []
+                old = old_snapshots.get((new.name, new.version_get_or_generate()))
+                if not old:
+                    continue
+                new.merge_intervals(old)
+                if new.is_forward_only:
+                    new.dev_intervals = new.intervals.copy()
 
             self.__missing_intervals = {
                 (snapshot.name, snapshot.version_get_or_generate()): missing
                 for snapshot, missing in missing_intervals(
-                    snapshots.values(),
+                    self.snapshots,
                     start=self._start,
                     end=self._end,
                     execution_time=self._execution_time,
                     restatements=self.restatements,
+                    is_dev=self.is_dev,
                     ignore_cron=True,
                 ).items()
             }
@@ -728,12 +727,6 @@ def _build_snapshots_and_dag(
             ignored_snapshot_names,
         )
 
-    def _has_paused_forward_only(self) -> bool:
-        for name, snapshot in self._snapshot_mapping.items():
-            if snapshot.is_paused_forward_only or self._is_forward_only_model(name):
-                return True
-        return False
-
     def _is_forward_only_model(self, model_name: str) -> bool:
         def _is_forward_only_expected(snapshot: Snapshot) -> bool:
             # Returns True if the snapshot is not categorized yet but is expected

sqlmesh/core/plan/evaluator.py

@@ -86,6 +86,8 @@ def evaluate(self, plan: Plan) -> None:
             }
             after_promote_snapshots = all_names - before_promote_snapshots
 
+        update_intervals_for_new_snapshots(plan.new_snapshots, self.state_sync)
+
         self._push(plan)
         self._restate(plan)
         self._backfill(plan, before_promote_snapshots)
@@ -377,3 +379,13 @@ def can_evaluate_before_promote(
     return not snapshot.is_paused_forward_only and not any(
         snapshots[p_id].is_paused_forward_only for p_id in snapshot.parents
     )
+
+
+def update_intervals_for_new_snapshots(
+    snapshots: t.Collection[Snapshot], state_sync: StateSync
+) -> None:
+    for snapshot in state_sync.refresh_snapshot_intervals(snapshots):
+        if snapshot.is_forward_only:
+            snapshot.dev_intervals = snapshot.intervals.copy()
+            for start, end in snapshot.dev_intervals:
+                state_sync.add_interval(snapshot, start, end, is_dev=True)

sqlmesh/core/scheduler.py

@@ -15,7 +15,6 @@
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotEvaluator,
-    SnapshotIdLike,
     earliest_start_date,
     missing_intervals,
 )
@@ -44,7 +43,7 @@ class Scheduler:
     The scheduler comes equipped with a simple ThreadPoolExecutor based evaluation engine.
 
     Args:
-        snapshots: A collection of snapshots/ids.
+        snapshots: A collection of snapshots.
         snapshot_evaluator: The snapshot evaluator to execute queries.
         state_sync: The state sync to pull saved snapshots.
         max_workers: The maximum number of parallel queries to run.
@@ -53,15 +52,15 @@ class Scheduler:
 
     def __init__(
         self,
-        snapshots: t.Iterable[SnapshotIdLike],
+        snapshots: t.Iterable[Snapshot],
         snapshot_evaluator: SnapshotEvaluator,
         state_sync: StateSync,
         max_workers: int = 1,
         console: t.Optional[Console] = None,
         notification_target_manager: t.Optional[NotificationTargetManager] = None,
     ):
         self.state_sync = state_sync
-        self.snapshots = self.state_sync.get_snapshots(snapshots)
+        self.snapshots = {s.snapshot_id: s for s in snapshots}
         self.snapshot_per_version = _resolve_one_snapshot_per_version(self.snapshots.values())
         self.snapshot_evaluator = snapshot_evaluator
         self.max_workers = max_workers
@@ -107,6 +106,8 @@ def batches(
         if selected_snapshots is not None:
             snapshots = [s for s in snapshots if s.name in selected_snapshots]
 
+        self.state_sync.refresh_snapshot_intervals(snapshots)
+
         return compute_interval_params(
             snapshots,
             start=start or earliest_start_date(snapshots),

sqlmesh/core/snapshot/definition.py

@@ -691,11 +691,7 @@ def missing_intervals(
         Returns:
             A list of all the missing intervals as epoch timestamps.
         """
-        intervals = (
-            self.dev_intervals
-            if is_dev and self.is_forward_only and self.is_paused
-            else self.intervals
-        )
+        intervals = self.dev_intervals if is_dev and self.is_paused_forward_only else self.intervals
 
         if self.is_symbolic or (self.is_seed and intervals):
             return []
@@ -775,7 +771,6 @@ def categorize_as(self, category: SnapshotChangeCategory) -> None:
         Args:
             category: The change category to assign to this snapshot.
         """
-        assert not self.intervals  # a snapshot that has been processed should not be recategorized
         is_forward_only = category in (
             SnapshotChangeCategory.FORWARD_ONLY,
             SnapshotChangeCategory.INDIRECT_NON_BREAKING,

sqlmesh/core/state_sync/base.py

@@ -275,6 +275,17 @@ def remove_interval(
             all_snapshots: All snapshots can be passed in to skip fetching matching snapshot versions.
         """
 
+    @abc.abstractmethod
+    def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.List[Snapshot]:
+        """Updates given snapshots with latest intervals from the state.
+
+        Args:
+            snapshots: The snapshots to refresh.
+
+        Returns:
+            The updated snapshots.
+        """
+
     @abc.abstractmethod
     def promote(self, environment: Environment, no_gaps: bool = False) -> PromotionResult:
         """Update the environment to reflect the current state.

sqlmesh/core/state_sync/engine_adapter.py

@@ -537,11 +537,12 @@ def remove_interval(
             snapshot_ids = ", ".join(str(s.snapshot_id) for s, _ in snapshot_intervals)
             logger.info("Removing interval for snapshots: %s", snapshot_ids)
 
-        self.engine_adapter.insert_append(
-            self.intervals_table,
-            _intervals_to_df(snapshot_intervals, is_dev=False, is_removed=True),
-            columns_to_types=self._interval_columns_to_types,
-        )
+        for is_dev in (True, False):
+            self.engine_adapter.insert_append(
+                self.intervals_table,
+                _intervals_to_df(snapshot_intervals, is_dev=is_dev, is_removed=True),
+                columns_to_types=self._interval_columns_to_types,
+            )
 
     @transactional()
     def compact_intervals(self) -> None:
@@ -558,6 +559,16 @@ def compact_intervals(self) -> None:
                 self.intervals_table, exp.column("id").isin(*interval_ids)
             )
 
+    def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.List[Snapshot]:
+        if not snapshots:
+            return []
+
+        _, intervals = self._get_snapshot_intervals(snapshots)
+        for s in snapshots:
+            s.intervals = []
+            s.dev_intervals = []
+        return Snapshot.hydrate_with_intervals_by_version(snapshots, intervals)
+
     def max_interval_end_for_environment(self, environment: str) -> t.Optional[int]:
         env = self._get_environment(environment)
         if not env:

sqlmesh/schedulers/airflow/plan.py

@@ -8,7 +8,10 @@
 from sqlmesh.core import scheduler
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment
-from sqlmesh.core.plan import can_evaluate_before_promote
+from sqlmesh.core.plan import (
+    can_evaluate_before_promote,
+    update_intervals_for_new_snapshots,
+)
 from sqlmesh.core.snapshot import SnapshotTableInfo
 from sqlmesh.core.state_sync import EngineAdapterStateSync, StateSync
 from sqlmesh.core.state_sync.base import DelegatingStateSync
@@ -94,6 +97,8 @@ def create_plan_dag_spec(
             "Make sure your code base is up to date and try re-creating the plan"
         )
 
+    update_intervals_for_new_snapshots(new_snapshots.values(), state_sync)
+
     if request.environment.end_at:
         end = request.environment.end_at
         unpaused_dt = None
@@ -104,14 +109,17 @@ def create_plan_dag_spec(
         unpaused_dt = end
 
     if request.restatements:
+        intervals_to_remove = [
+            (s, request.restatements[s.name])
+            for s in all_snapshots.values()
+            if s.name in request.restatements and s.snapshot_id not in new_snapshots
+        ]
         state_sync.remove_interval(
-            [
-                (s, request.restatements[s.name])
-                for s in all_snapshots.values()
-                if s.name in request.restatements and s.snapshot_id not in new_snapshots
-            ],
+            intervals_to_remove,
             remove_shared_versions=not request.is_dev,
         )
+        for s, interval in intervals_to_remove:
+            all_snapshots[s.snapshot_id].remove_interval(interval)
 
     if not request.skip_backfill:
         backfill_batches = scheduler.compute_interval_params(

sqlmesh/schedulers/airflow/state_sync.py

@@ -208,6 +208,19 @@ def remove_interval(
         """
         raise NotImplementedError("Removing intervals is not supported by the Airflow state sync.")
 
+    def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.List[Snapshot]:
+        """Updates given snapshots with latest intervals from the state.
+
+        Args:
+            snapshots: The snapshots to refresh.
+
+        Returns:
+            The updated snapshots.
+        """
+        raise NotImplementedError(
+            "Refreshing snapshot intervals is not supported by the Airflow state sync."
+        )
+
     def promote(self, environment: Environment, no_gaps: bool = False) -> PromotionResult:
         """Update the environment to reflect the current state.