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_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: