feat: batch expired snapshots

Author: eakmanrq

sqlmesh/core/config/janitor.py

@@ -1,14 +1,26 @@
 from __future__ import annotations
 
+import typing as t
 
 from sqlmesh.core.config.base import BaseConfig
+from sqlmesh.utils.pydantic import field_validator
 
 
 class JanitorConfig(BaseConfig):
     """The configuration for the janitor.
 
     Args:
         warn_on_delete_failure: Whether to warn instead of erroring if the janitor fails to delete the expired environment schema / views.
+        expired_snapshots_batch_size: Maximum number of expired snapshots to clean in a single batch.
     """
 
     warn_on_delete_failure: bool = False
+    expired_snapshots_batch_size: t.Optional[int] = None
+
+    @field_validator("expired_snapshots_batch_size", mode="before")
+    @classmethod
+    def _validate_batch_size(cls, value: int) -> int:
+        batch_size = int(value)
+        if batch_size <= 0:
+            raise ValueError("expired_snapshots_batch_size must be greater than 0")
+        return batch_size

sqlmesh/core/context.py

@@ -109,6 +109,7 @@
     StateSync,
     cleanup_expired_views,
 )
+from sqlmesh.core.state_sync.common import delete_expired_snapshots
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
     ModelTextTestResult,
@@ -2852,19 +2853,14 @@ def _run_janitor(self, ignore_ttl: bool = False) -> None:
         # Clean up expired environments by removing their views and schemas
         self._cleanup_environments(current_ts=current_ts)
 
-        cleanup_targets = self.state_sync.get_expired_snapshots(
-            ignore_ttl=ignore_ttl, current_ts=current_ts
-        )
-
-        # Remove the expired snapshots tables
-        self.snapshot_evaluator.cleanup(
-            target_snapshots=cleanup_targets,
-            on_complete=self.console.update_cleanup_progress,
+        delete_expired_snapshots(
+            self.state_sync,
+            self.snapshot_evaluator,
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+            console=self.console,
+            batch_size=self.config.janitor.expired_snapshots_batch_size,
         )
-
-        # Delete the expired snapshot records from the state sync
-        self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl, current_ts=current_ts)
-
         self.state_sync.compact_intervals()
 
     def _cleanup_environments(self, current_ts: t.Optional[int] = None) -> None:

sqlmesh/core/state_sync/base.py

@@ -72,6 +72,50 @@ def _schema_version_validator(cls, v: t.Any) -> int:
 SCHEMA_VERSION: int = MIN_SCHEMA_VERSION + len(MIGRATIONS) - 1
 
 
+class BatchBoundary(PydanticModel):
+    updated_ts: int
+    name: str
+    identifier: str
+
+    def to_upper_batch_boundary(self) -> UpperBatchBoundary:
+        return UpperBatchBoundary(
+            updated_ts=self.updated_ts,
+            name=self.name,
+            identifier=self.identifier,
+        )
+
+    def to_lower_batch_boundary(self, batch_size: int) -> LowerBatchBoundary:
+        return LowerBatchBoundary(
+            updated_ts=self.updated_ts,
+            name=self.name,
+            identifier=self.identifier,
+            batch_size=batch_size,
+        )
+
+
+class UpperBatchBoundary(BatchBoundary):
+    @classmethod
+    def include_all_boundary(cls) -> UpperBatchBoundary:
+        # 9999-12-31T23:59:59.999Z in epoch milliseconds
+        return UpperBatchBoundary(updated_ts=253_402_300_799_999, name="", identifier="")
+
+
+class LowerBatchBoundary(BatchBoundary):
+    batch_size: int
+
+    @classmethod
+    def init_batch_boundary(cls, batch_size: int) -> LowerBatchBoundary:
+        return LowerBatchBoundary(updated_ts=0, name="", identifier="", batch_size=batch_size)
+
+
+class ExpiredSnapshotBatch(PydanticModel):
+    """A batch of expired snapshots to be cleaned up."""
+
+    expired_snapshot_ids: t.Set[SnapshotId]
+    cleanup_tasks: t.List[SnapshotTableCleanupTask]
+    batch_boundary: BatchBoundary
+
+
 class PromotionResult(PydanticModel):
     added: t.List[SnapshotTableInfo]
     removed: t.List[SnapshotTableInfo]
@@ -315,15 +359,23 @@ def export(self, environment_names: t.Optional[t.List[str]] = None) -> StateStre
 
     @abc.abstractmethod
     def get_expired_snapshots(
-        self, current_ts: t.Optional[int] = None, ignore_ttl: bool = False
-    ) -> t.List[SnapshotTableCleanupTask]:
-        """Aggregates the id's of the expired snapshots and creates a list of table cleanup tasks.
+        self,
+        *,
+        batch_boundary: BatchBoundary,
+        current_ts: t.Optional[int] = None,
+        ignore_ttl: bool = False,
+    ) -> t.Optional[ExpiredSnapshotBatch]:
+        """Returns a single batch of expired snapshots ordered by (updated_ts, name, identifier).
 
-        Expired snapshots are snapshots that have exceeded their time-to-live
-        and are no longer in use within an environment.
+        Args:
+            current_ts: Timestamp used to evaluate expiration.
+            ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
+            batch_boundary: If provided, gets snapshot relative to the given boundary.
+                If lower boundary then snapshots later than that will be returned (exclusive).
+                If upper boundary then snapshots earlier than that will be returned (inclusive).
 
         Returns:
-           The list of table cleanup tasks.
+            A batch describing expired snapshots or None if no snapshots are pending cleanup.
         """
 
     @abc.abstractmethod
@@ -363,7 +415,9 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
 
     @abc.abstractmethod
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
+        self,
+        ignore_ttl: bool = False,
+        upper_batch_boundary: t.Optional[UpperBatchBoundary] = None,
     ) -> None:
         """Removes expired snapshots.
 
@@ -373,6 +427,8 @@ def delete_expired_snapshots(
         Args:
             ignore_ttl: Ignore the TTL on the snapshot when considering it expired. This has the effect of deleting
                 all snapshots that are not referenced in any environment
+            upper_batch_boundary: The upper boundary to delete expired snapshots till (inclusive). If not provided,
+                deletes all expired snapshots.
         """
 
     @abc.abstractmethod

sqlmesh/core/state_sync/cache.py

@@ -11,7 +11,7 @@
     SnapshotInfoLike,
 )
 from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
-from sqlmesh.core.state_sync.base import DelegatingStateSync, StateSync
+from sqlmesh.core.state_sync.base import DelegatingStateSync, StateSync, UpperBatchBoundary
 from sqlmesh.utils.date import TimeLike, now_timestamp
 
 
@@ -108,11 +108,15 @@ def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         self.state_sync.delete_snapshots(snapshot_ids)
 
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
+        self,
+        ignore_ttl: bool = False,
+        upper_batch_boundary: t.Optional[UpperBatchBoundary] = None,
     ) -> None:
-        current_ts = current_ts or now_timestamp()
         self.snapshot_cache.clear()
-        self.state_sync.delete_expired_snapshots(current_ts=current_ts, ignore_ttl=ignore_ttl)
+        self.state_sync.delete_expired_snapshots(
+            upper_batch_boundary=upper_batch_boundary,
+            ignore_ttl=ignore_ttl,
+        )
 
     def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
         for snapshot_intervals in snapshots_intervals:

sqlmesh/core/state_sync/common.py

@@ -14,14 +14,16 @@
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlmesh.core.environment import Environment, EnvironmentStatements
 from sqlmesh.utils.errors import SQLMeshError
-from sqlmesh.core.snapshot import Snapshot
+from sqlmesh.core.snapshot import Snapshot, SnapshotEvaluator
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter.base import EngineAdapter
-    from sqlmesh.core.state_sync.base import Versions
+    from sqlmesh.core.state_sync.base import Versions, ExpiredSnapshotBatch, StateReader, StateSync
 
 logger = logging.getLogger(__name__)
 
+EXPIRED_SNAPSHOT_DEFAULT_BATCH_SIZE = 200
+
 
 def cleanup_expired_views(
     default_adapter: EngineAdapter,
@@ -215,3 +217,94 @@ def __iter__(self) -> t.Iterator[StateStreamContents]:
                 yield EnvironmentsChunk(environments)
 
         return _StateStream()
+
+
+def for_each_expired_snapshot_batch(
+    state_reader: StateReader,
+    *,
+    current_ts: int,
+    ignore_ttl: bool = False,
+    callback: t.Callable[[ExpiredSnapshotBatch], None],
+    batch_size: t.Optional[int] = None,
+) -> int:
+    """Iterate over expired snapshot batches and invoke callback for each batch.
+
+    Args:
+        state_reader: StateReader instance to query expired snapshots from.
+        current_ts: Timestamp used to evaluate expiration.
+        ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
+        callback: Function to invoke for each batch of expired snapshots.
+        batch_size: Maximum number of snapshots to fetch per batch.
+    """
+    from sqlmesh.core.state_sync.base import LowerBatchBoundary
+
+    batch_size = batch_size if batch_size is not None else EXPIRED_SNAPSHOT_DEFAULT_BATCH_SIZE
+    batch_boundary = LowerBatchBoundary.init_batch_boundary(batch_size=batch_size)
+    num_expired_snapshots = 0
+
+    while True:
+        batch = state_reader.get_expired_snapshots(
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+            batch_boundary=batch_boundary,
+        )
+
+        if not batch:
+            return num_expired_snapshots
+
+        callback(batch)
+
+        batch_boundary = batch.batch_boundary.to_lower_batch_boundary(batch_size=batch_size)
+        num_expired_snapshots += len(batch.expired_snapshot_ids)
+
+
+def delete_expired_snapshots(
+    state_sync: StateSync,
+    snapshot_evaluator: SnapshotEvaluator,
+    *,
+    current_ts: int,
+    ignore_ttl: bool = False,
+    batch_size: t.Optional[int] = None,
+    console: t.Optional[Console] = None,
+) -> None:
+    """Delete all expired snapshots in batches.
+
+    This helper function encapsulates the logic for deleting expired snapshots in batches,
+    eliminating code duplication across different use cases.
+
+    Args:
+        state_sync: StateSync instance to query and delete expired snapshots from.
+        snapshot_evaluator: SnapshotEvaluator instance to clean up tables associated with snapshots.
+        current_ts: Timestamp used to evaluate expiration.
+        ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
+        batch_size: Maximum number of snapshots to fetch per batch.
+        console: Optional console for reporting progress.
+
+    Returns:
+        The total number of deleted expired snapshots.
+    """
+
+    def process_batch(batch: ExpiredSnapshotBatch) -> None:
+        logger.info(
+            "Processing batch of size %s and max_updated_ts of %s",
+            len(batch.expired_snapshot_ids),
+            batch.batch_boundary.updated_ts,
+        )
+        snapshot_evaluator.cleanup(
+            target_snapshots=batch.cleanup_tasks,
+            on_complete=console.update_cleanup_progress if console else None,
+        )
+        state_sync.delete_expired_snapshots(
+            upper_batch_boundary=batch.batch_boundary.to_upper_batch_boundary(),
+            ignore_ttl=ignore_ttl,
+        )
+        logger.info("Cleaned up expired snapshots batch")
+
+    num_deleted = for_each_expired_snapshot_batch(
+        state_sync,
+        current_ts=current_ts,
+        ignore_ttl=ignore_ttl,
+        callback=process_batch,
+        batch_size=batch_size,
+    )
+    logger.info("Cleaned up %s expired snapshots", num_deleted)

sqlmesh/core/state_sync/db/facade.py

@@ -35,17 +35,19 @@
     SnapshotInfoLike,
     SnapshotIntervals,
     SnapshotNameVersion,
-    SnapshotTableCleanupTask,
     SnapshotTableInfo,
     start_date,
 )
 from sqlmesh.core.snapshot.definition import (
     Interval,
 )
 from sqlmesh.core.state_sync.base import (
+    ExpiredSnapshotBatch,
     PromotionResult,
     StateSync,
     Versions,
+    BatchBoundary,
+    UpperBatchBoundary,
 )
 from sqlmesh.core.state_sync.common import (
     EnvironmentsChunk,
@@ -261,26 +263,37 @@ def invalidate_environment(self, name: str, protect_prod: bool = True) -> None:
         self.environment_state.invalidate_environment(name, protect_prod)
 
     def get_expired_snapshots(
-        self, current_ts: t.Optional[int] = None, ignore_ttl: bool = False
-    ) -> t.List[SnapshotTableCleanupTask]:
+        self,
+        *,
+        batch_boundary: BatchBoundary,
+        current_ts: t.Optional[int] = None,
+        ignore_ttl: bool = False,
+    ) -> t.Optional[ExpiredSnapshotBatch]:
         current_ts = current_ts or now_timestamp()
         return self.snapshot_state.get_expired_snapshots(
-            self.environment_state.get_environments(), current_ts=current_ts, ignore_ttl=ignore_ttl
+            environments=self.environment_state.get_environments(),
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+            batch_boundary=batch_boundary,
         )
 
     def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
         return self.environment_state.get_expired_environments(current_ts=current_ts)
 
     @transactional()
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
+        self,
+        ignore_ttl: bool = False,
+        upper_batch_boundary: t.Optional[UpperBatchBoundary] = None,
     ) -> None:
-        current_ts = current_ts or now_timestamp()
-        for expired_snapshot_ids, cleanup_targets in self.snapshot_state._get_expired_snapshots(
-            self.environment_state.get_environments(), ignore_ttl=ignore_ttl, current_ts=current_ts
-        ):
-            self.snapshot_state.delete_snapshots(expired_snapshot_ids)
-            self.interval_state.cleanup_intervals(cleanup_targets, expired_snapshot_ids)
+        upper_batch_boundary = upper_batch_boundary or UpperBatchBoundary.include_all_boundary()
+        batch = self.get_expired_snapshots(
+            ignore_ttl=ignore_ttl,
+            batch_boundary=upper_batch_boundary,
+        )
+        if batch and batch.expired_snapshot_ids:
+            self.snapshot_state.delete_snapshots(batch.expired_snapshot_ids)
+            self.interval_state.cleanup_intervals(batch.cleanup_tasks, batch.expired_snapshot_ids)
 
     @transactional()
     def delete_expired_environments(