Fix independent ownership of snapshot tables during cleanup

Signed-off-by: Iurie Virlan <38211890+IuraHD@users.noreply.github.com>

Author: IuraHD

sqlmesh/core/snapshot/definition.py

@@ -760,8 +760,10 @@ def hydrate_with_intervals_by_version(
             List of target snapshots with hydrated intervals.
         """
         intervals_by_name_version = defaultdict(list)
+        intervals_by_dev_version = defaultdict(list)
         for interval in intervals:
             intervals_by_name_version[(interval.name, interval.version)].append(interval)
+            intervals_by_dev_version[(interval.name, interval.dev_version)].append(interval)
 
         result = []
         for snapshot in snapshots:
@@ -771,6 +773,17 @@ def hydrate_with_intervals_by_version(
             for interval in snapshot_intervals:
                 snapshot.merge_intervals(interval)
 
+            for interval in intervals_by_dev_version.get((snapshot.name, snapshot.dev_version), []):
+                if interval.version != snapshot.version_get_or_generate():
+                    snapshot.merge_intervals(
+                        interval.copy(
+                            update={
+                                "intervals": [],
+                                "last_altered_ts": None,
+                                "pending_restatement_intervals": [],
+                            }
+                        )
+                    )
             result.append(snapshot)
 
         return result
@@ -1555,6 +1568,28 @@ def __getstate__(self) -> t.Dict[t.Any, t.Any]:
 class SnapshotTableCleanupTask(PydanticModel):
     snapshot: SnapshotTableInfo
     dev_table_only: bool
+    # Default preserves the behavior of tasks serialized before independent cleanup.
+    delete_dev_table: bool = True
+    # Physical objects and logical interval versions have different ownership.
+    # None retains the meaning of cleanup tasks persisted before these fields.
+    delete_prod_intervals: t.Optional[bool] = None
+    delete_dev_intervals: t.Optional[bool] = None
+
+    @property
+    def clears_prod_intervals(self) -> bool:
+        return (
+            not self.dev_table_only
+            if self.delete_prod_intervals is None
+            else self.delete_prod_intervals
+        )
+
+    @property
+    def clears_dev_intervals(self) -> bool:
+        return (
+            self.delete_dev_table
+            if self.delete_dev_intervals is None
+            else self.delete_dev_intervals
+        )
 
 
 SnapshotIdLike = t.Union[SnapshotId, SnapshotIdAndVersion, SnapshotTableInfo, Snapshot]

sqlmesh/core/snapshot/evaluator.py

@@ -543,17 +543,16 @@ def cleanup(
                 len(skipped),
                 ", ".join(f"{sid} (gateway={gw})" for sid, gw in skipped),
             )
-        snapshots_to_dev_table_only = {
-            t.snapshot.snapshot_id: t.dev_table_only for t in filtered_targets
-        }
+        cleanup_tasks_by_snapshot = {t.snapshot.snapshot_id: t for t in filtered_targets}
         with self.concurrent_context():
             errors, _ = concurrent_apply_to_snapshots(
                 [t.snapshot for t in filtered_targets],
                 lambda s: self._cleanup_snapshot(
                     s,
-                    snapshots_to_dev_table_only[s.snapshot_id],
+                    cleanup_tasks_by_snapshot[s.snapshot_id].dev_table_only,
                     self.get_adapter(s.model_gateway),
                     on_complete,
+                    delete_dev_table=cleanup_tasks_by_snapshot[s.snapshot_id].delete_dev_table,
                 ),
                 self.ddl_concurrent_tasks,
                 reverse_order=True,
@@ -1347,10 +1346,13 @@ def _cleanup_snapshot(
         dev_table_only: bool,
         adapter: EngineAdapter,
         on_complete: t.Optional[t.Callable[[str], None]],
+        delete_dev_table: bool = True,
     ) -> None:
         snapshot = snapshot.table_info
 
-        table_names = [(False, snapshot.table_name(is_deployable=False))]
+        table_names = (
+            [(False, snapshot.table_name(is_deployable=False))] if delete_dev_table else []
+        )
         if not dev_table_only:
             table_names.append((True, snapshot.table_name(is_deployable=True)))
 

sqlmesh/core/state_sync/db/interval.py

@@ -277,16 +277,35 @@ def _get_snapshot_intervals(
             t.Tuple[str, str, t.Optional[str], t.Optional[str]], SnapshotIntervals
         ] = {}
 
-        for where in (
-            snapshot_name_version_filter(
-                self.engine_adapter,
-                snapshots,
-                alias="intervals",
-                batch_size=self.SNAPSHOT_BATCH_SIZE,
+        filters = (
+            list(
+                snapshot_name_version_filter(
+                    self.engine_adapter,
+                    snapshots,
+                    alias="intervals",
+                    batch_size=self.SNAPSHOT_BATCH_SIZE,
+                )
             )
             if snapshots
             else [None]
-        ):
+        )
+        if snapshots:
+            dev_keys = [
+                SnapshotNameVersion(name=s.name, version=s.dev_version)
+                for s in snapshots
+                if hasattr(s, "dev_version")
+            ]
+            filters.extend(
+                where.and_(exp.column("is_dev", table="intervals"))
+                for where in snapshot_name_version_filter(
+                    self.engine_adapter,
+                    dev_keys,
+                    alias="intervals",
+                    version_column_name="dev_version",
+                    batch_size=self.SNAPSHOT_BATCH_SIZE,
+                )
+            )
+        for where in filters:
             rows = fetchall(self.engine_adapter, query.where(where))
             for (
                 interval_id,
@@ -301,6 +320,8 @@ def _get_snapshot_intervals(
                 is_pending_restatement,
                 last_altered_ts,
             ) in rows:
+                if interval_id in interval_ids:
+                    continue
                 interval_ids.add(interval_id)
                 merge_key = (name, version, dev_version, identifier)
                 # Pending restatement intervals are merged by name and version
@@ -420,7 +441,7 @@ def _delete_intervals_by_dev_version(self, targets: t.List[SnapshotTableCleanupT
         dev_keys_to_delete = [
             SnapshotNameVersion(name=t.snapshot.name, version=t.snapshot.dev_version)
             for t in targets
-            if t.dev_table_only
+            if t.clears_dev_intervals
         ]
         if not dev_keys_to_delete:
             return
@@ -436,7 +457,7 @@ def _delete_intervals_by_dev_version(self, targets: t.List[SnapshotTableCleanupT
 
     def _delete_intervals_by_version(self, targets: t.List[SnapshotTableCleanupTask]) -> None:
         """Deletes intervals for snapshot versions that are no longer used."""
-        non_dev_keys_to_delete = [t.snapshot for t in targets if not t.dev_table_only]
+        non_dev_keys_to_delete = [t.snapshot for t in targets if t.clears_prod_intervals]
         if not non_dev_keys_to_delete:
             return
 
@@ -446,7 +467,9 @@ def _delete_intervals_by_version(self, targets: t.List[SnapshotTableCleanupTask]
             alias=None,
             batch_size=self.SNAPSHOT_BATCH_SIZE,
         ):
-            self.engine_adapter.delete_from(self.intervals_table, where)
+            self.engine_adapter.delete_from(
+                self.intervals_table, where.and_(exp.column("is_dev").not_())
+            )
 
 
 def _intervals_to_df(

sqlmesh/core/state_sync/db/snapshot.py

@@ -37,7 +37,6 @@
 )
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, TimeLike, to_timestamp
-from sqlmesh.utils import unique
 
 if t.TYPE_CHECKING:
     import pandas as pd
@@ -225,12 +224,6 @@ def get_expired_snapshots(
         if not expired_candidates:
             return None
 
-        def _is_snapshot_used(snapshot: SnapshotIdAndVersion) -> bool:
-            return (
-                snapshot.snapshot_id in promoted_snapshot_ids
-                or snapshot.snapshot_id not in expired_candidates
-            )
-
         # Extract cursor values from last row for pagination
         last_row = rows[-1]
         last_row_boundary = RowBoundary(
@@ -244,51 +237,66 @@ def _is_snapshot_used(snapshot: SnapshotIdAndVersion) -> bool:
             end=last_row_boundary,
         )
 
-        unique_expired_versions = unique(expired_candidates.values())
         expired_snapshot_ids: t.Set[SnapshotId] = set()
         cleanup_tasks: t.List[SnapshotTableCleanupTask] = []
-
-        snapshots = self._get_snapshots_with_same_version(unique_expired_versions)
-
+        # A version identifies interval state, not a physical object. The same
+        # version may exist in different schemas or gateways. Load the table
+        # metadata for peers as well as candidates before determining ownership.
+        peers = self.get_snapshots_by_names(
+            {s.name for s in expired_candidates}, exclude_expired=False
+        )
+        full_snapshots = self._get_snapshots([s.snapshot_id for s in peers])
+        snapshots = list(full_snapshots.values())
         snapshots_by_version = defaultdict(set)
         snapshots_by_dev_version = defaultdict(set)
+        table_owners = defaultdict(set)
+
+        def table_key(snapshot: Snapshot, deployable: bool) -> t.Tuple[t.Optional[str], str]:
+            return (
+                snapshot.model_gateway,
+                snapshot.table_name(is_deployable=deployable),
+            )
+
         for s in snapshots:
             snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
             snapshots_by_dev_version[(s.name, s.dev_version)].add(s.snapshot_id)
+            if s.is_model and not s.is_symbolic:
+                for deployable in (True, False):
+                    table_owners[table_key(s, deployable)].add(s.snapshot_id)
 
-        expired_snapshots = [s for s in snapshots if not _is_snapshot_used(s)]
-        all_expired_snapshot_ids = {s.snapshot_id for s in expired_snapshots}
-
-        cleanup_targets: t.List[t.Tuple[SnapshotId, bool]] = []
+        expired_snapshots = [
+            full_snapshots[sid] for sid in expired_candidates if sid in full_snapshots
+        ]
         for snapshot in expired_snapshots:
-            shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
-            shared_version_snapshots.discard(snapshot.snapshot_id)
-
-            shared_dev_version_snapshots = snapshots_by_dev_version[
-                (snapshot.name, snapshot.dev_version)
-            ]
-            shared_dev_version_snapshots.discard(snapshot.snapshot_id)
-
-            if not shared_dev_version_snapshots:
-                dev_table_only = bool(shared_version_snapshots)
-                cleanup_targets.append((snapshot.snapshot_id, dev_table_only))
-
-        snapshot_ids_to_cleanup = [snapshot_id for snapshot_id, _ in cleanup_targets]
-        full_snapshots = self._get_snapshots(snapshot_ids_to_cleanup)
-        for snapshot_id, dev_table_only in cleanup_targets:
-            if snapshot_id in full_snapshots:
+            prod_owners = snapshots_by_version[(snapshot.name, snapshot.version)]
+            dev_owners = snapshots_by_dev_version[(snapshot.name, snapshot.dev_version)]
+            prod_owners.discard(snapshot.snapshot_id)
+            dev_owners.discard(snapshot.snapshot_id)
+            delete_prod = delete_dev = False
+            if snapshot.is_model and not snapshot.is_symbolic:
+                prod_table_owners = table_owners[table_key(snapshot, True)]
+                dev_table_owners = table_owners[table_key(snapshot, False)]
+                prod_table_owners.discard(snapshot.snapshot_id)
+                dev_table_owners.discard(snapshot.snapshot_id)
+                delete_prod = not prod_table_owners
+                delete_dev = not dev_table_owners
+            # An interval-only task is required when the last logical version
+            # expires but a physical object is retained by a different version.
+            if delete_prod or delete_dev or not prod_owners or not dev_owners:
                 cleanup_tasks.append(
                     SnapshotTableCleanupTask(
-                        snapshot=full_snapshots[snapshot_id].table_info,
-                        dev_table_only=dev_table_only,
+                        snapshot=snapshot.table_info,
+                        dev_table_only=not delete_prod,
+                        delete_dev_table=delete_dev,
+                        delete_prod_intervals=(
+                            None if (not prod_owners) == delete_prod else not prod_owners
+                        ),
+                        delete_dev_intervals=(
+                            None if (not dev_owners) == delete_dev else not dev_owners
+                        ),
                     )
                 )
-                expired_snapshot_ids.add(snapshot_id)
-                all_expired_snapshot_ids.discard(snapshot_id)
-
-        # Add any remaining expired snapshots that don't require cleanup
-        if all_expired_snapshot_ids:
-            expired_snapshot_ids.update(all_expired_snapshot_ids)
+            expired_snapshot_ids.add(snapshot.snapshot_id)
 
         if expired_snapshot_ids or cleanup_tasks:
             return ExpiredSnapshotBatch(