feat: add --environment flag to sqlmesh janitor for scoped cleanup

Signed-off-by: Michael Day <michael.day@cloudkitchens.com>
Signed-off-by: mday-io <mdaytn@gmail.com>

Author: mday-io

sqlmesh/cli/main.py

@@ -629,24 +629,36 @@ 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",
+    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.",
 )
 @click.option(
     "--force-delete",
     is_flag=True,
     help="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.",
 )
+@click.option(
+    "--environment",
+    "-e",
+    default=None,
+    help="Scope cleanup to a single expired environment. Global snapshot and interval compaction are skipped.",
+)
 @click.pass_context
 @error_handler
 @cli_analytics
-def janitor(ctx: click.Context, ignore_ttl: bool, force_delete: bool, **kwargs: t.Any) -> None:
+def janitor(
+    ctx: click.Context,
+    ignore_ttl: bool,
+    force_delete: bool,
+    environment: t.Optional[str],
+    **kwargs: t.Any,
+) -> None:
     """
     Run the janitor process on-demand.
 
     The janitor cleans up old environments and expired snapshots.
     """
-    ctx.obj.run_janitor(ignore_ttl, force_delete=force_delete, **kwargs)
+    ctx.obj.run_janitor(ignore_ttl, force_delete=force_delete, environment=environment, **kwargs)
 
 
 @cli.command("destroy")

sqlmesh/core/context.py

@@ -887,12 +887,20 @@ def _has_environment_changed() -> bool:
         return completion_status
 
     @python_api_analytics
-    def run_janitor(self, ignore_ttl: bool, force_delete: bool = False) -> bool:
+    def run_janitor(
+        self,
+        ignore_ttl: bool,
+        force_delete: bool = False,
+        environment: t.Optional[str] = None,
+    ) -> bool:
+        if environment is not None:
+            environment = Environment.sanitize_name(environment)
+
         success = False
 
         if self.console.start_cleanup(ignore_ttl):
             try:
-                self._run_janitor(ignore_ttl, force_delete=force_delete)
+                self._run_janitor(ignore_ttl, force_delete=force_delete, environment=environment)
                 success = True
             finally:
                 self.console.stop_cleanup(success=success)
@@ -1825,7 +1833,7 @@ def invalidate_environment(self, name: str, sync: bool = False) -> None:
         name = Environment.sanitize_name(name)
         self.state_sync.invalidate_environment(name)
         if sync:
-            self._cleanup_environments()
+            self._cleanup_environments(name=name)
             self.console.log_success(f"Environment '{name}' deleted.")
         else:
             self.console.log_success(f"Environment '{name}' invalidated.")
@@ -2896,27 +2904,35 @@ def _destroy(self) -> bool:
 
         return True
 
-    def _run_janitor(self, ignore_ttl: bool = False, force_delete: bool = False) -> None:
+    def _run_janitor(
+        self,
+        ignore_ttl: bool = False,
+        force_delete: bool = False,
+        environment: t.Optional[str] = None,
+    ) -> None:
         current_ts = now_timestamp()
         failures: t.List[str] = []
 
         # Clean up expired environments by removing their views and schemas
         failures.extend(
-            self._cleanup_environments(current_ts=current_ts, force_delete=force_delete)
+            self._cleanup_environments(
+                current_ts=current_ts, force_delete=force_delete, name=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,
+        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()
+            self.state_sync.compact_intervals()
 
         if failures:
             failure_string = "\n  - ".join(failures)
@@ -2929,15 +2945,23 @@ def _run_janitor(self, ignore_ttl: bool = False, force_delete: bool = False) ->
                 raise SQLMeshError(summary)
 
     def _cleanup_environments(
-        self, current_ts: t.Optional[int] = None, force_delete: bool = False
+        self,
+        current_ts: t.Optional[int] = None,
+        force_delete: bool = False,
+        name: t.Optional[str] = None,
     ) -> t.List[str]:
         current_ts = current_ts or now_timestamp()
         failures: t.List[str] = []
 
         expired_environments_summaries = self.state_sync.get_expired_environments(
-            current_ts=current_ts
+            current_ts=current_ts, name=name
         )
 
+        if name is not None and not expired_environments_summaries:
+            self.console.log_warning(
+                f"Environment '{name}' is not expired or does not exist. Nothing to clean up."
+            )
+
         for expired_env_summary in expired_environments_summaries:
             expired_env = self.state_reader.get_environment(expired_env_summary.name)
 
@@ -2954,7 +2978,7 @@ def _cleanup_environments(
         # we want to retry on the next janitor pass if drops failed, unless
         # force_delete is set in which case we purge state records regardless
         if not failures or force_delete:
-            self.state_sync.delete_expired_environments(current_ts=current_ts)
+            self.state_sync.delete_expired_environments(current_ts=current_ts, name=name)
         return failures
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:

sqlmesh/core/state_sync/base.py

@@ -321,10 +321,17 @@ def get_expired_snapshots(
         """
 
     @abc.abstractmethod
-    def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
+    def get_expired_environments(
+        self, current_ts: int, name: t.Optional[str] = None
+    ) -> t.List[EnvironmentSummary]:
         """Returns the expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.
+
+        Args:
+            current_ts: The current timestamp in milliseconds used to determine expiration.
+            name: If provided, only the environment with this name is considered.
+
         Returns:
             The list of environment summaries to remove.
         """
@@ -436,12 +443,16 @@ def finalize(self, environment: Environment) -> None:
 
     @abc.abstractmethod
     def delete_expired_environments(
-        self, current_ts: t.Optional[int] = None
+        self, current_ts: t.Optional[int] = None, name: t.Optional[str] = None
     ) -> t.List[EnvironmentSummary]:
         """Removes expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.
 
+        Args:
+            current_ts: The current timestamp in milliseconds. Defaults to now.
+            name: If provided, only the environment with this name is deleted.
+
         Returns:
             The list of removed environments.
         """

sqlmesh/core/state_sync/db/environment.py

@@ -167,31 +167,46 @@ def finalize(self, environment: Environment) -> None:
             where=environment_filter,
         )
 
-    def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
+    def get_expired_environments(
+        self, current_ts: int, name: t.Optional[str] = None
+    ) -> t.List[EnvironmentSummary]:
         """Returns the expired environments.
 
         Expired environments are environments that have exceeded their time-to-live value.
+
+        Args:
+            current_ts: The current timestamp in milliseconds used to determine expiration.
+            name: If provided, only the environment with this name is considered.
+
         Returns:
             The list of environment summaries to remove.
         """
-        return self._fetch_environment_summaries(
-            where=self._create_expiration_filter_expr(current_ts)
-        )
+        where: exp.Expression = self._create_expiration_filter_expr(current_ts)
+        if name is not None:
+            where = exp.and_(where, exp.column("name").eq(name))
+        return self._fetch_environment_summaries(where=where)
 
     def delete_expired_environments(
-        self, current_ts: t.Optional[int] = None
+        self, current_ts: t.Optional[int] = None, name: t.Optional[str] = None
     ) -> t.List[EnvironmentSummary]:
         """Deletes expired environments.
 
+        Args:
+            current_ts: The current timestamp in milliseconds. Defaults to now.
+            name: If provided, only the environment with this name is deleted.
+
         Returns:
             A list of deleted environments.
         """
         current_ts = current_ts or now_timestamp()
-        expired_environments = self.get_expired_environments(current_ts=current_ts)
+        expired_environments = self.get_expired_environments(current_ts=current_ts, name=name)
 
+        where: exp.Expression = self._create_expiration_filter_expr(current_ts)
+        if name is not None:
+            where = exp.and_(where, exp.column("name").eq(name))
         self.engine_adapter.delete_from(
             self.environments_table,
-            where=self._create_expiration_filter_expr(current_ts),
+            where=where,
         )
 
         # Delete the expired environments' corresponding environment statements

sqlmesh/core/state_sync/db/facade.py

@@ -276,8 +276,10 @@ def get_expired_snapshots(
             batch_range=batch_range,
         )
 
-    def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
-        return self.environment_state.get_expired_environments(current_ts=current_ts)
+    def get_expired_environments(
+        self, current_ts: int, name: t.Optional[str] = None
+    ) -> t.List[EnvironmentSummary]:
+        return self.environment_state.get_expired_environments(current_ts=current_ts, name=name)
 
     @transactional()
     def delete_expired_snapshots(
@@ -297,10 +299,10 @@ def delete_expired_snapshots(
 
     @transactional()
     def delete_expired_environments(
-        self, current_ts: t.Optional[int] = None
+        self, current_ts: t.Optional[int] = None, name: t.Optional[str] = None
     ) -> t.List[EnvironmentSummary]:
         current_ts = current_ts or now_timestamp()
-        return self.environment_state.delete_expired_environments(current_ts=current_ts)
+        return self.environment_state.delete_expired_environments(current_ts=current_ts, name=name)
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.snapshot_state.delete_snapshots(snapshot_ids)