Fix: Fix enviroment / snapshot cleanup order in janitor

Author: VaggelisD

sqlmesh/core/context.py

@@ -2353,19 +2353,23 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         # Clean up expired environments by removing their views and schemas
         self._cleanup_environments()
 
-        # Identify and delete expired snapshots
-        cleanup_targets = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+        expired_snapshot_ids, cleanup_targets = self.state_sync.get_expired_snapshots(
+            ignore_ttl=ignore_ttl
+        )
 
         # 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)
+
         self.state_sync.compact_intervals()
 
     def _cleanup_environments(self) -> None:
-        expired_environments = self.state_sync.delete_expired_environments()
+        expired_environments, filter_expr = self.state_sync.get_expired_environments()
 
         cleanup_expired_views(
             default_adapter=self.engine_adapter,
@@ -2374,6 +2378,7 @@ def _cleanup_environments(self) -> None:
             warn_on_delete_failure=self.config.janitor.warn_on_delete_failure,
             console=self.console,
         )
+        self.state_sync.delete_environments(expired_environments, filter_expr)
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:
         connection_name = connection_name.capitalize()

sqlmesh/core/state_sync/base.py

@@ -5,7 +5,7 @@
 import logging
 import pkgutil
 import typing as t
-
+from sqlglot import exp
 from sqlglot import __version__ as SQLGLOT_VERSION
 
 from sqlmesh import migrations
@@ -303,8 +303,23 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         """
 
     @abc.abstractmethod
-    def delete_expired_snapshots(
+    def get_expired_snapshots(
         self, ignore_ttl: bool = False
+    ) -> t.Tuple[t.Set[SnapshotId], 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 set of expired snapshot ids.
+            The list of table cleanup tasks.
+        """
+
+    @abc.abstractmethod
+    def delete_expired_snapshots(
+        self,
+        ignore_ttl: bool = False,
     ) -> t.List[SnapshotTableCleanupTask]:
         """Removes expired snapshots.
 
@@ -316,7 +331,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
@@ -385,6 +400,23 @@ def finalize(self, environment: Environment) -> None:
             environment: The target environment to finalize.
         """
 
+    @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 the environments specified by the arguments..
+
+        Returns:
+            The list of removed environments.
+        """
+
     @abc.abstractmethod
     def delete_expired_environments(self) -> t.List[Environment]:
         """Removes expired environments.

sqlmesh/core/state_sync/db/environment.py

@@ -162,11 +162,12 @@ 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) -> t.Tuple[t.List[Environment], exp.LTE]:
+        """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(
@@ -183,6 +184,14 @@ def delete_expired_environments(self) -> t.List[Environment]:
         )
         environments = [self._environment_from_row(r) for r in rows]
 
+        return environments, filter_expr
+
+    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
+        """Deletes expired environments.
+
+        Returns:
+            A list of deleted environments.
+        """
         self.engine_adapter.delete_from(
             self.environments_table,
             where=filter_expr,
@@ -198,7 +207,15 @@ def delete_expired_environments(self) -> t.List[Environment]:
                 where=exp.or_(*expired_environments),
             )
 
-        return environments
+    def delete_expired_environments(self) -> t.List[Environment]:
+        """Deletes expired environments.
+
+        Returns:
+            A list of deleted environments.
+        """
+        expired_environments, filter_expr = self.get_expired_environments()
+        self.delete_environments(expired_environments, filter_expr)
+        return expired_environments
 
     def get_environments(self) -> t.List[Environment]:
         """Fetches all environments.

sqlmesh/core/state_sync/db/facade.py

@@ -273,20 +273,35 @@ def unpause_snapshots(
     def invalidate_environment(self, name: str) -> None:
         self.environment_state.invalidate_environment(name)
 
+    def get_expired_snapshots(
+        self, ignore_ttl: bool = False
+    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
+        return self.snapshot_state.get_expired_snapshots(
+            self.environment_state.get_environments(), ignore_ttl=ignore_ttl
+        )
+
+    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
+        return self.environment_state.get_expired_environments()
+
     @transactional()
     def delete_expired_snapshots(
         self, ignore_ttl: bool = False
     ) -> t.List[SnapshotTableCleanupTask]:
-        expired_snapshot_ids, cleanup_targets = self.snapshot_state.delete_expired_snapshots(
-            self.environment_state.get_environments(), ignore_ttl=ignore_ttl
-        )
+        expired_snapshot_ids, cleanup_targets = self.get_expired_snapshots(ignore_ttl=ignore_ttl)
+
+        self.snapshot_state.delete_expired_snapshots(expired_snapshot_ids=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()
 
+    @transactional()
+    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
+        return self.environment_state.delete_environments(environments, filter_expr)
+
     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,16 +193,17 @@ def unpause_snapshots(
         if unrestorable_snapshots:
             self._update_snapshots(unrestorable_snapshots, unrestorable=True)
 
-    def delete_expired_snapshots(
+    def get_expired_snapshots(
         self, environments: t.Iterable[Environment], ignore_ttl: bool = False
     ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
-        """Deletes expired snapshots.
+        """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)
 
@@ -269,10 +270,16 @@ 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_expired_snapshots(self, expired_snapshot_ids: t.Set[SnapshotId]) -> None:
+        """Deletes expired snapshots.
+
+        Args:
+            expired_snapshot_ids: The list of expired snapshot ids.
+        """
+        self.delete_snapshots(expired_snapshot_ids)
+
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         """Deletes snapshots.
 

tests/core/test_integration.py

@@ -5609,3 +5609,82 @@ def test_plan_environment_statements_doesnt_cause_extra_diff(tmp_path: Path):
 
     # second plan - nothing has changed so should report no changes
     assert not ctx.plan(auto_apply=True, no_prompts=True).has_changes
+
+
+def test_janitor_cleanup_order(mocker: MockerFixture, tmp_path: Path):
+    def setup_scenario():
+        models_dir = tmp_path / "models"
+
+        if not models_dir.exists():
+            models_dir.mkdir()
+
+        model1_path = models_dir / "model1.sql"
+
+        with open(model1_path, "w") as f:
+            f.write("MODEL(name test.model1, kind FULL); SELECT 1 AS col")
+
+        config = Config(
+            model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        )
+        ctx = Context(paths=[tmp_path], config=config)
+
+        ctx.plan("dev", no_prompts=True, auto_apply=True)
+
+        model1_snapshot = ctx.get_snapshot("test.model1")
+
+        # Delete the model file to cause a snapshot expiration
+        model1_path.unlink()
+
+        ctx.load()
+
+        ctx.plan("dev", no_prompts=True, auto_apply=True)
+
+        # Invalidate the environment to cause an environment cleanup
+        ctx.invalidate_environment("dev")
+
+        try:
+            ctx._run_janitor(ignore_ttl=True)
+        except:
+            pass
+
+        return ctx, model1_snapshot
+
+    # Case 1: Assume that the snapshot cleanup yields an error, the snapshot records
+    # should still exist in the state sync so the next janitor can retry
+    mocker.patch(
+        "sqlmesh.core.snapshot.evaluator.SnapshotEvaluator.cleanup",
+        side_effect=Exception("snapshot cleanup error"),
+    )
+    ctx, model1_snapshot = setup_scenario()
+
+    # - Check that the snapshot record exists in the state sync
+    state_snapshot = ctx.state_sync.state_sync.get_snapshots([model1_snapshot.snapshot_id])
+    assert state_snapshot
+
+    # - Run the janitor again, this time it should succeed
+    mocker.patch("sqlmesh.core.snapshot.evaluator.SnapshotEvaluator.cleanup")
+    ctx._run_janitor(ignore_ttl=True)
+
+    # - Check that the snapshot record does not exist in the state sync anymore
+    state_snapshot = ctx.state_sync.state_sync.get_snapshots([model1_snapshot.snapshot_id])
+    assert not state_snapshot
+
+    # Case 2: Assume that the view cleanup yields an error, the enviroment
+    # record should still exist
+    mocker.patch(
+        "sqlmesh.core.context.cleanup_expired_views", side_effect=Exception("view cleanup error")
+    )
+    ctx, model1_snapshot = setup_scenario()
+
+    views = ctx.fetchdf("FROM duckdb_views() SELECT * EXCLUDE(sql) WHERE NOT internal")
+    assert views.empty
+
+    # - Check that the environment record exists in the state sync
+    assert ctx.state_sync.get_environment("dev")
+
+    # - Run the janitor again, this time it should succeed
+    mocker.patch("sqlmesh.core.context.cleanup_expired_views")
+    ctx._run_janitor(ignore_ttl=True)
+
+    # - Check that the environment record does not exist in the state sync anymore
+    assert not ctx.state_sync.get_environment("dev")