Fix: Fix enviroment / snapshot cleanup order in janitor

Author: VaggelisD

sqlmesh/core/context.py

@@ -117,7 +117,7 @@
 from sqlmesh.core.user import User
 from sqlmesh.utils import UniqueKeyDict, Verbosity
 from sqlmesh.utils.dag import DAG
-from sqlmesh.utils.date import TimeLike, now_ds, to_timestamp, format_tz_datetime
+from sqlmesh.utils.date import TimeLike, now_ds, to_timestamp, format_tz_datetime, now_timestamp
 from sqlmesh.utils.errors import (
     CircuitBreakerError,
     ConfigError,
@@ -2350,22 +2350,30 @@ def _context_diff(
         )
 
     def _run_janitor(self, ignore_ttl: bool = False) -> None:
+        current_ts = now_timestamp()
+
         # Clean up expired environments by removing their views and schemas
-        self._cleanup_environments()
+        self._cleanup_environments(current_ts=current_ts)
 
-        # Identify and delete expired snapshots
-        cleanup_targets = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+        cleanup_targets = self.state_sync.get_expired_snapshots(
+            ignore_ttl=ignore_ttl, current_ts=current_ts
+        )
 
         # Remove the expired snapshots tables
         self.snapshot_evaluator.cleanup(
             target_snapshots=cleanup_targets,
             on_complete=self.console.update_cleanup_progress,
         )
 
+        # Delete the expired snapshot records from the state sync
+        self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl, current_ts=current_ts)
+
         self.state_sync.compact_intervals()
 
-    def _cleanup_environments(self) -> None:
-        expired_environments = self.state_sync.delete_expired_environments()
+    def _cleanup_environments(self, current_ts: t.Optional[int] = None) -> None:
+        current_ts = current_ts or now_timestamp()
+
+        expired_environments = self.state_sync.get_expired_environments(current_ts=current_ts)
 
         cleanup_expired_views(
             default_adapter=self.engine_adapter,
@@ -2375,6 +2383,8 @@ def _cleanup_environments(self) -> None:
             console=self.console,
         )
 
+        self.state_sync.delete_expired_environments(current_ts=current_ts)
+
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:
         connection_name = connection_name.capitalize()
         try:

sqlmesh/core/state_sync/base.py

@@ -281,6 +281,28 @@ def export(self, environment_names: t.Optional[t.List[str]] = None) -> StateStre
             environment_names: An optional list of environment names to export. If not specified, all environments will be exported.
         """
 
+    @abc.abstractmethod
+    def get_expired_snapshots(
+        self, current_ts: int, ignore_ttl: bool = False
+    ) -> t.List[SnapshotTableCleanupTask]:
+        """Aggregates the id's of the expired snapshots and creates a list of table cleanup tasks.
+
+        Expired snapshots are snapshots that have exceeded their time-to-live
+        and are no longer in use within an environment.
+
+        Returns:
+           The list of table cleanup tasks.
+        """
+
+    @abc.abstractmethod
+    def get_expired_environments(self, current_ts: int) -> t.List[Environment]:
+        """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.
+        """
+
 
 class StateSync(StateReader, abc.ABC):
     """Abstract base class for snapshot and environment state management."""
@@ -309,7 +331,7 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
 
     @abc.abstractmethod
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False
+        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
     ) -> t.List[SnapshotTableCleanupTask]:
         """Removes expired snapshots.
 
@@ -321,7 +343,7 @@ def delete_expired_snapshots(
                 all snapshots that are not referenced in any environment
 
         Returns:
-            The list of table cleanup tasks.
+            The list of snapshot table cleanup tasks.
         """
 
     @abc.abstractmethod
@@ -391,7 +413,9 @@ def finalize(self, environment: Environment) -> None:
         """
 
     @abc.abstractmethod
-    def delete_expired_environments(self) -> t.List[Environment]:
+    def delete_expired_environments(
+        self, current_ts: t.Optional[int] = None
+    ) -> t.List[Environment]:
         """Removes expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.

sqlmesh/core/state_sync/cache.py

@@ -111,10 +111,13 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.state_sync.delete_snapshots(snapshot_ids)
 
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False
+        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
     ) -> t.List[SnapshotTableCleanupTask]:
+        current_ts = current_ts or now_timestamp()
         self.snapshot_cache.clear()
-        return self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+        return self.state_sync.delete_expired_snapshots(
+            current_ts=current_ts, ignore_ttl=ignore_ttl
+        )
 
     def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
         for snapshot_intervals in snapshots_intervals:

sqlmesh/core/state_sync/db/environment.py

@@ -162,43 +162,51 @@ def finalize(self, environment: Environment) -> None:
             where=environment_filter,
         )
 
-    def delete_expired_environments(self) -> t.List[Environment]:
-        """Deletes expired environments.
+    def get_expired_environments(self, current_ts: int) -> t.List[Environment]:
+        """Returns the expired environments.
 
+        Expired environments are environments that have exceeded their time-to-live value.
         Returns:
-            A list of deleted environments.
+            The list of environments to remove, the filter to remove environments.
         """
-        now_ts = now_timestamp()
-        filter_expr = exp.LTE(
-            this=exp.column("expiration_ts"),
-            expression=exp.Literal.number(now_ts),
-        )
-
         rows = fetchall(
             self.engine_adapter,
             self._environments_query(
-                where=filter_expr,
+                where=self._create_expiration_filter_expr(current_ts),
                 lock_for_update=True,
             ),
         )
-        environments = [self._environment_from_row(r) for r in rows]
+        expired_environments = [self._environment_from_row(r) for r in rows]
+
+        return expired_environments
+
+    def delete_expired_environments(
+        self, current_ts: t.Optional[int] = None
+    ) -> t.List[Environment]:
+        """Deletes expired environments.
+
+        Returns:
+            A list of deleted environments.
+        """
+        current_ts = current_ts or now_timestamp()
+        expired_environments = self.get_expired_environments(current_ts=current_ts)
 
         self.engine_adapter.delete_from(
             self.environments_table,
-            where=filter_expr,
+            where=self._create_expiration_filter_expr(current_ts),
         )
 
         # Delete the expired environments' corresponding environment statements
-        if expired_environments := [
+        if expired_environments_exprs := [
             exp.EQ(this=exp.column("environment_name"), expression=exp.Literal.string(env.name))
-            for env in environments
+            for env in expired_environments
         ]:
             self.engine_adapter.delete_from(
                 self.environment_statements_table,
-                where=exp.or_(*expired_environments),
+                where=exp.or_(*expired_environments_exprs),
             )
 
-        return environments
+        return expired_environments
 
     def get_environments(self) -> t.List[Environment]:
         """Fetches all environments.
@@ -308,6 +316,17 @@ def _environments_query(
             return query.lock(copy=False)
         return query
 
+    def _create_expiration_filter_expr(self, current_ts: int) -> exp.Expression:
+        """Creates a SQLGlot filter expression to find expired environments.
+
+        Args:
+            current_ts: The current timestamp.
+        """
+        return exp.LTE(
+            this=exp.column("expiration_ts"),
+            expression=exp.Literal.number(current_ts),
+        )
+
 
 def _environment_to_df(environment: Environment) -> pd.DataFrame:
     return pd.DataFrame(

sqlmesh/core/state_sync/db/facade.py

@@ -60,7 +60,7 @@
 from sqlmesh.core.state_sync.db.snapshot import SnapshotState
 from sqlmesh.core.state_sync.db.version import VersionState
 from sqlmesh.core.state_sync.db.migrator import StateMigrator
-from sqlmesh.utils.date import TimeLike, to_timestamp, time_like_to_str
+from sqlmesh.utils.date import TimeLike, to_timestamp, time_like_to_str, now_timestamp
 from sqlmesh.utils.errors import ConflictingPlanError, SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -273,19 +273,36 @@ def unpause_snapshots(
     def invalidate_environment(self, name: str) -> None:
         self.environment_state.invalidate_environment(name)
 
+    def get_expired_snapshots(
+        self, current_ts: int, ignore_ttl: bool = False
+    ) -> t.List[SnapshotTableCleanupTask]:
+        return self.snapshot_state.get_expired_snapshots(
+            self.environment_state.get_environments(), current_ts=current_ts, ignore_ttl=ignore_ttl
+        )
+
+    def get_expired_environments(self, current_ts: int) -> t.List[Environment]:
+        return self.environment_state.get_expired_environments(current_ts=current_ts)
+
     @transactional()
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False
+        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
     ) -> t.List[SnapshotTableCleanupTask]:
-        expired_snapshot_ids, cleanup_targets = self.snapshot_state.delete_expired_snapshots(
-            self.environment_state.get_environments(), ignore_ttl=ignore_ttl
+        current_ts = current_ts or now_timestamp()
+        expired_snapshot_ids, cleanup_targets = self.snapshot_state._get_expired_snapshots(
+            self.environment_state.get_environments(), ignore_ttl=ignore_ttl, current_ts=current_ts
         )
+
+        self.snapshot_state.delete_snapshots(expired_snapshot_ids)
         self.interval_state.cleanup_intervals(cleanup_targets, expired_snapshot_ids)
+
         return cleanup_targets
 
     @transactional()
-    def delete_expired_environments(self) -> t.List[Environment]:
-        return self.environment_state.delete_expired_environments()
+    def delete_expired_environments(
+        self, current_ts: t.Optional[int] = None
+    ) -> t.List[Environment]:
+        current_ts = current_ts or now_timestamp()
+        return self.environment_state.delete_expired_environments(current_ts=current_ts)
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.snapshot_state.delete_snapshots(snapshot_ids)

sqlmesh/core/state_sync/db/snapshot.py

@@ -193,19 +193,34 @@ def unpause_snapshots(
         if unrestorable_snapshots:
             self._update_snapshots(unrestorable_snapshots, unrestorable=True)
 
-    def delete_expired_snapshots(
-        self, environments: t.Iterable[Environment], ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
-        """Deletes expired snapshots.
+    def get_expired_snapshots(
+        self,
+        environments: t.Iterable[Environment],
+        current_ts: int,
+        ignore_ttl: bool = False,
+    ) -> t.List[SnapshotTableCleanupTask]:
+        """Aggregates the id's of the expired snapshots and creates a list of table cleanup tasks.
 
-        Args:
-            ignore_ttl: Whether to ignore the TTL of the 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 and cleanup targets.
+            The set of expired snapshot ids.
+            The list of table cleanup tasks.
         """
-        current_ts = now_timestamp(minute_floor=False)
+        _, cleanup_targets = self._get_expired_snapshots(
+            environments=environments,
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+        )
+        return cleanup_targets
 
+    def _get_expired_snapshots(
+        self,
+        environments: t.Iterable[Environment],
+        current_ts: int,
+        ignore_ttl: bool = False,
+    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
         expired_query = exp.select("name", "identifier", "version").from_(self.snapshots_table)
 
         if not ignore_ttl:
@@ -269,8 +284,6 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
                         )
                     )
 
-        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:

tests/core/state_sync/test_state_sync.py

@@ -1255,7 +1255,7 @@ def test_delete_expired_snapshots_promoted(
     env.snapshots_ = []
     state_sync.promote(env)
 
-    now_timestamp_mock = mocker.patch("sqlmesh.core.state_sync.db.snapshot.now_timestamp")
+    now_timestamp_mock = mocker.patch("sqlmesh.core.state_sync.db.facade.now_timestamp")
     now_timestamp_mock.return_value = now_timestamp() + 11000
 
     assert state_sync.delete_expired_snapshots() == [

tests/core/test_context.py

@@ -775,7 +775,7 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
     adapter_mock.dialect = "duckdb"
     state_sync_mock = mocker.MagicMock()
 
-    state_sync_mock.delete_expired_environments.return_value = [
+    state_sync_mock.get_expired_environments.return_value = [
         Environment(
             name="test_environment",
             suffix_target=EnvironmentSuffixTarget.TABLE,
@@ -798,6 +798,8 @@ 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 = []
+
     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