fix(janitor): scope snapshot cleanup for invalidated envs

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

Author: mday-io

docs/reference/cli.md

@@ -294,11 +294,15 @@ Usage: sqlmesh invalidate [OPTIONS] ENVIRONMENT
   of the janitor process.
 
 Options:
-  -s, --sync  Wait for the environment to be deleted before returning. If not
-              specified, the environment will be deleted asynchronously by the
-              janitor process. This option requires a connection to the data
-              warehouse.
-  --help      Show this message and exit.
+  -s, --sync          Wait for the environment to be deleted before returning.
+                      If not specified, the environment will be deleted
+                      asynchronously by the janitor process. This option
+                      requires a connection to the data warehouse.
+  --cleanup-snapshots
+                      After invalidating, synchronously delete unreferenced
+                      physical snapshot tables formerly referenced by this
+                      environment.
+  --help              Show this message and exit.
 ```
 
 ## janitor
@@ -313,14 +317,16 @@ Usage: sqlmesh janitor [OPTIONS]
 Options:
   --ignore-ttl      Cleanup snapshots that are not referenced in any
                     environment, regardless of when they're set to expire. Has
-                    no effect when --environment is specified.
+                    When --environment is specified, cleanup is scoped to
+                    snapshots formerly referenced by that environment.
   --force-delete    Delete expired environment and snapshot state records even
                     when the physical table or view drops fail. Any objects
                     that could not be dropped become orphaned and must be
                     removed manually.
   -e, --environment TEXT
-                    Scope cleanup to a single expired environment. Global
-                    snapshot and interval compaction are skipped.
+                    Scope cleanup to a single expired environment. With
+                    --ignore-ttl, snapshot cleanup is scoped to snapshots
+                    formerly referenced by this environment.
   --help            Show this message and exit.
 ```
 

sqlmesh/cli/main.py

@@ -623,7 +623,7 @@ def run(ctx: click.Context, environment: t.Optional[str] = None, **kwargs: t.Any
 @click.option(
     "--cleanup-snapshots",
     is_flag=True,
-    help="After invalidating, immediately delete physical snapshot tables that are exclusively owned by this environment (not referenced by any other environment). Cleanup runs synchronously regardless of --sync.",
+    help="After invalidating, synchronously delete unreferenced physical snapshot tables formerly referenced by this environment.",
 )
 @click.pass_context
 @error_handler
@@ -638,7 +638,7 @@ def invalidate(ctx: click.Context, environment: str, **kwargs: t.Any) -> None:
 @click.option(
     "--ignore-ttl",
     is_flag=True,
-    help="Cleanup snapshots that are not referenced in any environment, regardless of when they're set to expire. Has no effect when --environment is specified.",
+    help="Cleanup snapshots that are not referenced in any environment, regardless of when they're set to expire. When --environment is specified, cleanup is scoped to snapshots formerly referenced by that environment.",
 )
 @click.option(
     "--force-delete",
@@ -650,7 +650,7 @@ def invalidate(ctx: click.Context, environment: str, **kwargs: t.Any) -> None:
     "--environment",
     "-e",
     default=None,
-    help="Scope cleanup to a single expired environment. Global snapshot and interval compaction are skipped.",
+    help="Scope cleanup to a single expired environment. With --ignore-ttl, snapshot cleanup is scoped to snapshots formerly referenced by this environment.",
 )
 @click.pass_context
 @error_handler

sqlmesh/core/context.py

@@ -112,7 +112,6 @@
 from sqlmesh.core.janitor import (
     cleanup_expired_views,
     delete_expired_snapshots,
-    delete_snapshots_for_environment,
 )
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
@@ -1861,41 +1860,16 @@ def invalidate_environment(
             name: The name of the environment to invalidate.
             sync: If True, the call blocks until the environment is deleted. Otherwise, the environment will
                 be deleted asynchronously by the janitor process.
-            cleanup_snapshots: If True, immediately deletes physical snapshot tables that are exclusively
-                owned by this environment (not referenced by any other environment). Cleanup runs
-                synchronously regardless of --sync.
+            cleanup_snapshots: If True, immediately deletes unreferenced physical snapshot tables that were
+                formerly referenced by this environment. Cleanup runs synchronously regardless of sync.
         """
         name = Environment.sanitize_name(name)
-        sync = sync or cleanup_snapshots
-
-        target_snapshot_ids: t.Set[SnapshotId] = set()
-        if cleanup_snapshots:
-            # Capture snapshot IDs before invalidation so we can scope the cleanup afterwards.
-            env = self.state_sync.get_environment(name)
-            if env is None:
-                logger.warning("Environment '%s' does not exist; skipping snapshot cleanup.", name)
-                return
-            target_snapshot_ids = {s.snapshot_id for s in env.snapshots}
-
         self.state_sync.invalidate_environment(name)
-
-        if sync:
+        if cleanup_snapshots:
+            self._run_janitor(ignore_ttl=True, environment=name)
+            self.console.log_success(f"Environment '{name}' deleted.")
+        elif sync:
             self._cleanup_environments(name=name)
-            if cleanup_snapshots and target_snapshot_ids:
-                failures = delete_snapshots_for_environment(
-                    self.state_sync,
-                    self.snapshot_evaluator,
-                    target_snapshot_ids,
-                    console=self.console,
-                )
-                if failures:
-                    summary = "\n".join(failures)
-                    if self.config.janitor.warn_on_delete_failure:
-                        self.console.log_warning(
-                            f"Snapshot cleanup completed with failures:\n{summary}"
-                        )
-                    else:
-                        raise SQLMeshError(f"Snapshot cleanup completed with failures:\n{summary}")
             self.console.log_success(f"Environment '{name}' deleted.")
         else:
             self.console.log_success(f"Environment '{name}' invalidated.")
@@ -3030,6 +3004,16 @@ def _run_janitor(
         current_ts = now_timestamp()
         failures: t.List[str] = []
 
+        target_snapshot_ids: t.Set[SnapshotId] = set()
+        if environment is not None and ignore_ttl:
+            expired_environments = self.state_sync.get_expired_environments(
+                current_ts=current_ts, name=environment
+            )
+            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}
+
         # Clean up expired environments by removing their views and schemas
         failures.extend(
             self._cleanup_environments(
@@ -3050,6 +3034,25 @@ def _run_janitor(
                 )
             )
             self.state_sync.compact_intervals()
+        elif ignore_ttl and target_snapshot_ids and not self.state_reader.get_environment(
+            environment
+        ):
+            self.console.log_warning(
+                "Scoped snapshot cleanup will permanently delete unreferenced physical snapshot "
+                f"tables formerly referenced by environment '{environment}'."
+            )
+            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,
+                    target_snapshot_ids=target_snapshot_ids,
+                )
+            )
 
         if failures:
             failure_string = "\n  - ".join(failures)

sqlmesh/core/janitor.py

@@ -8,7 +8,7 @@
 from sqlmesh.core.console import Console
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.environment import Environment
-from sqlmesh.core.snapshot import SnapshotEvaluator, SnapshotId
+from sqlmesh.core.snapshot import SnapshotEvaluator, SnapshotIdLike
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.core.state_sync.common import (
     logger,
@@ -127,6 +127,7 @@ def delete_expired_snapshots(
     ignore_ttl: bool = False,
     force_delete: bool = False,
     batch_size: t.Optional[int] = None,
+    target_snapshot_ids: t.Optional[t.Collection[SnapshotIdLike]] = None,
     console: t.Optional[Console] = None,
 ) -> t.List[str]:
     """Delete all expired snapshots in batches.
@@ -153,6 +154,7 @@ def delete_expired_snapshots(
         current_ts=current_ts,
         ignore_ttl=ignore_ttl,
         batch_size=batch_size,
+        target_snapshot_ids=target_snapshot_ids,
     ):
         end_info = (
             f"updated_ts={batch.batch_range.end.updated_ts}"
@@ -184,6 +186,7 @@ def delete_expired_snapshots(
                         end=batch.batch_range.end,
                     ),
                     ignore_ttl=ignore_ttl,
+                    target_snapshot_ids=target_snapshot_ids,
                 )
                 logger.info("Cleaned up expired snapshots batch")
                 num_expired_snapshots += len(batch.expired_snapshot_ids)
@@ -193,72 +196,3 @@ def delete_expired_snapshots(
                 failures.append(message)
     logger.info("Cleaned up %s expired snapshots", num_expired_snapshots)
     return failures
-
-
-def delete_snapshots_for_environment(
-    state_sync: StateSync,
-    snapshot_evaluator: SnapshotEvaluator,
-    target_snapshot_ids: t.Collection[SnapshotId],
-    *,
-    force_delete: bool = False,
-    console: t.Optional[Console] = None,
-) -> t.List[str]:
-    """Delete snapshots that are exclusively owned by a specific (now-deleted) environment.
-
-    This performs a scoped cleanup: only the provided snapshot IDs are considered for deletion,
-    and only those that are not referenced by any remaining active environment will be removed.
-
-    Args:
-        state_sync: StateSync instance to query and delete snapshot state from.
-        snapshot_evaluator: SnapshotEvaluator instance to clean up physical tables.
-        target_snapshot_ids: The snapshot IDs to consider for deletion (typically from the
-            environment that was just invalidated/deleted).
-        force_delete: If True, delete snapshot state records even when physical table cleanup fails.
-        console: Optional console for reporting progress.
-
-    Returns:
-        List of failure messages encountered during cleanup.
-    """
-    if not target_snapshot_ids:
-        return []
-
-    failures: t.List[str] = []
-    batch = state_sync.get_expired_snapshots(
-        ignore_ttl=True,
-        batch_range=ExpiredBatchRange.all_batch_range(),
-        target_snapshot_ids=target_snapshot_ids,
-    )
-    if batch is None:
-        return failures
-
-    logger.info(
-        "Cleaning up %s snapshots exclusively owned by invalidated environment",
-        len(batch.expired_snapshot_ids),
-    )
-
-    cleanup_succeeded = True
-    if batch.cleanup_tasks:
-        try:
-            snapshot_evaluator.cleanup(
-                target_snapshots=batch.cleanup_tasks,
-                on_complete=console.update_cleanup_progress if console else None,
-            )
-        except Exception as failed_drops:
-            message = f"Failed to clean up: {failed_drops}"
-            logger.warning(message)
-            failures.append(message)
-            cleanup_succeeded = False
-
-    if cleanup_succeeded or force_delete:
-        try:
-            state_sync.delete_snapshots(batch.expired_snapshot_ids)
-            logger.info(
-                "Cleaned up %s snapshots from invalidated environment",
-                len(batch.expired_snapshot_ids),
-            )
-        except Exception as e:
-            message = f"Failed to delete snapshot state records: {e}"
-            logger.warning(message)
-            failures.append(message)
-
-    return failures

sqlmesh/core/state_sync/common.py

@@ -16,6 +16,7 @@
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
+    SnapshotIdLike,
     SnapshotTableCleanupTask,
     SnapshotTableInfo,
 )
@@ -288,6 +289,7 @@ def iter_expired_snapshot_batches(
     current_ts: int,
     ignore_ttl: bool = False,
     batch_size: t.Optional[int] = None,
+    target_snapshot_ids: t.Optional[t.Collection[SnapshotIdLike]] = None,
 ) -> t.Iterator[ExpiredSnapshotBatch]:
     """Yields expired snapshot batches.
 
@@ -306,6 +308,7 @@ def iter_expired_snapshot_batches(
             current_ts=current_ts,
             ignore_ttl=ignore_ttl,
             batch_range=batch_range,
+            target_snapshot_ids=target_snapshot_ids,
         )
 
         if batch is None:

tests/core/integration/test_aux_commands.py

@@ -482,7 +482,7 @@ def test_invalidating_environment(sushi_context: Context):
 
 
 def test_invalidate_environment_cleanup_snapshots_scoped(tmp_path: Path):
-    """Test that --cleanup-snapshots only deletes snapshots exclusively owned by the invalidated env."""
+    """Test that --cleanup-snapshots only deletes unreferenced snapshots from the invalidated env."""
     models_dir = tmp_path / "models"
     models_dir.mkdir()
     (models_dir / "model1.sql").write_text("MODEL(name test.model1, kind FULL); SELECT 1 AS col")
@@ -522,8 +522,9 @@ def test_invalidate_environment_cleanup_snapshots_scoped(tmp_path: Path):
     assert ctx.state_sync.get_environment("prod") is not None
 
 
-def test_invalidate_environment_cleanup_snapshots_exclusive(tmp_path: Path):
-    """Test that --cleanup-snapshots deletes snapshots exclusively owned by the invalidated env."""
+def test_invalidate_environment_cleanup_snapshots_warns_and_drops_physical_tables(
+    tmp_path: Path, mocker: MockerFixture
+):
     models_dir = tmp_path / "models"
     models_dir.mkdir()
     (models_dir / "model1.sql").write_text("MODEL(name test.model1, kind FULL); SELECT 1 AS col")
@@ -533,22 +534,67 @@ def test_invalidate_environment_cleanup_snapshots_exclusive(tmp_path: Path):
         config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")),
     )
 
-    # Apply model1 to dev only (not prod). These snapshots will be exclusively owned by dev.
     ctx.plan("dev", no_prompts=True, auto_apply=True)
+    snapshot = ctx.get_snapshot("test.model1")
+    assert snapshot is not None
+    snapshot_ids = [snapshot.snapshot_id]
+    physical_table_names = [
+        snapshot.table_name(is_deployable=False),
+        snapshot.table_name(is_deployable=True),
+    ]
+    assert any(ctx.engine_adapter.table_exists(table_name) for table_name in physical_table_names)
 
-    dev_env = ctx.state_sync.get_environment("dev")
-    assert dev_env is not None
-    dev_snapshot_ids = {s.snapshot_id for s in dev_env.snapshots}
-    assert dev_snapshot_ids
+    warning_mock = mocker.patch.object(ctx.console, "log_warning")
 
     ctx.invalidate_environment("dev", cleanup_snapshots=True)
 
-    # The dev environment record should be gone.
+    warning_mock.assert_any_call(
+        "Scoped snapshot cleanup will permanently delete unreferenced physical snapshot tables "
+        "formerly referenced by environment 'dev'."
+    )
     assert ctx.state_sync.get_environment("dev") is None
+    assert not ctx.state_sync.get_snapshots(snapshot_ids)
+    assert not any(ctx.engine_adapter.table_exists(table_name) for table_name in physical_table_names)
+
+
+def test_janitor_environment_ignore_ttl_cleans_only_scoped_snapshots(
+    tmp_path: Path, mocker: MockerFixture
+):
+    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")
 
-    # All dev-exclusive snapshots should have been deleted.
-    remaining_snapshots = ctx.state_sync.get_snapshots(list(dev_snapshot_ids))
-    assert not remaining_snapshots
+    ctx = Context(
+        paths=[tmp_path],
+        config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")),
+    )
+
+    ctx.plan("dev_a", no_prompts=True, auto_apply=True)
+    dev_a_snapshot = ctx.get_snapshot("test.model1")
+    assert dev_a_snapshot is not None
+
+    model_path.write_text("MODEL(name test.model1, kind FULL); SELECT 2 AS col")
+    ctx.load()
+    ctx.plan("dev_b", no_prompts=True, auto_apply=True)
+    dev_b_snapshot = ctx.get_snapshot("test.model1")
+    assert dev_b_snapshot is not None
+    assert dev_a_snapshot.snapshot_id != dev_b_snapshot.snapshot_id
+
+    ctx.invalidate_environment("dev_a")
+    ctx.invalidate_environment("dev_b")
+    warning_mock = mocker.patch.object(ctx.console, "log_warning")
+
+    ctx.run_janitor(ignore_ttl=True, environment="dev_a")
+
+    warning_mock.assert_any_call(
+        "Scoped snapshot cleanup will permanently delete unreferenced physical snapshot tables "
+        "formerly referenced by environment 'dev_a'."
+    )
+    assert ctx.state_sync.get_environment("dev_a") is None
+    assert ctx.state_sync.get_environment("dev_b") is not None
+    assert not ctx.state_sync.get_snapshots([dev_a_snapshot.snapshot_id])
+    assert ctx.state_sync.get_snapshots([dev_b_snapshot.snapshot_id])
 
 
 @time_machine.travel("2023-01-08 15:00:00 UTC")