adjust intervals based on force rebuild at runtime

Author: izeigerman

sqlmesh/core/plan/builder.py

@@ -16,6 +16,7 @@
 )
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.environment import EnvironmentNamingInfo
+from sqlmesh.core.plan.common import should_force_rebuild
 from sqlmesh.core.plan.definition import (
     Plan,
     SnapshotMapping,
@@ -276,7 +277,7 @@ def build(self) -> Plan:
 
         self._check_destructive_changes(directly_modified)
         self._categorize_snapshots(dag, indirectly_modified)
-        self._adjust_new_snapshot_intervals()
+        self._adjust_snapshot_intervals()
 
         deployability_index = (
             DeployabilityIndex.create(
@@ -508,21 +509,22 @@ def _build_models_to_backfill(
             ).sorted
         }
 
-    def _adjust_new_snapshot_intervals(self) -> None:
-        old_snapshots = {
-            (old.name, old.version_get_or_generate()): old
-            for _, old in self._context_diff.modified_snapshots.values()
-        }
-
-        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:
+    def _adjust_snapshot_intervals(self) -> None:
+        for new, old in self._context_diff.modified_snapshots.values():
+            if not new.is_model or not old.is_model:
                 continue
-            new.merge_intervals(old)
-            if new.is_forward_only:
-                new.dev_intervals = new.intervals.copy()
+            is_same_version = old.version_get_or_generate() == new.version_get_or_generate()
+            if is_same_version and should_force_rebuild(old, new):
+                # If the difference between 2 snapshots requires a full rebuild,
+                # then clear the intervals for the new snapshot.
+                self._context_diff.snapshots[new.snapshot_id].intervals = []
+            elif new.snapshot_id in self._context_diff.new_snapshots:
+                new.intervals = []
+                new.dev_intervals = []
+                if is_same_version:
+                    new.merge_intervals(old)
+                    if new.is_forward_only:
+                        new.dev_intervals = new.intervals.copy()
 
     def _check_destructive_changes(self, directly_modified: t.Set[SnapshotId]) -> None:
         for s_id in sorted(directly_modified):
@@ -589,7 +591,7 @@ def _categorize_snapshots(
             forward_only = self._forward_only or self._is_forward_only_change(s_id)
             if forward_only and s_id.name in self._context_diff.modified_snapshots:
                 new, old = self._context_diff.modified_snapshots[s_id.name]
-                if _should_force_rebuild(old, new) or snapshot.is_seed:
+                if should_force_rebuild(old, new) or snapshot.is_seed:
                     # Breaking kind changes and seed changes can't be forward-only.
                     forward_only = False
 
@@ -614,7 +616,7 @@ def _categorize_snapshot(
         if self._context_diff.directly_modified(s_id.name):
             if self._auto_categorization_enabled:
                 new, old = self._context_diff.modified_snapshots[s_id.name]
-                if _should_force_rebuild(old, new):
+                if should_force_rebuild(old, new):
                     snapshot.categorize_as(SnapshotChangeCategory.BREAKING, False)
                     return
 
@@ -772,7 +774,7 @@ def _is_forward_only_change(self, s_id: SnapshotId) -> bool:
         if snapshot.name in self._context_diff.modified_snapshots:
             _, old = self._context_diff.modified_snapshots[snapshot.name]
             # If the model kind has changed in a breaking way, then we can't consider this to be a forward-only change.
-            if snapshot.is_model and _should_force_rebuild(old, snapshot):
+            if snapshot.is_model and should_force_rebuild(old, snapshot):
                 return False
         return (
             snapshot.is_model and snapshot.model.forward_only and bool(snapshot.previous_versions)
@@ -870,19 +872,3 @@ def _modified_and_added_snapshots(self) -> t.List[Snapshot]:
             if snapshot.name in self._context_diff.modified_snapshots
             or snapshot.snapshot_id in self._context_diff.added
         ]
-
-
-def _should_force_rebuild(old: Snapshot, new: Snapshot) -> bool:
-    if old.virtual_environment_mode != new.virtual_environment_mode:
-        # If the virtual environment mode has changed, then we need to rebuild
-        return True
-    if old.model.kind.name == new.model.kind.name:
-        # If the kind hasn't changed, then we don't need to rebuild
-        return False
-    if not old.is_incremental or not new.is_incremental:
-        # If either is not incremental, then we need to rebuild
-        return True
-    if old.model.partitioned_by == new.model.partitioned_by:
-        # If the partitioning hasn't changed, then we don't need to rebuild
-        return False
-    return True

sqlmesh/core/plan/common.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from sqlmesh.core.snapshot import Snapshot
+
+
+def should_force_rebuild(old: Snapshot, new: Snapshot) -> bool:
+    if old.virtual_environment_mode != new.virtual_environment_mode:
+        # If the virtual environment mode has changed, then we need to rebuild
+        return True
+    if old.model.kind.name == new.model.kind.name:
+        # If the kind hasn't changed, then we don't need to rebuild
+        return False
+    if not old.is_incremental or not new.is_incremental:
+        # If either is not incremental, then we need to rebuild
+        return True
+    if old.model.partitioned_by == new.model.partitioned_by:
+        # If the partitioning hasn't changed, then we don't need to rebuild
+        return False
+    return True

sqlmesh/core/plan/stages.py

@@ -1,7 +1,8 @@
 import typing as t
 
 from dataclasses import dataclass
-from sqlmesh.core.environment import EnvironmentStatements, EnvironmentNamingInfo
+from sqlmesh.core.environment import EnvironmentStatements, EnvironmentNamingInfo, Environment
+from sqlmesh.core.plan.common import should_force_rebuild
 from sqlmesh.core.plan.definition import EvaluatablePlan
 from sqlmesh.core.state_sync import StateReader
 from sqlmesh.core.scheduler import merged_missing_intervals, SnapshotToIntervals
@@ -230,8 +231,9 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
         all_selected_for_backfill_snapshots = {
             s.snapshot_id for s in snapshots.values() if plan.is_selected_for_backfill(s.name)
         }
+        existing_environment = self.state_reader.get_environment(plan.environment.name)
 
-        self._adjust_intervals(snapshots_by_name, plan)
+        self._adjust_intervals(snapshots_by_name, plan, existing_environment)
 
         deployability_index = DeployabilityIndex.create(snapshots, start=plan.start)
         deployability_index_for_creation = deployability_index
@@ -269,7 +271,7 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
                     missing_intervals_after_promote[snapshot] = intervals
 
         promoted_snapshots, demoted_snapshots, demoted_environment_naming_info = (
-            self._get_promoted_demoted_snapshots(plan)
+            self._get_promoted_demoted_snapshots(plan, existing_environment)
         )
 
         stages: t.List[PlanStage] = []
@@ -459,11 +461,10 @@ def _should_update_virtual_layer(snapshot: SnapshotTableInfo) -> bool:
         )
 
     def _get_promoted_demoted_snapshots(
-        self, plan: EvaluatablePlan
+        self, plan: EvaluatablePlan, existing_environment: t.Optional[Environment]
     ) -> t.Tuple[
         t.Set[SnapshotTableInfo], t.Set[SnapshotTableInfo], t.Optional[EnvironmentNamingInfo]
     ]:
-        existing_environment = self.state_reader.get_environment(plan.environment.name)
         if existing_environment:
             new_table_infos = {
                 table_info.name: table_info for table_info in plan.environment.promoted_snapshots
@@ -579,10 +580,40 @@ def _should_create(s: Snapshot) -> bool:
         return [s for s in snapshots.values() if _should_create(s)]
 
     def _adjust_intervals(
-        self, snapshots_by_name: t.Dict[str, Snapshot], plan: EvaluatablePlan
+        self,
+        snapshots_by_name: t.Dict[str, Snapshot],
+        plan: EvaluatablePlan,
+        existing_environment: t.Optional[Environment],
     ) -> None:
         # Make sure the intervals are up to date and restatements are reflected
         self.state_reader.refresh_snapshot_intervals(snapshots_by_name.values())
+
+        if existing_environment:
+            new_snapshot_ids = set()
+            new_snapshot_versions = set()
+            for s in snapshots_by_name.values():
+                if s.is_model:
+                    new_snapshot_ids.add(s.snapshot_id)
+                    new_snapshot_versions.add(s.name_version)
+            # Only compare to old snapshots that share the same version as the new snapshots
+            old_snapshot_ids = {
+                s.snapshot_id
+                for s in existing_environment.snapshots
+                if s.is_model
+                and s.name_version in new_snapshot_versions
+                and s.snapshot_id not in new_snapshot_ids
+            }
+            if old_snapshot_ids:
+                old_snapshots = self.state_reader.get_snapshots(old_snapshot_ids)
+                for old in old_snapshots.values():
+                    new = snapshots_by_name.get(old.name)
+                    if not new or old.version != new.version:
+                        continue
+                    if should_force_rebuild(old, new):
+                        # If the difference between 2 snapshots requires a full rebuild,
+                        # then clear the intervals for the new snapshot.
+                        new.intervals = []
+
         for new_snapshot in plan.new_snapshots:
             if new_snapshot.is_forward_only:
                 # Forward-only snapshots inherit intervals in dev because of cloning

sqlmesh/core/snapshot/definition.py

@@ -893,19 +893,15 @@ def merge_intervals(self, other: t.Union[Snapshot, SnapshotIntervals]) -> None:
         Args:
             other: The target snapshot to inherit intervals from.
         """
-        if self.is_no_rebuild or self.virtual_environment_mode.is_full or not self.is_paused:
-            # If the virtual environment mode is not full we can only merge prod intervals if this snapshot
-            # is currently promoted in production or if it's forward-only / metadata / indirect non-breaking.
-            # Otherwise, we want to ignore any existing intervals and backfill this snapshot from scratch.
-            effective_from_ts = self.normalized_effective_from_ts or 0
-            apply_effective_from = effective_from_ts > 0 and self.identifier != other.identifier
-            for start, end in other.intervals:
-                # If the effective_from is set, then intervals that come after it must come from
-                # the current snapshots.
-                if apply_effective_from and start < effective_from_ts:
-                    end = min(end, effective_from_ts)
-                if not apply_effective_from or end <= effective_from_ts:
-                    self.add_interval(start, end)
+        effective_from_ts = self.normalized_effective_from_ts or 0
+        apply_effective_from = effective_from_ts > 0 and self.identifier != other.identifier
+        for start, end in other.intervals:
+            # If the effective_from is set, then intervals that come after it must come from
+            # the current snapshots.
+            if apply_effective_from and start < effective_from_ts:
+                end = min(end, effective_from_ts)
+            if not apply_effective_from or end <= effective_from_ts:
+                self.add_interval(start, end)
 
         if self.dev_version == other.dev_version:
             # Merge dev intervals if the dev versions match which would mean

tests/core/test_integration.py

@@ -2694,7 +2694,11 @@ def test_restatement_plan_across_environments_snapshot_with_shared_version(
     assert isinstance(previous_kind, IncrementalByTimeRangeKind)
 
     model = model.copy(
-        update={"kind": IncrementalUnmanagedKind(), "physical_version": "pinned_version_12345"}
+        update={
+            "kind": IncrementalUnmanagedKind(),
+            "physical_version": "pinned_version_12345",
+            "partitioned_by_": [exp.column("event_date")],
+        }
     )
     context.upsert_model(model)
     context.plan("prod", auto_apply=True, no_prompts=True)