Revise to contain gateway info in SnapshotTableCleanupTask

Author: themisvaltinos

sqlmesh/core/context.py

@@ -2299,22 +2299,15 @@ def _context_diff(
         )
 
     def _run_janitor(self, ignore_ttl: bool = False) -> None:
-        # Get expired environments and removes their views and schemas
+        # Clean up expired environments by removing their views and schemas
         self._cleanup_environments()
 
-        # Get expired snapshots and corresponding gateways per snapshot when applied
-        expired_snapshots_ids, cleanup_targets, snapshot_gateways = (
-            self.state_sync.get_expired_snapshots(ignore_ttl=ignore_ttl)
-        )
-
-        # Clean up snapshots and intervals from the state sync
-        self.state_sync.delete_snapshots(expired_snapshots_ids)
-        self.state_sync.cleanup_intervals(cleanup_targets, expired_snapshots_ids)
+        # Identify and delete expired snapshots
+        cleanup_targets = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
 
         # Remove the expired snapshots tables
         self.snapshot_evaluator.cleanup(
             target_snapshots=cleanup_targets,
-            snapshot_gateways=snapshot_gateways,
             on_complete=self.console.update_cleanup_progress,
         )
 

sqlmesh/core/snapshot/definition.py

@@ -1338,6 +1338,7 @@ def __getstate__(self) -> t.Dict[t.Any, t.Any]:
 class SnapshotTableCleanupTask(PydanticModel):
     snapshot: SnapshotTableInfo
     dev_table_only: bool
+    gateway: t.Optional[str] = None
 
 
 SnapshotIdLike = t.Union[SnapshotId, SnapshotTableInfo, Snapshot]

sqlmesh/core/snapshot/evaluator.py

@@ -426,7 +426,6 @@ def migrate(
     def cleanup(
         self,
         target_snapshots: t.Iterable[SnapshotTableCleanupTask],
-        snapshot_gateways: t.Optional[t.Dict[str, str]] = None,
         on_complete: t.Optional[t.Callable[[str], None]] = None,
     ) -> None:
         """Cleans up the given snapshots by removing its table
@@ -438,16 +437,15 @@ def cleanup(
         snapshots_to_dev_table_only = {
             t.snapshot.snapshot_id: t.dev_table_only for t in target_snapshots
         }
+        snapshot_gateways = {t.snapshot.snapshot_id: t.gateway for t in target_snapshots}
 
         with self.concurrent_context():
             concurrent_apply_to_snapshots(
                 [t.snapshot for t in target_snapshots],
                 lambda s: self._cleanup_snapshot(
                     s,
                     snapshots_to_dev_table_only[s.snapshot_id],
-                    self.get_adapter(
-                        snapshot_gateways.get(s.snapshot_id.name) if snapshot_gateways else None
-                    ),
+                    self._get_adapter(snapshot_gateways[s.snapshot_id]),
                     on_complete,
                 ),
                 self.ddl_concurrent_tasks,

sqlmesh/core/state_sync/base.py

@@ -7,7 +7,6 @@
 import typing as t
 
 from sqlglot import __version__ as SQLGLOT_VERSION
-from sqlglot import exp
 
 from sqlmesh import migrations
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
@@ -320,23 +319,6 @@ def delete_expired_snapshots(
             The list of table cleanup tasks.
         """
 
-    @abc.abstractmethod
-    def get_expired_snapshots(
-        self, ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask], t.Dict[str, str]]:
-        """Gets expired snapshots.
-
-        Expired snapshots are snapshots that have exceeded their time-to-live
-        and are no longer in use within an environment.
-
-        Args:
-            ignore_ttl: Ignore the TTL on the snapshot when considering it expired. This has the effect of deleting
-                all snapshots that are not referenced in any environment
-
-        Returns:
-            A tuple of expired snapshot IDs, cleanup targets and gateway per snapshot dictionary.
-        """
-
     @abc.abstractmethod
     def invalidate_environment(self, name: str) -> None:
         """Invalidates the target environment by setting its expiration timestamp to now.
@@ -345,14 +327,6 @@ def invalidate_environment(self, name: str) -> None:
             name: The name of the environment to invalidate.
         """
 
-    @abc.abstractmethod
-    def cleanup_intervals(
-        self,
-        cleanup_targets: t.List[SnapshotTableCleanupTask],
-        expired_snapshot_ids: t.Set[SnapshotId],
-    ) -> None:
-        """Cleans up intervals."""
-
     @abc.abstractmethod
     def remove_intervals(
         self,
@@ -421,24 +395,6 @@ def delete_expired_environments(self) -> t.List[Environment]:
             The list of removed environments.
         """
 
-    @abc.abstractmethod
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
-        """Returns the expired environments.
-
-        Expired environments are environments that have exceeded their time-to-live value.
-
-        Returns:
-            The list of environments to remove, the filter to remove environments.
-        """
-
-    @abc.abstractmethod
-    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
-        """Removes environments.
-
-        Returns:
-            The list of removed environments.
-        """
-
     @abc.abstractmethod
     def unpause_snapshots(
         self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike

sqlmesh/core/state_sync/db/facade.py

@@ -288,33 +288,6 @@ def delete_expired_snapshots(
     def delete_expired_environments(self) -> t.List[Environment]:
         return self.environment_state.delete_expired_environments()
 
-    @transactional()
-    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
-        self.environment_state.delete_environments(
-            environments=environments, filter_expr=filter_expr
-        )
-
-    @transactional()
-    def cleanup_intervals(
-        self,
-        cleanup_targets: t.List[SnapshotTableCleanupTask],
-        expired_snapshot_ids: t.Set[SnapshotId],
-    ) -> None:
-        self.interval_state.cleanup_intervals(cleanup_targets, expired_snapshot_ids)
-
-    def get_expired_snapshots(
-        self, ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask], t.Dict[str, str]]:
-        expired_snapshot_ids, cleanup_targets, snapshot_gateways = (
-            self.snapshot_state.get_expired_snapshots(
-                self.environment_state.get_environments(), ignore_ttl=ignore_ttl
-            )
-        )
-        return expired_snapshot_ids, cleanup_targets, snapshot_gateways
-
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
-        return self.environment_state.get_expired_environments()
-
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.snapshot_state.delete_snapshots(snapshot_ids)
 

sqlmesh/core/state_sync/db/snapshot.py

@@ -205,26 +205,6 @@ def delete_expired_snapshots(
             A tuple of expired snapshot IDs and cleanup targets.
         """
 
-        expired_snapshot_ids, cleanup_targets, _ = self.get_expired_snapshots(
-            environments=environments, ignore_ttl=ignore_ttl
-        )
-
-        if expired_snapshot_ids:
-            self.delete_snapshots(expired_snapshot_ids)
-
-        return expired_snapshot_ids, cleanup_targets
-
-    def get_expired_snapshots(
-        self, environments: t.Iterable[Environment], ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask], t.Dict[str, str]]:
-        """Gets the expired snapshots.
-
-        Args:
-            ignore_ttl: Whether to ignore the TTL of the snapshots.
-
-        Returns:
-            A tuple of expired snapshot IDs, cleanup targets and gateway per snapshot dictionary.
-        """
         current_ts = now_timestamp(minute_floor=False)
 
         expired_query = exp.select("name", "identifier", "version").from_(self.snapshots_table)
@@ -241,7 +221,7 @@ def get_expired_snapshots(
             for name, identifier, version in fetchall(self.engine_adapter, expired_query)
         }
         if not expired_candidates:
-            return set(), [], {}
+            return set(), []
 
         promoted_snapshot_ids = {
             snapshot.snapshot_id
@@ -261,7 +241,6 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
         )
         cleanup_targets = []
         expired_snapshot_ids = set()
-        snapshot_gateways: t.Dict[str, str] = {}
         for versions_batch in version_batches:
             snapshots = self._get_snapshots_with_same_version(versions_batch)
 
@@ -275,9 +254,6 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
             expired_snapshot_ids.update([s.snapshot_id for s in expired_snapshots])
 
             for snapshot in expired_snapshots:
-                if (node := snapshot.raw_snapshot.get("node")) and (gateway := node.get("gateway")):
-                    snapshot_gateways[snapshot.snapshot_id.name] = gateway
-
                 shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
                 shared_version_snapshots.discard(snapshot.snapshot_id)
 
@@ -291,10 +267,14 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
                         SnapshotTableCleanupTask(
                             snapshot=snapshot.full_snapshot.table_info,
                             dev_table_only=bool(shared_version_snapshots),
+                            gateway=snapshot.raw_snapshot.get("node", {}).get("gateway", None),
                         )
                     )
 
-        return expired_snapshot_ids, cleanup_targets, snapshot_gateways
+        if expired_snapshot_ids:
+            self.delete_snapshots(expired_snapshot_ids)
+
+        return expired_snapshot_ids, cleanup_targets
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         """Deletes snapshots.

sqlmesh/schedulers/airflow/state_sync.py

@@ -3,7 +3,6 @@
 import logging
 import typing as t
 
-from sqlglot import exp
 from sqlmesh.core.console import Console
 from sqlmesh.core.environment import Environment, EnvironmentStatements
 from sqlmesh.core.snapshot import (
@@ -294,41 +293,6 @@ def delete_expired_environments(self) -> t.List[Environment]:
             "Deleting expired environments is not supported by the Airflow state sync."
         )
 
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
-        """Returns the expired environments.
-
-        Expired environments are environments that have exceeded their time-to-live value.
-
-        """
-        raise NotImplementedError(
-            "Getting expired environments is not supported by the Airflow state sync."
-        )
-
-    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
-        """Removes environments.
-
-        Returns:
-            The list of removed environments.
-        """
-        raise NotImplementedError(
-            "Deleting environments is not supported by the Airflow state sync."
-        )
-
-    def get_expired_snapshots(
-        self, ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask], t.Dict[str, str]]:
-        """Gets expired snapshots.
-
-        Expired snapshots are snapshots that have exceeded their time-to-live
-        and are no longer in use within an environment.
-
-        Returns:
-            A tuple of expired snapshot IDs, cleanup targets and gateway per snapshot dictionary.
-        """
-        raise NotImplementedError(
-            "Getting expired snapshots is not supported by the Airflow state sync."
-        )
-
     def unpause_snapshots(
         self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
@@ -354,17 +318,6 @@ def compact_intervals(self) -> None:
             "Compacting intervals is not supported by the Airflow state sync."
         )
 
-    def cleanup_intervals(
-        self,
-        cleanup_targets: t.List[SnapshotTableCleanupTask],
-        expired_snapshot_ids: t.Set[SnapshotId],
-    ) -> None:
-        """Cleans up intervals."""
-
-        raise NotImplementedError(
-            "Cleaning up intervals is not supported by the Airflow state sync."
-        )
-
     def update_auto_restatements(
         self, next_auto_restatement_ts: t.Dict[SnapshotNameVersion, t.Optional[int]]
     ) -> None:

tests/core/test_context.py

@@ -803,7 +803,6 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
 
     sushi_context._engine_adapters = {sushi_context.config.default_gateway: adapter_mock}
     sushi_context._state_sync = state_sync_mock
-    state_sync_mock.get_expired_snapshots.return_value = (set({}), [], {})
     sushi_context._run_janitor()
     # Assert that the schemas are dropped just twice for the schema based environment
     # Make sure that external model schemas/tables are not dropped

tests/core/test_snapshot_evaluator.py

@@ -3970,13 +3970,15 @@ def test_multiple_engine_cleanup(snapshot: Snapshot, adapters, make_snapshot):
         f"sqlmesh__test_schema.test_schema__test_model__{snapshot_2.version}",
     )
 
-    snapshot_gateways = {snapshot.name: "default", snapshot_2.name: "secondary"}
     evaluator.cleanup(
         [
-            SnapshotTableCleanupTask(snapshot=snapshot.table_info, dev_table_only=True),
-            SnapshotTableCleanupTask(snapshot=snapshot_2.table_info, dev_table_only=True),
+            SnapshotTableCleanupTask(
+                snapshot=snapshot.table_info, dev_table_only=True, gateway="default"
+            ),
+            SnapshotTableCleanupTask(
+                snapshot=snapshot_2.table_info, dev_table_only=True, gateway="secondary"
+            ),
         ],
-        snapshot_gateways,
     )
 
     # The clean up will happen using the specific gateway the model was created with