diff --git a/misc/subscribe-sdk/python/.gitignore b/misc/subscribe-sdk/python/.gitignore new file mode 100644 index 0000000000000..14fa7b6820abb --- /dev/null +++ b/misc/subscribe-sdk/python/.gitignore @@ -0,0 +1,8 @@ +# Python caches, builds, and virtualenvs. +__pycache__/ +*.pyc +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +*.egg-info/ +.venv/ diff --git a/misc/subscribe-sdk/python/README.md b/misc/subscribe-sdk/python/README.md new file mode 100644 index 0000000000000..165cef7ae9891 --- /dev/null +++ b/misc/subscribe-sdk/python/README.md @@ -0,0 +1,66 @@ +# materialize-subscribe + +A correct-by-construction Python client for consuming Materialize +[`SUBSCRIBE`](https://materialize.com/docs/sql/subscribe/). + +```python +from materialize_subscribe import Subscribe, SubscribeClient + +with SubscribeClient.connect("postgres://materialize@localhost:6875") as client: + stream = client.subscribe( + Subscribe.object("winning_bids").envelope_upsert(["id"]) + ) + for batch in stream: + for change in batch.updates: + ... # apply `change` to your sink + # Persist `batch.resume_token.encode()` atomically with the applied + # changes for exactly-once state; hand it back to `client.resume`. +``` + +* The unit of consumption is a `ConsistentBatch`, never a bare row. +* Resuming takes an opaque `ResumeToken`; the `AS OF frontier - 1` arithmetic + lives inside it. +* Failures surface as a small, typed `SubscribeError` hierarchy. +* For a lower-level stream of decoded changes and progress markers, use + `subscribe_raw`. The batcher is built on it. + +## Cohorts + +To stay consistent across several views at once, a `Cohort` releases every +member only up to their shared minimum frontier, so each moment is a genuine +cross-view snapshot. + +```python +from materialize_subscribe import Cohort, Subscribe + +with Cohort.connect( + "postgres://materialize@localhost:6875", + [ + ("orders", Subscribe.object("orders")), + ("inventory", Subscribe.object("inventory")), + ], +) as cohort: + for moment in cohort: + for view in moment.views: + ... # view.name, view.updates at the joint frontier + # moment.resume_token checkpoints the whole cohort at once. +``` + +The protocol core has no runtime dependencies. Install the `client` extra to +connect: + +``` +pip install 'materialize-subscribe[client]' +``` + +See the [top-level README](../README.md) for the protocol this encodes and the +guarantees it provides. + +## Development + +``` +pip install -e '.[dev,client]' +pytest +mypy materialize_subscribe +ruff check . +``` diff --git a/misc/subscribe-sdk/python/materialize_subscribe/__init__.py b/misc/subscribe-sdk/python/materialize_subscribe/__init__.py new file mode 100644 index 0000000000000..6a2d84cae4ee7 --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/__init__.py @@ -0,0 +1,118 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A correct-by-construction client for consuming Materialize ``SUBSCRIBE``. + +Consuming ``SUBSCRIBE`` durably requires a specific protocol: buffer until a +progress message closes a timestamp, apply the closed batch and persist its +frontier atomically, and resume with ``SNAPSHOT = false AS OF frontier - 1``. +Every step has a failure mode that is invisible in testing. This package +encodes the protocol so callers cannot get it wrong: + +* The unit of consumption is a :class:`ConsistentBatch`, never a bare row. +* Resuming takes an opaque :class:`ResumeToken`; the ``- 1`` boundary + arithmetic lives inside it. +* Failures surface as a small, typed :class:`SubscribeError` hierarchy. + +The client is three layers, each usable on its own: :class:`RawStream` (via +:meth:`SubscribeClient.subscribe_raw`) is the decoded timestamped substrate; the +consistency engine buffers and releases below a frontier; a +:class:`ConsistentBatch` (one view) or :class:`CohortMoment` (a :class:`Cohort` +of N views) is the closed unit. A single view is the cohort of one. + +Example:: + + from materialize_subscribe import Subscribe, SubscribeClient + + with SubscribeClient.connect("postgres://materialize@localhost:6875") as client: + stream = client.subscribe( + Subscribe.object("winning_bids").envelope_upsert(["id"]) + ) + for batch in stream: + for change in batch.updates: + ... # apply `change` to your sink + # Persist `batch.resume_token.encode()` atomically with the applied + # changes for exactly-once state; hand it back to `client.resume`. +""" + +from __future__ import annotations + +from .batch import Batcher, ConsistentBatch +from .client import BatchStream, RawStream, SubscribeClient +from .cohort import Cohort, CohortMoment, ViewChanges +from .envelope import ( + Change, + Data, + Decoder, + Delete, + Diff, + DiffEnvelope, + Envelope, + KeyViolation, + Progress, + StreamMessage, + Upsert, + UpsertEnvelope, +) +from .errors import ( + BufferOverflow, + CohortLagExceeded, + CompactionHorizon, + DependencyDropped, + FatalError, + InvalidToken, + ProtocolError, + SchemaMismatch, + SubscribeError, + TransientError, +) +from .statement import Subscribe +from .token import CohortToken, ResumeToken + +__all__ = [ + "Batcher", + "BatchStream", + "BufferOverflow", + "Change", + "Cohort", + "CohortLagExceeded", + "CohortMoment", + "CohortToken", + "CompactionHorizon", + "ConsistentBatch", + "Data", + "Decoder", + "Delete", + "DependencyDropped", + "Diff", + "DiffEnvelope", + "Envelope", + "FatalError", + "InvalidToken", + "KeyViolation", + "Progress", + "ProtocolError", + "RawStream", + "ResumeToken", + "SchemaMismatch", + "StreamMessage", + "Subscribe", + "SubscribeClient", + "SubscribeError", + "TransientError", + "Upsert", + "UpsertEnvelope", + "ViewChanges", +] diff --git a/misc/subscribe-sdk/python/materialize_subscribe/batch.py b/misc/subscribe-sdk/python/materialize_subscribe/batch.py new file mode 100644 index 0000000000000..f30c7d790864d --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/batch.py @@ -0,0 +1,140 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The progress-driven batcher: the state machine that turns a stream of +changes and progress markers into *consistent, closed* batches. + +This is the part every hand-rolled subscribe consumer gets wrong. The rule it +enforces: + +* Buffer changes as they arrive. +* When a progress marker advances the frontier to ``F``, every change with a + timestamp strictly below ``F`` is now final. Emit exactly those as one batch, + tagged with frontier ``F`` and a resume token for ``F``. +* Never emit a change before its timestamp is closed, and never emit the same + timestamp twice. + +A consumer that applies each emitted batch and persists its token atomically +gets exactly-once *state*: after any crash, resuming from the last token neither +drops nor duplicates data. + +A single view is the one-member case of the cohort engine (see +:mod:`materialize_subscribe.cohort`): the release frontier is the view's own +progress frontier. Both share the :class:`ReleaseBuffer` buffer-and-consolidate +core. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional + +from .engine import ReleaseBuffer +from .envelope import Change, Data, Envelope, Progress, StreamMessage +from .errors import ProtocolError +from .token import ResumeToken + + +@dataclass(frozen=True) +class ConsistentBatch: + """A batch of changes for a contiguous, now-closed range of timestamps. + + Every change in ``updates`` has a timestamp strictly below ``frontier``, and + no future batch will contain a timestamp below ``frontier``. ``resume_token`` + checkpoints exactly this position. ``updates`` is consolidated to its net + effect within the batch, ordered by first appearance. + """ + + updates: List[Change] + frontier: int + resume_token: ResumeToken + is_snapshot: bool + + def is_empty(self) -> bool: + """Whether the batch carries no changes. + + Empty batches still advance the frontier and carry a fresh token, which + lets a consumer keep its checkpoint moving during idle periods so it + does not age out of the source's retained history. + """ + return not self.updates + + +class Batcher: + """Accumulates stream messages and emits :class:`ConsistentBatch` values as + timestamps close. + + Feed every decoded message to :meth:`push`. It returns a batch when a + progress marker closes one or more timestamps, and ``None`` otherwise. + """ + + def __init__( + self, fingerprint: str, envelope: Envelope, with_snapshot: bool + ) -> None: + self._fingerprint = fingerprint + self._with_snapshot = with_snapshot + self._buffer = ReleaseBuffer(envelope) + self._last_frontier: Optional[int] = None + self._first_frontier: Optional[int] = None + self._snapshot_emitted = False + + def push(self, message: StreamMessage) -> Optional[ConsistentBatch]: + """Feeds one decoded message. Returns a batch when the frontier + advances.""" + if isinstance(message, Data): + self._buffer.push_data(message.timestamp, message.change) + return None + + if isinstance(message, Progress): + return self._advance(message.frontier) + + raise ProtocolError(f"unexpected stream message: {message!r}") + + def _advance(self, frontier: int) -> Optional[ConsistentBatch]: + if frontier == self._last_frontier: + # No advance: nothing new is closed. A regression, by contrast, is + # caught by `observe_progress` below. + return None + self._buffer.observe_progress(frontier) + + if self._first_frontier is None: + self._first_frontier = frontier + + updates = self._buffer.release_below(frontier) + + # The snapshot is the data closed once we advance past the first frontier + # (the effective AS OF). Mark the first such batch. + # + # NOTE: this relies on the server emitting a leading Progress(as_of) + # before any snapshot rows (it does, unconditionally, when PROGRESS is + # set). That progress sets first_frontier; the snapshot rows sit at + # as_of and are closed by the next, higher progress. A transport that + # delivered snapshot data before any progress would mistag it. + is_snapshot = ( + self._with_snapshot + and not self._snapshot_emitted + and frontier != self._first_frontier + ) + if is_snapshot: + self._snapshot_emitted = True + + self._last_frontier = frontier + + return ConsistentBatch( + updates=updates, + frontier=frontier, + resume_token=ResumeToken(frontier=frontier, fingerprint=self._fingerprint), + is_snapshot=is_snapshot, + ) diff --git a/misc/subscribe-sdk/python/materialize_subscribe/client.py b/misc/subscribe-sdk/python/materialize_subscribe/client.py new file mode 100644 index 0000000000000..f9418caa44dcd --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/client.py @@ -0,0 +1,474 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The pgwire transport: connect to Materialize and drive a subscription. + +The transport is deliberately thin. It runs the ``DECLARE``/``FETCH`` loop and +hands every row to the tested :class:`Decoder`. All protocol judgement lives in +the decoder and the consistency engine, not here. + +A background thread drains the cursor continuously into a *bounded* buffer, +rather than fetching only when the consumer asks. This is deliberate: +Materialize buffers a subscription's unread output in ``environmentd`` without +bound, so a consumer that stops fetching pushes an unbounded cost onto the +server. Draining continuously keeps that buffer on the client, where it is +bounded and fails loud (:class:`BufferOverflow`) if the consumer cannot keep up. +The client falls over, never the server. + +``psycopg`` (v3) is imported lazily so the protocol core can be used and tested +without the driver installed. +""" + +from __future__ import annotations + +import queue as _queue +import threading +from typing import Any, Callable, List, Optional + +from .batch import Batcher, ConsistentBatch +from .engine import ReleaseBuffer # noqa: F401 - re-exported for cohort convenience +from .envelope import Decoder, Envelope, StreamMessage +from .errors import ( + BufferOverflow, + CompactionHorizon, + DependencyDropped, + FatalError, + SchemaMismatch, + SubscribeError, + TransientError, +) +from .statement import Subscribe +from .token import ResumeToken + +# The cursor name for the subscription. One subscription owns its connection, so +# a fixed name is safe. +_CURSOR = "mz_subscribe_cursor" + +# Default client-side buffer capacity, in decoded messages. A slow consumer that +# falls this far behind trips :class:`BufferOverflow`. +# +# TODO: make this configurable per subscription once we have real workloads to +# size it against. +DEFAULT_BUFFER_CAPACITY = 1 << 16 + +# How long a consumer blocks between checks that its producers are still alive, +# and how long ``close`` waits for a drain thread to wind down. +_POLL_SECONDS = 0.5 +_JOIN_SECONDS = 2.0 + +#: A message wrapper: transforms a decoded message before it enters the buffer. +#: Single-view uses identity; a cohort tags each message with its member index. +Wrap = Callable[[StreamMessage], Any] + + +def _identity(message: StreamMessage) -> Any: + return message + + +class _Buffer: + """The bounded, thread-safe hand-off between drain threads and the consumer. + + One or more drain threads ``put`` decoded messages; the consumer ``recv``s + them. When the buffer fills, ``put`` reports the overflow instead of blocking + (blocking would backpressure the server). Termination is observed by every + producer thread having exited: the consumer then surfaces the first recorded + error, or a clean end. + """ + + def __init__(self, capacity: int) -> None: + self._queue: _queue.Queue[Any] = _queue.Queue(maxsize=max(1, capacity)) + self._stop = threading.Event() + self._lock = threading.Lock() + self._error: Optional[SubscribeError] = None + self._threads: List[threading.Thread] = [] + + def set_threads(self, threads: List[threading.Thread]) -> None: + self._threads = threads + + def should_stop(self) -> bool: + return self._stop.is_set() + + def request_stop(self) -> None: + """Signals every drain to stop, without waiting.""" + self._stop.set() + + def put(self, item: Any) -> bool: + """Enqueues ``item``. Returns ``False`` if the buffer is full.""" + try: + self._queue.put_nowait(item) + return True + except _queue.Full: + return False + + def record_error(self, error: SubscribeError) -> None: + """Records the first terminal error and signals every drain to stop.""" + with self._lock: + if self._error is None: + self._error = error + self._stop.set() + + def recv(self) -> Any: + """Returns the next buffered item. + + Raises the terminal :class:`SubscribeError` if a drain failed, or + ``StopIteration`` on a clean end (every producer exited and the buffer + drained). + """ + while True: + try: + return self._queue.get(timeout=_POLL_SECONDS) + except _queue.Empty: + if not any(t.is_alive() for t in self._threads): + with self._lock: + error = self._error + if error is not None: + raise error + raise StopIteration + + def close(self) -> None: + """Stops every drain and waits briefly for it to release its cursor.""" + self._stop.set() + for thread in self._threads: + thread.join(timeout=_JOIN_SECONDS) + + +class SubscribeClient: + """A connection to Materialize for consuming subscriptions. + + One connection backs one subscription: starting a subscription hands the + connection to a background drain thread. Transaction-mode connection poolers + (such as PgBouncer) break the ``DECLARE``/``FETCH`` cursor and are not + supported. + + Usable as a context manager, which closes the active subscription (and its + connection) on exit. + """ + + def __init__(self, connection: Any, buffer_capacity: int) -> None: + self._conn: Optional[Any] = connection + self._buffer_capacity = buffer_capacity + self._stream: Optional[Any] = None + + @classmethod + def connect( + cls, conninfo: str, buffer_capacity: int = DEFAULT_BUFFER_CAPACITY + ) -> "SubscribeClient": + """Connects to Materialize using a libpq-style connection string. + + ``buffer_capacity`` is the client-side buffer size in messages; a + consumer that falls that far behind trips :class:`BufferOverflow`. + """ + psycopg = _import_psycopg() + try: + connection = psycopg.connect(conninfo, autocommit=True) + except psycopg.Error as exc: + raise _classify(exc) from exc + return cls(connection, buffer_capacity) + + def subscribe(self, subscribe: Subscribe) -> "BatchStream": + """Starts a new subscription, taking the initial snapshot.""" + raw = self._open_raw(subscribe, subscribe.to_sql_initial()) + stream = BatchStream( + raw, + Batcher(subscribe.fingerprint(), subscribe.envelope, with_snapshot=True), + ) + self._stream = stream + return stream + + def resume(self, subscribe: Subscribe, token: ResumeToken) -> "BatchStream": + """Resumes a subscription from ``token``, skipping the snapshot. + + Raises :class:`SchemaMismatch` if ``subscribe`` does not match the query + the token was taken against. + """ + _check_fingerprint(subscribe, token) + raw = self._open_raw(subscribe, subscribe.to_sql_resume(token)) + stream = BatchStream( + raw, + Batcher(subscribe.fingerprint(), subscribe.envelope, with_snapshot=False), + ) + self._stream = stream + return stream + + def subscribe_raw(self, subscribe: Subscribe) -> "RawStream": + """Starts a subscription and hands back the *raw* decoded stream: the + timestamped changes and progress markers, before any batching. + + This is the composable substrate. Most callers want :meth:`subscribe`, + which layers consistent batching on top. Reach for the raw stream to + build a different consistency policy of your own. + """ + raw = self._open_raw(subscribe, subscribe.to_sql_initial()) + self._stream = raw + return raw + + def resume_raw(self, subscribe: Subscribe, token: ResumeToken) -> "RawStream": + """Resumes a raw stream from ``token``. See :meth:`subscribe_raw`.""" + _check_fingerprint(subscribe, token) + raw = self._open_raw(subscribe, subscribe.to_sql_resume(token)) + self._stream = raw + return raw + + def close(self) -> None: + """Closes the active subscription and its connection.""" + if self._stream is not None: + self._stream.close() + self._stream = None + elif self._conn is not None: + self._conn.close() + self._conn = None + + def __enter__(self) -> "SubscribeClient": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() + + def _open_raw(self, subscribe: Subscribe, subscribe_sql: str) -> "RawStream": + if self._conn is None: + raise SubscribeError( + "this client already owns a subscription; use one client per " + "subscription" + ) + # Ownership of the connection transfers to the drain thread. + conn = self._conn + self._conn = None + _declare_cursor(conn, subscribe_sql) + buffer = _spawn_drain( + conn, + subscribe.envelope, + subscribe.is_bounded(), + _identity, + self._buffer_capacity, + ) + return RawStream(buffer) + + +class RawStream: + """The raw decoded stream for one subscription: :data:`StreamMessage` values + in arrival order, before any batching or consistency policy. + + This is layer one, the composable substrate the batcher and cohort build on. + Iterate to pull messages. Usable as a context manager, which stops the drain + on exit. + """ + + def __init__(self, buffer: _Buffer) -> None: + self._buffer = buffer + + def __iter__(self) -> "RawStream": + return self + + def __next__(self) -> StreamMessage: + message: StreamMessage = self._buffer.recv() + return message + + def close(self) -> None: + self._buffer.close() + + def __enter__(self) -> "RawStream": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() + + +class BatchStream: + """An iterator over the :class:`ConsistentBatch` values of one subscription. + + Iterate to pull batches. Each is complete: a consumer that applies it and + persists its ``resume_token`` atomically achieves exactly-once state. A + bounded (``UP TO``) subscription ends iteration when it drains; an unbounded + one blocks until the next batch is ready. This is a raw stream plus the + :class:`Batcher` consistency engine. + """ + + def __init__(self, raw: RawStream, batcher: Batcher) -> None: + self._raw = raw + self._batcher = batcher + + def __iter__(self) -> "BatchStream": + return self + + def __next__(self) -> ConsistentBatch: + for message in self._raw: + batch = self._batcher.push(message) + if batch is not None: + return batch + raise StopIteration + + def close(self) -> None: + self._raw.close() + + def __enter__(self) -> "BatchStream": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() + + +def _check_fingerprint(subscribe: Subscribe, token: ResumeToken) -> None: + """Rejects a resume whose query shape no longer matches the checkpoint.""" + fingerprint = subscribe.fingerprint() + if fingerprint != token.fingerprint: + raise SchemaMismatch(expected=token.fingerprint, actual=fingerprint) + + +def _declare_cursor(conn: Any, subscribe_sql: str) -> None: + """Opens the subscription cursor inside a transaction on ``conn``. + + A cursor must live inside a transaction, which also holds the read so the + frontier does not advance out from under a slow reader mid-batch. On failure + this rolls back and closes the connection, since ownership was already taken + from the caller. + """ + cursor = conn.cursor() + try: + cursor.execute("BEGIN") + cursor.execute(f"DECLARE {_CURSOR} CURSOR FOR {subscribe_sql}") + except Exception as exc: # noqa: BLE001 - re-raised as a typed error + try: + cursor.execute("ROLLBACK") + except Exception: # noqa: BLE001 - best-effort cleanup + pass + try: + conn.close() + except Exception: # noqa: BLE001 - best-effort cleanup + pass + raise _classify(exc) from exc + + +def _spawn_drain( + conn: Any, envelope: Envelope, bounded: bool, wrap: Wrap, capacity: int +) -> _Buffer: + """Starts one drain thread over ``conn`` feeding a fresh bounded buffer.""" + buffer = _Buffer(capacity) + thread = threading.Thread( + target=_run_drain, + args=(conn, envelope, bounded, wrap, buffer), + daemon=True, + ) + buffer.set_threads([thread]) + thread.start() + return buffer + + +def _run_drain( + conn: Any, envelope: Envelope, bounded: bool, wrap: Wrap, buffer: _Buffer +) -> None: + """The background drain: fetch, decode, and enqueue until the subscription + ends, the consumer stops it, or the buffer overflows. + + Owns ``conn``: on exit it rolls back (releasing the cursor so the server + stops buffering for a dead reader) and closes the connection. + """ + fetch_sql = f"FETCH ALL {_CURSOR} WITH (timeout = '1s')" + cursor = conn.cursor() + decoder: Optional[Decoder] = None + try: + while not buffer.should_stop(): + try: + cursor.execute(fetch_sql) + rows: List[Any] = cursor.fetchall() + except Exception as exc: # noqa: BLE001 - re-raised as a typed error + buffer.record_error(_classify(exc)) + return + + if rows and decoder is None: + decoder = Decoder([col.name for col in cursor.description], envelope) + for row in rows: + assert decoder is not None + try: + message = decoder.decode(row) + except SubscribeError as exc: + buffer.record_error(exc) + return + # Fail loud rather than block the drain: a full buffer means the + # consumer fell behind, and blocking would backpressure the + # server into unbounded buffering. + if not buffer.put(wrap(message)): + buffer.record_error(BufferOverflow()) + return + + # A fetch that returned no rows on a bounded subscription means it + # has drained. An unbounded one just idled; loop and fetch again. + if not rows and bounded: + return + finally: + try: + cursor.execute("ROLLBACK") + except Exception: # noqa: BLE001 - best-effort cleanup + pass + try: + conn.close() + except Exception: # noqa: BLE001 - best-effort cleanup + pass + + +def _import_psycopg() -> Any: + try: + import psycopg + except ImportError as exc: # pragma: no cover - environment dependent + raise SubscribeError( + "the psycopg package is required to connect; install " + "materialize-subscribe with the [client] extra" + ) from exc + return psycopg + + +def is_compaction_horizon(message: str) -> bool: + """Whether a server error message denotes the compaction horizon. + + Matched by text because Materialize does not yet expose a dedicated SQLSTATE + for it. Both the current and the historical (pre-#34712) wording are handled + so the classification survives server upgrades. This is the *only* place the + SDK matches an error message. + """ + return ( + "could not find a valid timestamp" in message + or "not valid for all inputs" in message + ) + + +def classify_server_message(message: str) -> SubscribeError: + """Maps a server error message to a typed error. + + Pure, so it is unit-tested and kept in lockstep with the Rust SDK. + """ + if is_compaction_horizon(message): + return CompactionHorizon() + if "dropped" in message and "dependenc" in message: + return DependencyDropped(message) + return FatalError(message) + + +def _classify(exc: Exception) -> SubscribeError: + """Classifies a driver exception into the SDK's taxonomy. + + A server error carries a SQLSTATE; anything without one is a connection-level + failure and therefore retryable. This discriminator mirrors the Rust SDK, so + the two agree on what is retryable. + """ + if isinstance(exc, SubscribeError): + return exc + + diag = getattr(exc, "diag", None) + message = str((diag and diag.message_primary) or exc) + + if getattr(exc, "sqlstate", None) is not None: + return classify_server_message(message) + + # No SQLSTATE means a connection-level failure: retryable. + return TransientError(message) diff --git a/misc/subscribe-sdk/python/materialize_subscribe/cohort.py b/misc/subscribe-sdk/python/materialize_subscribe/cohort.py new file mode 100644 index 0000000000000..1119a743ad955 --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/cohort.py @@ -0,0 +1,302 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-view consistency: subscribe to several views at once and observe them +at one shared, consistent moment. + +Materialize gives every object a place on one logical timeline, so an +``mz_timestamp`` means the same instant across views. A cohort exploits that: it +runs one ``SUBSCRIBE`` per view (each on its own connection) and releases every +view only up to the *minimum* closed frontier across all of them. That minimum +is the latest instant final for every view at once, so the changes it hands back +form a genuine cross-view snapshot, never a mix of a newer view with a staler +one. + +This is the same engine as the single-view :class:`Batcher`, generalized. A +single view is the cohort of one: its minimum frontier is its only frontier. The +cost of the guarantee is that a laggard member holds the joint moment back, and +the leading members' changes buffer in memory until it catches up. + +Dynamic membership (adding, dropping, merging, or splitting members of a live +cohort) is out of scope here. A cohort's membership is fixed for its life. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Any, List, Optional, Tuple + +from .client import ( + _JOIN_SECONDS, + DEFAULT_BUFFER_CAPACITY, + Wrap, + _Buffer, + _classify, + _declare_cursor, + _import_psycopg, + _run_drain, +) +from .engine import ReleaseBuffer +from .envelope import Change, Data, Envelope +from .errors import CohortLagExceeded, ProtocolError, SchemaMismatch +from .statement import Subscribe +from .token import CohortToken + +# Default cohort lag budget: the most changes the cohort buffers across all +# members while waiting for the slowest to advance the joint frontier. Exceeding +# it raises :class:`CohortLagExceeded` rather than growing memory without bound. +# +# TODO: make this configurable per cohort once there are workloads to size it. +DEFAULT_COHORT_LAG_BUDGET = 1 << 20 + + +@dataclass(frozen=True) +class ViewChanges: + """The changes one view contributes to a :class:`CohortMoment`.""" + + name: str + updates: List[Change] + + +@dataclass(frozen=True) +class CohortMoment: + """A consistent moment across every view in a cohort. + + Every view's ``updates`` are closed at the same ``frontier``, so applying the + whole moment moves a downstream store from one cross-view-consistent state to + the next. Persist ``resume_token`` atomically with the effects for + exactly-once state across the cohort. + """ + + views: List[ViewChanges] + frontier: int + resume_token: CohortToken + is_snapshot: bool + + def is_empty(self) -> bool: + """Whether no view carries any change. Empty moments still advance the + frontier and carry a fresh token.""" + return all(not view.updates for view in self.views) + + +class _CohortMember: + """One member's buffered state within a cohort.""" + + def __init__(self, name: str, fingerprint: str, envelope: Envelope) -> None: + self.name = name + self.fingerprint = fingerprint + self.buffer = ReleaseBuffer(envelope) + self.first_frontier: Optional[int] = None + + +class CohortEngine: + """The pure state machine behind a cohort: N per-member buffers released at + their shared minimum frontier. + + Kept separate from the transport so it can be exhaustively unit-tested + without a server. + """ + + def __init__( + self, + members: List[Tuple[str, str, Envelope]], + with_snapshot: bool, + lag_budget: int = DEFAULT_COHORT_LAG_BUDGET, + ) -> None: + self._members = [ + _CohortMember(name, fingerprint, envelope) + for name, fingerprint, envelope in members + ] + self._with_snapshot = with_snapshot + self._last_release: Optional[int] = None + self._snapshot_complete = False + self._lag_budget = lag_budget + + def push_data(self, idx: int, timestamp: int, change: Change) -> None: + """Buffers a change from member ``idx``. + + Raises :class:`CohortLagExceeded` if the cohort is already holding + ``lag_budget`` changes. This is the laggard backstop: a member that + stalls pins the joint frontier, so its peers' changes would otherwise + buffer without bound. The client fails loud instead. + """ + buffered = sum(m.buffer.pending() for m in self._members) + if buffered >= self._lag_budget: + raise CohortLagExceeded(buffered=buffered, limit=self._lag_budget) + self._members[idx].buffer.push_data(timestamp, change) + + def push_progress(self, idx: int, frontier: int) -> Optional[CohortMoment]: + """Records progress from member ``idx``, and emits a joint moment if the + minimum frontier across all members advanced.""" + member = self._members[idx] + member.buffer.observe_progress(frontier) + if member.first_frontier is None: + member.first_frontier = frontier + + # The joint frontier is the minimum across members, defined only once + # every member has reported at least one progress. Until then a silent + # member could still emit at an earlier timestamp, so nothing is safe to + # release. + frontiers = [m.buffer.frontier for m in self._members] + if any(f is None for f in frontiers): + return None + release = min(f for f in frontiers if f is not None) + + # Emit only when the joint frontier actually advances. + if self._last_release is not None and release <= self._last_release: + return None + + # Every member has reported, so every first frontier is set. The joint + # snapshot is complete once the release frontier passes the last member's + # AS OF. + max_first = max( + m.first_frontier for m in self._members if m.first_frontier is not None + ) + + views = [ + ViewChanges(name=m.name, updates=m.buffer.release_below(release)) + for m in self._members + ] + + was_complete = self._snapshot_complete + if not was_complete and release > max_first: + self._snapshot_complete = True + is_snapshot = ( + self._with_snapshot and not was_complete and self._snapshot_complete + ) + + self._last_release = release + fingerprints = [m.fingerprint for m in self._members] + + return CohortMoment( + views=views, + frontier=release, + resume_token=CohortToken(frontier=release, members=fingerprints), + is_snapshot=is_snapshot, + ) + + +class Cohort: + """A live cohort of subscriptions, yielding consistent cross-view moments. + + Each member runs its own ``SUBSCRIBE`` on its own connection, so a cohort of + N views holds N connections. All members drain into one shared bounded + buffer, so if the consumer falls behind, the cohort fails loud with + :class:`BufferOverflow` rather than letting the server buffer. + + Iterate to pull :class:`CohortMoment` values. Usable as a context manager, + which stops every member on exit. + """ + + def __init__(self, buffer: _Buffer, engine: CohortEngine) -> None: + self._buffer = buffer + self._engine = engine + + @classmethod + def connect(cls, conninfo: str, members: List[Tuple[str, Subscribe]]) -> "Cohort": + """Starts a fresh cohort over ``members``, each a ``(name, subscription)`` + pair. All members connect with ``conninfo`` and take their initial + snapshot. Names identify a member's changes in each moment, and need not + match the object name.""" + return cls._start(conninfo, members, None) + + @classmethod + def resume( + cls, + conninfo: str, + members: List[Tuple[str, Subscribe]], + token: CohortToken, + ) -> "Cohort": + """Resumes a cohort from ``token``, reconstructing the exact joint cut. + + Raises :class:`SchemaMismatch` if the members no longer match the token + (different shapes or order). + """ + return cls._start(conninfo, members, token) + + @classmethod + def _start( + cls, + conninfo: str, + members: List[Tuple[str, Subscribe]], + token: Optional[CohortToken], + ) -> "Cohort": + if not members: + raise ProtocolError("a cohort needs at least one member") + if token is not None: + fingerprints = [sub.fingerprint() for _, sub in members] + if fingerprints != token.members: + raise SchemaMismatch( + expected=",".join(token.members), + actual=",".join(fingerprints), + ) + + psycopg = _import_psycopg() + buffer = _Buffer(DEFAULT_BUFFER_CAPACITY) + threads: List[threading.Thread] = [] + specs: List[Tuple[str, str, Envelope]] = [] + try: + for idx, (name, sub) in enumerate(members): + sql = ( + sub.to_sql_initial() + if token is None + else sub.to_sql_resume_at(token.as_of()) + ) + conn = psycopg.connect(conninfo, autocommit=True) + _declare_cursor(conn, sql) + thread = threading.Thread( + target=_run_drain, + args=(conn, sub.envelope, sub.is_bounded(), _tagger(idx), buffer), + daemon=True, + ) + thread.start() + threads.append(thread) + specs.append((name, sub.fingerprint(), sub.envelope)) + except Exception as exc: # noqa: BLE001 - re-raised as a typed error + buffer.request_stop() + for thread in threads: + thread.join(timeout=_JOIN_SECONDS) + raise _classify(exc) from exc + + buffer.set_threads(threads) + return cls(buffer, CohortEngine(specs, with_snapshot=token is None)) + + def __iter__(self) -> "Cohort": + return self + + def __next__(self) -> CohortMoment: + while True: + idx, message = self._buffer.recv() + if isinstance(message, Data): + self._engine.push_data(idx, message.timestamp, message.change) + else: + moment = self._engine.push_progress(idx, message.frontier) + if moment is not None: + return moment + + def close(self) -> None: + self._buffer.close() + + def __enter__(self) -> "Cohort": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() + + +def _tagger(index: int) -> Wrap: + """A wrapper that tags every message with its member index.""" + return lambda message: (index, message) diff --git a/misc/subscribe-sdk/python/materialize_subscribe/engine.py b/misc/subscribe-sdk/python/materialize_subscribe/engine.py new file mode 100644 index 0000000000000..ebebfe8ce5f11 --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/engine.py @@ -0,0 +1,165 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The consistency engine: buffer timestamped changes, then release everything +strictly below a frontier as one consolidated batch. + +This one mechanism underlies both the single-view batcher and the multi-view +cohort. A single view is the one-member case: its release frontier is its own +progress frontier. A cohort of N views feeds N buffers and releases every buffer +below the *minimum* frontier across them, the latest timestamp closed for every +member at once. + +Whatever the release frontier, the released set is *consolidated* before it +leaves the engine: within one batch the net effect per row is computed and rows +that cancel to nothing are dropped. A batch is therefore a clean net delta at its +frontier, not a replay of intra-window churn. Consolidation is scoped to a single +batch, never across batches, so a consumer that applies batches in order still +sees every intermediate settled state. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from .envelope import ( + Change, + Delete, + Diff, + Envelope, + KeyViolation, + Upsert, + UpsertEnvelope, +) +from .errors import ProtocolError + + +class ReleaseBuffer: + """A per-member buffer of timestamped changes awaiting release. + + Holds changes until a frontier closes them, enforcing the two ordering + invariants the server guarantees: data never arrives below the frontier, and + the frontier never regresses. A violation is a :class:`ProtocolError`. + """ + + def __init__(self, envelope: Envelope) -> None: + self._envelope = envelope + self._buffered: List[Tuple[int, Change]] = [] + self._frontier: Optional[int] = None + + @property + def frontier(self) -> Optional[int]: + """The highest frontier observed so far.""" + return self._frontier + + def pending(self) -> int: + """The number of changes currently buffered, before consolidation. Used + by the cohort to bound how much a laggard can make its peers buffer.""" + return len(self._buffered) + + def push_data(self, timestamp: int, change: Change) -> None: + """Buffers a change. Raises if its timestamp is below the observed + frontier, which the server promises never happens.""" + if self._frontier is not None and timestamp < self._frontier: + raise ProtocolError( + f"change at timestamp {timestamp} arrived after the frontier " + f"advanced to {self._frontier}" + ) + self._buffered.append((timestamp, change)) + + def observe_progress(self, frontier: int) -> None: + """Records a progress frontier. Raises if it regresses.""" + if self._frontier is not None and frontier < self._frontier: + raise ProtocolError( + f"frontier went backwards from {self._frontier} to {frontier}" + ) + self._frontier = frontier + + def release_below(self, release_frontier: int) -> List[Change]: + """Drains and consolidates every buffered change with a timestamp + strictly below ``release_frontier``. Changes at or above it stay + buffered.""" + closed: List[Change] = [] + still_open: List[Tuple[int, Change]] = [] + for timestamp, change in self._buffered: + if timestamp < release_frontier: + closed.append(change) + else: + still_open.append((timestamp, change)) + self._buffered = still_open + return consolidate(self._envelope, closed) + + +def consolidate(envelope: Envelope, changes: List[Change]) -> List[Change]: + """Collapses a batch of changes to its net effect, preserving + first-appearance order. + + Diff changes sum their multiplicities per row; upsert changes keep only the + last operation per key. + """ + if isinstance(envelope, UpsertEnvelope): + return _consolidate_upsert(changes) + return _consolidate_diff(changes) + + +def _consolidate_diff(changes: List[Change]) -> List[Change]: + """Sums signed multiplicities per row, dropping rows whose net diff is + zero.""" + order: List[Any] = [] + totals: Dict[Any, List[Any]] = {} + for change in changes: + # A diff-envelope stream only ever carries `Diff`. + if isinstance(change, Diff): + key = _freeze(change.row) + entry = totals.get(key) + if entry is None: + totals[key] = [change.row, change.diff] + order.append(key) + else: + entry[1] += change.diff + return [ + Diff(row=totals[key][0], diff=totals[key][1]) + for key in order + if totals[key][1] != 0 + ] + + +def _consolidate_upsert(changes: List[Change]) -> List[Change]: + """Keeps the last operation observed for each key, in first-seen key order. + A window that upserts then deletes a key nets to the delete, and vice + versa.""" + order: List[Any] = [] + last: Dict[Any, Change] = {} + for change in changes: + # An upsert-envelope stream never carries a bare diff. + if isinstance(change, (Upsert, Delete, KeyViolation)): + key = _freeze(change.key) + if key not in last: + order.append(key) + last[key] = change + return [last[key] for key in order] + + +def _freeze(value: Any) -> Any: + """Turns a row (or nested value) into a hashable key for consolidation. + + Lists and dicts become tuples; scalars pass through. Rows only compare equal + when their frozen forms do, which is exactly per-cell value equality. + """ + if isinstance(value, (list, tuple)): + return tuple(_freeze(v) for v in value) + if isinstance(value, dict): + return tuple(sorted((k, _freeze(v)) for k, v in value.items())) + return value diff --git a/misc/subscribe-sdk/python/materialize_subscribe/envelope.py b/misc/subscribe-sdk/python/materialize_subscribe/envelope.py new file mode 100644 index 0000000000000..7045bb1de1fb9 --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/envelope.py @@ -0,0 +1,217 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Turning raw ``SUBSCRIBE`` rows into typed changes and progress markers. + +Decoding is by column *name*, not position, because the column layout differs +between envelopes (the upsert envelope interleaves key, ``mz_state``, and value +columns) and because it makes the decoder robust to the exact ordering the +server chooses. + +The decoder accepts either the pgwire text encoding (``"t"``/``"f"``, numeric +strings) or the native Python values that psycopg produces (``bool``, ``int``, +``Decimal``); the metadata columns are coerced either way. Payload and key +values are passed through unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List, Optional, Sequence, Union + +from .errors import ProtocolError + +#: A row of column values. ``None`` is SQL ``NULL``. +Row = List[Any] + + +@dataclass(frozen=True) +class DiffEnvelope: + """Raw differential updates: each change carries a signed multiplicity.""" + + +@dataclass(frozen=True) +class UpsertEnvelope: + """Server-side upsert compaction keyed by the given columns.""" + + key: Sequence[str] + + +#: How the server shapes the update stream. Fixed for a subscription's lifetime. +Envelope = Union[DiffEnvelope, UpsertEnvelope] + + +@dataclass(frozen=True) +class Diff: + """A differential update. + + ``diff`` is the signed multiplicity: positive means ``diff`` copies were + inserted, negative means ``|diff|`` were retracted. The magnitude is + preserved rather than collapsed to +/-1. + """ + + row: Row + diff: int + + +@dataclass(frozen=True) +class Upsert: + """An upsert-envelope insert or update: ``value`` is the new row for ``key``.""" + + key: Row + value: Row + + +@dataclass(frozen=True) +class Delete: + """An upsert-envelope delete: ``key`` no longer has a value.""" + + key: Row + + +@dataclass(frozen=True) +class KeyViolation: + """An upsert-envelope key violation (best-effort per the server).""" + + key: Row + + +#: A decoded change to the subscribed relation. +Change = Union[Diff, Upsert, Delete, KeyViolation] + + +@dataclass(frozen=True) +class Data: + """A change occurring at ``timestamp``.""" + + timestamp: int + change: Change + + +@dataclass(frozen=True) +class Progress: + """A progress marker: no more changes below ``frontier`` will arrive.""" + + frontier: int + + +#: A decoded message from the stream. +StreamMessage = Union[Data, Progress] + +_COL_TIMESTAMP = "mz_timestamp" +_COL_PROGRESSED = "mz_progressed" +_COL_DIFF = "mz_diff" +_COL_STATE = "mz_state" + +_STATE_UPSERT = "upsert" +_STATE_DELETE = "delete" +_STATE_KEY_VIOLATION = "key_violation" + + +class Decoder: + """Decodes raw ``SUBSCRIBE`` rows into :data:`StreamMessage` values. + + Built once from the result's column names (the SDK always subscribes + ``WITH (PROGRESS)``), then reused for every row. + """ + + def __init__(self, columns: Sequence[str], envelope: Envelope) -> None: + self._columns = list(columns) + + self._timestamp_idx = self._require(_COL_TIMESTAMP) + self._progressed_idx = self._require(_COL_PROGRESSED) + + if isinstance(envelope, UpsertEnvelope): + self._state_idx: Optional[int] = self._require(_COL_STATE) + self._key_idxs = [self._require_key(k) for k in envelope.key] + meta = {self._timestamp_idx, self._progressed_idx, self._state_idx} + meta.update(self._key_idxs) + self._value_idxs = [i for i in range(len(self._columns)) if i not in meta] + self._diff_idx: Optional[int] = None + else: + self._diff_idx = self._require(_COL_DIFF) + self._state_idx = None + self._key_idxs = [] + meta = {self._timestamp_idx, self._progressed_idx, self._diff_idx} + self._payload_idxs = [i for i in range(len(self._columns)) if i not in meta] + + def _require(self, name: str) -> int: + try: + return self._columns.index(name) + except ValueError: + raise ProtocolError(f"result is missing the {name} column") from None + + def _require_key(self, name: str) -> int: + try: + return self._columns.index(name) + except ValueError: + raise ProtocolError( + f"upsert key column {name} is not in the result" + ) from None + + def decode(self, row: Sequence[Any]) -> StreamMessage: + """Decodes a single raw row.""" + timestamp = _to_int(row[self._timestamp_idx], _COL_TIMESTAMP) + + if _to_bool(row[self._progressed_idx]): + return Progress(frontier=timestamp) + + if self._state_idx is not None: + key = [row[i] for i in self._key_idxs] + state = row[self._state_idx] + state = str(state) if state is not None else None + if state == _STATE_UPSERT: + change: Change = Upsert( + key=key, value=[row[i] for i in self._value_idxs] + ) + elif state == _STATE_DELETE: + change = Delete(key=key) + elif state == _STATE_KEY_VIOLATION: + change = KeyViolation(key=key) + else: + raise ProtocolError(f"unexpected mz_state value: {state!r}") + else: + assert self._diff_idx is not None + change = Diff( + row=[row[i] for i in self._payload_idxs], + diff=_to_int(row[self._diff_idx], _COL_DIFF), + ) + + return Data(timestamp=timestamp, change=change) + + +def _to_int(value: Any, column: str) -> int: + """Coerces a metadata cell to ``int``, accepting text or native numbers.""" + if value is None: + raise ProtocolError(f"{column} was NULL") + try: + return int(value) + except (ValueError, TypeError) as exc: + raise ProtocolError(f"{column} {value!r} is not an integer: {exc}") from exc + + +def _to_bool(value: Any) -> bool: + """Coerces ``mz_progressed`` to ``bool``, accepting text or native values.""" + if value is None: + return False + if isinstance(value, bool): + return value + if isinstance(value, str): + if value == "t": + return True + if value == "f": + return False + raise ProtocolError(f"unexpected mz_progressed value: {value!r}") + return bool(value) diff --git a/misc/subscribe-sdk/python/materialize_subscribe/errors.py b/misc/subscribe-sdk/python/materialize_subscribe/errors.py new file mode 100644 index 0000000000000..239760842cead --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/errors.py @@ -0,0 +1,127 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The error taxonomy consumers match on to decide how to recover.""" + +from __future__ import annotations + + +class SubscribeError(Exception): + """Base class for every error raised by the SDK. + + The ``is_retryable`` property captures the common "reconnect and resume" + versus "give up" split so callers rarely need to catch every subclass. + """ + + #: Whether reconnecting and resuming from the last token is a sensible + #: automatic response. Overridden to ``True`` only by :class:`TransientError`. + is_retryable: bool = False + + +class CompactionHorizon(SubscribeError): + """The resume point is older than the object's retained history. + + The gap cannot be filled. Recover by re-snapshotting or by accepting a gap, + per the caller's policy, or widen ``RETAIN HISTORY`` on the subscribed + object to make a larger resume window available. + """ + + def __init__(self) -> None: + super().__init__( + "resume point is older than the subscribed object's retained " + "history (compaction horizon); re-snapshot or widen RETAIN HISTORY" + ) + + +class DependencyDropped(SubscribeError): + """A dependency of the subscribed object (or its cluster) was dropped.""" + + +class SchemaMismatch(SubscribeError): + """A resumed query does not match the one the checkpoint was taken against. + + Resuming would silently mix incompatible results, so it is refused. + """ + + def __init__(self, expected: str, actual: str) -> None: + self.expected = expected + self.actual = actual + super().__init__( + f"resumed query does not match the checkpoint " + f"(expected {expected}, got {actual})" + ) + + +class TransientError(SubscribeError): + """A transient connection or protocol interruption; retryable.""" + + is_retryable = True + + +class BufferOverflow(SubscribeError): + """The consumer fell too far behind and the client-side buffer filled. + + The client drains the server continuously so the server never buffers + unboundedly, which means a slow consumer must fail here rather than push + that cost onto Materialize. Recover by consuming faster, widening the + buffer, or resuming from the last token once caught up. + """ + + def __init__(self) -> None: + super().__init__( + "client-side buffer overflowed: the consumer fell behind the " + "subscription. Consume faster or widen the buffer, then resume " + "from the last token" + ) + + +class CohortLagExceeded(SubscribeError): + """A cohort member lagged so far behind that the changes buffered for its + peers, waiting on it, exceeded the lag budget. + + Consistency requires holding a leading member's changes until the slowest + catches up, so a stalled member would otherwise grow memory without bound. + Recover by investigating the slow view, widening the budget, or resuming + from the last token once it recovers. + """ + + def __init__(self, buffered: int, limit: int) -> None: + self.buffered = buffered + self.limit = limit + super().__init__( + f"cohort lag budget exceeded: {buffered} changes buffered waiting " + f"on a lagging member (budget {limit}). Investigate the slow view " + f"or widen the budget" + ) + + +class InvalidToken(SubscribeError): + """A resume token could not be decoded.""" + + +class ProtocolError(SubscribeError): + """The server sent a stream that violates the subscribe protocol. + + Indicates a bug, not a recoverable condition. + """ + + +class FatalError(SubscribeError): + """A non-retryable server error. + + An auth failure, a malformed query, or a dataflow error that poisons the + stream (Materialize repeats such an error until the data or view changes, + so retrying will not help). + """ diff --git a/misc/subscribe-sdk/python/materialize_subscribe/statement.py b/misc/subscribe-sdk/python/materialize_subscribe/statement.py new file mode 100644 index 0000000000000..bb7c023586191 --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/statement.py @@ -0,0 +1,151 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Building the ``SUBSCRIBE`` statement, including the initial-versus-resume +``SNAPSHOT`` and ``AS OF`` handling that must line up with the resume token.""" + +from __future__ import annotations + +from typing import Iterable, Optional + +from .envelope import DiffEnvelope, Envelope, UpsertEnvelope +from .token import ResumeToken + +# FNV-1a 64-bit constants. Matches the Rust SDK byte-for-byte so a fingerprint +# (and therefore a resume token) is interchangeable between the two. +_FNV_OFFSET = 0xCBF29CE484222325 +_FNV_PRIME = 0x100000001B3 +_U64_MASK = 0xFFFFFFFFFFFFFFFF + + +class Subscribe: + """A subscription specification: what to read, how to shape it, and where to + stop. + + ``PROGRESS`` is always requested. It is the primitive the batcher relies on + to know when a timestamp is closed, and requesting it unconditionally keeps + the column layout stable. + + Build one with :meth:`object` or :meth:`query`, then chain + :meth:`envelope_upsert` and :meth:`up_to`. + """ + + def __init__(self, relation_is_query: bool, relation: str) -> None: + self._relation_is_query = relation_is_query + self._relation = relation + self._envelope: Envelope = DiffEnvelope() + self._up_to: Optional[int] = None + + @classmethod + def object(cls, name: str) -> "Subscribe": + """Subscribes to a catalog object (table, view, source, or index). + + ``name`` is passed through verbatim, so a qualified name must already be + a valid SQL object reference (e.g. ``my_db.public.orders``). + """ + return cls(relation_is_query=False, relation=name) + + @classmethod + def query(cls, sql: str) -> "Subscribe": + """Subscribes to the result of a ``SELECT``.""" + return cls(relation_is_query=True, relation=sql) + + def envelope_upsert(self, key: Iterable[str]) -> "Subscribe": + """Switches to the upsert envelope keyed by the given columns.""" + self._envelope = UpsertEnvelope(key=list(key)) + return self + + def up_to(self, timestamp: int) -> "Subscribe": + """Bounds the subscription: it terminates once the frontier reaches + ``timestamp`` (exclusive), yielding a deterministic window.""" + self._up_to = timestamp + return self + + @property + def envelope(self) -> Envelope: + """The envelope this subscription uses.""" + return self._envelope + + def is_bounded(self) -> bool: + """Whether the subscription is bounded by an ``UP TO``.""" + return self._up_to is not None + + def fingerprint(self) -> str: + """A stable identifier for this subscription's shape (relation plus + envelope), embedded in resume tokens so a changed query is caught on + resume rather than silently mixing incompatible results.""" + h = _FNV_OFFSET + if self._relation_is_query: + h = _fnv_update(h, b"query:") + else: + h = _fnv_update(h, b"object:") + h = _fnv_update(h, self._relation.encode("utf-8")) + if isinstance(self._envelope, UpsertEnvelope): + h = _fnv_update(h, b"|upsert:") + for k in self._envelope.key: + h = _fnv_update(h, k.encode("utf-8")) + h = _fnv_update(h, b",") + else: + h = _fnv_update(h, b"|diff") + return f"{h:016x}" + + def to_sql_initial(self) -> str: + """The SQL for an initial subscription: takes the snapshot and lets the + server choose the ``AS OF``.""" + return self._to_sql(snapshot=True, as_of=None) + + def to_sql_resume(self, token: ResumeToken) -> str: + """The SQL for resuming from ``token``: no snapshot, with ``AS OF`` set + to ``token.as_of()`` so emission begins exactly where it left off.""" + return self.to_sql_resume_at(token.as_of()) + + def to_sql_resume_at(self, as_of: int) -> str: + """The SQL for resuming at an explicit ``as_of``: no snapshot, emitting + only timestamps strictly greater than ``as_of``. Used by a cohort, where + every member resumes at the same joint ``AS OF``.""" + return self._to_sql(snapshot=False, as_of=as_of) + + def _to_sql(self, snapshot: bool, as_of: Optional[int]) -> str: + relation = ( + self._relation if not self._relation_is_query else f"({self._relation})" + ) + sql = f"SUBSCRIBE {relation}" + + if isinstance(self._envelope, UpsertEnvelope): + keys = ", ".join(_quote_ident(k) for k in self._envelope.key) + sql += f" ENVELOPE UPSERT (KEY ({keys}))" + + snapshot_sql = "true" if snapshot else "false" + sql += f" WITH (PROGRESS, SNAPSHOT = {snapshot_sql})" + + if as_of is not None: + sql += f" AS OF {as_of}" + if self._up_to is not None: + sql += f" UP TO {self._up_to}" + return sql + + +def _quote_ident(ident: str) -> str: + """Double-quotes a SQL identifier, escaping embedded quotes, so a key column + name cannot break out of the ``KEY (...)`` clause.""" + escaped = ident.replace('"', '""') + return f'"{escaped}"' + + +def _fnv_update(hash_: int, data: bytes) -> int: + for b in data: + hash_ ^= b + hash_ = (hash_ * _FNV_PRIME) & _U64_MASK + return hash_ diff --git a/misc/subscribe-sdk/python/materialize_subscribe/token.py b/misc/subscribe-sdk/python/materialize_subscribe/token.py new file mode 100644 index 0000000000000..97937cd807e8f --- /dev/null +++ b/misc/subscribe-sdk/python/materialize_subscribe/token.py @@ -0,0 +1,158 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The opaque checkpoint that makes resuming gap-free.""" + +from __future__ import annotations + +import base64 +import json +from dataclasses import dataclass +from typing import List + +from .errors import InvalidToken + +# Current on-the-wire token format. Bumped only if the encoded shape changes, so +# a stored token from an older SDK is rejected with a clear error rather than +# misinterpreted. Matches the Rust SDK so tokens are cross-compatible. +_TOKEN_FORMAT = 1 + + +@dataclass(frozen=True) +class ResumeToken: + """An opaque, serializable checkpoint marking how far a subscription has + been durably consumed. + + A token records a *closed* frontier: every update with a timestamp strictly + below :attr:`frontier` has been delivered. Resuming re-subscribes with + ``SNAPSHOT = false AS OF frontier - 1``, which the server reads as "emit + updates with timestamp strictly greater than ``frontier - 1``", i.e. + timestamp ``>= frontier``. That is exactly the set not yet delivered, so the + resume is gap-free and overlap-free. + + The ``- 1`` is the single most error-prone part of the subscribe protocol. + It lives here, in :meth:`as_of`, so callers never compute it. + """ + + frontier: int + """The closed frontier: every update below this timestamp is delivered.""" + + fingerprint: str + """The query fingerprint this token was taken against.""" + + def as_of(self) -> int: + """The ``AS OF`` value for a gap-free ``SNAPSHOT = false`` resume. + + Saturates at zero so a token for the very first frontier resumes from + the beginning rather than going negative. + """ + return max(0, self.frontier - 1) + + def encode(self) -> str: + """Encodes the token to a compact, URL-safe string for durable storage. + + The encoding is opaque: callers persist and return the string unchanged. + """ + payload = { + "format": _TOKEN_FORMAT, + "frontier": self.frontier, + "fingerprint": self.fingerprint, + } + raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") + + @classmethod + def decode(cls, encoded: str) -> ResumeToken: + """Decodes a token previously produced by :meth:`encode`.""" + try: + padded = encoded + "=" * (-len(encoded) % 4) + raw = base64.urlsafe_b64decode(padded.encode("ascii")) + payload = json.loads(raw) + except (ValueError, TypeError) as exc: + raise InvalidToken(f"malformed token: {exc}") from exc + + if not isinstance(payload, dict) or payload.get("format") != _TOKEN_FORMAT: + raise InvalidToken( + f"unsupported token format {payload.get('format') if isinstance(payload, dict) else '?'} " + f"(this SDK understands {_TOKEN_FORMAT})" + ) + try: + return cls( + frontier=int(payload["frontier"]), + fingerprint=str(payload["fingerprint"]), + ) + except (KeyError, ValueError, TypeError) as exc: + raise InvalidToken(f"malformed token: {exc}") from exc + + +@dataclass(frozen=True) +class CohortToken: + """An opaque, serializable checkpoint for a whole cohort of subscriptions. + + A cohort is released at one shared frontier: the minimum closed frontier + across its members. This token records that single joint frontier plus each + member's fingerprint, in member order. Resuming re-subscribes every member + with ``SNAPSHOT = false AS OF frontier - 1``, reconstructing the exact joint + cut, and refuses if the members no longer match. + + The encoded shape matches the Rust SDK so a cohort token is cross-compatible. + """ + + frontier: int + """The joint closed frontier shared by every member.""" + + members: List[str] + """The member fingerprints, in the order the cohort was created.""" + + def as_of(self) -> int: + """The ``AS OF`` value every member uses for a gap-free resume. + + Saturates at zero, like :meth:`ResumeToken.as_of`. + """ + return max(0, self.frontier - 1) + + def encode(self) -> str: + """Encodes the token to a compact, URL-safe string for durable + storage.""" + payload = { + "format": _TOKEN_FORMAT, + "frontier": self.frontier, + "members": list(self.members), + } + raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") + + @classmethod + def decode(cls, encoded: str) -> CohortToken: + """Decodes a token previously produced by :meth:`encode`.""" + try: + padded = encoded + "=" * (-len(encoded) % 4) + raw = base64.urlsafe_b64decode(padded.encode("ascii")) + payload = json.loads(raw) + except (ValueError, TypeError) as exc: + raise InvalidToken(f"malformed cohort token: {exc}") from exc + + if not isinstance(payload, dict) or payload.get("format") != _TOKEN_FORMAT: + raise InvalidToken( + f"unsupported token format {payload.get('format') if isinstance(payload, dict) else '?'} " + f"(this SDK understands {_TOKEN_FORMAT})" + ) + try: + return cls( + frontier=int(payload["frontier"]), + members=[str(m) for m in payload["members"]], + ) + except (KeyError, ValueError, TypeError) as exc: + raise InvalidToken(f"malformed cohort token: {exc}") from exc diff --git a/misc/subscribe-sdk/python/pyproject.toml b/misc/subscribe-sdk/python/pyproject.toml new file mode 100644 index 0000000000000..ae7ea87421ab9 --- /dev/null +++ b/misc/subscribe-sdk/python/pyproject.toml @@ -0,0 +1,54 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[project] +name = "materialize-subscribe" +version = "0.1.0" +description = "Correct-by-construction client for consuming Materialize SUBSCRIBE." +readme = "README.md" +requires-python = ">=3.9" +license = "Apache-2.0" +authors = [{ name = "Materialize Inc", email = "support@materialize.com" }] +keywords = ["materialize", "subscribe", "cdc", "streaming", "database"] + +# The protocol core has no runtime dependencies; the driver is an extra so the +# core can be used and tested without it. +dependencies = [] + +[project.optional-dependencies] +client = ["psycopg[binary]>=3.2"] +dev = ["pytest>=8.0", "mypy>=1.10", "ruff>=0.6"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["materialize_subscribe"] + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.ruff] +target-version = "py39" + +[tool.mypy] +strict = true + +# psycopg is an optional extra imported lazily, so it may be absent when +# type-checking the core. +[[tool.mypy.overrides]] +module = "psycopg" +ignore_missing_imports = true diff --git a/misc/subscribe-sdk/python/tests/test_batch.py b/misc/subscribe-sdk/python/tests/test_batch.py new file mode 100644 index 0000000000000..668f5e16ad6ad --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_batch.py @@ -0,0 +1,144 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from materialize_subscribe import ( + Batcher, + Data, + Diff, + DiffEnvelope, + Progress, + ProtocolError, +) + + +def insert(id_: str) -> Diff: + return Diff(row=[id_], diff=1) + + +def data(timestamp: int, change: Diff) -> Data: + return Data(timestamp=timestamp, change=change) + + +def batcher(fingerprint: str, with_snapshot: bool) -> Batcher: + return Batcher(fingerprint, DiffEnvelope(), with_snapshot=with_snapshot) + + +def test_buffers_until_the_timestamp_closes() -> None: + b = batcher("fp", with_snapshot=False) + assert b.push(data(5, insert("a"))) is None + assert b.push(data(5, insert("b"))) is None + batch = b.push(Progress(frontier=6)) + assert batch is not None + assert batch.updates == [insert("a"), insert("b")] + assert batch.frontier == 6 + assert batch.resume_token.frontier == 6 + + +def test_closes_a_timestamp_with_mixed_retractions_and_multiplicities() -> None: + # A single timestamp can hold inserts, retractions, and multiplicities + # beyond one. Distinct rows all survive consolidation, in first-seen order. + b = batcher("fp", with_snapshot=False) + insert3 = Diff(row=["a"], diff=3) + retract = Diff(row=["b"], diff=-1) + b.push(data(5, insert3)) + b.push(data(5, retract)) + b.push(data(5, insert("c"))) + batch = b.push(Progress(frontier=6)) + assert batch is not None + assert batch.updates == [insert3, retract, insert("c")] + + +def test_batch_is_consolidated_to_its_net_effect() -> None: + # Repeated changes to the same row within a batch collapse: `a` inserted + # twice and retracted once nets to +1; `b` inserted then retracted vanishes. + b = batcher("fp", with_snapshot=False) + b.push(data(5, insert("a"))) + b.push(data(5, insert("a"))) + b.push(data(6, insert("b"))) + b.push(data(6, Diff(row=["a"], diff=-1))) + b.push(data(6, Diff(row=["b"], diff=-1))) + batch = b.push(Progress(frontier=7)) + assert batch is not None + assert batch.updates == [Diff(row=["a"], diff=1)] + + +def test_data_at_the_frontier_stays_open() -> None: + b = batcher("fp", with_snapshot=False) + b.push(data(5, insert("a"))) + b.push(data(6, insert("b"))) + batch = b.push(Progress(frontier=6)) + assert batch is not None and batch.updates == [insert("a")] + batch = b.push(Progress(frontier=7)) + assert batch is not None and batch.updates == [insert("b")] + + +def test_empty_batches_still_advance_the_token() -> None: + b = batcher("fp", with_snapshot=False) + batch = b.push(Progress(frontier=10)) + assert batch is not None + assert batch.is_empty() + assert batch.frontier == 10 + assert batch.resume_token.frontier == 10 + + +def test_duplicate_frontier_emits_nothing() -> None: + b = batcher("fp", with_snapshot=False) + assert b.push(Progress(frontier=10)) is not None + assert b.push(Progress(frontier=10)) is None + + +def test_snapshot_batch_is_the_first_past_the_initial_frontier() -> None: + b = batcher("fp", with_snapshot=True) + leading = b.push(Progress(frontier=100)) + assert leading is not None and not leading.is_snapshot and leading.is_empty() + b.push(data(100, insert("snap1"))) + b.push(data(100, insert("snap2"))) + snap = b.push(Progress(frontier=101)) + assert snap is not None and snap.is_snapshot + assert snap.updates == [insert("snap1"), insert("snap2")] + b.push(data(101, insert("live"))) + live = b.push(Progress(frontier=102)) + assert live is not None and not live.is_snapshot + + +def test_resume_never_marks_a_snapshot() -> None: + b = batcher("fp", with_snapshot=False) + b.push(Progress(frontier=100)) + b.push(data(100, insert("x"))) + batch = b.push(Progress(frontier=101)) + assert batch is not None and not batch.is_snapshot + + +def test_late_data_below_the_frontier_is_a_protocol_error() -> None: + b = batcher("fp", with_snapshot=False) + b.push(Progress(frontier=10)) + with pytest.raises(ProtocolError): + b.push(data(5, insert("late"))) + + +def test_regressing_frontier_is_a_protocol_error() -> None: + b = batcher("fp", with_snapshot=False) + b.push(Progress(frontier=10)) + with pytest.raises(ProtocolError): + b.push(Progress(frontier=9)) + + +def test_token_carries_the_fingerprint() -> None: + b = batcher("my-fingerprint", with_snapshot=False) + batch = b.push(Progress(frontier=1)) + assert batch is not None + assert batch.resume_token.fingerprint == "my-fingerprint" diff --git a/misc/subscribe-sdk/python/tests/test_client.py b/misc/subscribe-sdk/python/tests/test_client.py new file mode 100644 index 0000000000000..5d9f8c34d3b32 --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_client.py @@ -0,0 +1,44 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from materialize_subscribe import ( + CompactionHorizon, + DependencyDropped, + FatalError, +) +from materialize_subscribe.client import classify_server_message, is_compaction_horizon + + +def test_classifies_compaction_horizon_wordings() -> None: + assert is_compaction_horizon("could not find a valid timestamp for the query") + assert is_compaction_horizon("Timestamp (123) is not valid for all inputs") + assert not is_compaction_horizon("some unrelated error") + + +def test_classifies_server_messages() -> None: + # Mirrors the Rust SDK's classify_server_messages test so the two taxonomies + # stay in lockstep. + assert isinstance( + classify_server_message("could not find a valid timestamp for the query"), + CompactionHorizon, + ) + assert isinstance( + classify_server_message("dependency winning_bids was dropped"), + DependencyDropped, + ) + assert isinstance( + classify_server_message('column "x" does not exist'), + FatalError, + ) diff --git a/misc/subscribe-sdk/python/tests/test_cohort.py b/misc/subscribe-sdk/python/tests/test_cohort.py new file mode 100644 index 0000000000000..06c59b958eb05 --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_cohort.py @@ -0,0 +1,134 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List + +import pytest + +from materialize_subscribe import CohortLagExceeded, Diff, DiffEnvelope, ProtocolError +from materialize_subscribe.cohort import CohortEngine + + +def diff(id_: str, d: int) -> Diff: + return Diff(row=[id_], diff=d) + + +def engine(names: List[str], lag_budget: int = 1 << 20) -> CohortEngine: + specs = [(name, f"fp-{name}", DiffEnvelope()) for name in names] + return CohortEngine(specs, with_snapshot=True, lag_budget=lag_budget) + + +def test_no_release_until_every_member_reports() -> None: + e = engine(["a", "b"]) + assert e.push_progress(0, 10) is None + moment = e.push_progress(1, 5) + assert moment is not None and moment.frontier == 5 + + +def test_releases_at_the_minimum_frontier() -> None: + e = engine(["a", "b"]) + e.push_data(0, 3, diff("a3", 1)) + e.push_data(0, 7, diff("a7", 1)) + e.push_data(1, 4, diff("b4", 1)) + e.push_progress(0, 10) + moment = e.push_progress(1, 6) + assert moment is not None and moment.frontier == 6 + assert moment.views[0].updates == [diff("a3", 1)] + assert moment.views[1].updates == [diff("b4", 1)] + + +def test_laggard_holds_the_joint_moment() -> None: + e = engine(["fast", "slow"]) + e.push_data(0, 5, diff("x", 1)) + assert e.push_progress(0, 100) is None + moment = e.push_progress(1, 6) + assert moment is not None and moment.frontier == 6 + assert moment.views[0].updates == [diff("x", 1)] + assert moment.views[1].updates == [] + + +def test_moment_carries_a_token_at_the_joint_frontier() -> None: + e = engine(["a", "b"]) + e.push_progress(0, 8) + moment = e.push_progress(1, 5) + assert moment is not None + assert moment.resume_token.frontier == 5 + assert moment.resume_token.members == ["fp-a", "fp-b"] + + +def test_per_view_changes_are_consolidated() -> None: + e = engine(["a", "b"]) + e.push_data(0, 3, diff("x", 1)) + e.push_data(0, 4, diff("x", -1)) + e.push_progress(0, 10) + moment = e.push_progress(1, 6) + assert moment is not None + assert moment.views[0].updates == [] + + +def test_snapshot_completes_once_the_last_member_passes_its_as_of() -> None: + e = engine(["a", "b"]) + assert e.push_progress(0, 5) is None + leading = e.push_progress(1, 8) + assert leading is not None and leading.frontier == 5 and not leading.is_snapshot + e.push_progress(0, 9) + done = e.push_progress(1, 9) + assert done is not None and done.frontier == 9 and done.is_snapshot + e.push_progress(0, 11) + live = e.push_progress(1, 11) + assert live is not None and not live.is_snapshot + + +def test_single_member_cohort_matches_a_batcher() -> None: + e = engine(["only"]) + leading = e.push_progress(0, 100) + assert leading is not None and leading.frontier == 100 and not leading.is_snapshot + e.push_data(0, 100, diff("snap", 1)) + snap = e.push_progress(0, 101) + assert snap is not None and snap.is_snapshot + assert snap.views[0].updates == [diff("snap", 1)] + + +def test_protocol_errors_surface_from_members() -> None: + e = engine(["a", "b"]) + e.push_progress(0, 10) + with pytest.raises(ProtocolError): + e.push_data(0, 5, diff("late", 1)) + + +def test_a_laggard_past_the_lag_budget_fails_loud() -> None: + # Budget of two buffered changes. The `slow` member never reports, so the + # joint frontier never advances and `fast`'s changes pile up. + e = engine(["fast", "slow"], lag_budget=2) + e.push_data(0, 1, diff("a", 1)) + e.push_data(0, 2, diff("b", 1)) + with pytest.raises(CohortLagExceeded): + e.push_data(0, 3, diff("c", 1)) + + +def test_buffered_changes_free_the_budget_once_released() -> None: + # Once the joint frontier advances and drains the buffer, there is room to + # buffer again. + e = engine(["a", "b"], lag_budget=2) + e.push_data(0, 1, diff("x", 1)) + e.push_progress(0, 2) + moment = e.push_progress(1, 2) + assert moment is not None and moment.frontier == 2 + assert moment.views[0].updates == [diff("x", 1)] + # Room to buffer again (data at or above each member's frontier of 2). + e.push_data(0, 2, diff("y", 1)) + e.push_data(1, 2, diff("z", 1)) + with pytest.raises(CohortLagExceeded): + e.push_data(0, 2, diff("w", 1)) diff --git a/misc/subscribe-sdk/python/tests/test_engine.py b/misc/subscribe-sdk/python/tests/test_engine.py new file mode 100644 index 0000000000000..b91fbe25fe7fc --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_engine.py @@ -0,0 +1,88 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from materialize_subscribe import Delete, Diff, DiffEnvelope, ProtocolError, Upsert +from materialize_subscribe.engine import ReleaseBuffer, consolidate + + +def diff(id_: str, d: int) -> Diff: + return Diff(row=[id_], diff=d) + + +def upsert(id_: str, value: str) -> Upsert: + return Upsert(key=[id_], value=[value]) + + +def test_diff_consolidation_sums_multiplicities_per_row() -> None: + out = consolidate(DiffEnvelope(), [diff("a", 3), diff("b", -1), diff("a", 2)]) + assert out == [diff("a", 5), diff("b", -1)] + + +def test_diff_consolidation_drops_net_zero_rows() -> None: + out = consolidate(DiffEnvelope(), [diff("a", 1), diff("b", 1), diff("a", -1)]) + assert out == [diff("b", 1)] + + +def test_diff_consolidation_preserves_first_seen_order() -> None: + out = consolidate(DiffEnvelope(), [diff("c", 1), diff("a", 1), diff("b", 1)]) + assert out == [diff("c", 1), diff("a", 1), diff("b", 1)] + + +def test_upsert_consolidation_keeps_last_op_per_key() -> None: + from materialize_subscribe import UpsertEnvelope + + env = UpsertEnvelope(key=["id"]) + out = consolidate(env, [upsert("1", "v1"), upsert("2", "x"), upsert("1", "v2")]) + assert out == [upsert("1", "v2"), upsert("2", "x")] + + +def test_upsert_consolidation_collapses_upsert_then_delete() -> None: + from materialize_subscribe import UpsertEnvelope + + env = UpsertEnvelope(key=["id"]) + out = consolidate(env, [upsert("1", "v"), Delete(key=["1"])]) + assert out == [Delete(key=["1"])] + + +def test_release_holds_data_at_or_above_the_frontier() -> None: + buf = ReleaseBuffer(DiffEnvelope()) + buf.push_data(5, diff("a", 1)) + buf.push_data(6, diff("b", 1)) + assert buf.release_below(6) == [diff("a", 1)] + assert buf.release_below(7) == [diff("b", 1)] + + +def test_release_consolidates_across_the_window() -> None: + buf = ReleaseBuffer(DiffEnvelope()) + buf.push_data(5, diff("a", 1)) + buf.push_data(6, diff("a", -1)) + buf.observe_progress(7) + assert buf.release_below(7) == [] + + +def test_data_below_the_frontier_is_rejected() -> None: + buf = ReleaseBuffer(DiffEnvelope()) + buf.observe_progress(10) + with pytest.raises(ProtocolError): + buf.push_data(5, diff("late", 1)) + + +def test_regressing_frontier_is_rejected() -> None: + buf = ReleaseBuffer(DiffEnvelope()) + buf.observe_progress(10) + with pytest.raises(ProtocolError): + buf.observe_progress(9) diff --git a/misc/subscribe-sdk/python/tests/test_envelope.py b/misc/subscribe-sdk/python/tests/test_envelope.py new file mode 100644 index 0000000000000..625986a1885c7 --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_envelope.py @@ -0,0 +1,98 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from materialize_subscribe import ( + Data, + Decoder, + Delete, + Diff, + DiffEnvelope, + KeyViolation, + Progress, + ProtocolError, + Upsert, + UpsertEnvelope, +) + + +def test_diff_decoder_reads_progress() -> None: + decoder = Decoder( + ["mz_timestamp", "mz_progressed", "mz_diff", "id", "name"], DiffEnvelope() + ) + # Text encoding (mz_progressed = "t"). + assert decoder.decode(["100", "t", None, None, None]) == Progress(frontier=100) + # Native psycopg values (mz_progressed = True) decode the same way. + assert decoder.decode([100, True, None, None, None]) == Progress(frontier=100) + + +def test_diff_decoder_preserves_multiplicity_and_sign() -> None: + decoder = Decoder( + ["mz_timestamp", "mz_progressed", "mz_diff", "id"], DiffEnvelope() + ) + assert decoder.decode(["7", "f", "3", "42"]) == Data( + timestamp=7, change=Diff(row=["42"], diff=3) + ) + assert decoder.decode(["7", "f", "-2", "42"]) == Data( + timestamp=7, change=Diff(row=["42"], diff=-2) + ) + + +def test_upsert_decoder_handles_all_states() -> None: + columns = ["id", "mz_state", "mz_timestamp", "name", "email", "mz_progressed"] + decoder = Decoder(columns, UpsertEnvelope(key=["id"])) + + assert decoder.decode(["1", "upsert", "50", "alice", "a@x", "f"]) == Data( + timestamp=50, change=Upsert(key=["1"], value=["alice", "a@x"]) + ) + assert decoder.decode(["1", "delete", "51", None, None, "f"]) == Data( + timestamp=51, change=Delete(key=["1"]) + ) + assert decoder.decode(["1", "key_violation", "52", None, None, "f"]) == Data( + timestamp=52, change=KeyViolation(key=["1"]) + ) + + +def test_upsert_decoder_reads_progress() -> None: + columns = ["id", "mz_state", "mz_timestamp", "name", "mz_progressed"] + decoder = Decoder(columns, UpsertEnvelope(key=["id"])) + assert decoder.decode([None, None, "200", None, "t"]) == Progress(frontier=200) + + +def test_multi_column_key_is_projected_in_order() -> None: + columns = ["region", "id", "mz_state", "mz_timestamp", "total", "mz_progressed"] + decoder = Decoder(columns, UpsertEnvelope(key=["region", "id"])) + assert decoder.decode(["us", "9", "upsert", "5", "100", "f"]) == Data( + timestamp=5, change=Upsert(key=["us", "9"], value=["100"]) + ) + + +def test_missing_metadata_column_is_rejected() -> None: + with pytest.raises(ProtocolError): + Decoder(["mz_timestamp", "mz_diff", "id"], DiffEnvelope()) + + +def test_unknown_key_column_is_rejected() -> None: + columns = ["id", "mz_state", "mz_timestamp", "mz_progressed"] + with pytest.raises(ProtocolError): + Decoder(columns, UpsertEnvelope(key=["nonexistent"])) + + +def test_unexpected_state_is_rejected() -> None: + columns = ["id", "mz_state", "mz_timestamp", "mz_progressed"] + decoder = Decoder(columns, UpsertEnvelope(key=["id"])) + with pytest.raises(ProtocolError): + decoder.decode(["1", "bogus", "5", "f"]) diff --git a/misc/subscribe-sdk/python/tests/test_statement.py b/misc/subscribe-sdk/python/tests/test_statement.py new file mode 100644 index 0000000000000..a9920d2ec67b2 --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_statement.py @@ -0,0 +1,91 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from materialize_subscribe import ResumeToken, Subscribe + + +def test_initial_takes_snapshot_and_omits_as_of() -> None: + assert ( + Subscribe.object("orders").to_sql_initial() + == "SUBSCRIBE orders WITH (PROGRESS, SNAPSHOT = true)" + ) + + +def test_resume_disables_snapshot_and_sets_as_of_to_frontier_minus_one() -> None: + sub = Subscribe.object("orders") + token = ResumeToken(frontier=100, fingerprint=sub.fingerprint()) + assert ( + sub.to_sql_resume(token) + == "SUBSCRIBE orders WITH (PROGRESS, SNAPSHOT = false) AS OF 99" + ) + + +def test_query_relation_is_parenthesized() -> None: + assert ( + Subscribe.query("SELECT id FROM orders WHERE paid").to_sql_initial() + == "SUBSCRIBE (SELECT id FROM orders WHERE paid) WITH (PROGRESS, SNAPSHOT = true)" + ) + + +def test_upsert_envelope_precedes_with_clause() -> None: + assert ( + Subscribe.object("orders").envelope_upsert(["id"]).to_sql_initial() + == 'SUBSCRIBE orders ENVELOPE UPSERT (KEY ("id")) WITH (PROGRESS, SNAPSHOT = true)' + ) + + +def test_multi_column_key_and_up_to() -> None: + sql = ( + Subscribe.object("sales") + .envelope_upsert(["region", "id"]) + .up_to(500) + .to_sql_initial() + ) + assert sql == ( + 'SUBSCRIBE sales ENVELOPE UPSERT (KEY ("region", "id")) ' + "WITH (PROGRESS, SNAPSHOT = true) UP TO 500" + ) + + +def test_key_identifiers_are_quote_escaped() -> None: + sql = Subscribe.object("t").envelope_upsert(['weird"name']).to_sql_initial() + assert 'KEY ("weird""name")' in sql + + +def test_fingerprint_is_stable_and_distinguishes_shape() -> None: + a = Subscribe.object("orders") + assert a.fingerprint() == Subscribe.object("orders").fingerprint() + assert a.fingerprint() != Subscribe.object("other").fingerprint() + assert ( + a.fingerprint() + != Subscribe.object("orders").envelope_upsert(["id"]).fingerprint() + ) + assert ( + Subscribe.object("orders").envelope_upsert(["id"]).fingerprint() + != Subscribe.object("orders").envelope_upsert(["id", "region"]).fingerprint() + ) + # UP TO does not change the shape, so it does not change the fingerprint. + assert a.fingerprint() == Subscribe.object("orders").up_to(9).fingerprint() + + +def test_fingerprint_matches_cross_language_vectors() -> None: + # These exact values are also asserted by the Rust SDK, so both compute the + # same fingerprint (and therefore interchangeable tokens) for the same + # subscription shape. + assert Subscribe.object("orders").fingerprint() == "f44d7c35ad6f569c" + assert ( + Subscribe.object("orders").envelope_upsert(["id"]).fingerprint() + == "a9392428ac0ed8c9" + ) diff --git a/misc/subscribe-sdk/python/tests/test_token.py b/misc/subscribe-sdk/python/tests/test_token.py new file mode 100644 index 0000000000000..bec50a6d4b27a --- /dev/null +++ b/misc/subscribe-sdk/python/tests/test_token.py @@ -0,0 +1,73 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import base64 +import json + +import pytest + +from materialize_subscribe import CohortToken, InvalidToken, ResumeToken + + +def test_as_of_is_frontier_minus_one_and_saturates() -> None: + assert ResumeToken(frontier=0, fingerprint="q").as_of() == 0 + assert ResumeToken(frontier=1, fingerprint="q").as_of() == 0 + assert ResumeToken(frontier=5, fingerprint="q").as_of() == 4 + + +def test_encode_decode_round_trips() -> None: + token = ResumeToken(frontier=42, fingerprint="fingerprint-abc") + decoded = ResumeToken.decode(token.encode()) + assert decoded == token + assert decoded.frontier == 42 + assert decoded.fingerprint == "fingerprint-abc" + + +def test_encoding_is_stable_and_cross_language() -> None: + # This exact string is also asserted by the Rust SDK's test suite, so a + # token minted by either SDK decodes in the other. If this changes, the two + # SDKs have diverged and tokens are no longer interchangeable. + encoded = ResumeToken(frontier=42, fingerprint="x").encode() + assert encoded == "eyJmb3JtYXQiOjEsImZyb250aWVyIjo0MiwiZmluZ2VycHJpbnQiOiJ4In0" + + +def test_decode_rejects_garbage() -> None: + with pytest.raises(InvalidToken): + ResumeToken.decode("!!! not base64 !!!") + with pytest.raises(InvalidToken): + ResumeToken.decode(base64.urlsafe_b64encode(b"not json").decode()) + + +def test_decode_rejects_unknown_format() -> None: + future = json.dumps({"format": 999, "frontier": 1, "fingerprint": "q"}).encode() + encoded = base64.urlsafe_b64encode(future).decode().rstrip("=") + with pytest.raises(InvalidToken): + ResumeToken.decode(encoded) + + +def test_cohort_token_round_trips_and_carries_members() -> None: + token = CohortToken(frontier=42, members=["fp-a", "fp-b"]) + decoded = CohortToken.decode(token.encode()) + assert decoded == token + assert decoded.frontier == 42 + assert decoded.as_of() == 41 + assert decoded.members == ["fp-a", "fp-b"] + + +def test_cohort_token_encoding_is_stable_and_cross_language() -> None: + # Asserted byte-for-byte by the Rust SDK too, so a cohort token minted by + # either SDK decodes in the other. + encoded = CohortToken(frontier=7, members=["x", "y"]).encode() + assert encoded == "eyJmb3JtYXQiOjEsImZyb250aWVyIjo3LCJtZW1iZXJzIjpbIngiLCJ5Il19"