Feat: Drop schemas for expired environments (#1177)

* drop schemas for expired environments

* add test

* fix location for dropping schemas

* airflow drop environment schemas

* remove the demotion of snapshots

* cleanup

* fix cleanup test

Author: eakmanrq

sqlmesh/core/context.py

@@ -1170,8 +1170,13 @@ def _load_configs(
 
     def _run_janitor(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()
-        for expired_environment in expired_environments:
-            self.snapshot_evaluator.demote(expired_environment.snapshots, expired_environment.name)
+        expired_schemas = {
+            snapshot.qualified_view_name.schema_for_environment(expired_environment.name)
+            for expired_environment in expired_environments
+            for snapshot in expired_environment.snapshots
+        }
+        for expired_schema in expired_schemas:
+            self.engine_adapter.drop_schema(expired_schema, ignore_if_not_exists=True, cascade=True)
 
         expired_snapshots = self.state_sync.delete_expired_snapshots()
         self.snapshot_evaluator.cleanup(expired_snapshots)

sqlmesh/engines/commands.py

@@ -114,8 +114,14 @@ def cleanup(
 ) -> None:
     if isinstance(command_payload, str):
         command_payload = CleanupCommandPayload.parse_raw(command_payload)
-    for environment in command_payload.environments:
-        evaluator.demote(environment.snapshots, environment.name)
+
+    expired_schemas = {
+        snapshot.qualified_view_name.schema_for_environment(environment.name)
+        for environment in command_payload.environments
+        for snapshot in environment.snapshots
+    }
+    for expired_schema in expired_schemas:
+        evaluator.adapter.drop_schema(expired_schema, ignore_if_not_exists=True, cascade=True)
     evaluator.cleanup(command_payload.snapshots)
 
 

tests/core/test_integration.py

@@ -711,6 +711,33 @@ def test_incremental_time_self_reference(
     )
 
 
+@pytest.mark.integration
+@pytest.mark.core_integration
+def test_invalidating_environment(sushi_context: Context):
+    def get_schemas() -> t.Set[str]:
+        return set(
+            sushi_context.state_sync.state_sync.engine_adapter.fetchdf("SHOW ALL TABLES")["schema"]  # type: ignore
+            .to_dict()
+            .values()
+        )
+
+    environment = "dev"
+    apply_to_environment(sushi_context, environment)
+    start_environment = sushi_context.state_sync.get_environment("dev")
+    assert start_environment is not None
+    start_schemas = get_schemas()
+    sushi_context.invalidate_environment("dev")
+    invalidate_environment = sushi_context.state_sync.get_environment("dev")
+    assert invalidate_environment is not None
+    schemas_prior_to_janitor = get_schemas()
+    assert invalidate_environment.expiration_ts < start_environment.expiration_ts  # type: ignore
+    assert start_schemas == schemas_prior_to_janitor
+    sushi_context._run_janitor()
+    schemas_after_janitor = get_schemas()
+    assert sushi_context.state_sync.get_environment("dev") is None
+    assert start_schemas - schemas_after_janitor == {"sushi__dev"}
+
+
 def initial_add(context: Context, environment: str):
     assert not context.state_reader.get_environment(environment)
 

tests/schedulers/airflow/operators/test_targets.py

@@ -1,4 +1,5 @@
 import typing as t
+from unittest.mock import call
 
 import pytest
 from airflow.exceptions import AirflowSkipException
@@ -134,7 +135,6 @@ def test_cleanup_target_execute(mocker: MockerFixture, make_snapshot: t.Callable
 
     context = Context(ti=task_instance_mock)  # type: ignore
 
-    evaluator_demote_mock = mocker.patch("sqlmesh.core.snapshot.evaluator.SnapshotEvaluator.demote")
     evaluator_cleanup_mock = mocker.patch(
         "sqlmesh.core.snapshot.evaluator.SnapshotEvaluator.cleanup"
     )
@@ -143,9 +143,12 @@ def test_cleanup_target_execute(mocker: MockerFixture, make_snapshot: t.Callable
 
     target = targets.SnapshotCleanupTarget()
 
-    target.execute(context, lambda: mocker.Mock(), "spark")
+    evaluator_adapter_mock = mocker.MagicMock()
+    target.execute(context, lambda: evaluator_adapter_mock, "spark")
 
-    evaluator_demote_mock.assert_called_once_with([snapshot.table_info], "test_env")
+    evaluator_adapter_mock.cursor().execute.assert_has_calls(
+        [call("DROP SCHEMA IF EXISTS default__test_env CASCADE")]
+    )
     evaluator_cleanup_mock.assert_called_once_with([snapshot.table_info])
 
     task_instance_mock.xcom_pull.assert_called_once_with(key="snapshot_cleanup_command")