PR Feedback 1

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,
@@ -2334,10 +2334,14 @@ 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()
 
-        _, cleanup_targets = self.state_sync.get_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(
@@ -2346,21 +2350,25 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         )
 
         # Delete the expired snapshot records from the state sync
-        self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+        self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl, current_ts=current_ts)
 
         self.state_sync.compact_intervals()
 
-    def _cleanup_environments(self) -> None:
-        cleanup_targets = self.state_sync.get_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,
             engine_adapters=self.engine_adapters,
-            environments=cleanup_targets.expired_environments,
+            environments=expired_environments,
             warn_on_delete_failure=self.config.janitor.warn_on_delete_failure,
             console=self.console,
         )
-        self.state_sync.delete_environments(cleanup_targets)
+        self.state_sync.delete_environments(
+            environments=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/environment.py

@@ -4,7 +4,6 @@
 import re
 import typing as t
 
-from sqlglot import exp
 from pydantic import Field
 
 from sqlmesh.core import constants as c
@@ -229,11 +228,6 @@ class EnvironmentStatements(PydanticModel):
     jinja_macros: t.Optional[JinjaMacroRegistry] = None
 
 
-class EnvironmentCleanupTask(PydanticModel):
-    expired_environments: t.List[Environment]
-    filter_expr: exp.LTE
-
-
 def execute_environment_statements(
     adapter: EngineAdapter,
     environment_statements: t.List[EnvironmentStatements],

sqlmesh/core/state_sync/base.py

@@ -8,12 +8,7 @@
 from sqlglot import __version__ as SQLGLOT_VERSION
 
 from sqlmesh import migrations
-from sqlmesh.core.environment import (
-    Environment,
-    EnvironmentNamingInfo,
-    EnvironmentStatements,
-    EnvironmentCleanupTask,
-)
+from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -308,7 +303,7 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
 
     @abc.abstractmethod
     def get_expired_snapshots(
-        self, ignore_ttl: bool = False
+        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.
 
@@ -322,8 +317,7 @@ def get_expired_snapshots(
 
     @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.
 
@@ -405,7 +399,7 @@ def finalize(self, environment: Environment) -> None:
         """
 
     @abc.abstractmethod
-    def get_expired_environments(self) -> EnvironmentCleanupTask:
+    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.
@@ -414,15 +408,18 @@ def get_expired_environments(self) -> EnvironmentCleanupTask:
         """
 
     @abc.abstractmethod
-    def delete_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
-        """Removes the environments specified by the arguments..
+    def delete_environments(self, environments: t.List[Environment], current_ts: int) -> None:
+        """Removes the environments specified by the arguments.
 
-        Returns:
-            The list of removed environments.
+        Args:
+            environments: The environments to remove.
+            current_ts: The current timestamp.
         """
 
     @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

@@ -12,7 +12,7 @@
     fetchall,
     fetchone,
 )
-from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentCleanupTask
+from sqlmesh.core.environment import Environment, EnvironmentStatements
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, time_like_to_str
 from sqlmesh.utils.errors import SQLMeshError
@@ -162,62 +162,57 @@ def finalize(self, environment: Environment) -> None:
             where=environment_filter,
         )
 
-    def get_expired_environments(self) -> EnvironmentCleanupTask:
+    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.
         """
-        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_filter_expr(current_ts),
                 lock_for_update=True,
             ),
         )
         expired_environments = [self._environment_from_row(r) for r in rows]
 
-        return EnvironmentCleanupTask(
-            expired_environments=expired_environments, filter_expr=filter_expr
-        )
+        return expired_environments
 
-    def delete_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
+    def delete_environments(self, environments: t.List[Environment], current_ts: int) -> None:
         """Deletes expired environments.
 
         Returns:
             A list of deleted environments.
         """
         self.engine_adapter.delete_from(
             self.environments_table,
-            where=cleanup_targets.filter_expr,
+            where=self._create_filter_expr(current_ts),
         )
 
         # Delete the expired environments' corresponding environment statements
         if expired_environments := [
             exp.EQ(this=exp.column("environment_name"), expression=exp.Literal.string(env.name))
-            for env in cleanup_targets.expired_environments
+            for env in environments
         ]:
             self.engine_adapter.delete_from(
                 self.environment_statements_table,
                 where=exp.or_(*expired_environments),
             )
 
-    def delete_expired_environments(self) -> t.List[Environment]:
+    def delete_expired_environments(
+        self, current_ts: t.Optional[int] = None
+    ) -> t.List[Environment]:
         """Deletes expired environments.
 
         Returns:
             A list of deleted environments.
         """
-        cleanup_targets = self.get_expired_environments()
-        self.delete_environments(cleanup_targets)
-        return cleanup_targets.expired_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]:
         """Fetches all environments.
@@ -321,6 +316,12 @@ def _environments_query(
             return query.lock(copy=False)
         return query
 
+    def _create_filter_expr(self, current_ts: int) -> exp.Expression:
+        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

@@ -26,7 +26,7 @@
 
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentCleanupTask
+from sqlmesh.core.environment import Environment, EnvironmentStatements
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -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__)
@@ -274,33 +274,39 @@ def invalidate_environment(self, name: str) -> None:
         self.environment_state.invalidate_environment(name)
 
     def get_expired_snapshots(
-        self, ignore_ttl: bool = False
+        self, current_ts: int, 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
+            self.environment_state.get_environments(), current_ts=current_ts, ignore_ttl=ignore_ttl
         )
 
-    def get_expired_environments(self) -> EnvironmentCleanupTask:
-        return self.environment_state.get_expired_environments()
+    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.get_expired_snapshots(ignore_ttl=ignore_ttl)
+        current_ts = current_ts or now_timestamp()
+        expired_snapshot_ids, cleanup_targets = self.get_expired_snapshots(
+            ignore_ttl=ignore_ttl, current_ts=current_ts
+        )
 
-        self.snapshot_state.delete_expired_snapshots(expired_snapshot_ids=expired_snapshot_ids)
+        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)
 
     @transactional()
-    def delete_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
-        return self.environment_state.delete_environments(cleanup_targets=cleanup_targets)
+    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

@@ -194,7 +194,10 @@ def unpause_snapshots(
             self._update_snapshots(unrestorable_snapshots, unrestorable=True)
 
     def get_expired_snapshots(
-        self, environments: t.Iterable[Environment], ignore_ttl: bool = False
+        self,
+        environments: t.Iterable[Environment],
+        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.
 
@@ -205,8 +208,6 @@ def get_expired_snapshots(
             The set of expired snapshot ids.
             The list of table cleanup tasks.
         """
-        current_ts = now_timestamp(minute_floor=False)
-
         expired_query = exp.select("name", "identifier", "version").from_(self.snapshots_table)
 
         if not ignore_ttl:
@@ -272,14 +273,6 @@ def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
 
         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/state_sync/test_state_sync.py

@@ -1255,7 +1255,9 @@ 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.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() == [