feat: support stream configuration during on-demand stream creation

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:

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)

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")

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,
         )

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

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()

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:

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: