From 16f26e7bddecba1f298174592733e9c492fec0b7 Mon Sep 17 00:00:00 2001 From: Mehul Arora Date: Sat, 12 Sep 2026 20:04:55 -0400 Subject: [PATCH 1/2] feat: `stream_config` on append, read, and sessions sent as `s2-stream-config` header --- examples/docs/streams.py | 20 +++ src/s2_sdk/_append_session.py | 5 + src/s2_sdk/_mappers.py | 5 + src/s2_sdk/_ops.py | 42 ++++++- src/s2_sdk/_producer.py | 3 + src/s2_sdk/_s2s/_append_session.py | 13 +- src/s2_sdk/_s2s/_read_session.py | 12 +- src/s2_sdk/_types.py | 8 ++ tests/test_stream_config_header.py | 196 +++++++++++++++++++++++++++++ 9 files changed, 297 insertions(+), 7 deletions(-) create mode 100644 tests/test_stream_config_header.py diff --git a/examples/docs/streams.py b/examples/docs/streams.py index ea0f5a1..9006a33 100644 --- a/examples/docs/streams.py +++ b/examples/docs/streams.py @@ -18,6 +18,7 @@ ReadLimit, Record, SeqNum, + StreamConfig, TailOffset, Timestamp, ) @@ -112,6 +113,25 @@ async def check_tail_example(stream): # ANCHOR_END: check-tail +async def auto_create_config_example(stream): + # ANCHOR: auto-create-config + # Applied only if the stream is created by this call; ignored if it exists. + # Unset fields inherit the basin's default stream configuration. + stream_config = StreamConfig(retention_policy=3600) + + await stream.append( + AppendInput(records=[Record(body=b"hello")], stream_config=stream_config) + ) + await stream.read(start=SeqNum(0), stream_config=stream_config) + + # Sessions send the config each time they connect. + async with stream.append_session(stream_config=stream_config) as session: + await session.submit(AppendInput(records=[Record(body=b"hello")])) + async with stream.read_session(start=SeqNum(0), stream_config=stream_config): + pass + # ANCHOR_END: auto-create-config + + async def read_session_example(stream): # ANCHOR: read-session async with stream.read_session(start=SeqNum(0)) as session: diff --git a/src/s2_sdk/_append_session.py b/src/s2_sdk/_append_session.py index 6a4a54f..d3fafbe 100644 --- a/src/s2_sdk/_append_session.py +++ b/src/s2_sdk/_append_session.py @@ -17,6 +17,7 @@ AppendInput, Compression, Retry, + StreamConfig, metered_bytes, ) from s2_sdk._validators import validate_append_input @@ -48,6 +49,7 @@ class AppendSession: "_permits", "_input_queue", "_retry", + "_stream_config", "_stream_name", "_task", "_unacked", @@ -62,12 +64,14 @@ def __init__( max_unacked_bytes: int, max_unacked_batches: int | None, encryption_key: str | None = None, + stream_config: StreamConfig | None = None, ) -> None: self._client = client self._stream_name = stream_name self._retry = retry self._compression = compression self._encryption_key = encryption_key + self._stream_config = stream_config self._permits = _AppendPermits(max_unacked_bytes, max_unacked_batches) self._input_queue: asyncio.Queue[AppendInput | None] = asyncio.Queue() @@ -134,6 +138,7 @@ async def _run(self) -> None: compression=self._compression, ack_timeout=self._client._request_timeout, encryption_key=self._encryption_key, + stream_config=self._stream_config, ): unacked = self._unacked.popleft() self._permits.release(unacked.metered_bytes) diff --git a/src/s2_sdk/_mappers.py b/src/s2_sdk/_mappers.py index c43b958..c9f6ed3 100644 --- a/src/s2_sdk/_mappers.py +++ b/src/s2_sdk/_mappers.py @@ -1,3 +1,4 @@ +import json from collections.abc import Iterable from datetime import datetime from typing import Any, Literal @@ -114,6 +115,10 @@ def stream_reconfiguration_to_json(config: StreamConfig) -> dict[str, Any]: return stream_config_to_json(config) or {} +def stream_config_header(config: StreamConfig) -> str: + return json.dumps(stream_config_to_json(config) or {}) + + def stream_config_from_json(data: dict[str, Any]) -> StreamConfig: retention_policy: int | Literal["infinite"] | None = None rp = data.get("retention_policy") diff --git a/src/s2_sdk/_ops.py b/src/s2_sdk/_ops.py index c28f40e..fe585f1 100644 --- a/src/s2_sdk/_ops.py +++ b/src/s2_sdk/_ops.py @@ -27,6 +27,7 @@ read_limit_params, read_start_params, stream_config_from_json, + stream_config_header, stream_config_to_json, stream_info_from_json, stream_reconfiguration_to_json, @@ -38,6 +39,7 @@ from s2_sdk._s2s._read_session import run_read_session from s2_sdk._types import ( _S2_ENCRYPTION_KEY_HEADER, + _S2_STREAM_CONFIG_HEADER, ONE_MIB, Compression, Endpoints, @@ -985,12 +987,17 @@ def name(self) -> str: return self._name def _request_headers( - self, headers: dict[str, str] | None = None + self, + headers: dict[str, str] | None = None, + stream_config: types.StreamConfig | None = None, ) -> dict[str, str] | None: - if self._encryption_key is None: + if self._encryption_key is None and stream_config is None: return headers merged = dict(headers or {}) - merged[_S2_ENCRYPTION_KEY_HEADER] = self._encryption_key + if self._encryption_key is not None: + merged[_S2_ENCRYPTION_KEY_HEADER] = self._encryption_key + if stream_config is not None: + merged[_S2_STREAM_CONFIG_HEADER] = stream_config_header(stream_config) return merged @fallible @@ -1031,7 +1038,8 @@ async def append(self, inp: types.AppendInput) -> types.AppendAck: { "content-type": "application/x-protobuf", "accept": "application/x-protobuf", - } + }, + inp.stream_config, ), ) ack = pb.AppendAck() @@ -1044,6 +1052,7 @@ def append_session( *, max_unacked_bytes: int = 5 * ONE_MIB, max_unacked_batches: int | None = None, + stream_config: types.StreamConfig | None = None, ) -> AppendSession: """Open a session for appending batches of records continuously. @@ -1054,6 +1063,10 @@ def append_session( batches before backpressure is applied. Default is 5 MiB. max_unacked_batches: Maximum number of unacknowledged batches before backpressure is applied. If ``None``, no limit is applied. + stream_config: Configuration to apply if the stream is created by + the session. Unset fields inherit the basin's default stream + configuration. Ignored if the stream already exists. Sent + whenever the session connects. Returns: An :class:`AppendSession` to use as an async context manager. @@ -1078,6 +1091,7 @@ def append_session( max_unacked_bytes=max_unacked_bytes, max_unacked_batches=max_unacked_batches, encryption_key=self._encryption_key, + stream_config=stream_config, ) @fallible @@ -1088,6 +1102,7 @@ def producer( match_seq_num: int | None = None, batching: types.Batching | None = None, max_unacked_bytes: int = 5 * ONE_MIB, + stream_config: types.StreamConfig | None = None, ) -> Producer: """Open a producer with per-record submit and auto-batching. @@ -1099,6 +1114,10 @@ def producer( values are used. See :class:`Batching`. max_unacked_bytes: Maximum total metered bytes of unacknowledged batches before backpressure is applied. Default is 5 MiB. + stream_config: Configuration to apply if the stream is created by + the producer. Unset fields inherit the basin's default stream + configuration. Ignored if the stream already exists. Sent + whenever the underlying session connects. Returns: A :class:`Producer` to use as an async context manager. @@ -1128,6 +1147,7 @@ def producer( match_seq_num=match_seq_num, max_unacked_bytes=max_unacked_bytes, batching=batching, + stream_config=stream_config, ) @fallible @@ -1140,6 +1160,7 @@ async def read( clamp_to_tail: bool = False, wait: int | None = None, ignore_command_records: bool = False, + stream_config: types.StreamConfig | None = None, ) -> types.ReadBatch: """Read a batch of records from a stream. @@ -1152,6 +1173,9 @@ async def read( exceeds the tail, instead of raising. wait: Number of seconds to wait for records before returning. ignore_command_records: Filter out command records from the batch. + stream_config: Configuration to apply if the stream is created on + read. Unset fields inherit the basin's default stream + configuration. Ignored if the stream already exists. Returns: A :class:`ReadBatch` containing sequenced records and an optional @@ -1173,7 +1197,9 @@ async def read( "GET", _stream_path(self.name, "/records"), params=params, - headers=self._request_headers({"accept": "application/x-protobuf"}), + headers=self._request_headers( + {"accept": "application/x-protobuf"}, stream_config + ), ) proto_batch = pb.ReadBatch() @@ -1196,6 +1222,7 @@ def read_session( clamp_to_tail: bool = False, wait: int | None = None, ignore_command_records: bool = False, + stream_config: types.StreamConfig | None = None, ) -> ReadSession: """Read batches of records from a stream continuously. @@ -1210,6 +1237,10 @@ def read_session( wait: Number of seconds to wait for new records when the tail is reached. ignore_command_records: Filter out command records from batches. + stream_config: Configuration to apply if the stream is created by + the session. Unset fields inherit the basin's default stream + configuration. Ignored if the stream already exists. Sent + whenever the session connects. Returns: A :class:`ReadSession` that yields batches of records. @@ -1232,6 +1263,7 @@ def read_session( wait, retry=self._retry, encryption_key=self._encryption_key, + stream_config=stream_config, ), ignore_command_records=ignore_command_records, ) diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 7a4d747..0cc32e5 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -24,6 +24,7 @@ IndexedAppendAck, Record, Retry, + StreamConfig, ) @@ -74,6 +75,7 @@ def __init__( max_unacked_bytes: int, batching: Batching, encryption_key: str | None = None, + stream_config: StreamConfig | None = None, ) -> None: self._session = AppendSession( client=client, @@ -83,6 +85,7 @@ def __init__( max_unacked_bytes=max_unacked_bytes, max_unacked_batches=None, encryption_key=encryption_key, + stream_config=stream_config, ) self._fencing_token = fencing_token self._match_seq_num = match_seq_num diff --git a/src/s2_sdk/_s2s/_append_session.py b/src/s2_sdk/_s2s/_append_session.py index fdef2a2..1337d0e 100644 --- a/src/s2_sdk/_s2s/_append_session.py +++ b/src/s2_sdk/_s2s/_append_session.py @@ -11,7 +11,11 @@ from s2_sdk._client import HttpClient from s2_sdk._exceptions import ReadTimeoutError, S2ClientError from s2_sdk._frame_signal import FrameSignal -from s2_sdk._mappers import append_ack_from_proto, append_input_to_proto +from s2_sdk._mappers import ( + append_ack_from_proto, + append_input_to_proto, + stream_config_header, +) from s2_sdk._retrier import ( AdvisedReconnectLimiter, Attempt, @@ -29,11 +33,13 @@ ) from s2_sdk._types import ( _S2_ENCRYPTION_KEY_HEADER, + _S2_STREAM_CONFIG_HEADER, AppendAck, AppendInput, AppendRetryPolicy, Compression, Retry, + StreamConfig, ) logger = logging.getLogger(__name__) @@ -72,6 +78,7 @@ async def run_append_session( compression: Compression, ack_timeout: float, encryption_key: str | None = None, + stream_config: StreamConfig | None = None, ) -> AsyncIterable[AppendAck]: input_queue: asyncio.Queue[AppendInput | None] = asyncio.Queue( maxsize=_QUEUE_MAX_SIZE @@ -112,6 +119,7 @@ async def retrying_inner(): ack_timeout, reconnect_limiter, encryption_key, + stream_config, ) if ( outcome is _AttemptOutcome.RECONNECT_ADVISED @@ -183,6 +191,7 @@ async def _run_attempt( ack_timeout: float, reconnect_limiter: AdvisedReconnectLimiter, encryption_key: str | None = None, + stream_config: StreamConfig | None = None, ) -> _AttemptOutcome: inflight_inputs = session_state.inflight_inputs headers = { @@ -191,6 +200,8 @@ async def _run_attempt( } if encryption_key is not None: headers[_S2_ENCRYPTION_KEY_HEADER] = encryption_key + if stream_config is not None: + headers[_S2_STREAM_CONFIG_HEADER] = stream_config_header(stream_config) ack_deadline_armed = asyncio.Event() advised_reconnect = asyncio.Event() diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index cac6384..766e9d1 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -8,7 +8,12 @@ import s2_sdk._generated.s2.v1.s2_pb2 as pb from s2_sdk._client import HttpClient 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._mappers import ( + read_batch_from_proto, + read_limit_params, + read_start_params, + stream_config_header, +) from s2_sdk._read_session import ( _ReadSessionBatch, _ReadSessionEvent, @@ -26,9 +31,11 @@ from s2_sdk._s2s._protocol import parse_error_info, read_messages from s2_sdk._types import ( _S2_ENCRYPTION_KEY_HEADER, + _S2_STREAM_CONFIG_HEADER, ReadLimit, Retry, SeqNum, + StreamConfig, TailOffset, Timestamp, metered_bytes, @@ -49,6 +56,7 @@ async def run_read_session( wait: int | None, retry: Retry, encryption_key: str | None = None, + stream_config: StreamConfig | None = None, ) -> AsyncGenerator[_ReadSessionEvent, None]: params = _build_read_params(start, limit, until_timestamp, clamp_to_tail, wait) max_retries = retry._max_retries() @@ -65,6 +73,8 @@ async def run_read_session( headers = {"content-type": "s2s/proto"} if encryption_key is not None: headers[_S2_ENCRYPTION_KEY_HEADER] = encryption_key + if stream_config is not None: + headers[_S2_STREAM_CONFIG_HEADER] = stream_config_header(stream_config) while True: if wait is not None: diff --git a/src/s2_sdk/_types.py b/src/s2_sdk/_types.py index 3baa75f..b81de14 100644 --- a/src/s2_sdk/_types.py +++ b/src/s2_sdk/_types.py @@ -12,6 +12,7 @@ ONE_MIB = 1024 * 1024 _S2_ENCRYPTION_KEY_HEADER = "s2-encryption-key" +_S2_STREAM_CONFIG_HEADER = "s2-stream-config" def _parse_scheme(url: str) -> str: @@ -166,6 +167,13 @@ class AppendInput: """Fencing token to match against the stream's current fencing token. If unset, no matching is performed. If set and mismatched, the append fails.""" + stream_config: StreamConfig | None = None + """Configuration to apply if the stream is created on append. Unset fields inherit the + basin's default stream configuration. Ignored if the stream already exists. + + Only used by :meth:`~S2Stream.append`. Sessions take ``stream_config`` when opened instead; + see :meth:`~S2Stream.append_session`.""" + @dataclass(slots=True) class StreamPosition: diff --git a/tests/test_stream_config_header.py b/tests/test_stream_config_header.py new file mode 100644 index 0000000..03fb1b5 --- /dev/null +++ b/tests/test_stream_config_header.py @@ -0,0 +1,196 @@ +import json +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import Any, cast +from unittest.mock import patch + +import s2_sdk._generated.s2.v1.s2_pb2 as pb +import s2_sdk._s2s._append_session as s2s_append_session +import s2_sdk._s2s._read_session as s2s_read_session +from s2_sdk._client import HttpClient +from s2_sdk._mappers import stream_config_header +from s2_sdk._ops import S2Stream +from s2_sdk._s2s._protocol import Message +from s2_sdk._types import ( + AppendInput, + Compression, + ReadLimit, + Record, + Retry, + SeqNum, + StorageClass, + StreamConfig, + Timestamping, + TimestampingMode, +) + +_CONFIG = StreamConfig(retention_policy=3600, delete_on_empty_min_age=300) +_HEADER = "s2-stream-config" + + +def _ack() -> pb.AppendAck: + return pb.AppendAck( + start=pb.StreamPosition(seq_num=0, timestamp=1), + end=pb.StreamPosition(seq_num=1, timestamp=1), + tail=pb.StreamPosition(seq_num=1, timestamp=1), + ) + + +def _batch() -> pb.ReadBatch: + return pb.ReadBatch(records=[pb.SequencedRecord(seq_num=0, timestamp=1, body=b"a")]) + + +def _response_body(method: str) -> bytes: + proto = _ack() if method == "POST" else _batch() + return proto.SerializeToString() + + +@dataclass(slots=True) +class _Response: + content: bytes + status_code: int = 200 + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + message = Message( + self.content, + terminal=False, + compression=Compression.NONE, + reconnect_advised=False, + ) + yield cast(bytes, message) + + def retire_connection(self) -> None: + pass + + +@dataclass(slots=True) +class _Client: + headers: list[dict[str, str] | None] = field(default_factory=list) + _request_timeout: float = 1.0 + + async def unary_request(self, method: str, path: str, **kwargs: Any) -> _Response: + self.headers.append(kwargs.get("headers")) + return _Response(_response_body(method)) + + @asynccontextmanager + async def streaming_request( + self, + method: str, + path: str, + *, + headers: dict[str, str] | None = None, + content: AsyncGenerator[bytes, None] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[_Response, None]: + self.headers.append(headers) + if content is not None: + async for _ in content: + pass + yield _Response(_response_body(method)) + + +def _decoded_messages(stream: AsyncIterator[bytes]) -> AsyncIterator[Message]: + return cast(AsyncIterator[Message], stream) + + +def _stream(client: _Client, encryption_key: str | None = None) -> S2Stream: + return S2Stream( + "stream", + cast(HttpClient, client), + retry=Retry(max_attempts=1), + compression=Compression.NONE, + encryption_key=encryption_key, + ) + + +def test_stream_config_header_uses_api_field_names() -> None: + config = StreamConfig( + storage_class=StorageClass.EXPRESS, + retention_policy="infinite", + timestamping=Timestamping(mode=TimestampingMode.CLIENT_REQUIRE, uncapped=True), + delete_on_empty_min_age=300, + ) + assert json.loads(stream_config_header(config)) == { + "storage_class": "express", + "retention_policy": {"infinite": {}}, + "timestamping": {"mode": "client-require", "uncapped": True}, + "delete_on_empty": {"min_age_secs": 300}, + } + assert stream_config_header(StreamConfig()) == "{}" + + +async def test_unary_append_and_read_send_header() -> None: + client = _Client() + stream = _stream(client, encryption_key="key") + + await stream.append(AppendInput(records=[Record(body=b"a")], stream_config=_CONFIG)) + await stream.read(start=SeqNum(0), stream_config=_CONFIG) + + assert len(client.headers) == 2 + for headers in client.headers: + assert headers is not None + assert headers["s2-encryption-key"] == "key" + assert json.loads(headers[_HEADER]) == { + "retention_policy": {"age": 3600}, + "delete_on_empty": {"min_age_secs": 300}, + } + + +async def test_unary_append_and_read_omit_header_when_unset() -> None: + client = _Client() + stream = _stream(client) + + await stream.append(AppendInput(records=[Record(body=b"a")])) + await stream.read(start=SeqNum(0)) + + assert all(_HEADER not in (headers or {}) for headers in client.headers) + + +async def test_read_session_sends_header() -> None: + client = _Client() + + with patch.object(s2s_read_session, "read_messages", new=_decoded_messages): + async for _ in s2s_read_session.run_read_session( + cast(HttpClient, client), + "stream", + SeqNum(0), + ReadLimit(count=1), + until_timestamp=None, + clamp_to_tail=False, + wait=None, + retry=Retry(max_attempts=1), + stream_config=_CONFIG, + ): + pass + + assert client.headers == [ + {"content-type": "s2s/proto", _HEADER: stream_config_header(_CONFIG)} + ] + + +async def test_append_session_sends_header() -> None: + client = _Client() + + async def inputs() -> AsyncIterator[AppendInput]: + yield AppendInput(records=[Record(body=b"a")]) + + with patch.object(s2s_append_session, "read_messages", new=_decoded_messages): + async for _ in s2s_append_session.run_append_session( + cast(HttpClient, client), + "stream", + inputs(), + Retry(max_attempts=1), + Compression.NONE, + ack_timeout=1.0, + stream_config=_CONFIG, + ): + pass + + assert client.headers == [ + { + "content-type": "s2s/proto", + "accept": "s2s/proto", + _HEADER: stream_config_header(_CONFIG), + } + ] From cce763f7a8dff160c6a8431266e1923087e62c49 Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:21:33 -0700 Subject: [PATCH 2/2] integration test instead of mock --- tests/test_stream_config_header.py | 196 ----------------------------- tests/test_stream_ops.py | 112 +++++++++++++++++ 2 files changed, 112 insertions(+), 196 deletions(-) delete mode 100644 tests/test_stream_config_header.py diff --git a/tests/test_stream_config_header.py b/tests/test_stream_config_header.py deleted file mode 100644 index 03fb1b5..0000000 --- a/tests/test_stream_config_header.py +++ /dev/null @@ -1,196 +0,0 @@ -import json -from collections.abc import AsyncGenerator, AsyncIterator -from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from typing import Any, cast -from unittest.mock import patch - -import s2_sdk._generated.s2.v1.s2_pb2 as pb -import s2_sdk._s2s._append_session as s2s_append_session -import s2_sdk._s2s._read_session as s2s_read_session -from s2_sdk._client import HttpClient -from s2_sdk._mappers import stream_config_header -from s2_sdk._ops import S2Stream -from s2_sdk._s2s._protocol import Message -from s2_sdk._types import ( - AppendInput, - Compression, - ReadLimit, - Record, - Retry, - SeqNum, - StorageClass, - StreamConfig, - Timestamping, - TimestampingMode, -) - -_CONFIG = StreamConfig(retention_policy=3600, delete_on_empty_min_age=300) -_HEADER = "s2-stream-config" - - -def _ack() -> pb.AppendAck: - return pb.AppendAck( - start=pb.StreamPosition(seq_num=0, timestamp=1), - end=pb.StreamPosition(seq_num=1, timestamp=1), - tail=pb.StreamPosition(seq_num=1, timestamp=1), - ) - - -def _batch() -> pb.ReadBatch: - return pb.ReadBatch(records=[pb.SequencedRecord(seq_num=0, timestamp=1, body=b"a")]) - - -def _response_body(method: str) -> bytes: - proto = _ack() if method == "POST" else _batch() - return proto.SerializeToString() - - -@dataclass(slots=True) -class _Response: - content: bytes - status_code: int = 200 - - async def aiter_bytes(self) -> AsyncIterator[bytes]: - message = Message( - self.content, - terminal=False, - compression=Compression.NONE, - reconnect_advised=False, - ) - yield cast(bytes, message) - - def retire_connection(self) -> None: - pass - - -@dataclass(slots=True) -class _Client: - headers: list[dict[str, str] | None] = field(default_factory=list) - _request_timeout: float = 1.0 - - async def unary_request(self, method: str, path: str, **kwargs: Any) -> _Response: - self.headers.append(kwargs.get("headers")) - return _Response(_response_body(method)) - - @asynccontextmanager - async def streaming_request( - self, - method: str, - path: str, - *, - headers: dict[str, str] | None = None, - content: AsyncGenerator[bytes, None] | None = None, - **kwargs: Any, - ) -> AsyncGenerator[_Response, None]: - self.headers.append(headers) - if content is not None: - async for _ in content: - pass - yield _Response(_response_body(method)) - - -def _decoded_messages(stream: AsyncIterator[bytes]) -> AsyncIterator[Message]: - return cast(AsyncIterator[Message], stream) - - -def _stream(client: _Client, encryption_key: str | None = None) -> S2Stream: - return S2Stream( - "stream", - cast(HttpClient, client), - retry=Retry(max_attempts=1), - compression=Compression.NONE, - encryption_key=encryption_key, - ) - - -def test_stream_config_header_uses_api_field_names() -> None: - config = StreamConfig( - storage_class=StorageClass.EXPRESS, - retention_policy="infinite", - timestamping=Timestamping(mode=TimestampingMode.CLIENT_REQUIRE, uncapped=True), - delete_on_empty_min_age=300, - ) - assert json.loads(stream_config_header(config)) == { - "storage_class": "express", - "retention_policy": {"infinite": {}}, - "timestamping": {"mode": "client-require", "uncapped": True}, - "delete_on_empty": {"min_age_secs": 300}, - } - assert stream_config_header(StreamConfig()) == "{}" - - -async def test_unary_append_and_read_send_header() -> None: - client = _Client() - stream = _stream(client, encryption_key="key") - - await stream.append(AppendInput(records=[Record(body=b"a")], stream_config=_CONFIG)) - await stream.read(start=SeqNum(0), stream_config=_CONFIG) - - assert len(client.headers) == 2 - for headers in client.headers: - assert headers is not None - assert headers["s2-encryption-key"] == "key" - assert json.loads(headers[_HEADER]) == { - "retention_policy": {"age": 3600}, - "delete_on_empty": {"min_age_secs": 300}, - } - - -async def test_unary_append_and_read_omit_header_when_unset() -> None: - client = _Client() - stream = _stream(client) - - await stream.append(AppendInput(records=[Record(body=b"a")])) - await stream.read(start=SeqNum(0)) - - assert all(_HEADER not in (headers or {}) for headers in client.headers) - - -async def test_read_session_sends_header() -> None: - client = _Client() - - with patch.object(s2s_read_session, "read_messages", new=_decoded_messages): - async for _ in s2s_read_session.run_read_session( - cast(HttpClient, client), - "stream", - SeqNum(0), - ReadLimit(count=1), - until_timestamp=None, - clamp_to_tail=False, - wait=None, - retry=Retry(max_attempts=1), - stream_config=_CONFIG, - ): - pass - - assert client.headers == [ - {"content-type": "s2s/proto", _HEADER: stream_config_header(_CONFIG)} - ] - - -async def test_append_session_sends_header() -> None: - client = _Client() - - async def inputs() -> AsyncIterator[AppendInput]: - yield AppendInput(records=[Record(body=b"a")]) - - with patch.object(s2s_append_session, "read_messages", new=_decoded_messages): - async for _ in s2s_append_session.run_append_session( - cast(HttpClient, client), - "stream", - inputs(), - Retry(max_attempts=1), - Compression.NONE, - ack_timeout=1.0, - stream_config=_CONFIG, - ): - pass - - assert client.headers == [ - { - "content-type": "s2s/proto", - "accept": "s2s/proto", - _HEADER: stream_config_header(_CONFIG), - } - ] diff --git a/tests/test_stream_ops.py b/tests/test_stream_ops.py index 48c69e4..f4ff9a0 100644 --- a/tests/test_stream_ops.py +++ b/tests/test_stream_ops.py @@ -26,6 +26,7 @@ S2Stream, SeqNum, SeqNumMismatchError, + StorageClass, StreamConfig, TailOffset, Timestamp, @@ -805,6 +806,117 @@ async def test_read_session_beyond_tail_errors(self, stream: S2Stream): assert exc_info.value.tail.seq_num == 0 +@pytest.mark.stream +class TestAutoCreateStreamConfig: + async def _run_with_stream_config( + self, + stream: S2Stream, + operation: str, + config: StreamConfig, + ) -> None: + if operation == "append": + await stream.append( + AppendInput( + records=[Record(body=b"auto-created")], stream_config=config + ) + ) + elif operation == "append_session": + async with stream.append_session(stream_config=config) as session: + ticket = await session.submit( + AppendInput(records=[Record(body=b"auto-created")]) + ) + await ticket + elif operation == "producer": + async with stream.producer(stream_config=config) as producer: + ticket = await producer.submit(Record(body=b"auto-created")) + await ticket + elif operation == "read": + with pytest.raises(ReadUnwrittenError): + await stream.read(start=SeqNum(0), stream_config=config) + elif operation == "read_session": + with pytest.raises(ReadUnwrittenError): + async with stream.read_session( + start=SeqNum(0), + limit=ReadLimit(count=1), + stream_config=config, + ) as session: + async for _ in session: + pass + else: + raise AssertionError(f"unsupported operation: {operation}") + + @pytest.mark.parametrize( + "operation", ["append", "append_session", "producer", "read", "read_session"] + ) + async def test_auto_create_config_overrides_basin_defaults_and_inherits_unset_fields( + self, + s2: S2, + basin: S2Basin, + stream_name: str, + operation: str, + ): + await s2.reconfigure_basin( + basin.name, + config=BasinConfig( + create_stream_on_append=True, + create_stream_on_read=True, + default_stream_config=StreamConfig( + storage_class=StorageClass.STANDARD, + retention_policy=7200, + ), + ), + ) + stream_config = StreamConfig( + retention_policy=3600, + timestamping=Timestamping( + mode=TimestampingMode.CLIENT_PREFER, + uncapped=True, + ), + delete_on_empty_min_age=3600, + ) + + await self._run_with_stream_config( + basin.stream(stream_name), operation, stream_config + ) + + actual = await basin.get_stream_config(stream_name) + assert actual.storage_class == StorageClass.STANDARD + assert actual.retention_policy == 3600 + assert actual.timestamping is not None + assert actual.timestamping.mode == TimestampingMode.CLIENT_PREFER + assert actual.timestamping.uncapped is True + assert actual.delete_on_empty_min_age == 3600 + + @pytest.mark.parametrize( + "operation", ["append", "append_session", "producer", "read", "read_session"] + ) + async def test_auto_create_config_is_ignored_for_existing_stream( + self, + basin: S2Basin, + stream_name: str, + operation: str, + ): + await basin.create_stream( + stream_name, + config=StreamConfig( + storage_class=StorageClass.EXPRESS, + retention_policy=7200, + ), + ) + await self._run_with_stream_config( + basin.stream(stream_name), + operation, + StreamConfig( + storage_class=StorageClass.STANDARD, + retention_policy=3600, + ), + ) + + actual = await basin.get_stream_config(stream_name) + assert actual.storage_class == StorageClass.EXPRESS + assert actual.retention_policy == 7200 + + @pytest.mark.stream @pytest.mark.parametrize("compression", [Compression.GZIP, Compression.ZSTD]) class TestCompression: