sqlmesh.core.state_sync.common

  1from __future__ import annotations
  2
  3import logging
  4import typing as t
  5from functools import wraps
  6import itertools
  7import abc
  8
  9from dataclasses import dataclass
 10
 11from pydantic_core.core_schema import ValidationInfo
 12from sqlglot import exp
 13
 14from sqlmesh.utils.pydantic import PydanticModel, field_validator
 15from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentNamingInfo
 16from sqlmesh.core.snapshot import (
 17    Snapshot,
 18    SnapshotId,
 19    SnapshotTableCleanupTask,
 20    SnapshotTableInfo,
 21)
 22
 23if t.TYPE_CHECKING:
 24    from sqlmesh.core.state_sync.base import Versions, StateReader
 25
 26logger = logging.getLogger(__name__)
 27
 28EXPIRED_SNAPSHOT_DEFAULT_BATCH_SIZE = 200
 29
 30
 31def transactional() -> t.Callable[[t.Callable], t.Callable]:
 32    def decorator(func: t.Callable) -> t.Callable:
 33        @wraps(func)
 34        def wrapper(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
 35            if not hasattr(self, "_transaction"):
 36                return func(self, *args, **kwargs)
 37
 38            with self._transaction():
 39                return func(self, *args, **kwargs)
 40
 41        return wrapper
 42
 43    return decorator
 44
 45
 46T = t.TypeVar("T")
 47
 48
 49def chunk_iterable(iterable: t.Iterable[T], size: int = 10) -> t.Iterable[t.Iterable[T]]:
 50    iterator = iter(iterable)
 51    for first in iterator:
 52        yield itertools.chain([first], itertools.islice(iterator, size - 1))
 53
 54
 55class EnvironmentWithStatements(PydanticModel):
 56    environment: Environment
 57    statements: t.List[EnvironmentStatements] = []
 58
 59
 60@dataclass
 61class VersionsChunk:
 62    versions: Versions
 63
 64
 65class SnapshotsChunk:
 66    def __init__(self, items: t.Iterator[Snapshot]):
 67        self.items = items
 68
 69    def __iter__(self) -> t.Iterator[Snapshot]:
 70        return self.items
 71
 72
 73class EnvironmentsChunk:
 74    def __init__(self, items: t.Iterator[EnvironmentWithStatements]):
 75        self.items = items
 76
 77    def __iter__(self) -> t.Iterator[EnvironmentWithStatements]:
 78        return self.items
 79
 80
 81StateStreamContents = t.Union[VersionsChunk, SnapshotsChunk, EnvironmentsChunk]
 82
 83
 84class StateStream(abc.ABC):
 85    """
 86    Represents a stream of state either going into the StateSync (perhaps loaded from a file)
 87    or out of the StateSync (perhaps being dumped to a file)
 88
 89    Iterating over the stream produces the following chunks:
 90
 91        VersionsChunk: The versions of the objects contained in this StateStream
 92        SnapshotsChunk: Is itself an iterator that streams Snapshot objects. Note that they should be fully populated with any relevant Intervals
 93        EnvironmentsChunk: Is itself an iterator emitting a stream of Environments with any EnvironmentStatements attached
 94
 95    The idea here is to give some structure to the stream and ensure that callers have the opportunity to process all its components while not
 96    needing to worry about the order they are emitted in
 97    """
 98
 99    @abc.abstractmethod
100    def __iter__(self) -> t.Iterator[StateStreamContents]:
101        pass
102
103    @classmethod
104    def from_iterators(
105        cls: t.Type["StateStream"],
106        versions: Versions,
107        snapshots: t.Iterator[Snapshot],
108        environments: t.Iterator[EnvironmentWithStatements],
109    ) -> "StateStream":
110        class _StateStream(cls):  # type: ignore
111            def __iter__(self) -> t.Iterator[StateStreamContents]:
112                yield VersionsChunk(versions)
113
114                yield SnapshotsChunk(snapshots)
115
116                yield EnvironmentsChunk(environments)
117
118        return _StateStream()
119
120
121class ExpiredBatchRange(PydanticModel):
122    start: RowBoundary
123    end: t.Union[RowBoundary, LimitBoundary]
124
125    @classmethod
126    def init_batch_range(cls, batch_size: int) -> ExpiredBatchRange:
127        return ExpiredBatchRange(
128            start=RowBoundary.lowest_boundary(),
129            end=LimitBoundary(batch_size=batch_size),
130        )
131
132    @classmethod
133    def all_batch_range(cls) -> ExpiredBatchRange:
134        return ExpiredBatchRange(
135            start=RowBoundary.lowest_boundary(),
136            end=RowBoundary.highest_boundary(),
137        )
138
139    @classmethod
140    def _expanded_tuple_comparison(
141        cls,
142        columns: t.List[exp.Column],
143        values: t.List[t.Union[exp.Literal, exp.Neg]],
144        operator: t.Type[exp.Expr],
145    ) -> exp.Condition:
146        """Generate expanded tuple comparison that works across all SQL engines.
147
148        Converts tuple comparisons like (a, b, c) OP (x, y, z) into an expanded form
149        that's compatible with all SQL engines, since native tuple comparisons have
150        inconsistent support across engines (especially DuckDB, MySQL, SQLite).
151
152        Repro of problem with DuckDB:
153            "SELECT * FROM VALUES(1,'2') as test(a,b) WHERE ((a, b) > (1, 'foo')) AND ((a, b) <= (10, 'baz'))"
154
155        Args:
156            columns: List of column expressions to compare
157            values: List of value expressions to compare against
158            operator: The comparison operator class (exp.GT, exp.GTE, exp.LT, exp.LTE)
159
160        Examples:
161            (a, b, c) > (x, y, z) expands to:
162                a > x OR (a = x AND b > y) OR (a = x AND b = y AND c > z)
163
164            (a, b, c) <= (x, y, z) expands to:
165                a < x OR (a = x AND b < y) OR (a = x AND b = y AND c <= z)
166
167            (a, b, c) >= (x, y, z) expands to:
168                a > x OR (a = x AND b > y) OR (a = x AND b = y AND c >= z)
169
170        Returns:
171            An expanded OR expression representing the tuple comparison
172        """
173        if operator not in (exp.GT, exp.GTE, exp.LT, exp.LTE):
174            raise ValueError(f"Unsupported operator: {operator}. Use GT, GTE, LT, or LTE.")
175
176        # For <= and >=, we use the strict operator for all but the last column
177        # e.g., (a, b) <= (x, y) becomes: a < x OR (a = x AND b <= y)
178        # For < and >, we use the strict operator throughout
179        # e.g., (a, b) > (x, y) becomes: a > x OR (a = x AND b > x)
180        strict_operator: t.Type[exp.Expr]
181        final_operator: t.Type[exp.Expr]
182
183        if operator in (exp.LTE, exp.GTE):
184            # For inclusive operators (<=, >=), use strict form for intermediate columns
185            # but keep inclusive form for the last column
186            strict_operator = exp.LT if operator == exp.LTE else exp.GT
187            final_operator = operator  # Keep LTE/GTE for last column
188        else:
189            # For strict operators (<, >), use them throughout
190            strict_operator = operator
191            final_operator = operator
192
193        conditions: t.List[exp.Expr] = []
194        for i in range(len(columns)):
195            # Build equality conditions for all columns before current
196            equality_conditions = [exp.EQ(this=columns[j], expression=values[j]) for j in range(i)]
197
198            # Use the final operator for the last column, strict for others
199            comparison_op = final_operator if i == len(columns) - 1 else strict_operator
200            comparison_condition = comparison_op(this=columns[i], expression=values[i])
201
202            if equality_conditions:
203                conditions.append(exp.and_(*equality_conditions, comparison_condition))
204            else:
205                conditions.append(comparison_condition)
206
207        return exp.or_(*conditions) if len(conditions) > 1 else t.cast(exp.Condition, conditions[0])
208
209    @property
210    def where_filter(self) -> exp.Condition:
211        # Use expanded tuple comparisons for cross-engine compatibility
212        # Native tuple comparisons like (a, b) > (x, y) don't work reliably across all SQL engines
213        columns = [
214            exp.column("updated_ts"),
215            exp.column("name"),
216            exp.column("identifier"),
217        ]
218        start_values = [
219            exp.Literal.number(self.start.updated_ts),
220            exp.Literal.string(self.start.name),
221            exp.Literal.string(self.start.identifier),
222        ]
223
224        start_condition = self._expanded_tuple_comparison(columns, start_values, exp.GT)
225
226        range_filter: exp.Condition
227        if isinstance(self.end, RowBoundary):
228            end_values = [
229                exp.Literal.number(self.end.updated_ts),
230                exp.Literal.string(self.end.name),
231                exp.Literal.string(self.end.identifier),
232            ]
233            end_condition = self._expanded_tuple_comparison(columns, end_values, exp.LTE)
234            range_filter = exp.and_(start_condition, end_condition)
235        else:
236            range_filter = start_condition
237        return range_filter
238
239
240class RowBoundary(PydanticModel):
241    updated_ts: int
242    name: str
243    identifier: str
244
245    @classmethod
246    def lowest_boundary(cls) -> RowBoundary:
247        return RowBoundary(updated_ts=0, name="", identifier="")
248
249    @classmethod
250    def highest_boundary(cls) -> RowBoundary:
251        # 9999-12-31T23:59:59.999Z in epoch milliseconds
252        return RowBoundary(updated_ts=253_402_300_799_999, name="", identifier="")
253
254
255class LimitBoundary(PydanticModel):
256    batch_size: int
257
258    @classmethod
259    def init_batch_boundary(cls, batch_size: int) -> LimitBoundary:
260        return LimitBoundary(batch_size=batch_size)
261
262
263class PromotionResult(PydanticModel):
264    added: t.List[SnapshotTableInfo]
265    removed: t.List[SnapshotTableInfo]
266    removed_environment_naming_info: t.Optional[EnvironmentNamingInfo]
267
268    @field_validator("removed_environment_naming_info")
269    def _validate_removed_environment_naming_info(
270        cls, v: t.Optional[EnvironmentNamingInfo], info: ValidationInfo
271    ) -> t.Optional[EnvironmentNamingInfo]:
272        if v and not info.data.get("removed"):
273            raise ValueError("removed_environment_naming_info must be None if removed is empty")
274        return v
275
276
277class ExpiredSnapshotBatch(PydanticModel):
278    """A batch of expired snapshots to be cleaned up."""
279
280    expired_snapshot_ids: t.Set[SnapshotId]
281    cleanup_tasks: t.List[SnapshotTableCleanupTask]
282    batch_range: ExpiredBatchRange
283
284
285def iter_expired_snapshot_batches(
286    state_reader: StateReader,
287    *,
288    current_ts: int,
289    ignore_ttl: bool = False,
290    batch_size: t.Optional[int] = None,
291) -> t.Iterator[ExpiredSnapshotBatch]:
292    """Yields expired snapshot batches.
293
294    Args:
295        state_reader: StateReader instance to query expired snapshots from.
296        current_ts: Timestamp used to evaluate expiration.
297        ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
298        batch_size: Maximum number of snapshots to fetch per batch.
299    """
300
301    batch_size = batch_size if batch_size is not None else EXPIRED_SNAPSHOT_DEFAULT_BATCH_SIZE
302    batch_range = ExpiredBatchRange.init_batch_range(batch_size=batch_size)
303
304    while True:
305        batch = state_reader.get_expired_snapshots(
306            current_ts=current_ts,
307            ignore_ttl=ignore_ttl,
308            batch_range=batch_range,
309        )
310
311        if batch is None:
312            return
313
314        yield batch
315
316        assert isinstance(batch.batch_range.end, RowBoundary), (
317            "Only RowBoundary is supported for pagination currently"
318        )
319        batch_range = ExpiredBatchRange(
320            start=batch.batch_range.end,
321            end=LimitBoundary(batch_size=batch_size),
322        )