From 7265a0471fbeae319564829d7698975e9d6da8b7 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 14:54:36 -0700 Subject: [PATCH 1/3] Added client Stream handles and Stream reads. `Client.create_stream` creates a Stream or gets the one that exists, and `Client.stream` binds one with no call. A `ClientStream` in a query result uses the client that made the query. A read yields Records, the Superseded items that withdraw earlier Records, and the close. --- temporalio/client/_client.py | 91 +++- temporalio/common/__init__.py | 6 + temporalio/common/_streams/__init__.py | 4 + temporalio/common/_streams/_access.py | 133 ++++++ temporalio/common/_streams/_convert.py | 66 +++ temporalio/common/_streams/_plugin.py | 9 +- temporalio/common/_streams/_reader.py | 316 ++++++++++++++ temporalio/common/_streams/_stream_client.py | 434 +++++++++++++++++++ temporalio/common/_streams/_writer.py | 122 ++++++ tests/contrib/streams/fake_extension.py | 170 ++++++++ tests/contrib/streams/test_facade.py | 211 +++++++++ tests/contrib/streams/test_reader.py | 81 ++++ tests/contrib/streams/test_stream_client.py | 84 ++++ 13 files changed, 1724 insertions(+), 3 deletions(-) create mode 100644 temporalio/common/_streams/_access.py create mode 100644 temporalio/common/_streams/_convert.py create mode 100644 temporalio/common/_streams/_reader.py create mode 100644 temporalio/common/_streams/_stream_client.py create mode 100644 temporalio/common/_streams/_writer.py create mode 100644 tests/contrib/streams/fake_extension.py create mode 100644 tests/contrib/streams/test_facade.py create mode 100644 tests/contrib/streams/test_reader.py create mode 100644 tests/contrib/streams/test_stream_client.py diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index f171bfa5c..569b17aef 100644 --- a/temporalio/client/_client.py +++ b/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: ( diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index c1104fce3..1bcd2f008 100644 --- a/temporalio/common/__init__.py +++ b/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", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 299ee3ced..8e4785be2 100644 --- a/temporalio/common/_streams/__init__.py +++ b/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", ] diff --git a/temporalio/common/_streams/_access.py b/temporalio/common/_streams/_access.py new file mode 100644 index 000000000..20bf6bd0a --- /dev/null +++ b/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) diff --git a/temporalio/common/_streams/_convert.py b/temporalio/common/_streams/_convert.py new file mode 100644 index 000000000..5d538bcbd --- /dev/null +++ b/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 diff --git a/temporalio/common/_streams/_plugin.py b/temporalio/common/_streams/_plugin.py index 1763ea9b8..e22c9ffac 100644 --- a/temporalio/common/_streams/_plugin.py +++ b/temporalio/common/_streams/_plugin.py @@ -159,9 +159,16 @@ def default_name(backends: Sequence[StreamBackend]) -> str: def configure_client( config: temporalio.client.ClientConfig, ) -> temporalio.client.ClientConfig: - """Check the client's backends, and apply the plugin configuration of each one.""" + """Check the client's backends, and apply the plugin configuration of each one. + + The client's query results bind each ``ClientStream`` in them to the client. + """ + # The client Stream module imports this one. + from temporalio.common._streams._stream_client import BindQueryResults + backends = check_backends(config.get("stream_backends", ()), "client") config["stream_backends"] = backends + config["interceptors"] = [*config.get("interceptors", []), BindQueryResults()] for backend in backends: config = backend._configure_client(config) return config diff --git a/temporalio/common/_streams/_reader.py b/temporalio/common/_streams/_reader.py new file mode 100644 index 000000000..321c52cb4 --- /dev/null +++ b/temporalio/common/_streams/_reader.py @@ -0,0 +1,316 @@ +"""Reads of a Stream, for Activities and clients.""" + +from __future__ import annotations + +import array +import asyncio +import dataclasses +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import Any, Generic, TypeVar + +import temporalio.converter +from temporalio.api.enums.v1 import StreamProducerKind +from temporalio.common._streams import _backend, _convert, _errors +from temporalio.common._streams._items import ( + Closed, + CloseStatus, + End, + Record, + Superseded, + producer_key, + withdraws, +) +from temporalio.common._streams.system import client as service_client +from temporalio.common._streams.system import models + +T = TypeVar("T") + +Item = Record[Any] | Superseded | Closed + +# The service waits at most 60 s per read, so a longer wait is a loop of reads. +_MAX_WAIT = timedelta(seconds=60) + +# An empty page that comes back in less than half of its wait is followed by a pause, so a +# Store that ignores the wait isn't asked again at once. The pause doubles up to the last one. +_FIRST_PAUSE = 0.05 +_LAST_PAUSE = 2.0 + + +@dataclasses.dataclass +class Position: + """Where a read resumes from. The read moves it past each page that it finished.""" + + next_offset: int | None + + +class Subscription(Generic[T]): + """The items of one read of a Stream, in order. Iterate it with ``async for``. + + It yields :py:class:`Record` items, :py:class:`Superseded` items that withdraw earlier + Records, and a :py:class:`Closed` item last when the Stream is closed. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + items: AsyncGenerator[Item, None], + position: Position, + options: _backend.CallOptions | None = None, + ) -> None: + """Create a read. Use a handle's ``subscribe`` instead.""" + self._items = items + self._position = position + # A client read's rpc_metadata and rpc_timeout, for each call of each step. + self._options = options + # The Producer and the offsets of each Record that this read gave, without values, so a + # Superseded item can list the offsets that it withdraws. + self._shown: dict[ + tuple[Any, ...], tuple[models.Producer, array.array[int]] + ] = {} + + @property + def next_offset(self) -> int | None: + """The offset to resume from, after the last item or the last page that the read got. + + Pass it as ``offset`` to a later read. After a ``Superseded`` item, it's that item's own + offset, since the first Record of its batch has the same offset and comes next. A read + from there gives the ``Superseded`` item again, which withdraws nothing new. After a + Record that failed to decode, it's past that Record. It's None before the first page of + a read from ``END`` or of a read with ``last``. + """ + return self._position.next_offset + + def __aiter__(self) -> Subscription[T]: + """The read is its own iterator.""" + return self + + async def __anext__(self) -> Record[T] | Superseded | Closed: + """The next item. + + Raises: + StreamError: A Record's value failed to decode with the Stream's payload type. The + read ends, and :py:attr:`next_offset` is past that Record. + """ + with _backend.using(self._options): + item = await self._items.__anext__() + if isinstance(item, Superseded): + # The Record of the same entry has the same offset and comes next. + self._position.next_offset = item.offset + else: + self._position.next_offset = max( + self._position.next_offset or 0, item.offset + 1 + ) + if ( + isinstance(item, Record) + and item.producer is not None + # Nothing withdraws a client's Records, so a long read of them keeps no offsets. + and item.producer.kind != StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT + ): + key = producer_key(item.producer) + _, runs = self._shown.setdefault(key, (item.producer, array.array("Q"))) + # A batch has consecutive offsets, so a run of them costs one first and last pair. + if runs and runs[-1] + 1 == item.offset: + runs[-1] = item.offset + else: + runs.extend((item.offset, item.offset)) + elif isinstance(item, Superseded): + item = dataclasses.replace( + item, withdrawn_offsets=self._withdraw(item.supersession) + ) + return item + + def _withdraw(self, supersession: models.Supersession) -> tuple[int, ...]: + offsets: list[int] = [] + for key, (producer, runs) in list(self._shown.items()): + if withdraws(supersession, producer): + for first, last in zip(runs[::2], runs[1::2]): + offsets.extend(range(first, last + 1)) + # A withdrawn Record stays withdrawn, so the read forgets it. + del self._shown[key] + return tuple(sorted(offsets)) + + async def aclose(self) -> None: + """End the read early, from the task that reads it. It makes no more calls. + + If another task is waiting in the read for new Records, this raises ``RuntimeError``. + Cancel that task instead. + """ + with _backend.using(self._options): + await self._items.aclose() + + +async def read( + service: service_client.StreamServiceClientExtension, + converter: temporalio.converter.DataConverter, + stream_id: str, + *, + position: Position, + payload_type: Any, + from_offset: int | End | None, + last: int | None = None, + wait_new_records: bool, + create_wait: timedelta | None, +) -> AsyncGenerator[Item, None]: + """Yield the items of a Stream in order, until a close. + + With ``wait_new_records`` False, stop at the newest Record. With ``last``, start with the + newest ``last`` Records. + """ + loop = asyncio.get_running_loop() + create_until = loop.time() + create_wait.total_seconds() if create_wait else None + + def request() -> models.ReadRecordsRequest: + create_left = ( + timedelta(seconds=max(create_until - loop.time(), 0)) + if create_until is not None + else None + ) + return models.ReadRecordsRequest( + stream=models.StreamRef(id=stream_id), + # END applies only to the first page. Later pages start at the last page's next offset. + from_offset=from_offset if isinstance(from_offset, int) else None, + start=models.StreamReadStart.StreamReadStartEnd + if from_offset is End.END + else models.StreamReadStart.StreamReadStartBeginning + if from_offset is None and not last + else None, + last=last, + create_wait=min(create_left, _backend.longest_wait(_MAX_WAIT)) + if create_left + else None, + ) + + while True: + pages = _pages(service, request(), wait_new_records) + try: + async for page in pages: + for record in page.records or []: + try: + item = await _item(converter, record, payload_type) + except Exception as err: + # Move next_offset past the Record, so a resume doesn't fail on it again. + # The error gives the offset, for a read of that Record. + offset = record.offset or 0 + position.next_offset = max( + position.next_offset or 0, offset + 1 + ) + raise _errors.StreamError( + f"Can't decode the Record at the offset {offset}: {err}" + ) from err + yield item + if isinstance(item, Closed): + return + if page.next_offset is not None: + # A resume starts where the page ended. + position.next_offset = max( + position.next_offset or 0, page.next_offset + ) + return + except _errors.StreamNotFoundError: + if create_until is None or loop.time() >= create_until: + raise + finally: + await pages.aclose() + + +async def _pages( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, + wait_new_records: bool, +) -> AsyncGenerator[models.ReadRecordsResponse, None]: + # A read that waits long-polls each page. A read with no wait reads pages until a page has + # no Record and doesn't advance. + if request.last: + # The newest Records come in one page that never waits. A waiting read then continues + # after them. + page = await _read_page(service, request) + yield page + if page.done or not wait_new_records or page.next_offset is None: + return + request = dataclasses.replace(request, last=None, from_offset=page.next_offset) + while True: + if wait_new_records: + page = await _read_until_records( + service, request, _backend.longest_wait(_MAX_WAIT) + ) + else: + page = await _read_page(service, request) + yield page + if page.done: + return + if not wait_new_records and ( + page.next_offset is None + # An empty page that didn't advance is at the newest Record. + or (not page.records and page.next_offset <= (request.from_offset or 0)) + ): + return + if page.next_offset is not None: + request = dataclasses.replace(request, from_offset=page.next_offset) + + +async def _read_page( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, +) -> models.ReadRecordsResponse: + return await (await service._read_records(request)) + + +async def _read_until_records( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, + deadline: timedelta, +) -> models.ReadRecordsResponse: + # Each read waits for what is left of the deadline, and an empty page moves the next read to + # its next offset. At the deadline the last page comes back, so the caller keeps its place. + loop = asyncio.get_running_loop() + end = loop.time() + deadline.total_seconds() + pause = _FIRST_PAUSE + while True: + wait = min(max(end - loop.time(), 0.001), _MAX_WAIT.total_seconds()) + request = dataclasses.replace(request, wait=timedelta(seconds=wait)) + asked = loop.time() + page = await _read_page(service, request) + if page.records or page.done: + return page + if page.next_offset is not None: + request = dataclasses.replace(request, from_offset=page.next_offset) + left = end - loop.time() + if left <= 0: + return page + if loop.time() - asked < wait / 2: + await asyncio.sleep(min(pause, left)) + pause = min(pause * 2, _LAST_PAUSE) + else: + pause = _FIRST_PAUSE + + +async def _item( + converter: temporalio.converter.DataConverter, + record: models.StreamRecord, + payload_type: Any, +) -> Item: + offset = record.offset or 0 + match record.item: + case models.StreamItemData(value=data): + value = _convert.decode(converter, data.body, payload_type) + return Record(data=value, offset=offset, producer=data.producer) + case models.StreamItemSuperseded(value=supersession): + return Superseded(offset=offset, supersession=supersession) + case models.StreamItemClose(value=close): + result = _convert.decode(converter, close.result, None) + failure = close.failure + # A close with no status takes it from its failure. + status = close.status or ( + CloseStatus.FAILED if failure is not None else CloseStatus.COMPLETED + ) + return Closed( + offset=offset, + result=result, + failure=_convert.failure(converter, failure) + if failure is not None + else None, + status=CloseStatus(status), + ) diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py new file mode 100644 index 000000000..210f9bf3e --- /dev/null +++ b/temporalio/common/_streams/_stream_client.py @@ -0,0 +1,434 @@ +"""Streams in client code.""" + +from __future__ import annotations + +import contextvars +import itertools +import os +import uuid +import weakref +from collections.abc import AsyncGenerator, Mapping +from datetime import timedelta +from types import TracebackType +from typing import Any, TypeVar + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.client +from temporalio.api.enums.v1 import StreamErrorReason, StreamProducerKind, StreamStatus +from temporalio.common._streams import _backend, _errors, _ids +from temporalio.common._streams._access import Access, ReadSpec, read_spec, subscription +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._defn import StreamDefinition, definition_of +from temporalio.common._streams._handle import ( + Creating, + StreamHandle, + StreamOptions, + exit_failure, +) +from temporalio.common._streams._items import End +from temporalio.common._streams._plugin import client_backend +from temporalio.common._streams._reader import Item, Position, Subscription +from temporalio.common._streams._writer import StreamWriter + +T = TypeVar("T") + +# The handles of one client and Stream write as one Producer. Each client gets its own key, +# since two clients can have two runtime objects, whose writers would send clashing sequences. +_process_key = str(uuid.uuid4()) +_clients = itertools.count(1) +_CLIENT_KEYS: weakref.WeakKeyDictionary[temporalio.client.Client, str] = ( + weakref.WeakKeyDictionary() +) + + +def _new_process_key() -> None: + # A forked child has the parent's key and count, so two children would give their clients + # the same Producer ids. + global _process_key, _clients + _process_key = str(uuid.uuid4()) + _clients = itertools.count(1) + _CLIENT_KEYS.clear() + + +# Windows has no fork. +_register_at_fork = getattr(os, "register_at_fork", None) +if _register_at_fork is not None: + _register_at_fork(after_in_child=_new_process_key) + +# The client whose query result the converter decodes. A ClientStream in the result takes it, so +# a reader uses a Stream that a Workflow hands out with no other call. +_decoding_client: contextvars.ContextVar[temporalio.client.Client | None] = ( + contextvars.ContextVar("__temporal_stream_decoding_client", default=None) +) + + +def _client_producer( + client: temporalio.client.Client, named: tuple[str, int] | None +) -> stream_v1.Producer: + if named is not None: + name, epoch = named + return stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT, + id=_ids.named_client_producer_id(name), + epoch=stream_v1.ProducerEpoch(attempt=epoch), + ) + key = _CLIENT_KEYS.get(client) + if key is None: + key = f"{_process_key}-{next(_clients)}" + _CLIENT_KEYS[client] = key + return stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT, + id=_ids.client_producer_id(client.identity, key), + epoch=stream_v1.ProducerEpoch(attempt=1), + ) + + +class ClientStream(StreamHandle[T]): + """A Stream in client code. The client can write and read it, and close a client's Stream. + + Get one from :py:meth:`temporalio.client.Client.create_stream` or + :py:meth:`temporalio.client.Client.stream`. A ``ClientStream`` that a query of a client + returns, such as a Workflow's ``stream`` query, uses that client. + + Each publish writes before it returns. All handles of one client and Stream write as one + Producer. :py:meth:`producer` returns a handle that writes as a named Producer. Use a handle + in ``async with`` to close the Stream at the end of the block. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + id: str, + workflow_id: str | None = None, + first_execution_run_id: str | None = None, + backend: str | None = None, + payload_type: Any = None, + ) -> None: + """Create a handle. Use :py:meth:`temporalio.client.Client.stream` instead.""" + super().__init__( + id=id, + workflow_id=workflow_id, + first_execution_run_id=first_execution_run_id, + backend=backend, + payload_type=payload_type, + ) + self._client = _decoding_client.get() + self._stream_backend: BackendCalls | None = None + self._writer: StreamWriter | None = None + self._closed = False + # The name and the epoch of a named Producer. + self._named: tuple[str, int] | None = None + + def _bind( + self, client: temporalio.client.Client, backend: BackendCalls + ) -> ClientStream[T]: + self._client = client + self._stream_backend = backend + return self + + def _access(self) -> Access: + client = self._client + if client is None: + raise _errors.StreamConfigError( + "This ClientStream came from a payload that no client query decoded, so it " + "has no client. Use Client.stream(definition, stream=handle)." + ) + if self._stream_backend is None: + _, self._stream_backend = client_backend(client, self._backend) + backend = self._stream_backend + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + writers=backend.writers(client.namespace, client.data_converter), + stream_id=self._id, + producer=_client_producer(client, self._named), + free=backend.free_writer, + ) + return self._writer + + return Access( + backend=backend, + converter=client.data_converter, + namespace=client.namespace, + stream_id=self._id, + writer=writer, + ) + + def producer(self, name: str, epoch: int) -> ClientStream[T]: + """Return this Stream as a handle that writes as the Producer ``name`` at ``epoch``. + + Use it for a writer that must fence its earlier instances, such as a server that + restarts. Keep ``epoch`` in durable storage, and raise it at each start. After a write + at a higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The + Records of the earlier epochs aren't withdrawn, since they're real messages. Run only + one live writer for each name. + + Args: + name: The Producer name, the same across restarts. + epoch: The writer's epoch, 1 or more. + """ + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + handle: ClientStream[T] = ClientStream( + id=self._id, + workflow_id=self._workflow_id, + first_execution_run_id=self._first_execution_run_id, + backend=self._backend, + payload_type=self._payload_type, + ) + handle._named = (name, epoch) + handle._client = self._client + handle._stream_backend = self._stream_backend + return handle + + async def publish( + self, + *values: T, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. ``rpc_metadata`` and ``rpc_timeout`` are as for + :py:meth:`get_offset`. + + Raises: + StreamNotFoundError: Nobody created the Stream. + """ + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await self._access().publish(values) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Subscription[T]: + """Read the Stream's items in order. + + The read ends after a ``Closed`` item, or at the newest Record when + ``wait_new_records`` is False. To stop after a quiet period, use your own timer, and + resume later with ``offset=`` set to the read's ``next_offset``. A ``Superseded`` item + withdraws earlier Records of a failed attempt. + + Args: + offset: Start at this offset, inclusive, or at ``common.END`` for new + Records only. Defaults to the first Record. + last: Start with the newest ``last`` Records, oldest first, such as for the first + view of a chat. Can't be used with ``offset``. A Record that a later + ``Superseded`` item withdraws doesn't count. At most 1000. + wait_new_records: Wait for new Records at the end of the Stream. + create_wait: How long to wait for a Stream that doesn't exist yet, such as the + Stream of a Workflow that just started. Defaults to no wait. + rpc_metadata: Headers of each call of the read. See :py:meth:`get_offset`. + rpc_timeout: The deadline of each call of the read, not of the whole read. A call + that waits for new Records returns before the deadline. + """ + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + return subscription( + lambda position: self._items(read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) + + async def _items( + self, read: ReadSpec, position: Position + ) -> AsyncGenerator[Item, None]: + items = self._access().items(read, position) + try: + async for item in items: + yield item + finally: + await items.aclose() + + async def get_offset( + self, + *, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> 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. + + Args: + rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so + they send none. + rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with + ``DEADLINE_EXCEEDED``. + """ + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await self._access().get_offset() + + def _check_close(self) -> None: + if self._workflow_id: + # Only the run chain can close a Workflow's Stream. The Store refuses a client's + # close too, but this check makes no call. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {self._id} from a client, since a Workflow owns it", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + ) + + async def close( + self, + *, + result: Any | None = None, + failure: BaseException | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> None: + """Write this handle's Records, then close the Stream. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`get_offset`. + + Raises: + StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. + """ + self._check_close() + self._closed = True + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + await self._access().close(result, failure) + + async def __aenter__(self) -> ClientStream[T]: + """Use the Stream in ``async with``, which closes it at the end. + + Raises: + StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. + """ + # Fail early, since the close at the end would raise and hide the block's own error. + self._check_close() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this + does nothing. + """ + if self._closed: + return + close, failure = exit_failure(exc) + if close: + await self.close(failure=failure) + + +def create( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str, + backend: str | None, + options: StreamOptions, +) -> Creating[ClientStream[T]]: + """Create a Stream that the client owns, or get the one that exists. + + See :py:meth:`temporalio.client.Client.create_stream`. + """ + + async def create() -> ClientStream[T]: + info = definition_of(definition) + backend_name, calls = client_backend(client, backend) + created = await calls.service(client.namespace).create_stream( + stream_id=id, retention=options.get("retention") + ) + if created.status == StreamStatus.STREAM_STATUS_CLOSED: + raise _errors.StreamClosedError( + f"The Stream {id} was created and closed before", + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + if created.status == StreamStatus.STREAM_STATUS_EXPIRED: + raise _errors.StreamExpiredError( + f"The Stream {id} expired", + reason=StreamErrorReason.STREAM_ERROR_REASON_EXPIRED, + ) + stream: ClientStream[T] = ClientStream( + id=id, backend=backend_name, payload_type=info.type + ) + return stream._bind(client, calls) + + return Creating(create) + + +def client_stream( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str | None, + stream: StreamHandle[Any] | None, + backend: str | None, +) -> ClientStream[T]: + """Get a Stream by its id, or from a handle that came in data. + + See :py:meth:`temporalio.client.Client.stream`. + """ + info = definition_of(definition) + if (id is None) == (stream is None): + raise TypeError("stream takes one of id= and stream=") + if stream is not None: + if backend and stream.backend and backend != stream.backend: + raise ValueError( + f"The handle is of the backend {stream.backend}, not {backend}" + ) + backend = backend or stream.backend + stream_id = stream.id if stream else id + if not stream_id: + # An empty id would fail only at the first call, with an unclear error. + raise ValueError("The Stream id is empty") + backend_name, calls = client_backend(client, backend) + handle: ClientStream[T] = ClientStream( + id=stream_id, + workflow_id=stream.workflow_id if stream else None, + first_execution_run_id=stream.first_execution_run_id if stream else None, + backend=backend_name, + payload_type=info.type, + ) + return handle._bind(client, calls) + + +class BindQueryResults(temporalio.client.Interceptor): + """Gives each ``ClientStream`` in a query result the client that made the query.""" + + def intercept_client( + self, next: temporalio.client.OutboundInterceptor + ) -> temporalio.client.OutboundInterceptor: + """Wrap the client's calls.""" + return _BindQueryResults(next) + + +class _BindQueryResults(temporalio.client.OutboundInterceptor): + def __init__(self, next: temporalio.client.OutboundInterceptor) -> None: + super().__init__(next) + # The innermost interceptor is the client's own, and only it has no next. + inner: Any = next + while hasattr(inner, "next"): + inner = inner.next + self._client: temporalio.client.Client = inner._client + + async def query_workflow(self, input: temporalio.client.QueryWorkflowInput) -> Any: + token = _decoding_client.set(self._client) + try: + return await super().query_workflow(input) + finally: + _decoding_client.reset(token) diff --git a/temporalio/common/_streams/_writer.py b/temporalio/common/_streams/_writer.py new file mode 100644 index 000000000..556407d92 --- /dev/null +++ b/temporalio/common/_streams/_writer.py @@ -0,0 +1,122 @@ +"""The writer of an Activity or a client Producer.""" + +from __future__ import annotations + +import asyncio +import weakref +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING, Any + +import temporalio.api.stream.v1 as stream_v1 +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.common._streams import _errors + +if TYPE_CHECKING: + from temporalio.common._streams.writer_system import client as writer_client + + +class StreamWriter: + """Writes the Records of one Producer to one Stream, through the generated Writer.""" + + def __init__( + self, + *, + writers: writer_client.StreamWriterClientExtension, + stream_id: str, + producer: stream_v1.Producer, + free: Callable[[int], None] | None = None, + previous_completed_task: stream_v1.WorkflowTaskEpoch | None = None, + failed_tasks: Sequence[int] = (), + ) -> None: + """Create a writer. It opens on the first write. + + ``writers`` encodes the values with the user's converter and codec. ``free`` releases + the crate's writer if Python frees this one without a close. + """ + self._writers = writers + self._stream_id = stream_id + self._producer = producer + self._free = free + self._release: weakref.finalize[Any, StreamWriter] | None = None + self._previous_completed_task = previous_completed_task + self._failed_tasks = list(failed_tasks) + self._writer: writer_client.Writer | None = None + self._open_lock = asyncio.Lock() + self._closed = False + + @property + def producer(self) -> stream_v1.Producer: + """The Producer that the writer writes as.""" + return self._producer + + async def withdraw_earlier_attempts(self) -> None: + """Open the writer of a new attempt, which withdraws the Records of earlier attempts. + + A Stream that doesn't exist has nothing to withdraw. A closed or expired Stream accepts + no write, so the attempt continues, and a later write raises the real error. + """ + try: + await self.open() + except ( + _errors.StreamNotFoundError, + _errors.StreamClosedError, + _errors.StreamExpiredError, + ): + pass + + async def open(self) -> writer_client.Writer: + """Open the crate's writer once. + + The open of a new attempt withdraws the Records of earlier attempts, even if it writes + nothing. + """ + from temporalio.common._streams.writer_system import models + + async with self._open_lock: + if self._writer is None: + producer = models._ProducerTransferTypeConverter().from_transfer_type( + self._producer, models.Producer + ) + previous = self._previous_completed_task + self._writer = await self._writers.open_writer( + stream=models.StreamRef(id=self._stream_id), + producer=producer, + previous_completed_task=models.WorkflowTaskEpoch( + started_event_id=previous.started_event_id + ) + if previous is not None + else None, + failed_tasks=self._failed_tasks, + ) + if self._free is not None: + # The crate keeps a writer for each open that nobody closed or freed. So + # free it when Python frees this writer. + self._release = weakref.finalize( + self, self._free, self._writer.writer_id + ) + return self._writer + + async def write(self, values: Sequence[Any]) -> int: + """Write the values as Records now, in order, and return the last offset.""" + from temporalio.common._streams.writer_system import models + + if self._closed: + raise _errors.StreamClosedError( + "The writer is closed", + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + writer = await self.open() + answer = await writer.write_records( + [models.AppendRecord(body=value) for value in values] + ) + return answer.last_offset or 0 + + async def close(self) -> None: + """Close the writer. The Stream stays open.""" + if self._closed: + return + self._closed = True + if self._release is not None: + self._release.detach() + if self._writer is not None: + await self._writer.close() diff --git a/tests/contrib/streams/fake_extension.py b/tests/contrib/streams/fake_extension.py new file mode 100644 index 000000000..c4945b81f --- /dev/null +++ b/tests/contrib/streams/fake_extension.py @@ -0,0 +1,170 @@ +"""An in-memory fake of the temporalio-streams runtime, for the facade tests.""" + +from __future__ import annotations + +import asyncio + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.api.streamservice.v1 as stream_pb +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.common import StreamBackend +from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.service import RPCStatusCode + + +def _error(code: RPCStatusCode, reason: StreamErrorReason.ValueType) -> Exception: + details = stream_v1.StreamErrorDetails(reason=reason).SerializeToString() + return StreamCallError(int(code), StreamErrorReason.Name(reason), details) + + +class _Stream: + def __init__(self) -> None: + self.records: list[stream_v1.StreamRecord] = [] + self.closed = False + self.changed = asyncio.Event() + + +class FakeExtension: + def __init__(self) -> None: + self.streams: dict[str, _Stream] = {} + self.writers: dict[int, tuple[str, stream_v1.Producer]] = {} + self.freed: list[int] = [] + self.namespaces: set[str] = set() + self.next_offset = 1 + + def _append(self, stream_id: str, record: stream_v1.StreamRecord) -> int: + stream = self.streams.get(stream_id) + if stream is None: + raise _error( + RPCStatusCode.NOT_FOUND, StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + ) + if stream.closed: + raise _error( + RPCStatusCode.FAILED_PRECONDITION, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + record.offset = self.next_offset + self.next_offset += 1 + stream.records.append(record) + stream.changed.set() + stream.changed = asyncio.Event() + return record.offset + + def free_writer(self, writer_id: int) -> None: + self.freed.append(writer_id) + + async def call(self, method: str, input: bytes) -> bytes: + if method == "CreateStream": + create = stream_pb.CreateStreamRequest.FromString(input) + self.namespaces.add(create.namespace) + self.streams.setdefault(create.stream_id, _Stream()) + return stream_pb.CreateStreamResponse( + stream=stream_v1.StreamRef(id=create.stream_id), created=True + ).SerializeToString() + if method == "OpenWriter": + open_writer = ext_pb.OpenWriterRequest.FromString(input) + writer_id = len(self.writers) + 1 + self.writers[writer_id] = (open_writer.stream.id, open_writer.producer) + return ext_pb.OpenWriterResponse(writer_id=writer_id).SerializeToString() + if method == "WriteRecords": + write = ext_pb.WriteRecordsRequest.FromString(input) + stream_id, producer = self.writers[write.writer_id] + offsets = [ + self._append( + stream_id, + stream_v1.StreamRecord( + data=stream_v1.StreamData(body=record.body, producer=producer) + ), + ) + for record in write.records + ] + return ext_pb.WriteRecordsResponse( + first_offset=offsets[0] if offsets else 0, + last_offset=offsets[-1] if offsets else 0, + ).SerializeToString() + if method == "CloseWriter": + return ext_pb.CloseWriterResponse( + last_offset=self.next_offset - 1 + ).SerializeToString() + if method == "CloseStream": + close = stream_pb.CloseStreamRequest.FromString(input) + self._append( + close.stream.id, + stream_v1.StreamRecord( + close=stream_v1.StreamClose( + producer=close.producer, + result=close.result if close.HasField("result") else None, + ) + ), + ) + self.streams[close.stream.id].closed = True + return b"" + if method == "DescribeStream": + describe = stream_pb.DescribeStreamRequest.FromString(input) + records = self.streams[describe.stream.id].records + return stream_pb.DescribeStreamResponse( + last_offset=records[-1].offset if records else 0 + ).SerializeToString() + if method == "ReadRecords": + return ( + await self._read(stream_pb.ReadRecordsRequest.FromString(input)) + ).SerializeToString() + raise _error( + RPCStatusCode.UNIMPLEMENTED, + StreamErrorReason.STREAM_ERROR_REASON_UNSUPPORTED, + ) + + async def _read( + self, request: stream_pb.ReadRecordsRequest + ) -> stream_pb.ReadRecordsResponse: + loop = asyncio.get_running_loop() + create_until = loop.time() + request.create_wait.ToTimedelta().total_seconds() + while request.stream.id not in self.streams: + if loop.time() >= create_until: + raise _error( + RPCStatusCode.NOT_FOUND, + StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND, + ) + await asyncio.sleep(0.01) + stream = self.streams[request.stream.id] + offset = request.from_offset or 1 + wait = request.wait.ToTimedelta().total_seconds() + until = loop.time() + wait + while True: + found = [record for record in stream.records if record.offset >= offset] + if found: + done = any(record.HasField("close") for record in found) + return stream_pb.ReadRecordsResponse( + records=found, next_offset=found[-1].offset + 1, done=done + ) + left = until - loop.time() + if left <= 0: + return stream_pb.ReadRecordsResponse(next_offset=offset) + try: + await asyncio.wait_for(stream.changed.wait(), left) + except asyncio.TimeoutError: + pass + + +class _FakeBackend(StreamBackend): + def __init__(self, calls: _ExtensionCalls) -> None: + super().__init__("fake-streams", backend_name="fake", default=True) + self._calls = calls + + def _backend_for( + self, + client: Client, # pyright: ignore[reportUnusedParameter] + ) -> _ExtensionCalls: + return self._calls + + +def fake_client(client: Client) -> tuple[Client, FakeExtension]: + """A copy of ``client`` whose one Stream backend, ``fake``, is a fake runtime.""" + extension = FakeExtension() + backend = _FakeBackend(_ExtensionCalls(extension)) + config = client.config() + config["stream_backends"] = [backend] + return Client(**config), extension diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py new file mode 100644 index 000000000..42068e79b --- /dev/null +++ b/tests/contrib/streams/test_facade.py @@ -0,0 +1,211 @@ +import asyncio +import gc +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +from temporalio import common +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import StreamErrorReason, StreamProducerKind +from temporalio.api.stream.v1 import StreamErrorDetails +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.common._streams import _errors, _stream_client +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError +from temporalio.service import RPCError, RPCStatusCode +from tests.contrib.streams.fake_extension import fake_client + + +@dataclass +class Turn: + kind: str + text: str = "" + + +def _details(reason: StreamErrorReason.ValueType) -> bytes: + return StreamErrorDetails(reason=reason).SerializeToString() + + +def test_runtime_errors_map_to_stream_errors(): + fenced = StreamErrorReason.STREAM_ERROR_REASON_FENCED + err = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "fenced", _details(fenced) + ) + ) + assert isinstance(err, common.StreamFencedError) + assert err.reason == fenced + # A reason with no class of its own gives the base class, with the reason. + sequence = StreamErrorReason.STREAM_ERROR_REASON_SEQUENCE + behind = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "behind", _details(sequence) + ) + ) + assert type(behind) is common.StreamError and behind.reason == sequence + other = _errors.from_call_error( + StreamCallError(int(RPCStatusCode.INTERNAL), "odd", b"") + ) + assert type(other) is common.StreamError + assert other.reason == StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def test_client_writes_reads_and_closes(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Turns, id="chat-42") + first = await stream.publish(Turn("model")) + second = await stream.publish(Turn("user", "hello")) + assert second == first + 1 + assert await stream.get_offset() == second + 1 + await stream.close(result="bye") + items = [item async for item in stream.subscribe()] + assert [type(item) for item in items] == [ + common.Record, + common.Record, + common.Closed, + ] + turn, token, closed = items + assert isinstance(turn, common.Record) and turn.data == Turn("model") + assert isinstance(token, common.Record) and token.data == Turn("user", "hello") + assert turn.producer is not None + assert turn.producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT + assert isinstance(closed, common.Closed) and closed.result == "bye" + assert extension.namespaces == {client.namespace} + + resumed = [item async for item in stream.subscribe(offset=token.offset)] + assert [item.offset for item in resumed] == [token.offset, closed.offset] + + +async def test_read_waits_for_create(client: Client): + client, _ = fake_client(client) + stream = client.stream(Tokens, id="wf-1:run-1:tokens") + with pytest.raises(common.StreamNotFoundError): + async for _ in stream.subscribe(): + pass + + async def create_later() -> None: + created = await client.create_stream(Tokens, id=stream.id) + await created.publish("five") + await created.close() + + task = asyncio.create_task(create_later()) + items = [item async for item in stream.subscribe(create_wait=timedelta(seconds=5))] + await task + assert isinstance(items[0], common.Record) and items[0].data == "five" + assert isinstance(items[-1], common.Closed) + + +def test_stream_errors_fail_the_workflow_not_the_task(): + closed = StreamErrorReason.STREAM_ERROR_REASON_CLOSED + err = common.StreamClosedError("closed", reason=closed) + assert isinstance(err, ApplicationError) + assert err.non_retryable and err.type == "StreamClosedError" + assert not common.StreamStorageError("down").non_retryable + mapped = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "closed", _details(closed) + ) + ) + assert mapped.code == RPCStatusCode.FAILED_PRECONDITION + # A Workflow await gets the error from a Local Activity, with the reason and no code. + returned = _errors.from_activity_error( + ApplicationError("closed", int(closed), type="StreamClosedError") + ) + assert isinstance(returned, common.StreamClosedError) + assert returned.reason == closed and returned.code is None + + +async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Tokens, id="two-writers") + other = client.stream(Tokens, id="two-writers") + await stream.publish("1") + await other.publish("2") + items = [item async for item in stream.subscribe(wait_new_records=False)] + producers = { + item.producer.id + for item in items + if isinstance(item, common.Record) and item.producer + } + assert len(producers) == 1 + # A handle that Python frees without a close frees its writer. So a server that makes a + # handle per request doesn't keep a writer for each. + del other + gc.collect() + assert len(extension.freed) == 1 + + +async def test_rpc_timeout_bounds_each_call_and_not_the_read(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Tokens, id="slow") + call = extension.call + + async def slow_describe(method: str, input: bytes) -> bytes: + if method == "DescribeStream": + await asyncio.sleep(1) + return await call(method, input) + + setattr(extension, "call", slow_describe) + with pytest.raises(RPCError) as err: + await stream.get_offset(rpc_timeout=timedelta(milliseconds=100)) + assert err.value.status == RPCStatusCode.DEADLINE_EXCEEDED + + async def publish_later() -> None: + await asyncio.sleep(1) + await stream.publish("late") + + # Each call of a waiting read returns before the deadline, so the read continues. + reader = stream.subscribe(rpc_timeout=timedelta(milliseconds=300)) + publisher = asyncio.create_task(publish_later()) + item = await anext(reader) + await reader.aclose() + await publisher + assert isinstance(item, common.Record) and item.data == "late" + + +class _CountingCodec(PayloadCodec): + def __init__(self) -> None: + self.calls: list[tuple[str, int]] = [] + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.calls.append(("encode", len(payloads))) + return list(payloads) + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.calls.append(("decode", len(payloads))) + return list(payloads) + + +async def test_codec_runs_once_for_each_write_and_page(client: Client): + # A remote codec costs a round trip for each call, so one call takes all the payloads of + # a message. + codec = _CountingCodec() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=codec) + client, _ = fake_client(Client(**config)) + stream = await client.create_stream(Tokens, id="coded") + await stream.publish(*(str(n) for n in range(5))) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + str(n) for n in range(5) + ] + assert codec.calls == [("encode", 5), ("decode", 5)] + + +async def test_forked_process_gets_its_own_client_producers(client: Client): + # A forked child has the parent's key and count, so the fork hook replaces them. + before = _stream_client._client_producer(client, None).id + _stream_client._new_process_key() + after = _stream_client._client_producer(client, None).id + assert before != after diff --git a/tests/contrib/streams/test_reader.py b/tests/contrib/streams/test_reader.py new file mode 100644 index 000000000..2a5653abf --- /dev/null +++ b/tests/contrib/streams/test_reader.py @@ -0,0 +1,81 @@ +from collections.abc import AsyncGenerator, Awaitable +from datetime import timedelta +from typing import Any + +from typing_extensions import override + +from temporalio.common._streams import _reader +from temporalio.common._streams._items import Record, Superseded +from temporalio.common._streams.system import client, models + + +class _Pages(client.StreamServiceClientExtension): + """Answers each read with the next page, and keeps the requests.""" + + def __init__(self, *pages: models.ReadRecordsResponse) -> None: + self.pages = list(pages) + self.requests: list[models.ReadRecordsRequest] = [] + + @override + def _create_nexus_client(self, *, service: str, endpoint: str | None) -> Any: + return self + + async def start_operation( + self, operation: str, input: Any, *, output_type: type[Any] | None = None + ) -> Awaitable[Any]: + assert operation == "ReadRecords" + assert output_type is models.ReadRecordsResponse + self.requests.append(input) + page = self.pages.pop(0) + + async def done() -> Any: + return page + + return done() + + +async def test_a_waiting_read_carries_the_offset_until_records_come(): + record = models.StreamRecord( + item=models.StreamItemData(models.StreamData(body=None)), offset=9 + ) + service = _Pages( + models.ReadRecordsResponse(next_offset=8), + models.ReadRecordsResponse(records=[record], next_offset=10), + ) + page = await _reader._read_until_records( + service, + models.ReadRecordsRequest(stream=models.StreamRef(id="wf:run"), from_offset=1), + timedelta(seconds=5), + ) + assert page.records and page.records[0].offset == 9 + first, second = service.requests + assert (first.from_offset, second.from_offset) == (1, 8) + assert second.wait is not None + assert timedelta() < second.wait <= timedelta(seconds=60) + + +async def test_superseded_lists_each_withdrawn_offset_once(): + kind = models.StreamProducerKind.StreamProducerKindActivity + failed = models.Producer( + id="act:run:1:1", kind=kind, epoch=models.ProducerEpoch(attempt=1) + ) + other = models.Producer( + id="act:run:2:2", kind=kind, epoch=models.ProducerEpoch(attempt=1) + ) + # A batch has consecutive offsets, and other Producers' Records come between batches. + written = [(10, failed), (11, failed), (12, failed), (13, other), (20, failed)] + supersession = models.Supersession( + producer_id=failed.id, producer_kind=kind, epoch=models.ProducerEpoch(attempt=2) + ) + + async def items() -> AsyncGenerator[_reader.Item, None]: + for offset, producer in written: + yield Record(data=None, offset=offset, producer=producer) + yield Superseded(offset=30, supersession=supersession) + yield Superseded(offset=31, supersession=supersession) + + read: _reader.Subscription[None] = _reader.Subscription( + items(), _reader.Position(0) + ) + superseded = [item async for item in read if isinstance(item, Superseded)] + assert [item.withdrawn_offsets for item in superseded] == [(10, 11, 12, 20), ()] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py new file mode 100644 index 000000000..be315b765 --- /dev/null +++ b/tests/contrib/streams/test_stream_client.py @@ -0,0 +1,84 @@ +import uuid +from dataclasses import dataclass + +import pytest + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.converter import DataConverter +from tests.contrib.streams.fake_extension import fake_client +from tests.helpers import new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def test_unknown_backend_raises_before_any_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(common.StreamConfigError, match="no Stream backend redis"): + await client.create_stream(Turns, id="chat-1", backend="redis") + with pytest.raises(common.StreamConfigError, match="It has fake"): + client.stream(Turns, id="chat-1", backend="redis") + assert not extension.streams and not extension.namespaces + stream = await client.create_stream(Turns, id="chat-1") + assert stream.backend == "fake" + + +async def test_client_without_backend_has_no_transport(client: Client): + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + client.stream(Turns, id="chat-1") + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + await client.create_stream(Turns, id="chat-1") + + +async def test_decoded_client_stream_has_no_client(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Turns, id="chat-2") + converter = DataConverter.default.payload_converter + decoded = converter.from_payload( + converter.to_payload(stream), common.ClientStream[Turn] + ) + assert decoded == stream + with pytest.raises(common.StreamConfigError, match="Client.stream"): + await decoded.get_offset() + + +@workflow.defn +class HandsOutStream: + @workflow.run + async def run(self) -> None: + pass + + @workflow.query + def stream(self) -> dict[str, str]: + # The proto JSON of a ref, as a Workflow Stream that a query returns. + return {"id": "handed-out"} + + +async def test_query_result_binds_its_client_stream(client: Client): + client, _ = fake_client(client) + created = await client.create_stream(Tokens, id="handed-out") + async with new_worker(client, HandsOutStream) as worker: + handle = await client.start_workflow( + HandsOutStream.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + stream = await handle.query("stream", result_type=common.ClientStream[str]) + # The ref names no backend, so the handle uses the default one. + assert stream.id == "handed-out" and stream.backend is None + await stream.publish("from the query result") + items = [item async for item in created.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "from the query result" + ] From 468ade2b4b3e1ee6317ed807114fe14fb53c6565 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 31 Oct 2026 01:56:08 -0700 Subject: [PATCH 2/3] Moved the client calls to the new bindings, and named handles `handle`. A closed or expired Stream fails the create, and a close sends only its failure, which sets the close status. A read skips an item kind that it doesn't know. --- temporalio/client/_client.py | 8 ++-- temporalio/common/_streams/_access.py | 18 ++++---- temporalio/common/_streams/_reader.py | 28 ++++++++---- temporalio/common/_streams/_stream_client.py | 46 ++++++++------------ temporalio/common/_streams/_writer.py | 15 +++---- tests/contrib/streams/fake_extension.py | 21 +++++---- tests/contrib/streams/test_facade.py | 13 +++--- tests/contrib/streams/test_reader.py | 39 +++++++++++------ 8 files changed, 102 insertions(+), 86 deletions(-) diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index 569b17aef..6102d6eaa 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -967,7 +967,7 @@ def stream( definition: type[StreamDefinition[_StreamPayloadT]], *, id: str | None = None, - stream: StreamHandle[Any] | None = None, + handle: 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. @@ -981,7 +981,7 @@ def stream( 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, + handle: 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``. @@ -989,13 +989,13 @@ def stream( 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. + TypeError: Neither or both of ``id`` and ``handle`` 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 + self, definition, id=id, handle=handle, backend=backend ) def get_workflow_handle_for( diff --git a/temporalio/common/_streams/_access.py b/temporalio/common/_streams/_access.py index 20bf6bd0a..77d34baca 100644 --- a/temporalio/common/_streams/_access.py +++ b/temporalio/common/_streams/_access.py @@ -10,7 +10,7 @@ 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._items import End, close_failure from temporalio.common._streams._reader import Item, Position, Subscription from temporalio.common._streams._writer import StreamWriter from temporalio.common._streams.system import models @@ -59,8 +59,8 @@ async def get_offset(self) -> int: 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)) + response = await service.describe_stream( + handle=models.StreamHandle(id=self.stream_id) ) return response.last_offset + 1 if response.last_offset else 0 @@ -71,15 +71,15 @@ async def close(self, result: Any | None, failure: BaseException | None) -> None 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), + # The generated close takes its Producer from a source, so the handle sends the request + # with its own. + await service._stream_service_close_stream( + models.CloseStreamRequest( + handle=models.StreamHandle(id=self.stream_id), producer=producer, result=result, - failure=failure, - status=models.StreamCloseStatus(status) if status is not None else None, + failure=close_failure(failure), ) ) diff --git a/temporalio/common/_streams/_reader.py b/temporalio/common/_streams/_reader.py index 321c52cb4..964bade2b 100644 --- a/temporalio/common/_streams/_reader.py +++ b/temporalio/common/_streams/_reader.py @@ -9,6 +9,7 @@ from datetime import timedelta from typing import Any, Generic, TypeVar +import temporalio.common import temporalio.converter from temporalio.api.enums.v1 import StreamProducerKind from temporalio.common._streams import _backend, _convert, _errors @@ -118,14 +119,14 @@ async def __anext__(self) -> Record[T] | Superseded | Closed: runs.extend((item.offset, item.offset)) elif isinstance(item, Superseded): item = dataclasses.replace( - item, withdrawn_offsets=self._withdraw(item.supersession) + item, withdrawn_offsets=self._withdraw(item.facts) ) return item - def _withdraw(self, supersession: models.Supersession) -> tuple[int, ...]: + def _withdraw(self, facts: models.StreamSuperseded) -> tuple[int, ...]: offsets: list[int] = [] for key, (producer, runs) in list(self._shown.items()): - if withdraws(supersession, producer): + if withdraws(facts, producer): for first, last in zip(runs[::2], runs[1::2]): offsets.extend(range(first, last + 1)) # A withdrawn Record stays withdrawn, so the read forgets it. @@ -169,7 +170,7 @@ def request() -> models.ReadRecordsRequest: else None ) return models.ReadRecordsRequest( - stream=models.StreamRef(id=stream_id), + handle=models.StreamHandle(id=stream_id), # END applies only to the first page. Later pages start at the last page's next offset. from_offset=from_offset if isinstance(from_offset, int) else None, start=models.StreamReadStart.StreamReadStartEnd @@ -190,6 +191,8 @@ def request() -> models.ReadRecordsRequest: for record in page.records or []: try: item = await _item(converter, record, payload_type) + if item is None: + continue except Exception as err: # Move next_offset past the Record, so a resume doesn't fail on it again. # The error gives the offset, for a read of that Record. @@ -255,7 +258,10 @@ async def _read_page( service: service_client.StreamServiceClientExtension, request: models.ReadRecordsRequest, ) -> models.ReadRecordsResponse: - return await (await service._read_records(request)) + # The bodies stay RawValue, so the read decodes each one with the Stream's payload type. + return await service._stream_service_read_records( + request, temporalio.common.RawValue + ) async def _read_until_records( @@ -289,16 +295,20 @@ async def _read_until_records( async def _item( converter: temporalio.converter.DataConverter, - record: models.StreamRecord, + record: models.StreamRecord[temporalio.common.RawValue], payload_type: Any, -) -> Item: +) -> Item | None: offset = record.offset or 0 match record.item: + case None: + # A newer backend's item kind that this SDK doesn't know is empty, so the read skips + # it. + return None case models.StreamItemData(value=data): value = _convert.decode(converter, data.body, payload_type) return Record(data=value, offset=offset, producer=data.producer) - case models.StreamItemSuperseded(value=supersession): - return Superseded(offset=offset, supersession=supersession) + case models.StreamItemSuperseded(value=facts): + return Superseded(offset=offset, facts=facts) case models.StreamItemClose(value=close): result = _convert.decode(converter, close.result, None) failure = close.failure diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py index 210f9bf3e..1508a2a3a 100644 --- a/temporalio/common/_streams/_stream_client.py +++ b/temporalio/common/_streams/_stream_client.py @@ -14,7 +14,7 @@ import temporalio.api.stream.v1 as stream_v1 import temporalio.client -from temporalio.api.enums.v1 import StreamErrorReason, StreamProducerKind, StreamStatus +from temporalio.api.enums.v1 import StreamProducerKind from temporalio.common._streams import _backend, _errors, _ids from temporalio.common._streams._access import Access, ReadSpec, read_spec, subscription from temporalio.common._streams._backend import BackendCalls @@ -134,7 +134,7 @@ def _access(self) -> Access: if client is None: raise _errors.StreamConfigError( "This ClientStream came from a payload that no client query decoded, so it " - "has no client. Use Client.stream(definition, stream=handle)." + "has no client. Use Client.stream(definition, handle=handle)." ) if self._stream_backend is None: _, self._stream_backend = client_backend(client, self._backend) @@ -279,8 +279,7 @@ def _check_close(self) -> None: # Only the run chain can close a Workflow's Stream. The Store refuses a client's # close too, but this check makes no call. raise _errors.StreamNotOwnerError( - f"Can't close the Stream {self._id} from a client, since a Workflow owns it", - reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + f"Can't close the Stream {self._id} from a client, since a Workflow owns it" ) async def close( @@ -350,19 +349,10 @@ def create( async def create() -> ClientStream[T]: info = definition_of(definition) backend_name, calls = client_backend(client, backend) - created = await calls.service(client.namespace).create_stream( - stream_id=id, retention=options.get("retention") + # A closed, expired or another owner's Stream fails the create. + await calls.service(client.namespace).create_stream( + definition, id=id, retention=options.get("retention") ) - if created.status == StreamStatus.STREAM_STATUS_CLOSED: - raise _errors.StreamClosedError( - f"The Stream {id} was created and closed before", - reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, - ) - if created.status == StreamStatus.STREAM_STATUS_EXPIRED: - raise _errors.StreamExpiredError( - f"The Stream {id} expired", - reason=StreamErrorReason.STREAM_ERROR_REASON_EXPIRED, - ) stream: ClientStream[T] = ClientStream( id=id, backend=backend_name, payload_type=info.type ) @@ -376,7 +366,7 @@ def client_stream( definition: type[StreamDefinition[T]], *, id: str | None, - stream: StreamHandle[Any] | None, + handle: StreamHandle[Any] | None, backend: str | None, ) -> ClientStream[T]: """Get a Stream by its id, or from a handle that came in data. @@ -384,27 +374,27 @@ def client_stream( See :py:meth:`temporalio.client.Client.stream`. """ info = definition_of(definition) - if (id is None) == (stream is None): - raise TypeError("stream takes one of id= and stream=") - if stream is not None: - if backend and stream.backend and backend != stream.backend: + if (id is None) == (handle is None): + raise TypeError("stream takes one of id= and handle=") + if handle is not None: + if backend and handle.backend and backend != handle.backend: raise ValueError( - f"The handle is of the backend {stream.backend}, not {backend}" + f"The handle is of the backend {handle.backend}, not {backend}" ) - backend = backend or stream.backend - stream_id = stream.id if stream else id + backend = backend or handle.backend + stream_id = handle.id if handle else id if not stream_id: # An empty id would fail only at the first call, with an unclear error. raise ValueError("The Stream id is empty") backend_name, calls = client_backend(client, backend) - handle: ClientStream[T] = ClientStream( + stream: ClientStream[T] = ClientStream( id=stream_id, - workflow_id=stream.workflow_id if stream else None, - first_execution_run_id=stream.first_execution_run_id if stream else None, + workflow_id=handle.workflow_id if handle else None, + first_execution_run_id=handle.first_execution_run_id if handle else None, backend=backend_name, payload_type=info.type, ) - return handle._bind(client, calls) + return stream._bind(client, calls) class BindQueryResults(temporalio.client.Interceptor): diff --git a/temporalio/common/_streams/_writer.py b/temporalio/common/_streams/_writer.py index 556407d92..c234ce882 100644 --- a/temporalio/common/_streams/_writer.py +++ b/temporalio/common/_streams/_writer.py @@ -8,15 +8,15 @@ from typing import TYPE_CHECKING, Any import temporalio.api.stream.v1 as stream_v1 -from temporalio.api.enums.v1 import StreamErrorReason from temporalio.common._streams import _errors if TYPE_CHECKING: from temporalio.common._streams.writer_system import client as writer_client + from temporalio.common._streams.writer_system.client import Writer class StreamWriter: - """Writes the Records of one Producer to one Stream, through the generated Writer.""" + """Writes the Records of one Producer to one Stream, through the writer client.""" def __init__( self, @@ -40,7 +40,7 @@ def __init__( self._release: weakref.finalize[Any, StreamWriter] | None = None self._previous_completed_task = previous_completed_task self._failed_tasks = list(failed_tasks) - self._writer: writer_client.Writer | None = None + self._writer: Writer | None = None self._open_lock = asyncio.Lock() self._closed = False @@ -64,7 +64,7 @@ async def withdraw_earlier_attempts(self) -> None: ): pass - async def open(self) -> writer_client.Writer: + async def open(self) -> Writer: """Open the crate's writer once. The open of a new attempt withdraws the Records of earlier attempts, even if it writes @@ -79,7 +79,7 @@ async def open(self) -> writer_client.Writer: ) previous = self._previous_completed_task self._writer = await self._writers.open_writer( - stream=models.StreamRef(id=self._stream_id), + handle=models.StreamHandle(id=self._stream_id), producer=producer, previous_completed_task=models.WorkflowTaskEpoch( started_event_id=previous.started_event_id @@ -101,10 +101,7 @@ async def write(self, values: Sequence[Any]) -> int: from temporalio.common._streams.writer_system import models if self._closed: - raise _errors.StreamClosedError( - "The writer is closed", - reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, - ) + raise _errors.StreamClosedError("The writer is closed") writer = await self.open() answer = await writer.write_records( [models.AppendRecord(body=value) for value in values] diff --git a/tests/contrib/streams/fake_extension.py b/tests/contrib/streams/fake_extension.py index c4945b81f..74a22d561 100644 --- a/tests/contrib/streams/fake_extension.py +++ b/tests/contrib/streams/fake_extension.py @@ -60,14 +60,19 @@ async def call(self, method: str, input: bytes) -> bytes: if method == "CreateStream": create = stream_pb.CreateStreamRequest.FromString(input) self.namespaces.add(create.namespace) - self.streams.setdefault(create.stream_id, _Stream()) + stream = self.streams.setdefault(create.stream_id, _Stream()) + if stream.closed: + raise _error( + RPCStatusCode.FAILED_PRECONDITION, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) return stream_pb.CreateStreamResponse( - stream=stream_v1.StreamRef(id=create.stream_id), created=True + handle=stream_v1.StreamHandle(id=create.stream_id), created=True ).SerializeToString() if method == "OpenWriter": open_writer = ext_pb.OpenWriterRequest.FromString(input) writer_id = len(self.writers) + 1 - self.writers[writer_id] = (open_writer.stream.id, open_writer.producer) + self.writers[writer_id] = (open_writer.handle.id, open_writer.producer) return ext_pb.OpenWriterResponse(writer_id=writer_id).SerializeToString() if method == "WriteRecords": write = ext_pb.WriteRecordsRequest.FromString(input) @@ -92,7 +97,7 @@ async def call(self, method: str, input: bytes) -> bytes: if method == "CloseStream": close = stream_pb.CloseStreamRequest.FromString(input) self._append( - close.stream.id, + close.handle.id, stream_v1.StreamRecord( close=stream_v1.StreamClose( producer=close.producer, @@ -100,11 +105,11 @@ async def call(self, method: str, input: bytes) -> bytes: ) ), ) - self.streams[close.stream.id].closed = True + self.streams[close.handle.id].closed = True return b"" if method == "DescribeStream": describe = stream_pb.DescribeStreamRequest.FromString(input) - records = self.streams[describe.stream.id].records + records = self.streams[describe.handle.id].records return stream_pb.DescribeStreamResponse( last_offset=records[-1].offset if records else 0 ).SerializeToString() @@ -122,14 +127,14 @@ async def _read( ) -> stream_pb.ReadRecordsResponse: loop = asyncio.get_running_loop() create_until = loop.time() + request.create_wait.ToTimedelta().total_seconds() - while request.stream.id not in self.streams: + while request.handle.id not in self.streams: if loop.time() >= create_until: raise _error( RPCStatusCode.NOT_FOUND, StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND, ) await asyncio.sleep(0.01) - stream = self.streams[request.stream.id] + stream = self.streams[request.handle.id] offset = request.from_offset or 1 wait = request.wait.ToTimedelta().total_seconds() until = loop.time() + wait diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index 42068e79b..0b9bc7cd7 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -108,22 +108,25 @@ async def create_later() -> None: def test_stream_errors_fail_the_workflow_not_the_task(): closed = StreamErrorReason.STREAM_ERROR_REASON_CLOSED - err = common.StreamClosedError("closed", reason=closed) + err = common.StreamClosedError("closed") assert isinstance(err, ApplicationError) assert err.non_retryable and err.type == "StreamClosedError" + assert err.reason == closed assert not common.StreamStorageError("down").non_retryable mapped = _errors.from_call_error( StreamCallError( int(RPCStatusCode.FAILED_PRECONDITION), "closed", _details(closed) ) ) - assert mapped.code == RPCStatusCode.FAILED_PRECONDITION - # A Workflow await gets the error from a Local Activity, with the reason and no code. + assert isinstance(mapped, common.StreamClosedError) and mapped.reason == closed + # A Workflow await gets the error from a Local Activity, with the reason in its details. returned = _errors.from_activity_error( - ApplicationError("closed", int(closed), type="StreamClosedError") + ApplicationError( + "closed", StreamErrorDetails(reason=closed), type="StreamClosedError" + ) ) assert isinstance(returned, common.StreamClosedError) - assert returned.reason == closed and returned.code is None + assert returned.reason == closed async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client): diff --git a/tests/contrib/streams/test_reader.py b/tests/contrib/streams/test_reader.py index 2a5653abf..2bbe63e81 100644 --- a/tests/contrib/streams/test_reader.py +++ b/tests/contrib/streams/test_reader.py @@ -1,6 +1,6 @@ -from collections.abc import AsyncGenerator, Awaitable +from collections.abc import AsyncGenerator, Mapping from datetime import timedelta -from typing import Any +from typing import Any, get_origin from typing_extensions import override @@ -17,21 +17,30 @@ def __init__(self, *pages: models.ReadRecordsResponse) -> None: self.requests: list[models.ReadRecordsRequest] = [] @override - def _create_nexus_client(self, *, service: str, endpoint: str | None) -> Any: + def _create_nexus_client(self, *, service: str, endpoint: str) -> Any: return self async def start_operation( - self, operation: str, input: Any, *, output_type: type[Any] | None = None - ) -> Awaitable[Any]: + self, + operation: str, + arg: Any, + *, + id: str, + result_type: type[Any] | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Any: + del id, rpc_metadata, rpc_timeout assert operation == "ReadRecords" - assert output_type is models.ReadRecordsResponse - self.requests.append(input) + assert get_origin(result_type) is models.ReadRecordsResponse + self.requests.append(arg) page = self.pages.pop(0) - async def done() -> Any: - return page + class Done: + async def result(self) -> Any: + return page - return done() + return Done() async def test_a_waiting_read_carries_the_offset_until_records_come(): @@ -44,7 +53,9 @@ async def test_a_waiting_read_carries_the_offset_until_records_come(): ) page = await _reader._read_until_records( service, - models.ReadRecordsRequest(stream=models.StreamRef(id="wf:run"), from_offset=1), + models.ReadRecordsRequest( + handle=models.StreamHandle(id="wf:run"), from_offset=1 + ), timedelta(seconds=5), ) assert page.records and page.records[0].offset == 9 @@ -64,15 +75,15 @@ async def test_superseded_lists_each_withdrawn_offset_once(): ) # A batch has consecutive offsets, and other Producers' Records come between batches. written = [(10, failed), (11, failed), (12, failed), (13, other), (20, failed)] - supersession = models.Supersession( + facts = models.StreamSuperseded( producer_id=failed.id, producer_kind=kind, epoch=models.ProducerEpoch(attempt=2) ) async def items() -> AsyncGenerator[_reader.Item, None]: for offset, producer in written: yield Record(data=None, offset=offset, producer=producer) - yield Superseded(offset=30, supersession=supersession) - yield Superseded(offset=31, supersession=supersession) + yield Superseded(offset=30, facts=facts) + yield Superseded(offset=31, facts=facts) read: _reader.Subscription[None] = _reader.Subscription( items(), _reader.Position(0) From adadfc4f8382969bd11fadecd908dc8b95343cca Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 1 Nov 2026 01:47:24 -0700 Subject: [PATCH 3/3] Gave `StreamHandle` the client calls, and dropped `ClientStream`. The proposed API has one handle type in every context. A handle passes each call to the calls of the context where it runs. In client code that's the client that made the handle or the client that decoded it. `Client.get_stream_handle` replaces `Client.stream`, and `publish` starts the write and returns its future. --- temporalio/client/_client.py | 31 +- temporalio/common/__init__.py | 3 - temporalio/common/_streams/__init__.py | 2 - temporalio/common/_streams/_access.py | 10 +- temporalio/common/_streams/_handle.py | 300 +++++++++++++- temporalio/common/_streams/_plugin.py | 8 +- temporalio/common/_streams/_stream_client.py | 412 +++++++++---------- temporalio/common/_streams/_writer.py | 101 ++++- tests/contrib/streams/test_facade.py | 4 +- tests/contrib/streams/test_stream_client.py | 93 ++++- 10 files changed, 675 insertions(+), 289 deletions(-) diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index 6102d6eaa..bdd12a239 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -112,7 +112,6 @@ if TYPE_CHECKING: from temporalio.common import ( - ClientStream, Creating, StreamBackend, StreamDefinition, @@ -320,7 +319,7 @@ def __init__( # Only a client with Stream backends loads the Streams package. from temporalio.common._streams import _plugin as _stream_plugin - config = _stream_plugin.configure_client(config) + config = _stream_plugin.configure_client(config, self) self._init_from_config(config) @@ -928,7 +927,7 @@ def create_stream( id: str, backend: str | None = None, **options: Unpack[StreamOptions], - ) -> Creating[ClientStream[_StreamPayloadT]]: + ) -> Creating[StreamHandle[_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`` @@ -962,18 +961,17 @@ def create_stream( self, definition, id=id, backend=backend, options=options ) - def stream( + def get_stream_handle( self, definition: type[StreamDefinition[_StreamPayloadT]], *, - id: str | None = None, - handle: StreamHandle[Any] | None = None, + id: str, 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. + ) -> StreamHandle[_StreamPayloadT]: + """Get a Stream by its id. It makes no call. - 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. + To create a Stream that the client owns, use :py:meth:`create_stream`. A handle that + came as data, such as from a Workflow's query, needs no call of this method. .. warning:: This API is experimental. @@ -981,21 +979,18 @@ def stream( Args: definition: The Stream definition class. id: The Stream id. - handle: 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``. + backend: The name of the backend. Defaults 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 ``handle`` are given. - ValueError: The id is empty, or ``backend`` is not the handle's backend. + ValueError: The id is empty. """ from temporalio.common._streams import _stream_client - return _stream_client.client_stream( - self, definition, id=id, handle=handle, backend=backend + return _stream_client.get_stream_handle( + self, definition, id=id, backend=backend ) def get_workflow_handle_for( diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 1bcd2f008..2d77f6c63 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -64,13 +64,11 @@ ) 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", @@ -1589,7 +1587,6 @@ 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", diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 8e4785be2..c803acdb7 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -26,11 +26,9 @@ ) 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", diff --git a/temporalio/common/_streams/_access.py b/temporalio/common/_streams/_access.py index 77d34baca..a68ed45c9 100644 --- a/temporalio/common/_streams/_access.py +++ b/temporalio/common/_streams/_access.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from collections.abc import AsyncGenerator, Callable, Sequence from dataclasses import dataclass from datetime import timedelta @@ -26,11 +27,14 @@ class Access: 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.""" + def publish(self, values: Sequence[Any]) -> asyncio.Future[int]: + """Start the write of the values as Records, in order, and return its future. + + The future gives the offset of the last Record. + """ if not values: raise TypeError("publish takes one or more values") - return await self.writer().write(values) + return self.writer().write(values) async def items( self, read: ReadSpec, position: Position diff --git a/temporalio/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py index 1a12d1ec3..deb0f39a2 100644 --- a/temporalio/common/_streams/_handle.py +++ b/temporalio/common/_streams/_handle.py @@ -2,20 +2,27 @@ from __future__ import annotations +import abc import asyncio +import contextvars import typing -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from datetime import timedelta from types import TracebackType -from typing import Any, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar import google.protobuf.json_format from typing_extensions import Self, TypedDict import temporalio.converter import temporalio.exceptions +from temporalio.common._streams import _errors from temporalio.common._streams.system import models +if TYPE_CHECKING: + from temporalio.common._streams._items import End + from temporalio.common._streams._reader import Subscription + T = TypeVar("T") @@ -171,12 +178,118 @@ def from_transfer_type( return _from_json(value, type_hint) +# The client whose data the converter decodes. A handle in the data takes it, so client code uses +# a Stream that a Workflow hands out with no other call. +decoding_client: contextvars.ContextVar[Any] = contextvars.ContextVar( + "__temporal_stream_decoding_client", default=None +) + + +class ContextCalls(abc.ABC): + """The Stream calls of a handle in one context, such as client code. + + The handle passes each call to the calls of the context where it runs. A call that the + context can't make raises :py:class:`StreamUnsupportedError`. + """ + + @abc.abstractmethod + def publish( + self, + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. See :py:meth:`StreamHandle.publish`.""" + + @abc.abstractmethod + def subscribe( + self, + handle: StreamHandle[Any], + *, + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read. See :py:meth:`StreamHandle.subscribe`.""" + + @abc.abstractmethod + async def get_offset( + self, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> int: + """Get the next offset. See :py:meth:`StreamHandle.get_offset`.""" + + @abc.abstractmethod + async def close( + self, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> None: + """Close the Stream. See :py:meth:`StreamHandle.close`.""" + + @abc.abstractmethod + async def enter(self, handle: StreamHandle[Any]) -> None: + """Start an ``async with`` block of the handle.""" + + @abc.abstractmethod + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """End an ``async with`` block of the handle. See :py:meth:`StreamHandle.__aexit__`.""" + + def publish_nowait( + self, handle: StreamHandle[Any], values: tuple[Any, ...] + ) -> None: + """Write without an await. See :py:meth:`StreamHandle.publish_nowait`.""" + del handle, values + raise _errors.StreamUnsupportedError( + "publish_nowait is for sync Workflow code. Elsewhere, use publish." + ) + + def close_nowait( + self, handle: StreamHandle[Any], result: Any, failure: BaseException | None + ) -> None: + """Close without an await. See :py:meth:`StreamHandle.close_nowait`.""" + del handle, result, failure + raise _errors.StreamUnsupportedError( + "close_nowait is for sync Workflow code. Elsewhere, use close." + ) + + def producer( + self, handle: StreamHandle[Any], name: str, epoch: int + ) -> StreamHandle[Any]: + """Return a handle of a named Producer. See :py:meth:`StreamHandle.producer`.""" + del handle, name, epoch + raise _errors.StreamUnsupportedError( + "Only client code writes as a named Producer" + ) + + class StreamHandle(Generic[T]): - """A handle on one Stream in Workflow, Activity and client code. ``T`` is the payload type. + """A handle on one Stream. ``T`` is the payload type. + + One handle type works in Workflow, Activity and client code. The SDK picks the write path + and the Producer from where the call runs. A call that the context can't make raises + :py:class:`StreamUnsupportedError`. - A handle is data. It's passed as a Workflow or Activity argument in the proto JSON form of - the ``temporal.api.stream.v1.StreamHandle`` message. The parameter's type, such as - ``StreamHandle[Turn]``, gives the payload type. + A handle is data. It's passed as a Workflow or Activity argument, or returned from a query, + in the proto JSON form of the ``temporal.api.stream.v1.StreamHandle`` message. The type + hint, such as ``StreamHandle[Turn]``, gives the payload type. + + In client code, a handle uses the client that made it, such as with + :py:meth:`temporalio.client.Client.get_stream_handle`, or the client that decoded it, such + as from a query, an Update or a Workflow result. A handle that no client decoded, such as + one that user code builds from JSON, raises :py:class:`StreamConfigError` at its first call. + All handles of one client and Stream write as one Producer, and :py:meth:`producer` returns + a handle that writes as a named Producer. A write error that no caller got raises at the + next call of the handle. .. warning:: This API is experimental. @@ -198,6 +311,10 @@ def __init__( self._backend = backend # The type that a read decodes each value with. None decodes with no type hint. self._payload_type = payload_type + # Client code binds a handle to a client and keeps its writer here. Workflow and Activity + # calls keep no state in the handle. + self._client: Any = decoding_client.get() + self._client_calls: ContextCalls | None = None @property def id(self) -> str: @@ -219,6 +336,177 @@ def backend(self) -> str | None: """Name of the backend that stores the Stream. None means the default one.""" return self._backend + def _calls(self) -> ContextCalls: + # The context modules import on first use, so Workflow code never imports the client's. + from temporalio.common._streams import _stream_client + + return _stream_client.calls(self) + + def publish( + self, + *values: T, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> asyncio.Future[int]: + """Start the write of one Record for each value, in order, and return its future. + + The write has its place in the Stream's order when the call returns, also when no + caller awaits the future. The writes of one handle keep the order of the calls. A cancel + of the future doesn't stop the write. A list is one value. In client code, the next call + of the handle raises a write error that no caller got. + + Args: + values: One value for each Record. + rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so + they send none. + rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with + ``DEADLINE_EXCEEDED``. + + Returns: + The future of the write. It gives the offset of the last Record, or raises + ``StreamError`` when the backend refused the write, such as + ``StreamNotFoundError`` for a Stream that nobody created. + + Raises: + StreamClosedError: This handle closed the Stream. + StreamConfigError: No client decoded the handle, or the client has no backend of + that name. + """ + return self._calls().publish(self, values, rpc_metadata, rpc_timeout) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Subscription[T]: + """Read the Stream's items in order. + + The read ends after a ``Closed`` item, or at the newest Record when + ``wait_new_records`` is False. To stop after a quiet period, use your own timer, and + resume later with ``offset=`` set to the read's ``next_offset``. A ``Superseded`` item + withdraws earlier Records of a failed attempt. + + Args: + offset: Start at this offset, inclusive, or at ``common.END`` for new + Records only. Defaults to the first Record. + last: Start with the newest ``last`` Records, oldest first, such as for the first + view of a chat. Can't be used with ``offset``. A Record that a later + ``Superseded`` item withdraws doesn't count. At most 1000. + wait_new_records: Wait for new Records at the end of the Stream. + create_wait: How long to wait for a Stream that doesn't exist yet, such as the + Stream of a Workflow that just started. Defaults to no wait. + rpc_metadata: Headers of each call of the read. See :py:meth:`publish`. + rpc_timeout: The deadline of each call of the read, not of the whole read. A call + that waits for new Records returns before the deadline. + """ + return self._calls().subscribe( + self, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + rpc_metadata=rpc_metadata, + rpc_timeout=rpc_timeout, + ) + + async def get_offset( + self, + *, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> 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. + ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + """ + return await self._calls().get_offset(self, rpc_metadata, rpc_timeout) + + async def close( + self, + *, + result: Any = None, + failure: BaseException | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> None: + """Write this handle's Records, then close the Stream. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + + Raises: + StreamNotOwnerError: The code doesn't own the Stream, such as a client with a + Workflow's Stream. Only the run chain closes a Workflow's Stream, and only a + client closes a client's Stream. + """ + await self._calls().close(self, result, failure, rpc_metadata, rpc_timeout) + + async def __aenter__(self) -> Self: + """Use the Stream in ``async with``, which closes it at the end. + + Raises: + StreamNotOwnerError: The code can't close the Stream. See :py:meth:`close`. + """ + await self._calls().enter(self) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this + does nothing. + """ + await self._calls().exit(self, exc) + + def publish_nowait(self, *values: T) -> None: + """Write one Record for each value, in order, without an await, in sync Workflow code. + + Raises: + StreamUnsupportedError: The code isn't Workflow code. + """ + self._calls().publish_nowait(self, values) + + def close_nowait( + self, *, result: Any = None, failure: BaseException | None = None + ) -> None: + """Close the Stream without an await, in sync Workflow code. + + Raises: + StreamUnsupportedError: The code isn't Workflow code. + """ + self._calls().close_nowait(self, result, failure) + + def producer(self, name: str, epoch: int) -> StreamHandle[T]: + """Return this Stream as a handle that writes as the Producer ``name`` at ``epoch``. + + Use it for a writer that must fence its earlier instances, such as a server that + restarts. Keep ``epoch`` in durable storage, and raise it at each start. After a write + at a higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The + Records of the earlier epochs aren't withdrawn, since they're real messages. Run only + one live writer for each name. + + Args: + name: The Producer name, the same across restarts. + epoch: The writer's epoch, 1 or more. + + Raises: + StreamUnsupportedError: The code isn't client code. + """ + return self._calls().producer(self, name, epoch) + def _fields(self) -> dict[str, str]: # The fields that name the Stream. They need no call, so a handle whose id is not known # yet still compares, hashes and prints. diff --git a/temporalio/common/_streams/_plugin.py b/temporalio/common/_streams/_plugin.py index e22c9ffac..51b2f85ce 100644 --- a/temporalio/common/_streams/_plugin.py +++ b/temporalio/common/_streams/_plugin.py @@ -157,18 +157,18 @@ def default_name(backends: Sequence[StreamBackend]) -> str: def configure_client( - config: temporalio.client.ClientConfig, + config: temporalio.client.ClientConfig, client: temporalio.client.Client ) -> temporalio.client.ClientConfig: """Check the client's backends, and apply the plugin configuration of each one. - The client's query results bind each ``ClientStream`` in them to the client. + The decodes of ``client`` bind each Stream handle in the data to it. """ # The client Stream module imports this one. - from temporalio.common._streams._stream_client import BindQueryResults + from temporalio.common._streams._stream_client import bind_decodes backends = check_backends(config.get("stream_backends", ()), "client") config["stream_backends"] = backends - config["interceptors"] = [*config.get("interceptors", []), BindQueryResults()] + config["data_converter"] = bind_decodes(config["data_converter"], client) for backend in backends: config = backend._configure_client(config) return config diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py index 1508a2a3a..7e95fd3f0 100644 --- a/temporalio/common/_streams/_stream_client.py +++ b/temporalio/common/_streams/_stream_client.py @@ -2,27 +2,33 @@ from __future__ import annotations -import contextvars +import asyncio +import dataclasses import itertools import os import uuid import weakref -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import timedelta -from types import TracebackType from typing import Any, TypeVar +from typing_extensions import Self + import temporalio.api.stream.v1 as stream_v1 import temporalio.client +import temporalio.converter +from temporalio.api.common.v1 import Payload from temporalio.api.enums.v1 import StreamProducerKind from temporalio.common._streams import _backend, _errors, _ids from temporalio.common._streams._access import Access, ReadSpec, read_spec, subscription from temporalio.common._streams._backend import BackendCalls from temporalio.common._streams._defn import StreamDefinition, definition_of from temporalio.common._streams._handle import ( + ContextCalls, Creating, StreamHandle, StreamOptions, + decoding_client, exit_failure, ) from temporalio.common._streams._items import End @@ -55,12 +61,6 @@ def _new_process_key() -> None: if _register_at_fork is not None: _register_at_fork(after_in_child=_new_process_key) -# The client whose query result the converter decodes. A ClientStream in the result takes it, so -# a reader uses a Stream that a Workflow hands out with no other call. -_decoding_client: contextvars.ContextVar[temporalio.client.Client | None] = ( - contextvars.ContextVar("__temporal_stream_decoding_client", default=None) -) - def _client_producer( client: temporalio.client.Client, named: tuple[str, int] | None @@ -83,68 +83,58 @@ def _client_producer( ) -class ClientStream(StreamHandle[T]): - """A Stream in client code. The client can write and read it, and close a client's Stream. +def calls(handle: StreamHandle[Any]) -> ClientCalls: + """The client calls of a handle, which the handle keeps. - Get one from :py:meth:`temporalio.client.Client.create_stream` or - :py:meth:`temporalio.client.Client.stream`. A ``ClientStream`` that a query of a client - returns, such as a Workflow's ``stream`` query, uses that client. + Raises: + StreamConfigError: No client decoded the handle, or the client has no backend of it. + """ + bound = handle._client_calls + if bound is None: + client = handle._client + if client is None: + # A lookup of some client would pick the wrong one when a process has two. So only + # a decode of a client binds a handle to it. + raise _errors.StreamConfigError( + f"The handle of the Stream {handle.id} has no client, since no client decoded " + "it. Use client.get_stream_handle(definition, id=handle.id)." + ) + _, backend = client_backend(client, handle.backend) + bound = ClientCalls(client, backend) + handle._client_calls = bound + assert isinstance(bound, ClientCalls) + return bound - Each publish writes before it returns. All handles of one client and Stream write as one - Producer. :py:meth:`producer` returns a handle that writes as a named Producer. Use a handle - in ``async with`` to close the Stream at the end of the block. - .. warning:: - This API is experimental. +class ClientCalls(ContextCalls): + """The Stream calls of one handle in client code. + + Each publish starts its write at the call. All handles of one client and Stream write as one + Producer, unless a handle is of a named Producer. """ def __init__( self, - *, - id: str, - workflow_id: str | None = None, - first_execution_run_id: str | None = None, - backend: str | None = None, - payload_type: Any = None, + client: temporalio.client.Client, + backend: BackendCalls, + named: tuple[str, int] | None = None, ) -> None: - """Create a handle. Use :py:meth:`temporalio.client.Client.stream` instead.""" - super().__init__( - id=id, - workflow_id=workflow_id, - first_execution_run_id=first_execution_run_id, - backend=backend, - payload_type=payload_type, - ) - self._client = _decoding_client.get() - self._stream_backend: BackendCalls | None = None + """Create the calls of a handle of ``client``. ``named`` is a Producer's name and epoch.""" + self._client = client + self._backend = backend + self._named = named self._writer: StreamWriter | None = None self._closed = False - # The name and the epoch of a named Producer. - self._named: tuple[str, int] | None = None - def _bind( - self, client: temporalio.client.Client, backend: BackendCalls - ) -> ClientStream[T]: - self._client = client - self._stream_backend = backend - return self - - def _access(self) -> Access: + def _access(self, handle: StreamHandle[Any]) -> Access: client = self._client - if client is None: - raise _errors.StreamConfigError( - "This ClientStream came from a payload that no client query decoded, so it " - "has no client. Use Client.stream(definition, handle=handle)." - ) - if self._stream_backend is None: - _, self._stream_backend = client_backend(client, self._backend) - backend = self._stream_backend + backend = self._backend def writer() -> StreamWriter: if self._writer is None: self._writer = StreamWriter( writers=backend.writers(client.namespace, client.data_converter), - stream_id=self._id, + stream_id=handle.id, producer=_client_producer(client, self._named), free=backend.free_writer, ) @@ -154,101 +144,59 @@ def writer() -> StreamWriter: backend=backend, converter=client.data_converter, namespace=client.namespace, - stream_id=self._id, + stream_id=handle.id, writer=writer, ) - def producer(self, name: str, epoch: int) -> ClientStream[T]: - """Return this Stream as a handle that writes as the Producer ``name`` at ``epoch``. - - Use it for a writer that must fence its earlier instances, such as a server that - restarts. Keep ``epoch`` in durable storage, and raise it at each start. After a write - at a higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The - Records of the earlier epochs aren't withdrawn, since they're real messages. Run only - one live writer for each name. - - Args: - name: The Producer name, the same across restarts. - epoch: The writer's epoch, 1 or more. - """ - if epoch < 1: - raise ValueError("A Producer's epoch is 1 or more") - handle: ClientStream[T] = ClientStream( - id=self._id, - workflow_id=self._workflow_id, - first_execution_run_id=self._first_execution_run_id, - backend=self._backend, - payload_type=self._payload_type, - ) - handle._named = (name, epoch) - handle._client = self._client - handle._stream_backend = self._stream_backend - return handle + def _raise_lost_write(self) -> None: + # Client code has no end, such as a run's, to raise a write error that no caller got. So + # the handle's next call raises it. + error = self._writer.unretrieved_error() if self._writer else None + if error is not None: + raise error - async def publish( + def publish( self, - *values: T, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> int: - """Write one Record for each value, in order, and return the last Record's offset. - - A list is one value. ``rpc_metadata`` and ``rpc_timeout`` are as for - :py:meth:`get_offset`. - - Raises: - StreamNotFoundError: Nobody created the Stream. - """ + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. The write's task takes the call options.""" + self._raise_lost_write() with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - return await self._access().publish(values) + return self._access(handle).publish(values) def subscribe( self, + handle: StreamHandle[Any], *, - offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> Subscription[T]: - """Read the Stream's items in order. - - The read ends after a ``Closed`` item, or at the newest Record when - ``wait_new_records`` is False. To stop after a quiet period, use your own timer, and - resume later with ``offset=`` set to the read's ``next_offset``. A ``Superseded`` item - withdraws earlier Records of a failed attempt. - - Args: - offset: Start at this offset, inclusive, or at ``common.END`` for new - Records only. Defaults to the first Record. - last: Start with the newest ``last`` Records, oldest first, such as for the first - view of a chat. Can't be used with ``offset``. A Record that a later - ``Superseded`` item withdraws doesn't count. At most 1000. - wait_new_records: Wait for new Records at the end of the Stream. - create_wait: How long to wait for a Stream that doesn't exist yet, such as the - Stream of a Workflow that just started. Defaults to no wait. - rpc_metadata: Headers of each call of the read. See :py:meth:`get_offset`. - rpc_timeout: The deadline of each call of the read, not of the whole read. A call - that waits for new Records returns before the deadline. - """ + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + self._raise_lost_write() read = read_spec( - self._payload_type, + handle._payload_type, offset=offset, last=last, wait_new_records=wait_new_records, create_wait=create_wait, ) return subscription( - lambda position: self._items(read, position), + lambda position: self._items(handle, read, position), read, _backend.options_of(rpc_metadata, rpc_timeout), ) async def _items( - self, read: ReadSpec, position: Position + self, handle: StreamHandle[Any], read: ReadSpec, position: Position ) -> AsyncGenerator[Item, None]: - items = self._access().items(read, position) + items = self._access(handle).items(read, position) try: async for item in items: yield item @@ -257,80 +205,68 @@ async def _items( async def get_offset( self, - *, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, ) -> 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. - - Args: - rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so - they send none. - rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with - ``DEADLINE_EXCEEDED``. - """ + """Get the next offset.""" + self._raise_lost_write() with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - return await self._access().get_offset() + return await self._access(handle).get_offset() - def _check_close(self) -> None: - if self._workflow_id: + def _check_close(self, handle: StreamHandle[Any]) -> None: + if handle.workflow_id: # Only the run chain can close a Workflow's Stream. The Store refuses a client's # close too, but this check makes no call. raise _errors.StreamNotOwnerError( - f"Can't close the Stream {self._id} from a client, since a Workflow owns it" + f"Can't close the Stream {handle.id} from a client, since a Workflow owns it" ) async def close( self, - *, - result: Any | None = None, - failure: BaseException | None = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, ) -> None: - """Write this handle's Records, then close the Stream. - - Readers get the Records before the close, then a ``Closed`` item with ``result`` or - ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as - canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`get_offset`. - - Raises: - StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. - """ - self._check_close() + """Write this handle's Records, then close the Stream.""" + self._check_close(handle) self._closed = True with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - await self._access().close(result, failure) - - async def __aenter__(self) -> ClientStream[T]: - """Use the Stream in ``async with``, which closes it at the end. + await self._access(handle).close(result, failure) + # The close runs first, so a lost write doesn't leave the Stream open. + self._raise_lost_write() - Raises: - StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. - """ + async def enter(self, handle: StreamHandle[Any]) -> None: + """Check that the client can close the Stream at the end of the block.""" # Fail early, since the close at the end would raise and hide the block's own error. - self._check_close() - return self + self._check_close(handle) - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - """Close the Stream, with the exception as its failure if there is one. - - A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as - ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this - does nothing. - """ + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """Close the Stream, unless this handle closed it already.""" if self._closed: return close, failure = exit_failure(exc) if close: - await self.close(failure=failure) + await self.close(handle, None, failure, {}, None) + + def producer( + self, handle: StreamHandle[Any], name: str, epoch: int + ) -> StreamHandle[Any]: + """Return a handle of the same Stream and client that writes as a named Producer.""" + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + named: StreamHandle[Any] = StreamHandle( + id=handle.id, + workflow_id=handle.workflow_id, + first_execution_run_id=handle.first_execution_run_id, + backend=handle.backend, + payload_type=handle._payload_type, + ) + named._client = self._client + named._client_calls = ClientCalls(self._client, self._backend, (name, epoch)) + return named def create( @@ -340,85 +276,107 @@ def create( id: str, backend: str | None, options: StreamOptions, -) -> Creating[ClientStream[T]]: +) -> Creating[StreamHandle[T]]: """Create a Stream that the client owns, or get the one that exists. See :py:meth:`temporalio.client.Client.create_stream`. """ - async def create() -> ClientStream[T]: + async def create() -> StreamHandle[T]: info = definition_of(definition) backend_name, calls = client_backend(client, backend) # A closed, expired or another owner's Stream fails the create. await calls.service(client.namespace).create_stream( definition, id=id, retention=options.get("retention") ) - stream: ClientStream[T] = ClientStream( - id=id, backend=backend_name, payload_type=info.type + return _bound( + client, calls, id=id, backend_name=backend_name, payload_type=info.type ) - return stream._bind(client, calls) return Creating(create) -def client_stream( +def get_stream_handle( client: temporalio.client.Client, definition: type[StreamDefinition[T]], *, - id: str | None, - handle: StreamHandle[Any] | None, + id: str, backend: str | None, -) -> ClientStream[T]: - """Get a Stream by its id, or from a handle that came in data. +) -> StreamHandle[T]: + """Get a Stream by its id. - See :py:meth:`temporalio.client.Client.stream`. + See :py:meth:`temporalio.client.Client.get_stream_handle`. """ info = definition_of(definition) - if (id is None) == (handle is None): - raise TypeError("stream takes one of id= and handle=") - if handle is not None: - if backend and handle.backend and backend != handle.backend: - raise ValueError( - f"The handle is of the backend {handle.backend}, not {backend}" - ) - backend = backend or handle.backend - stream_id = handle.id if handle else id - if not stream_id: + if not id: # An empty id would fail only at the first call, with an unclear error. raise ValueError("The Stream id is empty") backend_name, calls = client_backend(client, backend) - stream: ClientStream[T] = ClientStream( - id=stream_id, - workflow_id=handle.workflow_id if handle else None, - first_execution_run_id=handle.first_execution_run_id if handle else None, - backend=backend_name, - payload_type=info.type, + return _bound( + client, calls, id=id, backend_name=backend_name, payload_type=info.type ) - return stream._bind(client, calls) -class BindQueryResults(temporalio.client.Interceptor): - """Gives each ``ClientStream`` in a query result the client that made the query.""" +def _bound( + client: temporalio.client.Client, + backend: BackendCalls, + *, + id: str, + backend_name: str, + payload_type: Any, +) -> StreamHandle[Any]: + handle: StreamHandle[Any] = StreamHandle( + id=id, backend=backend_name, payload_type=payload_type + ) + handle._client = client + handle._client_calls = ClientCalls(client, backend) + return handle - def intercept_client( - self, next: temporalio.client.OutboundInterceptor - ) -> temporalio.client.OutboundInterceptor: - """Wrap the client's calls.""" - return _BindQueryResults(next) +class _BindingPayloadConverter( + temporalio.converter.PayloadConverter, temporalio.converter.WithSerializationContext +): + # Wraps the payload converter of a client with Stream backends. A query result, an Update or + # Workflow result, or any other decode of the client binds each handle in it to the client, + # so code uses a handle that came as data with no other call. + def __init__( + self, + inner: temporalio.converter.PayloadConverter, + client: weakref.ref[temporalio.client.Client], + ) -> None: + self.inner = inner + self._client = client -class _BindQueryResults(temporalio.client.OutboundInterceptor): - def __init__(self, next: temporalio.client.OutboundInterceptor) -> None: - super().__init__(next) - # The innermost interceptor is the client's own, and only it has no next. - inner: Any = next - while hasattr(inner, "next"): - inner = inner.next - self._client: temporalio.client.Client = inner._client + def to_payloads(self, values: Sequence[Any]) -> list[Payload]: + return self.inner.to_payloads(values) - async def query_workflow(self, input: temporalio.client.QueryWorkflowInput) -> Any: - token = _decoding_client.set(self._client) + def from_payloads( + self, payloads: Sequence[Payload], type_hints: list[type] | None = None + ) -> list[Any]: + token = decoding_client.set(self._client()) try: - return await super().query_workflow(input) + return self.inner.from_payloads(payloads, type_hints) finally: - _decoding_client.reset(token) + decoding_client.reset(token) + + def with_context(self, context: temporalio.converter.SerializationContext) -> Self: + inner = self.inner + if isinstance(inner, temporalio.converter.WithSerializationContext): + inner = inner.with_context(context) + return self if inner is self.inner else type(self)(inner, self._client) + + +def bind_decodes( + converter: temporalio.converter.DataConverter, client: temporalio.client.Client +) -> temporalio.converter.DataConverter: + """A copy of a client's data converter whose decodes bind each Stream handle to the client.""" + inner = converter.payload_converter + if isinstance(inner, _BindingPayloadConverter): + # A config that came from another client has that client's converter. + inner = inner.inner + bound = dataclasses.replace(converter) + # The client holds its converter, so the converter holds the client weakly. + object.__setattr__( + bound, "payload_converter", _BindingPayloadConverter(inner, weakref.ref(client)) + ) + return bound diff --git a/temporalio/common/_streams/_writer.py b/temporalio/common/_streams/_writer.py index c234ce882..ce05e166a 100644 --- a/temporalio/common/_streams/_writer.py +++ b/temporalio/common/_streams/_writer.py @@ -4,7 +4,7 @@ import asyncio import weakref -from collections.abc import Callable, Sequence +from collections.abc import Callable, Generator, Sequence from typing import TYPE_CHECKING, Any import temporalio.api.stream.v1 as stream_v1 @@ -15,6 +15,43 @@ from temporalio.common._streams.writer_system.client import Writer +class WriteFuture(asyncio.Future[int]): + """The future of a write. A caller that gets its result or its error sets ``retrieved``. + + A write error that no caller gets still raises, at the end of a Workflow run or an Activity + attempt, or at the next call of a client's handle. + """ + + retrieved = False + + def result(self) -> int: + """The last offset of the write.""" + self.retrieved = True + return super().result() + + def exception(self) -> BaseException | None: + """The error of the write, or None.""" + self.retrieved = True + return super().exception() + + def __await__(self) -> Generator[Any, None, int]: + """Wait for the write, and return its last offset.""" + # The await of asyncio's future doesn't call result(), so this one does. + if not self.done(): + self._asyncio_future_blocking = True + yield self + return self.result() + + __iter__ = __await__ + + +def unretrieved_error(future: asyncio.Future[int]) -> BaseException | None: + """The error of a failed write that no caller got, else None.""" + if not isinstance(future, WriteFuture) or future.retrieved or future.cancelled(): + return None + return asyncio.Future.exception(future) + + class StreamWriter: """Writes the Records of one Producer to one Stream, through the writer client.""" @@ -42,6 +79,11 @@ def __init__( self._failed_tasks = list(failed_tasks) self._writer: Writer | None = None self._open_lock = asyncio.Lock() + # The last write that started. Each write waits for the one before it, so the order is + # the order of the calls, even before a write's task first runs. + self._last: asyncio.Task[int] | None = None + # The writes that failed. The first error that no caller got raises later. + self._failed: list[WriteFuture] = [] self._closed = False @property @@ -96,12 +138,59 @@ async def open(self) -> Writer: ) return self._writer - async def write(self, values: Sequence[Any]) -> int: - """Write the values as Records now, in order, and return the last offset.""" - from temporalio.common._streams.writer_system import models + def write(self, values: Sequence[Any]) -> WriteFuture: + """Start the write of the values as Records, and return its future. + + Writes start in the order of the calls, also when nobody awaits a future. A cancel of + the future doesn't stop the write, since the write keeps its place in the order. The + future gives the last offset. + Raises: + StreamClosedError: The writer is closed. + """ if self._closed: raise _errors.StreamClosedError("The writer is closed") + loop = asyncio.get_running_loop() + task = loop.create_task(self._write_after(self._last, values)) + self._last = task + future = WriteFuture(loop=loop) + + def done(task: asyncio.Task[int]) -> None: + if self._last is task: + self._last = None + error = None if task.cancelled() else task.exception() + if future.cancelled(): + return + if task.cancelled(): + future.cancel() + elif error is not None: + future.set_exception(error) + self._failed.append(future) + else: + future.set_result(task.result()) + + task.add_done_callback(done) + return future + + def unretrieved_error(self) -> BaseException | None: + """The first write error that no caller got, or None. + + The caller owns the error, so the later errors of the same writes don't raise again. + """ + failed, self._failed = self._failed, [] + error = next((e for e in map(unretrieved_error, failed) if e is not None), None) + for future in failed: + future.retrieved = True + return error + + async def _write_after( + self, previous: asyncio.Task[int] | None, values: Sequence[Any] + ) -> int: + from temporalio.common._streams.writer_system import models + + if previous is not None: + # The earlier write's error belongs to its own caller. + await asyncio.wait([previous]) writer = await self.open() answer = await writer.write_records( [models.AppendRecord(body=value) for value in values] @@ -109,10 +198,12 @@ async def write(self, values: Sequence[Any]) -> int: return answer.last_offset or 0 async def close(self) -> None: - """Close the writer. The Stream stays open.""" + """Wait for the writes that started, then close the writer. The Stream stays open.""" if self._closed: return self._closed = True + if self._last is not None: + await asyncio.wait([self._last]) if self._release is not None: self._release.detach() if self._writer is not None: diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index 0b9bc7cd7..40ef3b9ad 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -89,7 +89,7 @@ async def test_client_writes_reads_and_closes(client: Client): async def test_read_waits_for_create(client: Client): client, _ = fake_client(client) - stream = client.stream(Tokens, id="wf-1:run-1:tokens") + stream = client.get_stream_handle(Tokens, id="wf-1:run-1:tokens") with pytest.raises(common.StreamNotFoundError): async for _ in stream.subscribe(): pass @@ -132,7 +132,7 @@ def test_stream_errors_fail_the_workflow_not_the_task(): async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client): client, extension = fake_client(client) stream = await client.create_stream(Tokens, id="two-writers") - other = client.stream(Tokens, id="two-writers") + other = client.get_stream_handle(Tokens, id="two-writers") await stream.publish("1") await other.publish("2") items = [item async for item in stream.subscribe(wait_new_records=False)] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index be315b765..e56210d49 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -1,3 +1,4 @@ +import asyncio import uuid from dataclasses import dataclass @@ -29,7 +30,7 @@ async def test_unknown_backend_raises_before_any_call(client: Client): with pytest.raises(common.StreamConfigError, match="no Stream backend redis"): await client.create_stream(Turns, id="chat-1", backend="redis") with pytest.raises(common.StreamConfigError, match="It has fake"): - client.stream(Turns, id="chat-1", backend="redis") + client.get_stream_handle(Turns, id="chat-1", backend="redis") assert not extension.streams and not extension.namespaces stream = await client.create_stream(Turns, id="chat-1") assert stream.backend == "fake" @@ -37,48 +38,102 @@ async def test_unknown_backend_raises_before_any_call(client: Client): async def test_client_without_backend_has_no_transport(client: Client): with pytest.raises(common.StreamConfigError, match="no Stream backend"): - client.stream(Turns, id="chat-1") + client.get_stream_handle(Turns, id="chat-1") with pytest.raises(common.StreamConfigError, match="no Stream backend"): await client.create_stream(Turns, id="chat-1") -async def test_decoded_client_stream_has_no_client(client: Client): +async def test_only_a_decode_of_the_client_binds_a_handle(client: Client): client, _ = fake_client(client) stream = await client.create_stream(Turns, id="chat-2") converter = DataConverter.default.payload_converter - decoded = converter.from_payload( - converter.to_payload(stream), common.ClientStream[Turn] + unbound = converter.from_payload( + converter.to_payload(stream), common.StreamHandle[Turn] ) - assert decoded == stream - with pytest.raises(common.StreamConfigError, match="Client.stream"): - await decoded.get_offset() + assert unbound == stream + with pytest.raises(common.StreamConfigError, match="get_stream_handle"): + await unbound.get_offset() + payloads = await client.data_converter.encode([stream]) + [bound] = await client.data_converter.decode(payloads, [common.StreamHandle[Turn]]) + assert await bound.get_offset() == 0 + # A client made from another client's config binds to itself. + other = Client(**client.config(active_config=True)) + [rebound] = await other.data_converter.decode(payloads, [common.StreamHandle[Turn]]) + assert rebound._client is other + + +async def test_publish_starts_the_write_at_the_call(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="ordered") + futures = [stream.publish(str(n)) for n in range(3)] + # A cancel of a future doesn't stop its write, which already has its place. + futures[1].cancel() + stream.publish("3") + await stream.close() + items = [item async for item in stream.subscribe()] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "0", + "1", + "2", + "3", + ] + assert [await futures[0], await futures[2]] == [items[0].offset, items[2].offset] + with pytest.raises(common.StreamClosedError): + stream.publish("4") + + +async def test_next_call_raises_a_write_error_that_no_caller_got(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="lost") + await client.get_stream_handle(Tokens, id="lost").close() + late = stream.publish("late") + await asyncio.wait([late]) + with pytest.raises(common.StreamClosedError): + await stream.get_offset() + # The call owns the error, so the next one makes its call. + assert await stream.get_offset() > 0 + + +async def test_sync_workflow_calls_raise_in_client_code(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="no-nowait") + with pytest.raises(common.StreamUnsupportedError, match="sync Workflow code"): + stream.publish_nowait("1") + with pytest.raises(common.StreamUnsupportedError, match="sync Workflow code"): + stream.close_nowait() @workflow.defn class HandsOutStream: @workflow.run - async def run(self) -> None: - pass + async def run(self) -> dict[str, str]: + return {"id": "handed-out"} @workflow.query def stream(self) -> dict[str, str]: - # The proto JSON of a ref, as a Workflow Stream that a query returns. + # The proto JSON of a handle, as a Workflow's Stream that a query returns. return {"id": "handed-out"} -async def test_query_result_binds_its_client_stream(client: Client): +async def test_query_and_workflow_results_bind_their_client(client: Client): client, _ = fake_client(client) created = await client.create_stream(Tokens, id="handed-out") async with new_worker(client, HandsOutStream) as worker: + # By name, since the method's own return type is the JSON. handle = await client.start_workflow( - HandsOutStream.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + "HandsOutStream", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + result_type=common.StreamHandle[str], ) - await handle.result() - stream = await handle.query("stream", result_type=common.ClientStream[str]) - # The ref names no backend, so the handle uses the default one. - assert stream.id == "handed-out" and stream.backend is None - await stream.publish("from the query result") + result = await handle.result() + queried = await handle.query("stream", result_type=common.StreamHandle[str]) + # The handle names no backend, so it uses the default one. + assert queried.id == "handed-out" and queried.backend is None + await queried.publish("from the query result") + await result.publish("from the Workflow result") items = [item async for item in created.subscribe(wait_new_records=False)] assert [item.data for item in items if isinstance(item, common.Record)] == [ - "from the query result" + "from the query result", + "from the Workflow result", ]