Revert to deleting first from state envs and snapshots

Author: themisvaltinos

sqlmesh/core/context.py

@@ -2295,14 +2295,15 @@ def _context_diff(
 
     def _run_janitor(self, ignore_ttl: bool = False) -> None:
         # Get expired environments and removes their views and schemas
-        expired_environments, filter_expr = self._cleanup_environments()
+        self._cleanup_environments()
 
         # Get expired snapshots and corresponding gateways per snapshot when applied
         expired_snapshots_ids, cleanup_targets, snapshot_gateways = (
             self.state_sync.get_expired_snapshots(ignore_ttl=ignore_ttl)
         )
 
-        # Clean up intervals from the state sync
+        # Clean up snapshots and intervals from the state sync
+        self.state_sync.delete_snapshots(expired_snapshots_ids)
         self.state_sync.cleanup_intervals(cleanup_targets, expired_snapshots_ids)
 
         # Remove the expired snapshots tables
@@ -2312,13 +2313,10 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
             on_complete=self.console.update_cleanup_progress,
         )
 
-        # Finally, remove the expired environments and snapshots from the state sync
-        self.state_sync.delete_environments(expired_environments, filter_expr)
-        self.state_sync.delete_snapshots(expired_snapshots_ids)
         self.state_sync.compact_intervals()
 
-    def _cleanup_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
-        expired_environments, filter_expr = self.state_sync.get_expired_environments()
+    def _cleanup_environments(self) -> None:
+        expired_environments = self.state_sync.delete_expired_environments()
 
         environment_snapshot_adapters: t.Dict[str, t.Dict[str, EngineAdapter]] = {}
         for environment in expired_environments:
@@ -2339,8 +2337,6 @@ def _cleanup_environments(self) -> t.Tuple[t.List[Environment], exp.LTE]:
             environment_snapshot_adapters=environment_snapshot_adapters,
         )
 
-        return expired_environments, filter_expr
-
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:
         connection_name = connection_name.capitalize()
         try:

tests/core/test_context.py

@@ -43,7 +43,6 @@
 from sqlmesh.utils.date import (
     make_inclusive_end,
     now,
-    now_timestamp,
     to_date,
     to_timestamp,
     yesterday_ds,
@@ -780,34 +779,28 @@ def test_janitor(sushi_context, mocker: MockerFixture) -> None:
     adapter_mock = mocker.MagicMock()
     adapter_mock.dialect = "duckdb"
     state_sync_mock = mocker.MagicMock()
-    now_ts = now_timestamp()
-    filter_expr = exp.LTE(
-        this=exp.column("expiration_ts"),
-        expression=exp.Literal.number(now_ts),
-    )
-    state_sync_mock.get_expired_environments.return_value = (
-        [
-            Environment(
-                name="test_environment",
-                suffix_target=EnvironmentSuffixTarget.TABLE,
-                snapshots=[x.table_info for x in sushi_context.snapshots.values()],
-                start_at="2022-01-01",
-                end_at="2022-01-01",
-                plan_id="test_plan_id",
-                previous_plan_id="test_plan_id",
-            ),
-            Environment(
-                name="test_environment",
-                suffix_target=EnvironmentSuffixTarget.SCHEMA,
-                snapshots=[x.table_info for x in sushi_context.snapshots.values()],
-                start_at="2022-01-01",
-                end_at="2022-01-01",
-                plan_id="test_plan_id",
-                previous_plan_id="test_plan_id",
-            ),
-        ],
-        filter_expr,
-    )
+
+    state_sync_mock.delete_expired_environments.return_value = [
+        Environment(
+            name="test_environment",
+            suffix_target=EnvironmentSuffixTarget.TABLE,
+            snapshots=[x.table_info for x in sushi_context.snapshots.values()],
+            start_at="2022-01-01",
+            end_at="2022-01-01",
+            plan_id="test_plan_id",
+            previous_plan_id="test_plan_id",
+        ),
+        Environment(
+            name="test_environment",
+            suffix_target=EnvironmentSuffixTarget.SCHEMA,
+            snapshots=[x.table_info for x in sushi_context.snapshots.values()],
+            start_at="2022-01-01",
+            end_at="2022-01-01",
+            plan_id="test_plan_id",
+            previous_plan_id="test_plan_id",
+        ),
+    ]
+
     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({}), [], {})