Wrap expired enviroment logic in pydantic class

Author: VaggelisD

sqlmesh/core/context.py

@@ -2351,16 +2351,16 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         self.state_sync.compact_intervals()
 
     def _cleanup_environments(self) -> None:
-        expired_environments, filter_expr = self.state_sync.get_expired_environments()
+        cleanup_targets = self.state_sync.get_expired_environments()
 
         cleanup_expired_views(
             default_adapter=self.engine_adapter,
             engine_adapters=self.engine_adapters,
-            environments=expired_environments,
+            environments=cleanup_targets.expired_environments,
             warn_on_delete_failure=self.config.janitor.warn_on_delete_failure,
             console=self.console,
         )
-        self.state_sync.delete_environments(expired_environments, filter_expr)
+        self.state_sync.delete_environments(cleanup_targets)
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:
         connection_name = connection_name.capitalize()

sqlmesh/core/environment.py

@@ -4,6 +4,7 @@
 import re
 import typing as t
 
+from sqlglot import exp
 from pydantic import Field
 
 from sqlmesh.core import constants as c
@@ -228,6 +229,11 @@ 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

@@ -5,11 +5,15 @@
 import logging
 import pkgutil
 import typing as t
-from sqlglot import exp
 from sqlglot import __version__ as SQLGLOT_VERSION
 
 from sqlmesh import migrations
-from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
+from sqlmesh.core.environment import (
+    Environment,
+    EnvironmentNamingInfo,
+    EnvironmentStatements,
+    EnvironmentCleanupTask,
+)
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -401,7 +405,7 @@ def finalize(self, environment: Environment) -> None:
         """
 
     @abc.abstractmethod
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
+    def get_expired_environments(self) -> EnvironmentCleanupTask:
         """Returns the expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.
@@ -410,7 +414,7 @@ def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
         """
 
     @abc.abstractmethod
-    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
+    def delete_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
         """Removes the environments specified by the arguments..
 
         Returns:

sqlmesh/core/state_sync/db/environment.py

@@ -12,7 +12,7 @@
     fetchall,
     fetchone,
 )
-from sqlmesh.core.environment import Environment, EnvironmentStatements
+from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentCleanupTask
 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,7 +162,7 @@ def finalize(self, environment: Environment) -> None:
             where=environment_filter,
         )
 
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
+    def get_expired_environments(self) -> EnvironmentCleanupTask:
         """Returns the expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.
@@ -182,25 +182,27 @@ def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
                 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 environments, filter_expr
+        return EnvironmentCleanupTask(
+            expired_environments=expired_environments, filter_expr=filter_expr
+        )
 
-    def delete_environments(self, environments: t.List[Environment], filter_expr: exp.LTE) -> None:
+    def delete_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
         """Deletes expired environments.
 
         Returns:
             A list of deleted environments.
         """
         self.engine_adapter.delete_from(
             self.environments_table,
-            where=filter_expr,
+            where=cleanup_targets.filter_expr,
         )
 
         # 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 environments
+            for env in cleanup_targets.expired_environments
         ]:
             self.engine_adapter.delete_from(
                 self.environment_statements_table,
@@ -213,9 +215,9 @@ def delete_expired_environments(self) -> t.List[Environment]:
         Returns:
             A list of deleted environments.
         """
-        expired_environments, filter_expr = self.get_expired_environments()
-        self.delete_environments(expired_environments, filter_expr)
-        return expired_environments
+        cleanup_targets = self.get_expired_environments()
+        self.delete_environments(cleanup_targets)
+        return cleanup_targets.expired_environments
 
     def get_environments(self) -> t.List[Environment]:
         """Fetches all environments.

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
+from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentCleanupTask
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -280,7 +280,7 @@ def get_expired_snapshots(
             self.environment_state.get_environments(), ignore_ttl=ignore_ttl
         )
 
-    def get_expired_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
+    def get_expired_environments(self) -> EnvironmentCleanupTask:
         return self.environment_state.get_expired_environments()
 
     @transactional()
@@ -299,8 +299,8 @@ 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_environments(self, cleanup_targets: EnvironmentCleanupTask) -> None:
+        return self.environment_state.delete_environments(cleanup_targets=cleanup_targets)
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.snapshot_state.delete_snapshots(snapshot_ids)

tests/core/test_context.py

@@ -30,7 +30,12 @@
 from sqlmesh.core.console import create_console, get_console
 from sqlmesh.core.dialect import parse, schema_
 from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
-from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
+from sqlmesh.core.environment import (
+    Environment,
+    EnvironmentNamingInfo,
+    EnvironmentStatements,
+    EnvironmentCleanupTask,
+)
 from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
 from sqlmesh.core.model import load_sql_based_model, model, SqlModel, Model
 from sqlmesh.core.model.cache import OptimizedQueryCache
@@ -782,8 +787,8 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
         expression=exp.Literal.number(now_ts),
     )
 
-    state_sync_mock.get_expired_environments.return_value = (
-        [
+    state_sync_mock.get_expired_environments.return_value = EnvironmentCleanupTask(
+        expired_environments=[
             Environment(
                 name="test_environment",
                 suffix_target=EnvironmentSuffixTarget.TABLE,
@@ -803,7 +808,7 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
                 previous_plan_id="test_plan_id",
             ),
         ],
-        filter_expr,
+        filter_expr=filter_expr,
     )
 
     sushi_context._engine_adapters = {sushi_context.config.default_gateway: adapter_mock}