PR Feedback 2

Author: VaggelisD

sqlmesh/core/context.py

@@ -2339,7 +2339,7 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         # Clean up expired environments by removing their views and schemas
         self._cleanup_environments(current_ts=current_ts)
 
-        _, cleanup_targets = self.state_sync.get_expired_snapshots(
+        cleanup_targets = self.state_sync.get_expired_snapshots(
             ignore_ttl=ignore_ttl, current_ts=current_ts
         )
 
@@ -2366,9 +2366,8 @@ def _cleanup_environments(self, current_ts: t.Optional[int] = None) -> None:
             warn_on_delete_failure=self.config.janitor.warn_on_delete_failure,
             console=self.console,
         )
-        self.state_sync.delete_environments(
-            environments=expired_environments, current_ts=current_ts
-        )
+
+        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()

sqlmesh/core/state_sync/base.py

@@ -275,6 +275,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."""
@@ -301,20 +323,6 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
             snapshot_ids: A list of snapshot like objects to delete.
         """
 
-    @abc.abstractmethod
-    def get_expired_snapshots(
-        self, current_ts: int, 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, current_ts: t.Optional[int] = None
@@ -398,24 +406,6 @@ def finalize(self, environment: Environment) -> None:
             environment: The target environment to finalize.
         """
 
-    @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.
-        """
-
-    @abc.abstractmethod
-    def delete_environments(self, environments: t.List[Environment], current_ts: int) -> None:
-        """Removes the environments specified by the arguments.
-
-        Args:
-            environments: The environments to remove.
-            current_ts: The current timestamp.
-        """
-
     @abc.abstractmethod
     def delete_expired_environments(
         self, current_ts: t.Optional[int] = None

sqlmesh/core/state_sync/db/environment.py

@@ -172,46 +172,40 @@ def get_expired_environments(self, current_ts: int) -> t.List[Environment]:
         rows = fetchall(
             self.engine_adapter,
             self._environments_query(
-                where=self._create_filter_expr(current_ts),
+                where=self._create_expiration_filter_expr(current_ts),
                 lock_for_update=True,
             ),
         )
         expired_environments = [self._environment_from_row(r) for r in rows]
 
         return expired_environments
 
-    def delete_environments(self, environments: t.List[Environment], current_ts: int) -> None:
+    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=self._create_filter_expr(current_ts),
+            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),
             )
 
-    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.delete_environments(expired_environments, current_ts=current_ts)
         return expired_environments
 
     def get_environments(self) -> t.List[Environment]:
@@ -316,7 +310,12 @@ def _environments_query(
             return query.lock(copy=False)
         return query
 
-    def _create_filter_expr(self, current_ts: int) -> exp.Expression:
+    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),

sqlmesh/core/state_sync/db/facade.py

@@ -275,7 +275,7 @@ def invalidate_environment(self, name: str) -> None:
 
     def get_expired_snapshots(
         self, current_ts: int, ignore_ttl: bool = False
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
+    ) -> t.List[SnapshotTableCleanupTask]:
         return self.snapshot_state.get_expired_snapshots(
             self.environment_state.get_environments(), current_ts=current_ts, ignore_ttl=ignore_ttl
         )
@@ -288,8 +288,8 @@ def delete_expired_snapshots(
         self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
     ) -> t.List[SnapshotTableCleanupTask]:
         current_ts = current_ts or now_timestamp()
-        expired_snapshot_ids, cleanup_targets = self.get_expired_snapshots(
-            ignore_ttl=ignore_ttl, current_ts=current_ts
+        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)
@@ -304,10 +304,6 @@ def delete_expired_environments(
         current_ts = current_ts or now_timestamp()
         return self.environment_state.delete_expired_environments(current_ts=current_ts)
 
-    @transactional()
-    def delete_environments(self, environments: t.List[Environment], current_ts: int) -> None:
-        self.environment_state.delete_environments(environments=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

@@ -198,7 +198,7 @@ def get_expired_snapshots(
         environments: t.Iterable[Environment],
         current_ts: int,
         ignore_ttl: bool = False,
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
+    ) -> 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
@@ -208,6 +208,19 @@ def get_expired_snapshots(
             The set of expired snapshot ids.
             The list of table cleanup tasks.
         """
+        _, 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:

tests/core/test_context.py

@@ -805,7 +805,7 @@ 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({}), [])
+    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