[Brook S5] Added client Stream handles and Stream reads.

temporalio/client/_client.py

@@ -13,13 +13,14 @@
     TYPE_CHECKING,
     Any,
     Concatenate,
+    TypeVar,
     cast,
     overload,
 )
 
 import nexusrpc
 from nexusrpc import OutputT
-from typing_extensions import Required, Self, TypedDict
+from typing_extensions import Required, Self, TypedDict, Unpack
 
 import temporalio.activity
 import temporalio.api.workflowservice.v1
@@ -110,11 +111,20 @@
 )
 
 if TYPE_CHECKING:
-    from temporalio.common import StreamBackend
+    from temporalio.common import (
+        ClientStream,
+        Creating,
+        StreamBackend,
+        StreamDefinition,
+        StreamHandle,
+        StreamOptions,
+    )
 
     from ._interceptor import Interceptor
     from ._plugin import Plugin
 
+_StreamPayloadT = TypeVar("_StreamPayloadT")
+
 
 class Client:
     """Client for accessing Temporal.
@@ -911,6 +921,83 @@ def get_workflow_handle(
             result_type=result_type,
         )
 
+    def create_stream(
+        self,
+        definition: type[StreamDefinition[_StreamPayloadT]],
+        *,
+        id: str,
+        backend: str | None = None,
+        **options: Unpack[StreamOptions],
+    ) -> Creating[ClientStream[_StreamPayloadT]]:
+        """Create a Stream that the client owns, or get the one that exists.
+
+        A create is create-or-get by Stream id. On a Stream that exists, it ignores ``options``
+        and returns that Stream. A Stream must exist before its first publish. The client closes
+        a Stream that it owns, and a Workflow only writes to it. Await the call to get the
+        handle, or use it in ``async with``, which closes the Stream at the end of the block::
+
+            async with client.create_stream(Turns, id="chat-42") as stream:
+                await stream.publish(Turn("user", "hi"))
+
+        .. warning::
+            This API is experimental.
+
+        Args:
+            definition: The Stream definition class.
+            id: The Stream id, which the caller picks, as for a Workflow id.
+            backend: The name of the backend. Defaults to the registered backend with
+                ``default=True``.
+            options: The options of a new Stream, such as ``retention``. See
+                :py:class:`temporalio.common.StreamOptions`.
+
+        Raises:
+            StreamConfigError: The client has no backend with that name, or no backend has
+                ``default=True`` and the call gives no backend. No call is made.
+            StreamClosedError: A Stream of that id was created and closed before.
+            StreamExpiredError: A Stream of that id expired.
+        """
+        from temporalio.common._streams import _stream_client
+
+        return _stream_client.create(
+            self, definition, id=id, backend=backend, options=options
+        )
+
+    def stream(
+        self,
+        definition: type[StreamDefinition[_StreamPayloadT]],
+        *,
+        id: str | None = None,
+        stream: StreamHandle[Any] | None = None,
+        backend: str | None = None,
+    ) -> ClientStream[_StreamPayloadT]:
+        """Get a Stream by its id, or from a handle that came in data, such as in a Record.
+
+        This method makes no call. To create a Stream that the client owns, use
+        :py:meth:`create_stream`. A Workflow hands out its Stream, such as through a query.
+
+        .. warning::
+            This API is experimental.
+
+        Args:
+            definition: The Stream definition class.
+            id: The Stream id.
+            stream: A handle of the Stream, such as one in a Record. The result keeps its id,
+                owner and backend.
+            backend: The name of the backend. Defaults to the handle's backend, or else to the
+                registered backend with ``default=True``.
+
+        Raises:
+            StreamConfigError: The client has no backend with that name, or no backend has
+                ``default=True`` and the call gives no backend.
+            TypeError: Neither or both of ``id`` and ``stream`` are given.
+            ValueError: The id is empty, or ``backend`` is not the handle's backend.
+        """
+        from temporalio.common._streams import _stream_client
+
+        return _stream_client.client_stream(
+            self, definition, id=id, stream=stream, backend=backend
+        )
+
     def get_workflow_handle_for(
         self,
         workflow: (

temporalio/common/__init__.py

@@ -63,11 +63,14 @@
         Superseded,
     )
     from temporalio.common._streams._plugin import StreamBackend
+    from temporalio.common._streams._reader import Subscription
+    from temporalio.common._streams._stream_client import ClientStream
 
 __all__ = [
     "ActivityIDConflictPolicy",
     "ActivityIDReusePolicy",
     "AutoUpgradeVersioningOverride",
+    "ClientStream",
     "CloseStatus",
     "Closed",
     "Creating",
@@ -119,6 +122,7 @@
     "StreamStorageError",
     "StreamTooLargeError",
     "StreamUnsupportedError",
+    "Subscription",
     "Superseded",
     "TypedSearchAttributes",
     "VersioningBehavior",
@@ -1585,6 +1589,7 @@ def wait_sync(self, timeout: float | None = None) -> None:
 # so this module imports the Stream names on first use.
 _STREAM_NAMES = frozenset(
     [
+        "ClientStream",
         "CloseStatus",
         "Closed",
         "Creating",
@@ -1607,6 +1612,7 @@ def wait_sync(self, timeout: float | None = None) -> None:
         "StreamStorageError",
         "StreamTooLargeError",
         "StreamUnsupportedError",
+        "Subscription",
         "Superseded",
         "defn",
     ]

temporalio/common/_streams/__init__.py

@@ -25,9 +25,12 @@
     Superseded,
 )
 from temporalio.common._streams._plugin import StreamBackend
+from temporalio.common._streams._reader import Subscription
+from temporalio.common._streams._stream_client import ClientStream
 
 __all__ = [
     "END",
+    "ClientStream",
     "CloseStatus",
     "Closed",
     "Creating",
@@ -49,6 +52,7 @@
     "StreamStorageError",
     "StreamTooLargeError",
     "StreamUnsupportedError",
+    "Subscription",
     "Superseded",
     "defn",
 ]

temporalio/common/_streams/_access.py

@@ -0,0 +1,133 @@
+"""Stream calls outside Workflow code, which Activity and client handles share."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Callable, Sequence
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any
+
+import temporalio.converter
+from temporalio.common._streams import _convert, _reader
+from temporalio.common._streams._backend import BackendCalls, CallOptions
+from temporalio.common._streams._items import End, close_status
+from temporalio.common._streams._reader import Item, Position, Subscription
+from temporalio.common._streams._writer import StreamWriter
+from temporalio.common._streams.system import models
+
+
+@dataclass
+class Access:
+    """The calls of one Stream handle, bound to a backend."""
+
+    backend: BackendCalls
+    converter: temporalio.converter.DataConverter
+    namespace: str
+    stream_id: str
+    writer: Callable[[], StreamWriter]
+
+    async def publish(self, values: Sequence[Any]) -> int:
+        """Write the values as Records now, in order, and return the offset of the last one."""
+        if not values:
+            raise TypeError("publish takes one or more values")
+        return await self.writer().write(values)
+
+    async def items(
+        self, read: ReadSpec, position: Position
+    ) -> AsyncGenerator[Item, None]:
+        """Read the Stream's items in order, and move ``position`` past each page."""
+        items = _reader.read(
+            self.backend.service(self.namespace, _convert.raw(self.converter)),
+            self.converter,
+            self.stream_id,
+            position=position,
+            payload_type=read.payload_type,
+            from_offset=read.offset,
+            last=read.last,
+            wait_new_records=read.wait_new_records,
+            create_wait=read.create_wait,
+        )
+        try:
+            async for item in items:
+                yield item
+        finally:
+            await items.aclose()
+
+    async def get_offset(self) -> int:
+        """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record.
+
+        Offsets are sparse and above zero, so ``offset=0`` reads from the start.
+        """
+        service = self.backend.service(self.namespace)
+        response = await (
+            await service.describe_stream(stream=models.StreamRef(id=self.stream_id))
+        )
+        return response.last_offset + 1 if response.last_offset else 0
+
+    async def close(self, result: Any | None, failure: BaseException | None) -> None:
+        """Write this Producer's Records, then close the Stream."""
+        writer = self.writer()
+        await writer.close()
+        producer = models._ProducerTransferTypeConverter().from_transfer_type(
+            writer.producer, models.Producer
+        )
+        status = close_status(failure)
+        service = self.backend.service(self.namespace, self.converter)
+        await (
+            await service.close_stream(
+                stream=models.StreamRef(id=self.stream_id),
+                producer=producer,
+                result=result,
+                failure=failure,
+                status=models.StreamCloseStatus(status) if status is not None else None,
+            )
+        )
+
+
+@dataclass
+class ReadSpec:
+    """What one read asks for."""
+
+    payload_type: Any
+    offset: int | End | None
+    last: int | None
+    wait_new_records: bool
+    create_wait: timedelta | None
+
+
+def read_spec(
+    payload_type: Any,
+    *,
+    offset: int | End | None,
+    last: int | None,
+    wait_new_records: bool,
+    create_wait: timedelta | None,
+) -> ReadSpec:
+    """The options of a read, and the type that it decodes each value with.
+
+    Raises:
+        ValueError: ``last`` is not above zero, or the call also gives ``offset``.
+    """
+    if last is not None and (last < 1 or offset is not None):
+        raise ValueError("last takes a count above zero and no offset")
+    return ReadSpec(
+        payload_type=payload_type,
+        offset=offset,
+        last=last,
+        wait_new_records=wait_new_records,
+        create_wait=create_wait,
+    )
+
+
+def subscription(
+    items: Callable[[Position], AsyncGenerator[Item, None]],
+    read: ReadSpec,
+    options: CallOptions | None = None,
+) -> Subscription[Any]:
+    """A read over the items that ``items`` gives."""
+    # A read from END or of the last Records has no resume offset until its first page.
+    if read.last is not None or read.offset is End.END:
+        position = Position(None)
+    else:
+        position = Position(read.offset or 0)
+    return Subscription(items(position), position, options)

temporalio/common/_streams/_convert.py

@@ -0,0 +1,66 @@
+"""Value conversion for Stream Records outside Workflow code."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+from typing import Any
+
+import temporalio.api.common.v1
+import temporalio.common
+import temporalio.converter
+import temporalio.exceptions
+
+
+class _RawPayloadConverter(temporalio.converter.PayloadConverter):
+    # The generated client runs the codec on Record bodies and keeps them as RawValue, so the
+    # handle converts each body with the Stream's payload type.
+    def to_payloads(
+        self, values: Sequence[Any]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        payloads = []
+        for value in values:
+            if not isinstance(value, temporalio.common.RawValue):
+                raise TypeError(f"expected a RawValue Record body, got {type(value)}")
+            payloads.append(value.payload)
+        return payloads
+
+    def from_payloads(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+        type_hints: list[type] | None = None,
+    ) -> list[Any]:
+        return [temporalio.common.RawValue(payload) for payload in payloads]
+
+
+def raw(
+    converter: temporalio.converter.DataConverter,
+) -> temporalio.converter.DataConverter:
+    """The converter with Record bodies kept as RawValue. Its codec still runs on them."""
+    return dataclasses.replace(converter, payload_converter_class=_RawPayloadConverter)
+
+
+def failure(
+    converter: temporalio.converter.DataConverter, failure: BaseException
+) -> BaseException:
+    """Convert the RawValue payloads of a read's close failure with the user's converter."""
+    if (
+        not isinstance(failure, temporalio.exceptions.FailureError)
+        or not failure.failure
+    ):
+        return failure
+    return converter.failure_converter.from_failure(
+        failure.failure, converter.payload_converter
+    )
+
+
+def decode(
+    converter: temporalio.converter.DataConverter, body: Any, type_hint: Any
+) -> Any:
+    """Convert a Record body that a read gave as RawValue, after the codec."""
+    if not isinstance(body, temporalio.common.RawValue):
+        return body
+    [value] = converter.payload_converter.from_payloads(
+        [body.payload], [type_hint] if type_hint else None
+    )
+    return value