fix(janitor): include previous finalized snapshots in scoped cleanup

When the invalidated environment is unfinalized, it still points at the
snapshots of its last finalized plan. Include them in the scoped cleanup
target, mirroring how SnapshotState.get_expired_snapshots treats them.

Also fold the duplicated global and scoped delete_expired_snapshots calls
into one, and revert an unrelated rewrite of the GitHub controller test.

Signed-off-by: mday-io <mdaytn@gmail.com>

Author: mday-io

sqlmesh/core/context.py

@@ -3185,7 +3185,17 @@ def _run_janitor(
             if expired_environments:
                 expired_env = self.state_reader.get_environment(expired_environments[0].name)
                 if expired_env:
-                    target_snapshot_ids = {s.snapshot_id for s in expired_env.snapshots}
+                    # An unfinalized environment may still point at the snapshots of its last
+                    # finalized plan, which would otherwise be left behind.
+                    expired_env_snapshots = (
+                        expired_env.snapshots
+                        if expired_env.finalized_ts is not None
+                        else [
+                            *expired_env.snapshots,
+                            *(expired_env.previous_finalized_snapshots or []),
+                        ]
+                    )
+                    target_snapshot_ids = {s.snapshot_id for s in expired_env_snapshots}
 
         # Clean up expired environments by removing their views and schemas
         failures.extend(
@@ -3194,28 +3204,19 @@ def _run_janitor(
             )
         )
 
-        if environment is None:
-            failures.extend(
-                delete_expired_snapshots(
-                    self.state_sync,
-                    self.snapshot_evaluator,
-                    current_ts=current_ts,
-                    ignore_ttl=ignore_ttl,
-                    force_delete=force_delete,
-                    console=self.console,
-                    batch_size=self.config.janitor.expired_snapshots_batch_size,
-                )
-            )
-            self.state_sync.compact_intervals()
-        elif (
-            ignore_ttl
-            and target_snapshot_ids
+        scoped_cleanup = (
+            environment is not None
+            and ignore_ttl
+            and bool(target_snapshot_ids)
             and not self.state_reader.get_environment(environment)
-        ):
+        )
+        if scoped_cleanup:
             self.console.log_warning(
                 "Scoped snapshot cleanup will permanently delete unreferenced physical snapshot "
                 f"tables formerly referenced by environment '{environment}'."
             )
+
+        if environment is None or scoped_cleanup:
             failures.extend(
                 delete_expired_snapshots(
                     self.state_sync,
@@ -3225,10 +3226,13 @@ def _run_janitor(
                     force_delete=force_delete,
                     console=self.console,
                     batch_size=self.config.janitor.expired_snapshots_batch_size,
-                    target_snapshot_ids=target_snapshot_ids,
+                    target_snapshot_ids=target_snapshot_ids if scoped_cleanup else None,
                 )
             )
 
+        if environment is None:
+            self.state_sync.compact_intervals()
+
         if failures:
             failure_string = "\n  - ".join(failures)
             summary = f"Janitor completed with failures:\n  {failure_string}"

tests/core/integration/test_aux_commands.py

@@ -599,6 +599,49 @@ def test_janitor_environment_ignore_ttl_cleans_only_scoped_snapshots(
     assert ctx.state_sync.get_snapshots([dev_b_snapshot.snapshot_id])
 
 
+def test_invalidate_environment_cleanup_snapshots_includes_previous_finalized(tmp_path: Path):
+    """An unfinalized environment's previous finalized snapshots are cleaned up too."""
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+    model_path = models_dir / "model1.sql"
+    model_path.write_text("MODEL(name test.model1, kind FULL); SELECT 1 AS col")
+
+    ctx = Context(
+        paths=[tmp_path],
+        config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")),
+    )
+
+    ctx.plan("dev", no_prompts=True, auto_apply=True)
+    old_snapshot = ctx.get_snapshot("test.model1")
+    assert old_snapshot is not None
+
+    model_path.write_text("MODEL(name test.model1, kind FULL); SELECT 2 AS col")
+    ctx.load()
+    ctx.plan("dev", no_prompts=True, auto_apply=True)
+    new_snapshot = ctx.get_snapshot("test.model1")
+    assert new_snapshot is not None
+    assert old_snapshot.snapshot_id != new_snapshot.snapshot_id
+
+    # Simulate a plan that failed partway: the environment is left unfinalized and still
+    # points at the snapshots from its last finalized plan.
+    dev_env = ctx.state_sync.get_environment("dev")
+    assert dev_env is not None
+    ctx.state_sync.state_sync.environment_state.update_environment(  # type: ignore
+        dev_env.copy(
+            update={
+                "finalized_ts": None,
+                "previous_finalized_snapshots": [old_snapshot.table_info],
+            }
+        )
+    )
+
+    ctx.invalidate_environment("dev", cleanup_snapshots=True)
+
+    assert ctx.state_sync.get_environment("dev") is None
+    assert not ctx.state_sync.get_snapshots([old_snapshot.snapshot_id])
+    assert not ctx.state_sync.get_snapshots([new_snapshot.snapshot_id])
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_evaluate_uncategorized_snapshot(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")

tests/integrations/github/cicd/test_github_controller.py

@@ -578,23 +578,21 @@ def test_try_invalidate_pr_environment(github_client, make_controller, mocker: M
     invalidate_controller = make_controller(
         "tests/fixtures/github/pull_request_synchronized.json", github_client
     )
-    invalidate_environment = mocker.patch.object(
-        invalidate_controller._context, "invalidate_environment"
-    )
+    invalidate_controller._context._state_sync = mocker.MagicMock()
     invalidate_controller.try_invalidate_pr_environment()
-    invalidate_environment.assert_called_once_with("hello_world_2")
+    invalidate_controller._context._state_sync.invalidate_environment.assert_called_once_with(
+        "hello_world_2"
+    )
 
     no_invalidate_controller = make_controller(
         "tests/fixtures/github/pull_request_synchronized.json",
         github_client,
         bot_config=GithubCICDBotConfig(invalidate_environment_after_deploy=False),
     )
-    no_invalidate_environment = mocker.patch.object(
-        no_invalidate_controller._context, "invalidate_environment"
-    )
+    no_invalidate_controller._context._state_sync = mocker.MagicMock()
     no_invalidate_controller.try_invalidate_pr_environment()
 
-    no_invalidate_environment.assert_not_called()
+    assert not no_invalidate_controller._context._state_sync.invalidate_environment.called
 
 
 def test_try_merge_pr(github_client, make_controller, mocker: MockerFixture):