use range instead of single boundary

Author: eakmanrq

sqlmesh/core/state_sync/base.py

@@ -11,7 +11,6 @@
 from sqlmesh import migrations
 from sqlmesh.core.environment import (
     Environment,
-    EnvironmentNamingInfo,
     EnvironmentStatements,
     EnvironmentSummary,
 )
@@ -21,17 +20,20 @@
     SnapshotIdLike,
     SnapshotIdAndVersionLike,
     SnapshotInfoLike,
-    SnapshotTableCleanupTask,
-    SnapshotTableInfo,
     SnapshotNameVersion,
     SnapshotIdAndVersion,
 )
 from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 from sqlmesh.utils import major_minor
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import SQLMeshError
-from sqlmesh.utils.pydantic import PydanticModel, ValidationInfo, field_validator
-from sqlmesh.core.state_sync.common import StateStream
+from sqlmesh.utils.pydantic import PydanticModel, field_validator
+from sqlmesh.core.state_sync.common import (
+    StateStream,
+    ExpiredSnapshotBatch,
+    PromotionResult,
+    ExpiredBatchRange,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -72,64 +74,6 @@ 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]
-    removed_environment_naming_info: t.Optional[EnvironmentNamingInfo]
-
-    @field_validator("removed_environment_naming_info")
-    def _validate_removed_environment_naming_info(
-        cls, v: t.Optional[EnvironmentNamingInfo], info: ValidationInfo
-    ) -> t.Optional[EnvironmentNamingInfo]:
-        if v and not info.data.get("removed"):
-            raise ValueError("removed_environment_naming_info must be None if removed is empty")
-        return v
-
-
 class StateReader(abc.ABC):
     """Abstract base class for read-only operations on snapshot and environment state."""
 
@@ -361,7 +305,7 @@ def export(self, environment_names: t.Optional[t.List[str]] = None) -> StateStre
     def get_expired_snapshots(
         self,
         *,
-        batch_boundary: BatchBoundary,
+        batch_range: ExpiredBatchRange,
         current_ts: t.Optional[int] = None,
         ignore_ttl: bool = False,
     ) -> t.Optional[ExpiredSnapshotBatch]:
@@ -370,9 +314,7 @@ def get_expired_snapshots(
         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).
+            batch_range: The range of the batch to fetch.
 
         Returns:
             A batch describing expired snapshots or None if no snapshots are pending cleanup.
@@ -418,7 +360,7 @@ def delete_expired_snapshots(
         self,
         ignore_ttl: bool = False,
         current_ts: t.Optional[int] = None,
-        upper_batch_boundary: t.Optional[UpperBatchBoundary] = None,
+        batch_range: t.Optional[ExpiredBatchRange] = None,
     ) -> None:
         """Removes expired snapshots.
 
@@ -429,8 +371,7 @@ def delete_expired_snapshots(
             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
             current_ts: Timestamp used to evaluate expiration.
-            upper_batch_boundary: The upper boundary to delete expired snapshots till (inclusive). If not provided,
-                deletes all expired snapshots.
+            batch_range: The range of snapshots to delete in this batch. If None, all expired snapshots are deleted.
         """
 
     @abc.abstractmethod

sqlmesh/core/state_sync/cache.py

@@ -11,7 +11,8 @@
     SnapshotInfoLike,
 )
 from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
-from sqlmesh.core.state_sync.base import DelegatingStateSync, StateSync, UpperBatchBoundary
+from sqlmesh.core.state_sync.base import DelegatingStateSync, StateSync
+from sqlmesh.core.state_sync.common import ExpiredBatchRange
 from sqlmesh.utils.date import TimeLike, now_timestamp
 
 
@@ -111,11 +112,11 @@ def delete_expired_snapshots(
         self,
         ignore_ttl: bool = False,
         current_ts: t.Optional[int] = None,
-        upper_batch_boundary: t.Optional[UpperBatchBoundary] = None,
+        batch_range: t.Optional[ExpiredBatchRange] = None,
     ) -> None:
         self.snapshot_cache.clear()
         self.state_sync.delete_expired_snapshots(
-            upper_batch_boundary=upper_batch_boundary,
+            batch_range=batch_range,
             ignore_ttl=ignore_ttl,
             current_ts=current_ts,
         )

sqlmesh/core/state_sync/common.py

@@ -7,18 +7,26 @@
 import abc
 
 from dataclasses import dataclass
+
+from pydantic_core.core_schema import ValidationInfo
 from sqlglot import exp
 
 from sqlmesh.core.console import Console
 from sqlmesh.core.dialect import schema_
-from sqlmesh.utils.pydantic import PydanticModel
-from sqlmesh.core.environment import Environment, EnvironmentStatements
+from sqlmesh.utils.pydantic import PydanticModel, field_validator
+from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentNamingInfo
 from sqlmesh.utils.errors import SQLMeshError
-from sqlmesh.core.snapshot import Snapshot, SnapshotEvaluator
+from sqlmesh.core.snapshot import (
+    Snapshot,
+    SnapshotEvaluator,
+    SnapshotId,
+    SnapshotTableCleanupTask,
+    SnapshotTableInfo,
+)
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter.base import EngineAdapter
-    from sqlmesh.core.state_sync.base import Versions, ExpiredSnapshotBatch, StateReader, StateSync
+    from sqlmesh.core.state_sync.base import Versions, StateReader, StateSync
 
 logger = logging.getLogger(__name__)
 
@@ -219,6 +227,109 @@ def __iter__(self) -> t.Iterator[StateStreamContents]:
         return _StateStream()
 
 
+class ExpiredBatchRange(PydanticModel):
+    start: RowBoundary
+    end: t.Union[RowBoundary, LimitBoundary]
+
+    @classmethod
+    def init_batch_range(cls, batch_size: int) -> ExpiredBatchRange:
+        return ExpiredBatchRange(
+            start=RowBoundary.lowest_boundary(),
+            end=LimitBoundary(batch_size=batch_size),
+        )
+
+    @classmethod
+    def all_batch_range(cls) -> ExpiredBatchRange:
+        return ExpiredBatchRange(
+            start=RowBoundary.lowest_boundary(),
+            end=RowBoundary.highest_boundary(),
+        )
+
+
+class RowBoundary(PydanticModel):
+    updated_ts: int
+    name: str
+    identifier: str
+
+    @classmethod
+    def lowest_boundary(cls) -> RowBoundary:
+        return RowBoundary(updated_ts=0, name="", identifier="")
+
+    @classmethod
+    def highest_boundary(cls) -> RowBoundary:
+        # 9999-12-31T23:59:59.999Z in epoch milliseconds
+        return RowBoundary(updated_ts=253_402_300_799_999, name="", identifier="")
+
+    # def to_upper_batch_boundary(self) -> EndRowBoundary:
+    #     return EndRowBoundary(
+    #         updated_ts=self.updated_ts,
+    #         name=self.name,
+    #         identifier=self.identifier,
+    #     )
+    #
+    # def to_lower_batch_boundary(self, batch_size: int) -> StartRowBoundary:
+    #     return StartRowBoundary(
+    #         updated_ts=self.updated_ts,
+    #         name=self.name,
+    #         identifier=self.identifier,
+    #         batch_size=batch_size,
+    #     )
+
+
+#
+#
+# class EndRowBoundary(RowBoundary):
+#     @classmethod
+#     def include_all_boundary(cls) -> EndRowBoundary:
+#         # 9999-12-31T23:59:59.999Z in epoch milliseconds
+#         return EndRowBoundary(updated_ts=253_402_300_799_999, name="", identifier="")
+#
+#     def to_start_batch_boundary(self) -> StartRowBoundary:
+#         return StartRowBoundary(
+#             updated_ts=self.updated_ts,
+#             name=self.name,
+#             identifier=self.identifier,
+#         )
+
+
+class LimitBoundary(PydanticModel):
+    batch_size: int
+
+    @classmethod
+    def init_batch_boundary(cls, batch_size: int) -> LimitBoundary:
+        return LimitBoundary(batch_size=batch_size)
+
+
+#
+# class StartRowBoundary(RowBoundary):
+#     @classmethod
+#     def init_batch_boundary(cls) -> StartRowBoundary:
+#         return StartRowBoundary(updated_ts=0, name="", identifier="")
+#
+
+
+class PromotionResult(PydanticModel):
+    added: t.List[SnapshotTableInfo]
+    removed: t.List[SnapshotTableInfo]
+    removed_environment_naming_info: t.Optional[EnvironmentNamingInfo]
+
+    @field_validator("removed_environment_naming_info")
+    def _validate_removed_environment_naming_info(
+        cls, v: t.Optional[EnvironmentNamingInfo], info: ValidationInfo
+    ) -> t.Optional[EnvironmentNamingInfo]:
+        if v and not info.data.get("removed"):
+            raise ValueError("removed_environment_naming_info must be None if removed is empty")
+        return v
+
+
+class ExpiredSnapshotBatch(PydanticModel):
+    """A batch of expired snapshots to be cleaned up."""
+
+    expired_snapshot_ids: t.Set[SnapshotId]
+    cleanup_tasks: t.List[SnapshotTableCleanupTask]
+    batch_range: ExpiredBatchRange
+
+
 def iter_expired_snapshot_batches(
     state_reader: StateReader,
     *,
@@ -234,24 +345,29 @@ def iter_expired_snapshot_batches(
         ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
         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)
+    batch_range = ExpiredBatchRange.init_batch_range(batch_size=batch_size)
 
     while True:
         batch = state_reader.get_expired_snapshots(
             current_ts=current_ts,
             ignore_ttl=ignore_ttl,
-            batch_boundary=batch_boundary,
+            batch_range=batch_range,
         )
 
         if batch is None:
             return
 
         yield batch
 
-        batch_boundary = batch.batch_boundary.to_lower_batch_boundary(batch_size=batch_size)
+        assert isinstance(batch.batch_range.end, RowBoundary), (
+            "Only RowBoundary is supported for pagination currently"
+        )
+        batch_range = ExpiredBatchRange(
+            start=batch.batch_range.end,
+            end=LimitBoundary(batch_size=batch_size),
+        )
 
 
 def delete_expired_snapshots(
@@ -286,17 +402,25 @@ def delete_expired_snapshots(
         ignore_ttl=ignore_ttl,
         batch_size=batch_size,
     ):
+        end_info = (
+            f"updated_ts={batch.batch_range.end.updated_ts}"
+            if isinstance(batch.batch_range.end, RowBoundary)
+            else f"limit={batch.batch_range.end.batch_size}"
+        )
         logger.info(
-            "Processing batch of size %s and max_updated_ts of %s",
+            "Processing batch of size %s with end %s",
             len(batch.expired_snapshot_ids),
-            batch.batch_boundary.updated_ts,
+            end_info,
         )
         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(),
+            batch_range=ExpiredBatchRange(
+                start=RowBoundary.lowest_boundary(),
+                end=batch.batch_range.end,
+            ),
             ignore_ttl=ignore_ttl,
         )
         logger.info("Cleaned up expired snapshots batch")