From 7265a0471fbeae319564829d7698975e9d6da8b7 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 14:54:36 -0700 Subject: [PATCH] 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" + ]