s2-sdk 0.1.0__py3-none-any.whl

s2_sdk/__init__.py

@@ -0,0 +1,129 @@
+import logging
+from importlib.metadata import version
+
+from s2_sdk._append_session import AppendSession, BatchSubmitTicket
+from s2_sdk._batching import append_inputs, append_record_batches
+from s2_sdk._exceptions import (
+    AppendConditionError,
+    FencingTokenMismatchError,
+    ReadUnwrittenError,
+    S2ClientError,
+    S2Error,
+    S2ServerError,
+    SeqNumMismatchError,
+)
+from s2_sdk._ops import S2, S2Basin, S2Stream
+from s2_sdk._producer import Producer, RecordSubmitTicket
+from s2_sdk._types import (
+    AccessTokenInfo,
+    AccessTokenScope,
+    AccountMetricSet,
+    Accumulation,
+    AppendAck,
+    AppendInput,
+    AppendRetryPolicy,
+    BasinConfig,
+    BasinInfo,
+    BasinMetricSet,
+    BasinScope,
+    Batching,
+    CommandRecord,
+    Compression,
+    Endpoints,
+    ExactMatch,
+    Gauge,
+    IndexedAppendAck,
+    Label,
+    MetricUnit,
+    Operation,
+    OperationGroupPermissions,
+    Page,
+    Permission,
+    PrefixMatch,
+    ReadBatch,
+    ReadLimit,
+    Record,
+    Retry,
+    Scalar,
+    SeqNum,
+    SequencedRecord,
+    StorageClass,
+    StreamConfig,
+    StreamInfo,
+    StreamMetricSet,
+    StreamPosition,
+    TailOffset,
+    Timeout,
+    TimeseriesInterval,
+    Timestamp,
+    Timestamping,
+    TimestampingMode,
+    metered_bytes,
+)
+
+__version__ = version("s2-sdk")
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+__all__ = [
+    "S2",
+    "S2Basin",
+    "S2Stream",
+    "AppendSession",
+    "BatchSubmitTicket",
+    "Producer",
+    "RecordSubmitTicket",
+    "append_record_batches",
+    "append_inputs",
+    "Retry",
+    "Timeout",
+    "Compression",
+    "AppendRetryPolicy",
+    "Batching",
+    "Endpoints",
+    "Record",
+    "AppendInput",
+    "AppendAck",
+    "IndexedAppendAck",
+    "StreamPosition",
+    "SeqNum",
+    "Timestamp",
+    "TailOffset",
+    "ReadBatch",
+    "ReadLimit",
+    "SequencedRecord",
+    "Page",
+    "CommandRecord",
+    "metered_bytes",
+    "StorageClass",
+    "TimestampingMode",
+    "Timestamping",
+    "StreamConfig",
+    "BasinConfig",
+    "BasinInfo",
+    "BasinScope",
+    "StreamInfo",
+    "ExactMatch",
+    "PrefixMatch",
+    "Permission",
+    "Operation",
+    "OperationGroupPermissions",
+    "AccessTokenScope",
+    "AccessTokenInfo",
+    "MetricUnit",
+    "TimeseriesInterval",
+    "AccountMetricSet",
+    "BasinMetricSet",
+    "StreamMetricSet",
+    "Scalar",
+    "Accumulation",
+    "Gauge",
+    "Label",
+    "S2Error",
+    "S2ClientError",
+    "S2ServerError",
+    "AppendConditionError",
+    "FencingTokenMismatchError",
+    "SeqNumMismatchError",
+    "ReadUnwrittenError",
+]

s2_sdk/_append_session.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import asyncio
+from collections import deque
+from typing import AsyncIterable, NamedTuple, Self
+
+from s2_sdk._client import HttpClient
+from s2_sdk._exceptions import S2ClientError
+from s2_sdk._s2s._append_session import run_append_session
+from s2_sdk._types import (
+    AppendAck,
+    AppendInput,
+    Compression,
+    Retry,
+    metered_bytes,
+)
+from s2_sdk._validators import validate_append_input
+
+
+class _UnackedBatch(NamedTuple):
+    ack_fut: asyncio.Future[AppendAck]
+    metered_bytes: int
+
+
+class AppendSession:
+    """Session for high-throughput appending with backpressure control.
+
+    Supports pipelining multiple :class:`AppendInput`\\ s while preserving
+    submission order.
+
+    Caution:
+        Returned by :meth:`S2Stream.append_session`. Do not instantiate directly.
+    """
+
+    __slots__ = (
+        "_closed",
+        "_client",
+        "_compression",
+        "_error",
+        "_permits",
+        "_queue",
+        "_retry",
+        "_stream_name",
+        "_task",
+        "_unacked",
+    )
+
+    def __init__(
+        self,
+        client: HttpClient,
+        stream_name: str,
+        retry: Retry,
+        compression: Compression,
+        max_unacked_bytes: int,
+        max_unacked_batches: int | None,
+    ) -> None:
+        self._client = client
+        self._stream_name = stream_name
+        self._retry = retry
+        self._compression = compression
+        self._permits = _AppendPermits(max_unacked_bytes, max_unacked_batches)
+
+        self._queue: asyncio.Queue[AppendInput | None] = asyncio.Queue()
+        self._unacked: deque[_UnackedBatch] = deque()
+        self._closed = False
+        self._error: BaseException | None = None
+
+        self._task = asyncio.get_running_loop().create_task(self._run())
+
+    async def submit(self, inp: AppendInput) -> BatchSubmitTicket:
+        """Submit a batch of records for appending.
+
+        Waits when backpressure limits are reached.
+        """
+        self._check_ready()
+        batch_bytes = metered_bytes(inp.records)
+        validate_append_input(len(inp.records), batch_bytes)
+
+        await self._permits.acquire(batch_bytes)
+        # Re-check after potentially waiting on backpressure.
+        try:
+            self._check_ready()
+        except BaseException:
+            self._permits.release(batch_bytes)
+            raise
+
+        ack_fut: asyncio.Future[AppendAck] = asyncio.get_running_loop().create_future()
+        self._unacked.append(_UnackedBatch(ack_fut, batch_bytes))
+        await self._queue.put(inp)
+        return BatchSubmitTicket(ack_fut)
+
+    def _check_ready(self) -> None:
+        if self._closed:
+            raise S2ClientError("AppendSession is closed")
+        if self._error is not None:
+            raise self._error
+
+    async def close(self) -> None:
+        """Close the session and wait for all submitted batches to be appended."""
+        if self._closed:
+            return
+        self._closed = True
+        await self._queue.put(None)
+        await self._task
+        if self._error is not None:
+            raise self._error
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
+        await self.close()
+        return False
+
+    async def _run(self) -> None:
+        try:
+            async for ack in run_append_session(
+                self._client,
+                self._stream_name,
+                self._input_iter(),
+                retry=self._retry,
+                compression=self._compression,
+                ack_timeout=self._client._request_timeout,
+            ):
+                self._resolve_next(ack)
+        except BaseException as e:
+            # Unwrap single-exception ExceptionGroups so callers see
+            # the original exception type (e.g. S2ServerError, SeqNumMismatchError).
+            exc = e
+            while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
+                exc = exc.exceptions[0]
+            self._fail_all(exc)
+
+    async def _input_iter(self) -> AsyncIterable[AppendInput]:
+        while True:
+            item = await self._queue.get()
+            if item is None:
+                return
+            yield item
+
+    def _resolve_next(self, ack: AppendAck) -> None:
+        unacked = self._unacked.popleft()
+        self._permits.release(unacked.metered_bytes)
+        unacked.ack_fut.set_result(ack)
+
+    def _fail_all(self, error: BaseException) -> None:
+        self._error = error
+        for unacked in self._unacked:
+            self._permits.release(unacked.metered_bytes)
+            if not unacked.ack_fut.done():
+                unacked.ack_fut.set_exception(error)
+        self._unacked.clear()
+        # Drain queue
+        while not self._queue.empty():
+            try:
+                self._queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+
+
+class BatchSubmitTicket:
+    """Awaitable that resolves to an :class:`AppendAck` once the batch is appended."""
+
+    __slots__ = ("_ack_fut",)
+
+    def __init__(self, ack_fut: asyncio.Future[AppendAck]) -> None:
+        self._ack_fut = ack_fut
+
+    def __await__(self):
+        return self._ack_fut.__await__()
+
+
+class _Semaphore:
+    __slots__ = ("_event", "_lock", "_value")
+
+    def __init__(self, value: int) -> None:
+        self._value = value
+        self._event = asyncio.Event()
+        self._event.set()
+        self._lock = asyncio.Lock()
+
+    async def acquire(self, n: int) -> None:
+        while True:
+            async with self._lock:
+                if self._value >= n:
+                    self._value -= n
+                    return
+                self._event.clear()
+            await self._event.wait()
+
+    def release(self, n: int) -> None:
+        self._value += n
+        self._event.set()
+
+
+class _AppendPermits:
+    __slots__ = ("_bytes", "_count")
+
+    def __init__(self, max_bytes: int, max_count: int | None = None) -> None:
+        self._bytes = _Semaphore(max_bytes)
+        self._count = _Semaphore(max_count) if max_count is not None else None
+
+    async def acquire(self, n_bytes: int) -> None:
+        if self._count is not None:
+            await self._count.acquire(1)
+        try:
+            await self._bytes.acquire(n_bytes)
+        except BaseException:
+            if self._count is not None:
+                self._count.release(1)
+            raise
+
+    def release(self, n_bytes: int) -> None:
+        self._bytes.release(n_bytes)
+        if self._count is not None:
+            self._count.release(1)

s2_sdk/_batching.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncIterable
+
+from s2_sdk._types import AppendInput, Batching, Record, metered_bytes
+from s2_sdk._validators import validate_batching
+
+
+class BatchAccumulator:
+    __slots__ = ("_batching", "_bytes", "_records")
+
+    def __init__(self, batching: Batching) -> None:
+        self._batching = batching
+        self._records: list[Record] = []
+        self._bytes = 0
+
+    def add(self, record: Record) -> None:
+        self._records.append(record)
+        self._bytes += metered_bytes((record,))
+
+    def take(self) -> list[Record]:
+        records = list(self._records)
+        self._records.clear()
+        self._bytes = 0
+        return records
+
+    def is_full(self) -> bool:
+        return (
+            len(self._records) >= self._batching.max_records
+            or self._bytes >= self._batching.max_bytes
+        )
+
+    def is_empty(self) -> bool:
+        return len(self._records) == 0
+
+    @property
+    def linger(self) -> float:
+        return self._batching.linger.total_seconds()
+
+
+async def append_record_batches(
+    records: AsyncIterable[Record],
+    *,
+    batching: Batching | None = None,
+) -> AsyncIterable[list[Record]]:
+    """Group records into batches based on count, bytes, and linger time."""
+    if batching is None:
+        batching = Batching()
+    validate_batching(batching.max_records, batching.max_bytes)
+    acc = BatchAccumulator(batching)
+    linger_secs = batching.linger.total_seconds()
+    aiter = records.__aiter__()
+
+    while True:
+        try:
+            record = await anext(aiter)
+        except StopAsyncIteration:
+            break
+
+        acc.add(record)
+        if acc.is_full():
+            yield acc.take()
+            continue
+
+        try:
+            while not acc.is_full():
+                if linger_secs > 0:
+                    record = await asyncio.wait_for(anext(aiter), timeout=linger_secs)
+                else:
+                    record = await anext(aiter)
+                acc.add(record)
+        except StopAsyncIteration:
+            pass
+        except TimeoutError:
+            pass
+
+        if not acc.is_empty():
+            yield acc.take()
+
+
+async def append_inputs(
+    records: AsyncIterable[Record],
+    *,
+    match_seq_num: int | None = None,
+    fencing_token: str | None = None,
+    batching: Batching | None = None,
+) -> AsyncIterable[AppendInput]:
+    """Group records into :class:`AppendInput` batches based on count, bytes, and linger time.
+
+    If ``match_seq_num`` is set, it applies to the first input and is auto-incremented for subsequent ones.
+    """
+    if batching is None:
+        batching = Batching()
+    async for batch in append_record_batches(records, batching=batching):
+        if not batch:
+            continue
+        append_input = AppendInput(
+            records=batch,
+            match_seq_num=match_seq_num,
+            fencing_token=fencing_token,
+        )
+        if match_seq_num is not None:
+            match_seq_num += len(batch)
+        yield append_input