From 4007c6e33d1ded165d97626e95a6771260b7d4ad Mon Sep 17 00:00:00 2001 From: Mehul Arora Date: Sat, 18 Jul 2026 15:20:09 -0400 Subject: [PATCH 1/5] feat: expose caught-up-to-tail signal on read sessions --- docs/source/api-reference.md | 7 + s2-specs | 2 +- src/s2_sdk/__init__.py | 4 + src/s2_sdk/_exceptions.py | 5 + src/s2_sdk/_ops.py | 39 +++--- src/s2_sdk/_read_session.py | 226 +++++++++++++++++++++++++++++++ src/s2_sdk/_s2s/_read_session.py | 71 ++++++---- tests/test_correctness.py | 76 ++++++++--- tests/test_read_session.py | 204 ++++++++++++++++++++++++++++ 9 files changed, 564 insertions(+), 70 deletions(-) create mode 100644 src/s2_sdk/_read_session.py create mode 100644 tests/test_read_session.py diff --git a/docs/source/api-reference.md b/docs/source/api-reference.md index 9ae8366..8d5e92a 100644 --- a/docs/source/api-reference.md +++ b/docs/source/api-reference.md @@ -22,6 +22,9 @@ .. autoclass:: BatchSubmitTicket() :members: +.. autoclass:: ReadSession() + :members: + .. autoclass:: Producer() :members: @@ -191,4 +194,8 @@ .. autoclass:: ReadUnwrittenError() :members: :show-inheritance: + +.. autoclass:: ReadSessionClosedError() + :members: + :show-inheritance: ``` diff --git a/s2-specs b/s2-specs index 329de93..973e0d9 160000 --- a/s2-specs +++ b/s2-specs @@ -1 +1 @@ -Subproject commit 329de93f7b240a4daef9edbeb98ced0699aab7d0 +Subproject commit 973e0d92166ee1386b71d4e7ade8bd0fc2aaf4b8 diff --git a/src/s2_sdk/__init__.py b/src/s2_sdk/__init__.py index 0950512..1af06a6 100644 --- a/src/s2_sdk/__init__.py +++ b/src/s2_sdk/__init__.py @@ -6,6 +6,7 @@ from s2_sdk._exceptions import ( AppendConditionError, FencingTokenMismatchError, + ReadSessionClosedError, ReadUnwrittenError, S2ClientError, S2Error, @@ -14,6 +15,7 @@ ) from s2_sdk._ops import S2, S2Basin, S2Stream from s2_sdk._producer import Producer, RecordSubmitTicket +from s2_sdk._read_session import ReadSession from s2_sdk._types import ( AccessTokenInfo, AccessTokenScope, @@ -98,6 +100,7 @@ "Timestamp", "TailOffset", "ReadBatch", + "ReadSession", "ReadLimit", "SequencedRecord", "Page", @@ -134,4 +137,5 @@ "FencingTokenMismatchError", "SeqNumMismatchError", "ReadUnwrittenError", + "ReadSessionClosedError", ] diff --git a/src/s2_sdk/_exceptions.py b/src/s2_sdk/_exceptions.py index d95a8b0..b31bdcc 100644 --- a/src/s2_sdk/_exceptions.py +++ b/src/s2_sdk/_exceptions.py @@ -16,6 +16,10 @@ class S2ClientError(S2Error): """Error originating from the client.""" +class ReadSessionClosedError(S2ClientError): + """Read session ended before becoming caught up.""" + + UNKNOWN_CODE = "unknown" @@ -126,6 +130,7 @@ def raise_for_416(body: dict, code: str) -> None: S2Error.__module__ = "s2_sdk" S2ClientError.__module__ = "s2_sdk" +ReadSessionClosedError.__module__ = "s2_sdk" S2ServerError.__module__ = "s2_sdk" AppendConditionError.__module__ = "s2_sdk" FencingTokenMismatchError.__module__ = "s2_sdk" diff --git a/src/s2_sdk/_ops.py b/src/s2_sdk/_ops.py index 6ddfc0e..82615c3 100644 --- a/src/s2_sdk/_ops.py +++ b/src/s2_sdk/_ops.py @@ -2,7 +2,7 @@ import uuid from collections.abc import AsyncIterator from datetime import datetime -from typing import Any, AsyncIterable, Self +from typing import Any, Self from urllib.parse import quote import s2_sdk._generated.s2.v1.s2_pb2 as pb @@ -33,6 +33,7 @@ tail_from_json, ) from s2_sdk._producer import Producer +from s2_sdk._read_session import ReadSession from s2_sdk._retrier import Retrier, http_retry_on, is_safe_to_retry_unary from s2_sdk._s2s._read_session import run_read_session from s2_sdk._types import ( @@ -1172,7 +1173,7 @@ async def read( return batch @fallible - async def read_session( + def read_session( self, *, start: types.SeqNum | types.Timestamp | types.TailOffset, @@ -1181,7 +1182,7 @@ async def read_session( clamp_to_tail: bool = False, wait: int | None = None, ignore_command_records: bool = False, - ) -> AsyncIterable[types.ReadBatch]: + ) -> ReadSession: """Read batches of records from a stream continuously. Args: @@ -1196,9 +1197,8 @@ async def read_session( reached. ignore_command_records: Filter out command records from batches. - Yields: - :class:`ReadBatch` — each containing a batch of records and an - optional tail position. + Returns: + A :class:`ReadSession` that yields batches of records. Note: Sessions without bounds (no ``limit`` or ``until_timestamp``) default @@ -1207,19 +1207,20 @@ async def read_session( ``wait`` makes a bounded session wait up to that many seconds for new records before ending. """ - async for batch in run_read_session( - self._client, - self.name, - start, - limit, - until_timestamp, - clamp_to_tail, - wait, - ignore_command_records, - retry=self._retry, - encryption_key=self._encryption_key, - ): - yield batch + return ReadSession( + run_read_session( + self._client, + self.name, + start, + limit, + until_timestamp, + clamp_to_tail, + wait, + ignore_command_records, + retry=self._retry, + encryption_key=self._encryption_key, + ) + ) def _s2_request_token() -> str: diff --git a/src/s2_sdk/_read_session.py b/src/s2_sdk/_read_session.py new file mode 100644 index 0000000..e71cf2b --- /dev/null +++ b/src/s2_sdk/_read_session.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterable, AsyncIterator, Awaitable +from dataclasses import dataclass, field +from typing import Self + +from s2_sdk._exceptions import ReadSessionClosedError, normalize_exception +from s2_sdk._types import ReadBatch, StreamPosition + + +@dataclass(slots=True) +class _ReadSessionUpdate: + batch: ReadBatch | None = None + caught_up_tail: StreamPosition | None = None + + +@dataclass(slots=True) +class _ReadDelivery: + batch: ReadBatch + caught_up_tail: StreamPosition | None + consumed: asyncio.Event = field(default_factory=asyncio.Event) + + +class _End: + pass + + +_END = _End() + + +class ReadSession(AsyncIterator[ReadBatch]): + """A continuous read returned by :meth:`S2Stream.read_session`. + + Use as an async context manager or call :meth:`close`. + """ + + __slots__ = ( + "_caught_up_tail", + "_closed", + "_deliveries", + "_error", + "_error_raised", + "_is_caught_up", + "_task", + "_updates", + "_waiters", + ) + + def __init__(self, updates: AsyncIterable[_ReadSessionUpdate]) -> None: + self._updates = updates + self._deliveries: asyncio.Queue[_ReadDelivery | _End] = asyncio.Queue() + self._task: asyncio.Task[None] | None = None + self._closed = False + self._error: BaseException | None = None + self._error_raised = False + self._is_caught_up = False + self._caught_up_tail: StreamPosition | None = None + self._waiters: set[asyncio.Future[StreamPosition]] = set() + + def is_caught_up(self) -> bool: + """Return whether all records through the latest reported tail were delivered. + + A later batch that does not reach a reported tail or a reconnect resets it. + Filtered command records count toward progress. Use + :meth:`S2Stream.check_tail` for the current stream tail. + """ + return self._is_caught_up + + def caught_up(self) -> Awaitable[StreamPosition]: + """Return an awaitable for the next caught-up tail. + + It resolves immediately if already caught up and remains pending across + retries. Keep consuming batches while waiting. Call again after falling + behind. + + Raises: + ReadSessionClosedError: The session ended before catching up. + S2Error: The read failed before catching up. + """ + loop = asyncio.get_running_loop() + waiter: asyncio.Future[StreamPosition] = loop.create_future() + if self._is_caught_up and self._caught_up_tail is not None: + waiter.set_result(_copy_position(self._caught_up_tail)) + elif self._closed: + waiter.set_exception(self._error or _session_closed_error()) + else: + self._waiters.add(waiter) + waiter.add_done_callback(self._waiters.discard) + self._ensure_started() + return waiter + + async def close(self) -> None: + """Close the session.""" + if self._closed: + return + if self._task is None: + self._finish() + self._signal_end() + return + current = asyncio.current_task() + cancellation_count = current.cancelling() if current is not None else 0 + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + finally: + if not self._closed: + self._finish() + self._signal_end() + if current is not None and current.cancelling() > cancellation_count: + raise asyncio.CancelledError + + async def aclose(self) -> None: + """Close the session.""" + await self.close() + + def __aiter__(self) -> Self: + self._ensure_started() + return self + + async def __anext__(self) -> ReadBatch: + self._ensure_started() + try: + item = await self._deliveries.get() + except asyncio.CancelledError: + await self.close() + raise + + if isinstance(item, _End): + self._deliveries.put_nowait(_END) + if self._error is not None and not self._error_raised: + self._error_raised = True + raise self._error + raise StopAsyncIteration + + if item.caught_up_tail is not None: + self._mark_caught_up(item.caught_up_tail) + item.consumed.set() + return item.batch + + async def __aenter__(self) -> Self: + self._ensure_started() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: + await self.close() + return False + + def _ensure_started(self) -> None: + if self._task is None and not self._closed: + self._task = asyncio.get_running_loop().create_task(self._run()) + + async def _run(self) -> None: + error: BaseException | None = None + cancelled = False + updates = aiter(self._updates) + try: + async for update in updates: + self._mark_behind() + batch = update.batch + if batch is None or not batch.records: + if update.caught_up_tail is not None: + self._mark_caught_up(update.caught_up_tail) + continue + + delivery = _ReadDelivery(batch, update.caught_up_tail) + await self._deliveries.put(delivery) + await delivery.consumed.wait() + except asyncio.CancelledError: + cancelled = True + except BaseException as exc: + error = normalize_exception(exc) + finally: + close = getattr(updates, "aclose", None) + if close is not None: + try: + await close() + except BaseException as exc: + if error is None and not cancelled: + error = normalize_exception(exc) + self._finish(error) + self._signal_end() + + def _mark_caught_up(self, tail: StreamPosition) -> None: + if self._closed: + return + self._is_caught_up = True + self._caught_up_tail = _copy_position(tail) + waiters = tuple(self._waiters) + self._waiters.clear() + for waiter in waiters: + if not waiter.done(): + waiter.set_result(_copy_position(tail)) + + def _mark_behind(self) -> None: + if not self._closed: + self._is_caught_up = False + + def _finish(self, error: BaseException | None = None) -> None: + if self._closed: + return + self._closed = True + self._error = error + if error is not None: + self._is_caught_up = False + waiters = tuple(self._waiters) + self._waiters.clear() + rejection = error or _session_closed_error() + for waiter in waiters: + if not waiter.done(): + waiter.set_exception(rejection) + + def _signal_end(self) -> None: + while not self._deliveries.empty(): + self._deliveries.get_nowait() + self._deliveries.put_nowait(_END) + + +def _copy_position(position: StreamPosition) -> StreamPosition: + return StreamPosition(position.seq_num, position.timestamp) + + +def _session_closed_error() -> ReadSessionClosedError: + return ReadSessionClosedError("Read session ended before catching up") diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index 3e01b76..5ef388c 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -2,12 +2,14 @@ import logging import math import time -from typing import Any, AsyncIterable +from collections.abc import AsyncIterator +from typing import Any import s2_sdk._generated.s2.v1.s2_pb2 as pb from s2_sdk._client import HttpClient from s2_sdk._exceptions import ReadTimeoutError from s2_sdk._mappers import read_batch_from_proto, read_limit_params, read_start_params +from s2_sdk._read_session import _ReadSessionUpdate from s2_sdk._retrier import Attempt, compute_backoff, http_retry_on from s2_sdk._s2s import _stream_records_path from s2_sdk._s2s._protocol import parse_error_info, read_messages @@ -17,6 +19,7 @@ ReadLimit, Retry, SeqNum, + StreamPosition, TailOffset, Timestamp, metered_bytes, @@ -38,7 +41,7 @@ async def run_read_session( ignore_command_records: bool, retry: Retry, encryption_key: str | None = None, -) -> AsyncIterable[ReadBatch]: +) -> AsyncIterator[_ReadSessionUpdate]: params = _build_read_params(start, limit, until_timestamp, clamp_to_tail, wait) max_retries = retry._max_retries() min_base_delay = retry.min_base_delay.total_seconds() @@ -90,37 +93,29 @@ async def run_read_session( if batch.tail is not None: last_tail_at = time.monotonic() - if not batch.records: - continue - - last_record = batch.records[-1] - params["seq_num"] = last_record.seq_num + 1 - params.pop("timestamp", None) - params.pop("tail_offset", None) - - if remaining_count is not None: - remaining_count = max(remaining_count - len(batch.records), 0) - params["count"] = remaining_count - if remaining_bytes is not None: - remaining_bytes = max( - remaining_bytes - metered_bytes(batch.records), 0 - ) - params["bytes"] = remaining_bytes - - if ignore_command_records: - batch = ReadBatch( - records=[ - r for r in batch.records if not r.is_command_record() - ], - tail=batch.tail, - ) - if batch.records: - yield batch + last_record = batch.records[-1] + params["seq_num"] = last_record.seq_num + 1 + params.pop("timestamp", None) + params.pop("tail_offset", None) + + if remaining_count is not None: + remaining_count = max( + remaining_count - len(batch.records), 0 + ) + params["count"] = remaining_count + if remaining_bytes is not None: + remaining_bytes = max( + remaining_bytes - metered_bytes(batch.records), 0 + ) + params["bytes"] = remaining_bytes + + yield _read_session_update(batch, ignore_command_records) return except Exception as e: if attempt.value < max_retries and http_retry_on(e): + yield _ReadSessionUpdate() backoff = compute_backoff( attempt.value, min_base_delay=min_base_delay, @@ -144,6 +139,26 @@ async def run_read_session( raise e +def _caught_up_tail(batch: ReadBatch) -> StreamPosition | None: + if batch.tail is None: + return None + if not batch.records or batch.records[-1].seq_num + 1 == batch.tail.seq_num: + return batch.tail + return None + + +def _read_session_update( + batch: ReadBatch, ignore_command_records: bool +) -> _ReadSessionUpdate: + caught_up_tail = _caught_up_tail(batch) + if ignore_command_records: + batch = ReadBatch( + records=[r for r in batch.records if not r.is_command_record()], + tail=batch.tail, + ) + return _ReadSessionUpdate(batch, caught_up_tail) + + def _build_read_params( start: SeqNum | Timestamp | TailOffset, limit: ReadLimit | None, diff --git a/tests/test_correctness.py b/tests/test_correctness.py index 5293ca2..8f1f5d1 100644 --- a/tests/test_correctness.py +++ b/tests/test_correctness.py @@ -3,7 +3,16 @@ import pytest -from s2_sdk import Batching, Record, Retry, S2Stream, SeqNum, SequencedRecord +from s2_sdk import ( + AppendInput, + Batching, + Record, + Retry, + S2Stream, + SeqNum, + SequencedRecord, + TailOffset, +) _NUM_RECORDS = 1024 _RECORD_IDX_HEADER = b"record-idx" @@ -29,27 +38,28 @@ async def read_records() -> None: last_seq_num: int | None = None num_records_read = 0 - async for batch in stream.read_session(start=SeqNum(0), wait=60): - for record in batch.records: - seq_num = record.seq_num - if last_seq_num is None: - assert seq_num == 0 - else: - assert seq_num == last_seq_num + 1 - last_seq_num = seq_num - - record_idx = _record_idx(record) - assert 0 <= record_idx < _NUM_RECORDS - assert record_idx <= next_record_idx - - if record_idx == next_record_idx: - next_record_idx = record_idx + 1 - num_records_read += 1 - - if next_record_idx == _NUM_RECORDS: - assert last_seq_num + 1 == num_records_read - assert num_records_read >= _NUM_RECORDS - return + async with stream.read_session(start=SeqNum(0), wait=60) as session: + async for batch in session: + for record in batch.records: + seq_num = record.seq_num + if last_seq_num is None: + assert seq_num == 0 + else: + assert seq_num == last_seq_num + 1 + last_seq_num = seq_num + + record_idx = _record_idx(record) + assert 0 <= record_idx < _NUM_RECORDS + assert record_idx <= next_record_idx + + if record_idx == next_record_idx: + next_record_idx = record_idx + 1 + num_records_read += 1 + + if next_record_idx == _NUM_RECORDS: + assert last_seq_num + 1 == num_records_read + assert num_records_read >= _NUM_RECORDS + return pytest.fail( "read session ended before all records were read: " @@ -72,6 +82,28 @@ async def append_records() -> None: task_group.create_task(append_records()) +@pytest.mark.correctness +@pytest.mark.asyncio +async def test_read_session_reports_caught_up_after_delivering_tail(stream: S2Stream): + await stream.append( + AppendInput(records=[Record(body=b"first"), Record(body=b"second")]) + ) + + async with stream.read_session(start=TailOffset(2), wait=60) as session: + caught_up = session.caught_up() + records = [] + + assert not session.is_caught_up() + while not session.is_caught_up(): + records.extend((await anext(session)).records) + + tail = await caught_up + assert [record.seq_num for record in records] == [ + tail.seq_num - 2, + tail.seq_num - 1, + ] + + def _indexed_record(idx: int) -> Record: return Record( body=b"", diff --git a/tests/test_read_session.py b/tests/test_read_session.py new file mode 100644 index 0000000..0c37479 --- /dev/null +++ b/tests/test_read_session.py @@ -0,0 +1,204 @@ +import asyncio +from collections.abc import AsyncIterator + +import pytest + +from s2_sdk import ReadSessionClosedError, S2ClientError +from s2_sdk._read_session import ReadSession, _ReadSessionUpdate +from s2_sdk._s2s._read_session import _read_session_update +from s2_sdk._types import ReadBatch, SequencedRecord, StreamPosition + + +def _record(seq_num: int, *, command: bool = False) -> SequencedRecord: + return SequencedRecord( + seq_num=seq_num, + body=b"command" if command else b"record", + headers=[(b"", b"fence")] if command else [], + timestamp=1, + ) + + +async def _updates(*updates: _ReadSessionUpdate) -> AsyncIterator[_ReadSessionUpdate]: + for update in updates: + yield update + + +async def test_caught_up_follows_batch_delivery(): + tail = StreamPosition(2, 1) + batch = ReadBatch([_record(0), _record(1)], tail) + session = ReadSession(_updates(_ReadSessionUpdate(batch, tail))) + pending = session.caught_up() + + await asyncio.sleep(0) + assert not session.is_caught_up() + assert isinstance(pending, asyncio.Future) + assert not pending.done() + + assert await anext(session) == batch + assert session.is_caught_up() + assert await pending == tail + assert await session.caught_up() == tail + await session.close() + + +async def test_heartbeat_waits_for_queued_batch(): + tail = StreamPosition(2, 1) + batch = ReadBatch([_record(0), _record(1)]) + session = ReadSession( + _updates( + _ReadSessionUpdate(batch), + _ReadSessionUpdate(ReadBatch([], tail), tail), + ) + ) + pending = session.caught_up() + + await asyncio.sleep(0) + assert not session.is_caught_up() + assert await anext(session) == batch + assert not session.is_caught_up() + + assert await pending == tail + assert session.is_caught_up() + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_filtered_command_counts_toward_caught_up(): + tail = StreamPosition(2, 1) + update = _read_session_update( + ReadBatch([_record(0), _record(1, command=True)], tail), + ignore_command_records=True, + ) + session = ReadSession(_updates(update)) + pending = session.caught_up() + + batch = await anext(session) + assert [record.seq_num for record in batch.records] == [0] + assert session.is_caught_up() + assert await pending == tail + await session.close() + + +async def test_caught_up_wait_survives_reconnect(): + resume = asyncio.Event() + tail = StreamPosition(3, 1) + + async def reconnecting() -> AsyncIterator[_ReadSessionUpdate]: + yield _ReadSessionUpdate() + await resume.wait() + yield _ReadSessionUpdate(ReadBatch([], tail), tail) + + session = ReadSession(reconnecting()) + pending = session.caught_up() + await asyncio.sleep(0) + + assert not session.is_caught_up() + assert isinstance(pending, asyncio.Future) + assert not pending.done() + resume.set() + assert await pending == tail + assert session.is_caught_up() + await session.close() + + +async def test_caught_up_wait_fails_when_session_ends(): + session = ReadSession(_updates()) + + with pytest.raises(ReadSessionClosedError): + await session.caught_up() + with pytest.raises(StopAsyncIteration): + await anext(session) + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_close_stops_the_read(): + stopped = asyncio.Event() + + async def updates() -> AsyncIterator[_ReadSessionUpdate]: + try: + yield _ReadSessionUpdate(ReadBatch([_record(0)])) + await asyncio.Event().wait() + finally: + stopped.set() + + session = ReadSession(updates()) + async with session: + await asyncio.sleep(0) + + assert stopped.is_set() + + +async def test_close_before_read_starts(): + session = ReadSession(_updates()) + aiter(session) + + await session.close() + + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_cancelled_caught_up_wait_is_removed(): + resume = asyncio.Event() + tail = StreamPosition(1, 1) + + async def updates() -> AsyncIterator[_ReadSessionUpdate]: + await resume.wait() + yield _ReadSessionUpdate(ReadBatch([], tail), tail) + + session = ReadSession(updates()) + pending = session.caught_up() + assert isinstance(pending, asyncio.Future) + pending.cancel() + await asyncio.sleep(0) + + assert not session._waiters + next_caught_up = session.caught_up() + resume.set() + assert await next_caught_up == tail + await session.close() + + +async def test_close_propagates_caller_cancellation(): + close_started = asyncio.Event() + + class Updates: + def __aiter__(self) -> "Updates": + return self + + async def __anext__(self) -> _ReadSessionUpdate: + await asyncio.Event().wait() + raise StopAsyncIteration + + async def aclose(self) -> None: + close_started.set() + await asyncio.Event().wait() + + session = ReadSession(Updates()) + aiter(session) + await asyncio.sleep(0) + closing = asyncio.create_task(session.close()) + await close_started.wait() + closing.cancel() + + with pytest.raises(asyncio.CancelledError): + await closing + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_read_error_rejects_wait_and_iteration(): + async def updates() -> AsyncIterator[_ReadSessionUpdate]: + raise ValueError("read failed") + yield _ReadSessionUpdate() + + session = ReadSession(updates()) + caught_up = session.caught_up() + + with pytest.raises(S2ClientError, match="read failed") as wait_error: + await caught_up + with pytest.raises(S2ClientError, match="read failed") as read_error: + await anext(session) + + assert read_error.value is wait_error.value From 61585cb5ac4da53db4119ab5133cc4dfb23fb7be Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:52:14 -0700 Subject: [PATCH 2/5] misc. polishes --- docs/source/api-reference.md | 4 - src/s2_sdk/__init__.py | 2 - src/s2_sdk/_append_session.py | 2 + src/s2_sdk/_client.py | 6 +- src/s2_sdk/_exceptions.py | 5 - src/s2_sdk/_ops.py | 4 +- src/s2_sdk/_producer.py | 2 + src/s2_sdk/_read_session.py | 243 +++++++++++++++++-------------- src/s2_sdk/_s2s/_read_session.py | 45 ++---- tests/test_read_session.py | 168 ++++++++++----------- 10 files changed, 239 insertions(+), 242 deletions(-) diff --git a/docs/source/api-reference.md b/docs/source/api-reference.md index 8d5e92a..d45607d 100644 --- a/docs/source/api-reference.md +++ b/docs/source/api-reference.md @@ -194,8 +194,4 @@ .. autoclass:: ReadUnwrittenError() :members: :show-inheritance: - -.. autoclass:: ReadSessionClosedError() - :members: - :show-inheritance: ``` diff --git a/src/s2_sdk/__init__.py b/src/s2_sdk/__init__.py index 1af06a6..fde2544 100644 --- a/src/s2_sdk/__init__.py +++ b/src/s2_sdk/__init__.py @@ -6,7 +6,6 @@ from s2_sdk._exceptions import ( AppendConditionError, FencingTokenMismatchError, - ReadSessionClosedError, ReadUnwrittenError, S2ClientError, S2Error, @@ -137,5 +136,4 @@ "FencingTokenMismatchError", "SeqNumMismatchError", "ReadUnwrittenError", - "ReadSessionClosedError", ] diff --git a/src/s2_sdk/_append_session.py b/src/s2_sdk/_append_session.py index 75750bd..5c7094f 100644 --- a/src/s2_sdk/_append_session.py +++ b/src/s2_sdk/_append_session.py @@ -28,6 +28,8 @@ class AppendSession: Supports pipelining multiple :class:`AppendInput`\\ s while preserving submission order. + Use it as an async context manager, or call :meth:`close` explicitly to close the session. + Caution: Returned by :meth:`S2Stream.append_session`. Do not instantiate directly. """ diff --git a/src/s2_sdk/_client.py b/src/s2_sdk/_client.py index b9d9241..892f0a2 100644 --- a/src/s2_sdk/_client.py +++ b/src/s2_sdk/_client.py @@ -5,7 +5,7 @@ import ssl import time from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field from importlib.metadata import version from typing import Any, NamedTuple @@ -423,10 +423,8 @@ async def close(self) -> None: self._closed = True if self._reaper_task is not None: self._reaper_task.cancel() - try: + with suppress(asyncio.CancelledError): await self._reaper_task - except asyncio.CancelledError: - pass for conns in self._hosts.values(): for pc in conns: await pc.close() diff --git a/src/s2_sdk/_exceptions.py b/src/s2_sdk/_exceptions.py index b31bdcc..d95a8b0 100644 --- a/src/s2_sdk/_exceptions.py +++ b/src/s2_sdk/_exceptions.py @@ -16,10 +16,6 @@ class S2ClientError(S2Error): """Error originating from the client.""" -class ReadSessionClosedError(S2ClientError): - """Read session ended before becoming caught up.""" - - UNKNOWN_CODE = "unknown" @@ -130,7 +126,6 @@ def raise_for_416(body: dict, code: str) -> None: S2Error.__module__ = "s2_sdk" S2ClientError.__module__ = "s2_sdk" -ReadSessionClosedError.__module__ = "s2_sdk" S2ServerError.__module__ = "s2_sdk" AppendConditionError.__module__ = "s2_sdk" FencingTokenMismatchError.__module__ = "s2_sdk" diff --git a/src/s2_sdk/_ops.py b/src/s2_sdk/_ops.py index 82615c3..b2700f9 100644 --- a/src/s2_sdk/_ops.py +++ b/src/s2_sdk/_ops.py @@ -1216,10 +1216,10 @@ def read_session( until_timestamp, clamp_to_tail, wait, - ignore_command_records, retry=self._retry, encryption_key=self._encryption_key, - ) + ), + ignore_command_records=ignore_command_records, ) diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 6014ce5..986b6d6 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -33,6 +33,8 @@ class Producer: Handles batching into :class:`AppendInput` automatically and uses an append session internally. + Use it as an async context manager, or call :meth:`close` explicitly to close the producer. + Caution: Returned by :meth:`S2Stream.producer`. Do not instantiate directly. """ diff --git a/src/s2_sdk/_read_session.py b/src/s2_sdk/_read_session.py index e71cf2b..e114ee9 100644 --- a/src/s2_sdk/_read_session.py +++ b/src/s2_sdk/_read_session.py @@ -1,94 +1,106 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterable, AsyncIterator, Awaitable +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable +from contextlib import suppress from dataclasses import dataclass, field from typing import Self -from s2_sdk._exceptions import ReadSessionClosedError, normalize_exception +from s2_sdk._exceptions import S2ClientError, normalize_exception from s2_sdk._types import ReadBatch, StreamPosition -@dataclass(slots=True) -class _ReadSessionUpdate: - batch: ReadBatch | None = None - caught_up_tail: StreamPosition | None = None +class _ReadSessionRetrying: + pass -@dataclass(slots=True) -class _ReadDelivery: +@dataclass(frozen=True, slots=True) +class _ReadSessionHeartbeat: + tail: StreamPosition + + +@dataclass(frozen=True, slots=True) +class _ReadSessionBatch: batch: ReadBatch - caught_up_tail: StreamPosition | None - consumed: asyncio.Event = field(default_factory=asyncio.Event) -class _End: - pass +_ReadSessionEvent = _ReadSessionRetrying | _ReadSessionHeartbeat | _ReadSessionBatch -_END = _End() +@dataclass(slots=True) +class _ReadDelivery: + batch: ReadBatch + caught_up_tail: StreamPosition | None + delivered: asyncio.Event = field(default_factory=asyncio.Event) class ReadSession(AsyncIterator[ReadBatch]): - """A continuous read returned by :meth:`S2Stream.read_session`. + """Async iterator that yields :class:`ReadBatch` values with caught-up tracking. + + Use it as an async context manager, or call :meth:`close` explicitly to close the session. - Use as an async context manager or call :meth:`close`. + Caution: + Returned by :meth:`S2Stream.read_session`. Do not instantiate directly. """ __slots__ = ( "_caught_up_tail", "_closed", - "_deliveries", + "_delivery_queue", "_error", - "_error_raised", - "_is_caught_up", + "_ignore_command_records", "_task", - "_updates", - "_waiters", + "_events", + "_caught_up_futs", ) - def __init__(self, updates: AsyncIterable[_ReadSessionUpdate]) -> None: - self._updates = updates - self._deliveries: asyncio.Queue[_ReadDelivery | _End] = asyncio.Queue() + def __init__( + self, + events: AsyncGenerator[_ReadSessionEvent, None], + *, + ignore_command_records: bool = False, + ) -> None: + self._events = events + self._ignore_command_records = ignore_command_records + self._delivery_queue: asyncio.Queue[_ReadDelivery | BaseException | None] = ( + asyncio.Queue(maxsize=1) + ) self._task: asyncio.Task[None] | None = None self._closed = False self._error: BaseException | None = None - self._error_raised = False - self._is_caught_up = False self._caught_up_tail: StreamPosition | None = None - self._waiters: set[asyncio.Future[StreamPosition]] = set() + self._caught_up_futs: set[asyncio.Future[StreamPosition]] = set() def is_caught_up(self) -> bool: - """Return whether all records through the latest reported tail were delivered. + """Whether this session has yielded all records with sequence numbers less than + the sequence number of the last observed tail position. - A later batch that does not reach a reported tail or a reconnect resets it. - Filtered command records count toward progress. Use - :meth:`S2Stream.check_tail` for the current stream tail. + The stream's tail may have advanced since this session last observed it. + Use :meth:`S2Stream.check_tail` if you need the current tail. """ - return self._is_caught_up + return self._caught_up_tail is not None def caught_up(self) -> Awaitable[StreamPosition]: - """Return an awaitable for the next caught-up tail. + """Return an awaitable that resolves to a tail position once this session is caught up to it. - It resolves immediately if already caught up and remains pending across - retries. Keep consuming batches while waiting. Call again after falling - behind. + See :meth:`is_caught_up` for the semantics of caught up. - Raises: - ReadSessionClosedError: The session ended before catching up. - S2Error: The read failed before catching up. + This awaitable only signals that the session is caught up. To yield batches, continue + iterating over the session. """ loop = asyncio.get_running_loop() - waiter: asyncio.Future[StreamPosition] = loop.create_future() - if self._is_caught_up and self._caught_up_tail is not None: - waiter.set_result(_copy_position(self._caught_up_tail)) + caught_up_fut: asyncio.Future[StreamPosition] = loop.create_future() + if self._caught_up_tail is not None: + caught_up_fut.set_result(_copy_position(self._caught_up_tail)) elif self._closed: - waiter.set_exception(self._error or _session_closed_error()) + caught_up_fut.set_exception( + self._error if self._error is not None else _session_closed_error() + ) else: - self._waiters.add(waiter) - waiter.add_done_callback(self._waiters.discard) + self._caught_up_futs.add(caught_up_fut) + caught_up_fut.add_done_callback(self._caught_up_futs.discard) self._ensure_started() - return waiter + return caught_up_fut async def close(self) -> None: """Close the session.""" @@ -96,53 +108,43 @@ async def close(self) -> None: return if self._task is None: self._finish() - self._signal_end() return current = asyncio.current_task() cancellation_count = current.cancelling() if current is not None else 0 self._task.cancel() - try: + with suppress(asyncio.CancelledError): await self._task - except asyncio.CancelledError: - pass - finally: - if not self._closed: - self._finish() - self._signal_end() + if not self._closed: + self._finish() if current is not None and current.cancelling() > cancellation_count: raise asyncio.CancelledError - async def aclose(self) -> None: - """Close the session.""" - await self.close() - def __aiter__(self) -> Self: self._ensure_started() return self async def __anext__(self) -> ReadBatch: - self._ensure_started() + session = self.__aiter__() try: - item = await self._deliveries.get() + delivery = await session._delivery_queue.get() except asyncio.CancelledError: await self.close() raise - if isinstance(item, _End): - self._deliveries.put_nowait(_END) - if self._error is not None and not self._error_raised: - self._error_raised = True - raise self._error + if delivery is None: + self._delivery_queue.put_nowait(None) raise StopAsyncIteration + if isinstance(delivery, BaseException): + self._delivery_queue.put_nowait(None) + raise delivery - if item.caught_up_tail is not None: - self._mark_caught_up(item.caught_up_tail) - item.consumed.set() - return item.batch + if delivery.caught_up_tail is not None: + self._mark_caught_up(delivery.caught_up_tail) + delivery.delivered.set() + return delivery.batch async def __aenter__(self) -> Self: - self._ensure_started() - return self + return self.__aiter__() async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: await self.close() @@ -155,72 +157,89 @@ def _ensure_started(self) -> None: async def _run(self) -> None: error: BaseException | None = None cancelled = False - updates = aiter(self._updates) + events = self._events try: - async for update in updates: - self._mark_behind() - batch = update.batch - if batch is None or not batch.records: - if update.caught_up_tail is not None: - self._mark_caught_up(update.caught_up_tail) - continue - - delivery = _ReadDelivery(batch, update.caught_up_tail) - await self._deliveries.put(delivery) - await delivery.consumed.wait() + async for event in events: + await self._handle_event(event) except asyncio.CancelledError: cancelled = True except BaseException as exc: error = normalize_exception(exc) finally: - close = getattr(updates, "aclose", None) - if close is not None: - try: - await close() - except BaseException as exc: - if error is None and not cancelled: - error = normalize_exception(exc) + try: + await events.aclose() + except BaseException as exc: + if error is None and not cancelled: + error = normalize_exception(exc) self._finish(error) - self._signal_end() + + async def _handle_event(self, event: _ReadSessionEvent) -> None: + if isinstance(event, _ReadSessionRetrying): + self._mark_not_caught_up() + elif isinstance(event, _ReadSessionHeartbeat): + self._mark_caught_up(event.tail) + else: + batch = event.batch + caught_up_tail = ( + batch.tail + if batch.tail is not None + and batch.records[-1].seq_num + 1 == batch.tail.seq_num + else None + ) + + if self._ignore_command_records: + batch = ReadBatch( + records=[r for r in batch.records if not r.is_command_record()], + tail=batch.tail, + ) + + if batch.records: + self._mark_not_caught_up() + delivery = _ReadDelivery(batch, caught_up_tail) + await self._delivery_queue.put(delivery) + await delivery.delivered.wait() + elif caught_up_tail is not None: + self._mark_caught_up(caught_up_tail) + else: + self._mark_not_caught_up() def _mark_caught_up(self, tail: StreamPosition) -> None: if self._closed: return - self._is_caught_up = True self._caught_up_tail = _copy_position(tail) - waiters = tuple(self._waiters) - self._waiters.clear() - for waiter in waiters: - if not waiter.done(): - waiter.set_result(_copy_position(tail)) + caught_up_futs = tuple(self._caught_up_futs) + self._caught_up_futs.clear() + for caught_up_fut in caught_up_futs: + if not caught_up_fut.done(): + caught_up_fut.set_result(_copy_position(tail)) - def _mark_behind(self) -> None: + def _mark_not_caught_up(self) -> None: if not self._closed: - self._is_caught_up = False + self._caught_up_tail = None def _finish(self, error: BaseException | None = None) -> None: if self._closed: return self._closed = True self._error = error - if error is not None: - self._is_caught_up = False - waiters = tuple(self._waiters) - self._waiters.clear() - rejection = error or _session_closed_error() - for waiter in waiters: - if not waiter.done(): - waiter.set_exception(rejection) + if self._error is not None: + self._caught_up_tail = None + caught_up_futs = tuple(self._caught_up_futs) + self._caught_up_futs.clear() + for caught_up_fut in caught_up_futs: + if not caught_up_fut.done(): + caught_up_fut.set_exception( + self._error if self._error is not None else _session_closed_error() + ) - def _signal_end(self) -> None: - while not self._deliveries.empty(): - self._deliveries.get_nowait() - self._deliveries.put_nowait(_END) + while not self._delivery_queue.empty(): + self._delivery_queue.get_nowait() + self._delivery_queue.put_nowait(self._error) def _copy_position(position: StreamPosition) -> StreamPosition: return StreamPosition(position.seq_num, position.timestamp) -def _session_closed_error() -> ReadSessionClosedError: - return ReadSessionClosedError("Read session ended before catching up") +def _session_closed_error() -> S2ClientError: + return S2ClientError("ReadSession is closed") diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index 5ef388c..e9e21fb 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -2,24 +2,27 @@ import logging import math import time -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator from typing import Any import s2_sdk._generated.s2.v1.s2_pb2 as pb from s2_sdk._client import HttpClient -from s2_sdk._exceptions import ReadTimeoutError +from s2_sdk._exceptions import ReadTimeoutError, S2ClientError from s2_sdk._mappers import read_batch_from_proto, read_limit_params, read_start_params -from s2_sdk._read_session import _ReadSessionUpdate +from s2_sdk._read_session import ( + _ReadSessionBatch, + _ReadSessionEvent, + _ReadSessionHeartbeat, + _ReadSessionRetrying, +) from s2_sdk._retrier import Attempt, compute_backoff, http_retry_on from s2_sdk._s2s import _stream_records_path from s2_sdk._s2s._protocol import parse_error_info, read_messages from s2_sdk._types import ( _S2_ENCRYPTION_KEY_HEADER, - ReadBatch, ReadLimit, Retry, SeqNum, - StreamPosition, TailOffset, Timestamp, metered_bytes, @@ -38,10 +41,9 @@ async def run_read_session( until_timestamp: int | None, clamp_to_tail: bool, wait: int | None, - ignore_command_records: bool, retry: Retry, encryption_key: str | None = None, -) -> AsyncIterator[_ReadSessionUpdate]: +) -> AsyncGenerator[_ReadSessionEvent, None]: params = _build_read_params(start, limit, until_timestamp, clamp_to_tail, wait) max_retries = retry._max_retries() min_base_delay = retry.min_base_delay.total_seconds() @@ -110,12 +112,17 @@ async def run_read_session( ) params["bytes"] = remaining_bytes - yield _read_session_update(batch, ignore_command_records) + if batch.records: + yield _ReadSessionBatch(batch) + elif batch.tail is not None: + yield _ReadSessionHeartbeat(batch.tail) + else: + raise S2ClientError("Read session received an empty batch without a tail") return except Exception as e: if attempt.value < max_retries and http_retry_on(e): - yield _ReadSessionUpdate() + yield _ReadSessionRetrying() backoff = compute_backoff( attempt.value, min_base_delay=min_base_delay, @@ -139,26 +146,6 @@ async def run_read_session( raise e -def _caught_up_tail(batch: ReadBatch) -> StreamPosition | None: - if batch.tail is None: - return None - if not batch.records or batch.records[-1].seq_num + 1 == batch.tail.seq_num: - return batch.tail - return None - - -def _read_session_update( - batch: ReadBatch, ignore_command_records: bool -) -> _ReadSessionUpdate: - caught_up_tail = _caught_up_tail(batch) - if ignore_command_records: - batch = ReadBatch( - records=[r for r in batch.records if not r.is_command_record()], - tail=batch.tail, - ) - return _ReadSessionUpdate(batch, caught_up_tail) - - def _build_read_params( start: SeqNum | Timestamp | TailOffset, limit: ReadLimit | None, diff --git a/tests/test_read_session.py b/tests/test_read_session.py index 0c37479..1ea8754 100644 --- a/tests/test_read_session.py +++ b/tests/test_read_session.py @@ -1,11 +1,16 @@ import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator import pytest -from s2_sdk import ReadSessionClosedError, S2ClientError -from s2_sdk._read_session import ReadSession, _ReadSessionUpdate -from s2_sdk._s2s._read_session import _read_session_update +from s2_sdk import S2ClientError +from s2_sdk._read_session import ( + ReadSession, + _ReadSessionBatch, + _ReadSessionEvent, + _ReadSessionHeartbeat, + _ReadSessionRetrying, +) from s2_sdk._types import ReadBatch, SequencedRecord, StreamPosition @@ -18,93 +23,92 @@ def _record(seq_num: int, *, command: bool = False) -> SequencedRecord: ) -async def _updates(*updates: _ReadSessionUpdate) -> AsyncIterator[_ReadSessionUpdate]: - for update in updates: - yield update +async def _events( + *events: _ReadSessionEvent, +) -> AsyncGenerator[_ReadSessionEvent, None]: + for event in events: + yield event -async def test_caught_up_follows_batch_delivery(): +async def test_batch_with_tail_marks_caught_up(): tail = StreamPosition(2, 1) batch = ReadBatch([_record(0), _record(1)], tail) - session = ReadSession(_updates(_ReadSessionUpdate(batch, tail))) - pending = session.caught_up() + session = ReadSession(_events(_ReadSessionBatch(batch))) + caught_up = session.caught_up() await asyncio.sleep(0) assert not session.is_caught_up() - assert isinstance(pending, asyncio.Future) - assert not pending.done() assert await anext(session) == batch assert session.is_caught_up() - assert await pending == tail + assert await caught_up == tail assert await session.caught_up() == tail await session.close() -async def test_heartbeat_waits_for_queued_batch(): +async def test_heartbeat_marks_caught_up_after_batch_without_tail(): tail = StreamPosition(2, 1) batch = ReadBatch([_record(0), _record(1)]) session = ReadSession( - _updates( - _ReadSessionUpdate(batch), - _ReadSessionUpdate(ReadBatch([], tail), tail), + _events( + _ReadSessionBatch(batch), + _ReadSessionHeartbeat(tail), ) ) - pending = session.caught_up() + caught_up = session.caught_up() await asyncio.sleep(0) assert not session.is_caught_up() + assert await anext(session) == batch assert not session.is_caught_up() - - assert await pending == tail + assert await caught_up == tail assert session.is_caught_up() with pytest.raises(StopAsyncIteration): await anext(session) -async def test_filtered_command_counts_toward_caught_up(): +async def test_filtered_command_record_still_counts_toward_caught_up(): tail = StreamPosition(2, 1) - update = _read_session_update( - ReadBatch([_record(0), _record(1, command=True)], tail), - ignore_command_records=True, + batch = ReadBatch([_record(0), _record(1, command=True)], tail) + session = ReadSession( + _events(_ReadSessionBatch(batch)), ignore_command_records=True ) - session = ReadSession(_updates(update)) - pending = session.caught_up() + caught_up = session.caught_up() batch = await anext(session) assert [record.seq_num for record in batch.records] == [0] assert session.is_caught_up() - assert await pending == tail + assert await caught_up == tail await session.close() -async def test_caught_up_wait_survives_reconnect(): - resume = asyncio.Event() +async def test_caught_up_resolves_after_session_retry(): + allow_heartbeat = asyncio.Event() tail = StreamPosition(3, 1) - async def reconnecting() -> AsyncIterator[_ReadSessionUpdate]: - yield _ReadSessionUpdate() - await resume.wait() - yield _ReadSessionUpdate(ReadBatch([], tail), tail) + async def retry_then_heartbeat() -> AsyncGenerator[_ReadSessionEvent, None]: + yield _ReadSessionRetrying() + await allow_heartbeat.wait() + yield _ReadSessionHeartbeat(tail) - session = ReadSession(reconnecting()) - pending = session.caught_up() - await asyncio.sleep(0) + session = ReadSession(retry_then_heartbeat()) + caught_up = session.caught_up() + await asyncio.sleep(0) assert not session.is_caught_up() - assert isinstance(pending, asyncio.Future) - assert not pending.done() - resume.set() - assert await pending == tail + + allow_heartbeat.set() + + assert await caught_up == tail assert session.is_caught_up() await session.close() -async def test_caught_up_wait_fails_when_session_ends(): - session = ReadSession(_updates()) +async def test_caught_up_fails_when_session_ends(): + session = ReadSession(_events()) - with pytest.raises(ReadSessionClosedError): + with pytest.raises(S2ClientError, match="ReadSession is closed"): await session.caught_up() with pytest.raises(StopAsyncIteration): await anext(session) @@ -112,17 +116,17 @@ async def test_caught_up_wait_fails_when_session_ends(): await anext(session) -async def test_close_stops_the_read(): +async def test_context_exit_stops_read_events(): stopped = asyncio.Event() - async def updates() -> AsyncIterator[_ReadSessionUpdate]: + async def read_events_until_stopped() -> AsyncGenerator[_ReadSessionEvent, None]: try: - yield _ReadSessionUpdate(ReadBatch([_record(0)])) + yield _ReadSessionBatch(ReadBatch([_record(0)])) await asyncio.Event().wait() finally: stopped.set() - session = ReadSession(updates()) + session = ReadSession(read_events_until_stopped()) async with session: await asyncio.sleep(0) @@ -130,8 +134,7 @@ async def updates() -> AsyncIterator[_ReadSessionUpdate]: async def test_close_before_read_starts(): - session = ReadSession(_updates()) - aiter(session) + session = ReadSession(_events()) await session.close() @@ -139,66 +142,63 @@ async def test_close_before_read_starts(): await anext(session) -async def test_cancelled_caught_up_wait_is_removed(): - resume = asyncio.Event() +async def test_cancelled_caught_up_await_does_not_affect_other_awaits(): + allow_heartbeat = asyncio.Event() tail = StreamPosition(1, 1) - async def updates() -> AsyncIterator[_ReadSessionUpdate]: - await resume.wait() - yield _ReadSessionUpdate(ReadBatch([], tail), tail) + async def delayed_heartbeat() -> AsyncGenerator[_ReadSessionEvent, None]: + await allow_heartbeat.wait() + yield _ReadSessionHeartbeat(tail) - session = ReadSession(updates()) - pending = session.caught_up() - assert isinstance(pending, asyncio.Future) - pending.cancel() + session = ReadSession(delayed_heartbeat()) + caught_up_1 = asyncio.ensure_future(session.caught_up()) + caught_up_2 = asyncio.ensure_future(session.caught_up()) await asyncio.sleep(0) + caught_up_1.cancel() - assert not session._waiters - next_caught_up = session.caught_up() - resume.set() - assert await next_caught_up == tail - await session.close() + with pytest.raises(asyncio.CancelledError): + await caught_up_1 + allow_heartbeat.set() + assert await caught_up_2 == tail + await session.close() -async def test_close_propagates_caller_cancellation(): - close_started = asyncio.Event() - class Updates: - def __aiter__(self) -> "Updates": - return self +async def test_close_does_not_swallow_caller_cancellation(): + read_cleanup_started = asyncio.Event() - async def __anext__(self) -> _ReadSessionUpdate: + async def blocked_read_events() -> AsyncGenerator[_ReadSessionEvent, None]: + try: await asyncio.Event().wait() - raise StopAsyncIteration - - async def aclose(self) -> None: - close_started.set() + finally: + read_cleanup_started.set() await asyncio.Event().wait() + yield _ReadSessionRetrying() - session = ReadSession(Updates()) + session = ReadSession(blocked_read_events()) aiter(session) await asyncio.sleep(0) - closing = asyncio.create_task(session.close()) - await close_started.wait() - closing.cancel() + close_task = asyncio.create_task(session.close()) + await read_cleanup_started.wait() + close_task.cancel() with pytest.raises(asyncio.CancelledError): - await closing + await close_task with pytest.raises(StopAsyncIteration): await anext(session) -async def test_read_error_rejects_wait_and_iteration(): - async def updates() -> AsyncIterator[_ReadSessionUpdate]: +async def test_read_error_fails_caught_up_and_iteration(): + async def failing_read_events() -> AsyncGenerator[_ReadSessionEvent, None]: raise ValueError("read failed") - yield _ReadSessionUpdate() + yield _ReadSessionRetrying() - session = ReadSession(updates()) + session = ReadSession(failing_read_events()) caught_up = session.caught_up() - with pytest.raises(S2ClientError, match="read failed") as wait_error: + with pytest.raises(S2ClientError, match="read failed") as caught_up_error: await caught_up with pytest.raises(S2ClientError, match="read failed") as read_error: await anext(session) - assert read_error.value is wait_error.value + assert read_error.value is caught_up_error.value From 3fb96beb3af254f0ca28d751320ea1a17d7729fb Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:54:15 -0700 Subject: [PATCH 3/5] cq fix --- src/s2_sdk/_s2s/_read_session.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index e9e21fb..59c0d1a 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -117,7 +117,9 @@ async def run_read_session( elif batch.tail is not None: yield _ReadSessionHeartbeat(batch.tail) else: - raise S2ClientError("Read session received an empty batch without a tail") + raise S2ClientError( + "Read session received an empty batch without a tail" + ) return except Exception as e: From 2932142db990b2d909376455babf562cf966b57d Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:02:44 -0700 Subject: [PATCH 4/5] clarify docstring --- src/s2_sdk/_read_session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/s2_sdk/_read_session.py b/src/s2_sdk/_read_session.py index e114ee9..3dec706 100644 --- a/src/s2_sdk/_read_session.py +++ b/src/s2_sdk/_read_session.py @@ -35,7 +35,7 @@ class _ReadDelivery: class ReadSession(AsyncIterator[ReadBatch]): - """Async iterator that yields :class:`ReadBatch` values with caught-up tracking. + """Async iterator that yields :class:`ReadBatch` values and tracks caught-up state. Use it as an async context manager, or call :meth:`close` explicitly to close the session. From 9f9f0d1bd31e1478a7957e01b6972c4b88c4b71c Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:50:37 -0700 Subject: [PATCH 5/5] return -> get --- src/s2_sdk/_read_session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/s2_sdk/_read_session.py b/src/s2_sdk/_read_session.py index 3dec706..9396ae0 100644 --- a/src/s2_sdk/_read_session.py +++ b/src/s2_sdk/_read_session.py @@ -81,7 +81,7 @@ def is_caught_up(self) -> bool: return self._caught_up_tail is not None def caught_up(self) -> Awaitable[StreamPosition]: - """Return an awaitable that resolves to a tail position once this session is caught up to it. + """Get an awaitable that resolves to a tail position once this session is caught up to it. See :meth:`is_caught_up` for the semantics of caught up.