Fix: When removing intervals also remove them for snapshots that have been deleted by the janitor process (#2489)

Author: izeigerman

sqlmesh/core/snapshot/definition.py

@@ -162,6 +162,10 @@ class SnapshotIntervals(PydanticModel, frozen=True):
     def snapshot_id(self) -> SnapshotId:
         return SnapshotId(name=self.name, identifier=self.identifier)
 
+    @property
+    def name_version(self) -> SnapshotNameVersion:
+        return SnapshotNameVersion(name=self.name, version=self.version)
+
 
 class SnapshotDataVersion(PydanticModel, frozen=True):
     fingerprint: SnapshotFingerprint

sqlmesh/core/state_sync/engine_adapter.py

@@ -645,26 +645,38 @@ def remove_interval(
         execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
+        intervals_to_remove: t.Sequence[
+            t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]
+        ] = snapshot_intervals
         if remove_shared_versions:
-            name_version_mapping = {
-                s.name_version: (s, interval) for s, interval in snapshot_intervals
-            }
-            all_snapshots = self._get_snapshots_with_same_version(
-                [s[0] for s in snapshot_intervals]
-            )
-            snapshot_intervals = [
-                (snapshot, name_version_mapping[snapshot.name_version][1])
+            name_version_mapping = {s.name_version: interval for s, interval in snapshot_intervals}
+            all_snapshots = []
+            for where in self._snapshot_name_version_filter(name_version_mapping, alias=None):
+                all_snapshots.extend(
+                    [
+                        SnapshotIntervals(
+                            name=r[0], identifier=r[1], version=r[2], intervals=[], dev_intervals=[]
+                        )
+                        for r in self._fetchall(
+                            exp.select("name", "identifier", "version")
+                            .from_(self.intervals_table)
+                            .where(where)
+                        )
+                    ]
+                )
+            intervals_to_remove = [
+                (snapshot, name_version_mapping[snapshot.name_version])
                 for snapshot in all_snapshots
             ]
 
         if logger.isEnabledFor(logging.INFO):
-            snapshot_ids = ", ".join(str(s.snapshot_id) for s, _ in snapshot_intervals)
+            snapshot_ids = ", ".join(str(s.snapshot_id) for s, _ in intervals_to_remove)
             logger.info("Removing interval for snapshots: %s", snapshot_ids)
 
         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),
+                _intervals_to_df(intervals_to_remove, is_dev=is_dev, is_removed=True),
                 columns_to_types=self._interval_columns_to_types,
             )
 
@@ -1262,7 +1274,9 @@ def _snapshot_id_filter(
                 )
 
     def _snapshot_name_version_filter(
-        self, snapshot_name_versions: t.Iterable[SnapshotNameVersionLike], alias: str = "snapshots"
+        self,
+        snapshot_name_versions: t.Iterable[SnapshotNameVersionLike],
+        alias: t.Optional[str] = "snapshots",
     ) -> t.Iterator[exp.Condition]:
         name_versions = sorted({(s.name, s.version) for s in snapshot_name_versions})
         batches = self._snapshot_batches(name_versions)
@@ -1305,7 +1319,7 @@ def _transaction(self) -> t.Iterator[None]:
 
 
 def _intervals_to_df(
-    snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+    snapshot_intervals: t.Sequence[t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]],
     is_dev: bool,
     is_removed: bool,
 ) -> pd.DataFrame:

tests/core/test_state_sync.py

@@ -330,6 +330,49 @@ def test_remove_interval(state_sync: EngineAdapterStateSync, make_snapshot: t.Ca
     ]
 
 
+def test_remove_interval_missing_snapshot(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+) -> None:
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            cron="@daily",
+            query=parse_one("select 1, ds"),
+        ),
+        version="a",
+    )
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="a",
+            cron="@daily",
+            query=parse_one("select 2::INT, '2022-01-01'::TEXT AS ds"),
+        ),
+        version="a",
+    )
+    # Only add snapshot_a to simulate that snapshot_b is missing
+    state_sync.push_snapshots([snapshot_a])
+    state_sync.add_interval(snapshot_a, "2020-01-01", "2020-01-10")
+    state_sync.add_interval(snapshot_b, "2020-01-11", "2020-01-30")
+
+    snapshots = state_sync.get_snapshots([snapshot_a, snapshot_b])
+    assert len(snapshots) == 1
+    assert snapshots[snapshot_a.snapshot_id].intervals == [
+        (to_timestamp("2020-01-01"), to_timestamp("2020-01-31")),
+    ]
+
+    state_sync.remove_interval(
+        [(snapshot_a, snapshot_a.inclusive_exclusive("2020-01-15", "2020-01-17"))],
+        remove_shared_versions=True,
+    )
+
+    snapshots = state_sync.get_snapshots([snapshot_a, snapshot_b])
+    assert len(snapshots) == 1
+    assert snapshots[snapshot_a.snapshot_id].intervals == [
+        (to_timestamp("2020-01-01"), to_timestamp("2020-01-15")),
+        (to_timestamp("2020-01-18"), to_timestamp("2020-01-31")),
+    ]
+
+
 def test_refresh_snapshot_intervals(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ) -> None: