From 7104162b608ab7e7077f17e2c4fe5ebc0927d581 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Thu, 5 Nov 2026 01:45:38 -0700 Subject: [PATCH] Added the client `Stream`, its typed creates and Stream reads. --- temporalio/client/__init__.py | 2 + temporalio/client/_client.py | 80 ++++- temporalio/client/_streams.py | 321 +++++++++++++++++++ temporalio/streams/__init__.py | 2 + temporalio/streams/_access.py | 134 ++++++++ temporalio/streams/_convert.py | 51 +++ temporalio/streams/_reader.py | 311 ++++++++++++++++++ temporalio/streams/_stream_client.py | 329 ++++++++++++++++++++ temporalio/streams/_writer.py | 213 +++++++++++++ tests/contrib/streams/fake_extension.py | 193 ++++++++++++ tests/contrib/streams/test_facade.py | 314 +++++++++++++++++++ tests/contrib/streams/test_reader.py | 115 +++++++ tests/contrib/streams/test_stream_client.py | 155 +++++++++ 13 files changed, 2219 insertions(+), 1 deletion(-) create mode 100644 temporalio/client/_streams.py create mode 100644 temporalio/streams/_access.py create mode 100644 temporalio/streams/_convert.py create mode 100644 temporalio/streams/_reader.py create mode 100644 temporalio/streams/_stream_client.py create mode 100644 temporalio/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/__init__.py b/temporalio/client/__init__.py index 2eef41a39d..dafe658fcb 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -185,6 +185,7 @@ ScheduleUpdate, ScheduleUpdateInput, ) +from ._streams import Stream from ._worker_versioning import ( BuildIdOp, BuildIdOpAddNewCompatible, @@ -397,4 +398,5 @@ "ParamType", "ReturnType", "SelfType", + "Stream", ] diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index 8cf2a8dc26..f9e9a4e025 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -13,6 +13,7 @@ TYPE_CHECKING, Any, Concatenate, + TypeVar, cast, overload, ) @@ -94,6 +95,7 @@ ScheduleBackfill, ScheduleHandle, ) +from ._streams import CreateStream, GetStream, Stream from ._worker_versioning import ( BuildIdOp, TaskReachabilityType, @@ -110,11 +112,13 @@ ) if TYPE_CHECKING: - from temporalio.streams import StreamBackend + from temporalio.streams import StreamBackend, StreamRef from ._interceptor import Interceptor from ._plugin import Plugin +_StreamPayloadT = TypeVar("_StreamPayloadT") + class Client: """Client for accessing Temporal. @@ -911,6 +915,80 @@ def get_workflow_handle( result_type=result_type, ) + @property + def create_stream(self) -> CreateStream[Any]: + """Create a Stream, which fails if a Stream of the id exists. + + Give the payload type as a type argument:: + + stream = await client.create_stream[Turn]("chat-42", retention=timedelta(days=1)) + async with stream: # optional, the block end closes it + await stream.publish(Turn("user", "hi")) + + A Stream must exist before its first publish, and the create calls the Store, so a bad + setting fails early. Code closes the Stream with ``close()``, or at the end of + ``async with``. The call takes the Stream id, an optional ``backend=`` name and the + options of :py:class:`temporalio.streams.StreamOptions`. See + :py:meth:`CreateStream.__call__`. + + .. warning:: + This API is experimental. + """ + return CreateStream(self, None, or_get=False) + + @property + def create_or_get_stream(self) -> CreateStream[Any]: + """Create a Stream, or get the one of the id that exists. + + It takes the arguments of :py:attr:`create_stream`. On a Stream that exists, it ignores + the options, so its retries are safe. It still fails on a closed or expired Stream. + + .. warning:: + This API is experimental. + """ + return CreateStream(self, None, or_get=True) + + @property + def get_stream(self) -> GetStream[Any]: + """Get a Stream by its id, with no call. + + Give the payload type as a type argument, as in ``client.get_stream[Turn]("chat-42")``. + Reads decode each Record with it, or as a JSON value with no type argument. For a ref + that came as data, such as from a Workflow's query, use :py:meth:`to_stream`. See + :py:meth:`GetStream.__call__`. + + .. warning:: + This API is experimental. + """ + return GetStream(self, None) + + def to_stream( + self, + ref: StreamRef[_StreamPayloadT], + *, + payload_type: type[_StreamPayloadT] | None = None, + ) -> Stream[_StreamPayloadT]: + """Get the Stream of a ref that came as data, such as from a query. It makes no call. + + The type argument of the ref is for type checkers only. Without ``payload_type``, a read + gives each Record as a JSON value, such as a dict for a dataclass. + + .. warning:: + This API is experimental. + + Args: + ref: The ref of the Stream. + payload_type: The type that reads decode each Record with. Defaults to no type, so + a read decodes each Record as a JSON value. + + Raises: + StreamConfigError: The client has no backend of the ref, or the ref has no backend + and no backend has ``default=True``. + """ + from temporalio.streams import _stream_client + + return _stream_client.to_stream(self, ref, payload_type=payload_type) + def get_workflow_handle_for( self, workflow: ( diff --git a/temporalio/client/_streams.py b/temporalio/client/_streams.py new file mode 100644 index 0000000000..a796d47b92 --- /dev/null +++ b/temporalio/client/_streams.py @@ -0,0 +1,321 @@ +"""Streams outside Workflow code. + +This module imports the Streams API of :py:mod:`temporalio.streams` on first use, so a client +with no Stream never imports it. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from datetime import timedelta +from types import TracebackType +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from typing_extensions import Self, Unpack + +if TYPE_CHECKING: + import temporalio.api.stream.v1 + import temporalio.client + from temporalio.streams import End, StreamOptions, StreamRef, Subscription + from temporalio.streams._stream_client import StreamCalls + +T = TypeVar("T") +_U = TypeVar("_U") + + +class Stream(Generic[T]): + """A Stream outside Workflow code, with its calls. ``T`` is the payload type. + + Get one with :py:attr:`temporalio.client.Client.create_stream`, + :py:attr:`temporalio.client.Client.create_or_get_stream`, + :py:attr:`temporalio.client.Client.get_stream`, or + :py:meth:`temporalio.client.Client.to_stream` for a ref that came as data. Its + :py:attr:`ref` is the data to pass on, such as to a Workflow. + + All Streams of one client and Stream id write as one Producer, and :py:meth:`producer` + returns a Stream that writes as a Producer with a name of its own. A write error that no + caller got raises at the next call of the Stream. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, ref: StreamRef[T], payload_type: Any, calls: StreamCalls + ) -> None: + """Create a Stream. Use a client's ``create_stream``, ``get_stream`` or ``to_stream``.""" + self._ref = ref + # The type that a read decodes each value with. None decodes with no type hint. + self._payload_type = payload_type + self._calls = calls + + @property + def ref(self) -> StreamRef[T]: + """The data that identifies the Stream, to pass on.""" + return self._ref + + @property + def proto_ref(self) -> temporalio.api.stream.v1.StreamRef: + """The ``StreamRef`` wire message of the Stream, for a field of a proto message.""" + from temporalio.streams import _ref + + return _ref.to_proto(self._ref) + + 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 Stream keep the order of the calls. A cancel + of the future doesn't stop the write. A list is one value. The next call of the Stream + 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 object closed the Stream. + """ + return self._calls.publish(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 when the Stream is closed and the read gave its last 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 ``streams.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 one + that a Workflow creates after it starts. 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. + """ + from temporalio.streams._access import read_spec + + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + return self._calls.subscribe(read, rpc_metadata, 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(rpc_metadata, rpc_timeout) + + async def close( + self, + *, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> None: + """Write the Records that this object published, then close the Stream. + + A close carries no value, so code with a last value publishes it first. Readers get the + Records before the close, then their read ends. Later writes fail with + ``StreamClosedError``. ``rpc_metadata`` and ``rpc_timeout`` are as for + :py:meth:`publish`. + + Raises: + StreamNotOwnerError: The code can't close the Stream. Only the run chain closes a + Stream bound to a Workflow, and it raises with no call. + """ + await self._calls.close(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() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, also when the block raised an exception or was canceled. + + Another ``BaseException``, such as ``KeyboardInterrupt``, leaves it open. If this + object already closed the Stream, this does nothing. + """ + await self._calls.exit(exc) + + def producer(self, name: str, epoch: int) -> Stream[T]: + """Return this Stream as a Stream 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:`temporalio.streams.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: + ValueError: The epoch is below 1. + """ + return Stream(self._ref, self._payload_type, self._calls.producer(name, epoch)) + + def __repr__(self) -> str: + """Show the ref.""" + return f"{type(self).__name__}({self._ref!r})" + + +# Python has no type arguments on a call, so a call that takes the payload type is an object +# whose subscript binds the type, as in client.get_stream[Turn]("chat-42"). The subscript passes +# the class on, so a read decodes each Record with it, and type checkers see Stream[Turn]. With +# no subscript, T is Any and a read gives JSON values. + + +class CreateStream(Generic[T]): + """The call of :py:attr:`temporalio.client.Client.create_stream` and + :py:attr:`temporalio.client.Client.create_or_get_stream`. + + Subscript it with the payload type, as in ``client.create_stream[Turn]("chat-42")``. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + client: temporalio.client.Client, + payload_type: type[T] | None, + *, + or_get: bool, + ) -> None: + """Create the call. Use the client's property instead.""" + self._client = client + self._payload_type = payload_type + self._or_get = or_get + + def __getitem__(self, payload_type: type[_U]) -> CreateStream[_U]: + """The call for Streams of ``payload_type``. + + It takes a class, such as ``Turn`` or ``list[str]``. Type checkers refuse a union, such + as ``str | None``, so use ``Any`` or a class that holds the value instead. + """ + return CreateStream(self._client, payload_type, or_get=self._or_get) + + async def __call__( + self, + id: str, + *, + backend: str | None = None, + **options: Unpack[StreamOptions], + ) -> Stream[T]: + """Create the Stream, and return it. + + Args: + 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.streams.StreamOptions`. + + Raises: + StreamAlreadyExistsError: ``create_stream`` of an id that exists. + StreamClosedError: A Stream of that id closed. + StreamExpiredError: A Stream of that id expired. + 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. + """ + from temporalio.streams import _stream_client + + return await _stream_client.create( + self._client, + self._payload_type, + id=id, + backend=backend, + options=options, + or_get=self._or_get, + ) + + +class GetStream(Generic[T]): + """The call of :py:attr:`temporalio.client.Client.get_stream`. + + Subscript it with the payload type, as in ``client.get_stream[Turn]("chat-42")``. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, client: temporalio.client.Client, payload_type: type[T] | None + ) -> None: + """Create the call. Use the client's property instead.""" + self._client = client + self._payload_type = payload_type + + def __getitem__(self, payload_type: type[_U]) -> GetStream[_U]: + """The call for Streams of ``payload_type``. It takes a class, as for ``create_stream``.""" + return GetStream(self._client, payload_type) + + def __call__(self, id: str, *, backend: str | None = None) -> Stream[T]: + """Get the Stream with no call. + + Args: + id: The Stream id. + 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. + ValueError: The id is empty. + """ + from temporalio.streams import _stream_client + + return _stream_client.get_stream( + self._client, self._payload_type, id=id, backend=backend + ) diff --git a/temporalio/streams/__init__.py b/temporalio/streams/__init__.py index e1fbd65560..a4c2a6a6e2 100644 --- a/temporalio/streams/__init__.py +++ b/temporalio/streams/__init__.py @@ -21,6 +21,7 @@ ) from temporalio.streams._items import END, End, ItemKind, Record, Superseded from temporalio.streams._plugin import StreamBackend +from temporalio.streams._reader import Subscription from temporalio.streams._ref import StreamOptions, StreamRef # nexgen generates these wire models. Refs, Records, Superseded items and errors carry them. @@ -80,6 +81,7 @@ "StreamSuperseded", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "Timestamp", "WorkflowTaskEpoch", diff --git a/temporalio/streams/_access.py b/temporalio/streams/_access.py new file mode 100644 index 0000000000..aeb8cdc2c3 --- /dev/null +++ b/temporalio/streams/_access.py @@ -0,0 +1,134 @@ +"""Stream calls outside Workflow code, which Activity and client Streams share.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, Callable, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import temporalio.converter +from temporalio.streams import _convert, _reader +from temporalio.streams._backend import BackendCalls, CallOptions +from temporalio.streams._items import End +from temporalio.streams._reader import Item, Position, Subscription +from temporalio.streams._writer import StreamWriter +from temporalio.streams.system import models + + +@dataclass +class Access: + """The calls of one Stream, bound to a backend.""" + + backend: BackendCalls + converter: temporalio.converter.DataConverter + namespace: str + stream_id: str + writer: Callable[[], StreamWriter] + + 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 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 service.send_describe_stream( + models.DescribeStreamRequest(ref=models.StreamRef(id=self.stream_id)) + ) + return response.last_offset + 1 if response.last_offset else 0 + + async def close(self) -> 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 + ) + service = self.backend.service(self.namespace, self.converter) + # The generated close takes its Producer from a source, so the Stream sends the request + # with its own. + await service.send_close_stream( + models.CloseStreamRequest( + ref=models.StreamRef(id=self.stream_id), producer=producer + ) + ) + + +@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 (not 1 <= last <= 1000 or offset is not None): + raise ValueError("last takes a count from 1 to 1000 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/streams/_convert.py b/temporalio/streams/_convert.py new file mode 100644 index 0000000000..8bf182a79e --- /dev/null +++ b/temporalio/streams/_convert.py @@ -0,0 +1,51 @@ +"""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 + + +class _RawPayloadConverter(temporalio.converter.PayloadConverter): + # The generated client runs the codec on Record bodies and keeps them as RawValue, so the + # Stream converts each body with its 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 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/streams/_reader.py b/temporalio/streams/_reader.py new file mode 100644 index 0000000000..f0f26d4ad5 --- /dev/null +++ b/temporalio/streams/_reader.py @@ -0,0 +1,311 @@ +"""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.common +import temporalio.converter +from temporalio.streams import _backend, _convert, _errors +from temporalio.streams._items import ( + End, + Record, + Superseded, + is_client, + producer_key, + withdraws, +) +from temporalio.streams.system import models + +T = TypeVar("T") + +Item = Record[Any] | Superseded + +# 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, and :py:class:`Superseded` items that withdraw earlier + Records. Each item's ``kind`` gives its type. A read of a closed Stream ends after the last + item. + + .. 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 Stream'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: + """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 not is_client(item.producer) + ): + 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.facts) + ) + return item + + def _withdraw(self, facts: models.StreamSuperseded) -> tuple[int, ...]: + offsets: list[int] = [] + for key, (producer, runs) in list(self._shown.items()): + 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. + 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: _backend.StreamServiceCalls, + 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 the Stream is closed and drained. + + 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 + pause = _FIRST_PAUSE + + 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( + ref=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) + 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. + 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 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: + left = create_until - loop.time() if create_until is not None else 0 + if left <= 0: + raise + # A Store that answers at once, with no wait for the create, isn't asked again at + # once. + await asyncio.sleep(min(pause, left)) + pause = min(pause * 2, _LAST_PAUSE) + finally: + await pages.aclose() + + +async def _pages( + service: _backend.StreamServiceCalls, + 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: _backend.StreamServiceCalls, + request: models.ReadRecordsRequest, +) -> models.ReadRecordsResponse: + # The bodies stay RawValue, so the read decodes each one with the Stream's payload type. + return await service.send_read_records(request, temporalio.common.RawValue) + + +async def _read_until_records( + service: _backend.StreamServiceCalls, + 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[temporalio.common.RawValue], + payload_type: Any, +) -> 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=facts): + return Superseded(offset=offset, facts=facts) diff --git a/temporalio/streams/_stream_client.py b/temporalio/streams/_stream_client.py new file mode 100644 index 0000000000..c92565aa9e --- /dev/null +++ b/temporalio/streams/_stream_client.py @@ -0,0 +1,329 @@ +"""Streams outside Workflow code.""" + +from __future__ import annotations + +import abc +import asyncio +import itertools +import os +import uuid +import weakref +from collections.abc import AsyncGenerator, Mapping +from datetime import timedelta +from typing import Any, TypeVar + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.client +import temporalio.nexus.system._support as _support +from temporalio.client._streams import Stream +from temporalio.streams import _backend, _errors, _ids +from temporalio.streams._access import Access, ReadSpec, subscription +from temporalio.streams._backend import BackendCalls +from temporalio.streams._plugin import client_backend +from temporalio.streams._reader import Item, Position, Subscription +from temporalio.streams._ref import StreamOptions, StreamRef +from temporalio.streams._writer import StreamWriter +from temporalio.streams.system import models + +T = TypeVar("T") + +# The Streams of one client and Stream id 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) + + +def _client_epoch(attempt: int) -> stream_v1.ProducerEpoch: + return stream_v1.ProducerEpoch(client=stream_v1.ClientEpoch(attempt=attempt)) + + +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( + id=_ids.named_client_producer_id(name), epoch=_client_epoch(epoch) + ) + key = _CLIENT_KEYS.get(client) + if key is None: + key = f"{_process_key}-{next(_clients)}" + _CLIENT_KEYS[client] = key + return stream_v1.Producer( + id=_ids.client_producer_id(client.identity, key), epoch=_client_epoch(1) + ) + + +class StreamCalls(abc.ABC): + """The calls of one :py:class:`Stream`, such as in client code.""" + + @abc.abstractmethod + def publish( + self, + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. See :py:meth:`Stream.publish`.""" + + @abc.abstractmethod + def subscribe( + self, + read: ReadSpec, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read. See :py:meth:`Stream.subscribe`.""" + + @abc.abstractmethod + async def get_offset( + self, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None + ) -> int: + """Get the next offset. See :py:meth:`Stream.get_offset`.""" + + @abc.abstractmethod + async def close( + self, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None + ) -> None: + """Close the Stream. See :py:meth:`Stream.close`.""" + + @abc.abstractmethod + async def enter(self) -> None: + """Start an ``async with`` block of the Stream.""" + + @abc.abstractmethod + async def exit(self, exc: BaseException | None) -> None: + """End an ``async with`` block of the Stream. See :py:meth:`Stream.__aexit__`.""" + + @abc.abstractmethod + def producer(self, name: str, epoch: int) -> StreamCalls: + """The calls of a Producer with a name of its own. See :py:meth:`Stream.producer`.""" + + +class ClientCalls(StreamCalls): + """The calls of one Stream in client code. + + Each publish starts its write at the call. All Streams of one client and Stream id write as + one Producer, unless a Stream is of a Producer with a name of its own. + """ + + def __init__( + self, + client: temporalio.client.Client, + backend: BackendCalls, + ref: StreamRef[Any], + named: tuple[str, int] | None = None, + ) -> None: + """Create the calls of ``ref``. ``named`` is a Producer's name and epoch.""" + self._client = client + self._backend = backend + self._ref = ref + self._named = named + self._writer: StreamWriter | None = None + self._closed = False + + def _access(self) -> Access: + client = self._client + backend = self._backend + stream_id = self._ref.id + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + writers=backend.writers(client.namespace, client.data_converter), + stream_id=stream_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=stream_id, + writer=writer, + ) + + 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 Stream's next call raises it. + error = self._writer.unretrieved_error() if self._writer else None + if error is not None: + raise error + + def publish( + self, + 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 self._access().publish(values) + + def subscribe( + self, + read: ReadSpec, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + self._raise_lost_write() + 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 + ) -> int: + """Get the next offset.""" + self._raise_lost_write() + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await self._access().get_offset() + + def _check_close(self) -> None: + if self._ref.owner is not None: + # Only the run chain closes a Stream bound to a Workflow. The Store refuses a + # client's close too, but this check makes no call. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {self._ref.id} from a client, since it is bound to a " + "Workflow" + ) + + async def close( + self, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None + ) -> None: + """Write the Records of these calls, then close the Stream.""" + self._check_close() + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + await self._access().close() + # Only a close that the Store took counts, so the end of async with tries a failed one + # again. + self._closed = True + # The close runs first, so a lost write doesn't leave the Stream open. + self._raise_lost_write() + + async def enter(self) -> 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() + + async def exit(self, exc: BaseException | None) -> None: + """Close the Stream, unless these calls closed it already.""" + if not self._closed and _support.exit_close(self, exc): + await self.close({}, None) + + def producer(self, name: str, epoch: int) -> StreamCalls: + """The calls of the same Stream and client as a Producer with a name of its own.""" + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + return ClientCalls(self._client, self._backend, self._ref, (name, epoch)) + + +async def create( + client: temporalio.client.Client, + payload_type: type[T] | None, + *, + id: str, + backend: str | None, + options: StreamOptions, + or_get: bool, +) -> Stream[T]: + """Create a Stream, or with ``or_get`` get the one that exists. + + See :py:attr:`temporalio.client.Client.create_stream`. + """ + if not id: + # An empty id would fail only at the Store, with an unclear error. + raise ValueError("The Stream id is empty") + backend_name, calls = client_backend(client, backend) + policy = ( + models.StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting + if or_get + else models.StreamIdConflictPolicy.StreamIdConflictPolicyFail + ) + # The generated create sends a new request id for each call. A closed, expired or another + # owner's Stream fails either create. + await calls.service(client.namespace).create_stream( + object, + id=id, + backend=backend_name, + retention=options.get("retention"), + id_conflict_policy=policy, + ) + ref: StreamRef[T] = StreamRef(id=id, backend=backend_name) + return _stream(client, calls, ref, payload_type) + + +def get_stream( + client: temporalio.client.Client, + payload_type: type[T] | None, + *, + id: str, + backend: str | None, +) -> Stream[T]: + """Get a Stream by its id. + + See :py:attr:`temporalio.client.Client.get_stream`. + """ + 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) + ref: StreamRef[T] = StreamRef(id=id, backend=backend_name) + return _stream(client, calls, ref, payload_type) + + +def to_stream( + client: temporalio.client.Client, + ref: StreamRef[T], + *, + payload_type: type[T] | None, +) -> Stream[T]: + """The Stream of a ref that came as data. + + See :py:meth:`temporalio.client.Client.to_stream`. + """ + _, calls = client_backend(client, ref.backend, hint="the ref's backend field") + return _stream(client, calls, ref, payload_type) + + +def _stream( + client: temporalio.client.Client, + backend: BackendCalls, + ref: StreamRef[T], + payload_type: Any, +) -> Stream[T]: + return Stream(ref, payload_type, ClientCalls(client, backend, ref)) diff --git a/temporalio/streams/_writer.py b/temporalio/streams/_writer.py new file mode 100644 index 0000000000..1a0c1fa0e0 --- /dev/null +++ b/temporalio/streams/_writer.py @@ -0,0 +1,213 @@ +"""The writer of an Activity or a client Producer.""" + +from __future__ import annotations + +import asyncio +import weakref +from collections.abc import Callable, Generator, Sequence +from typing import TYPE_CHECKING, Any + +import temporalio.api.stream.v1 as stream_v1 +from temporalio.streams import _errors + +if TYPE_CHECKING: + from temporalio.streams.writer_system import client as writer_client + from temporalio.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 Stream. + """ + + 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.""" + + 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 | 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 + 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: + """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.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( + ref=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 + + 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, since its Stream object closed it. + """ + if self._closed: + raise _errors.StreamClosedError( + f"Can't write to the Stream {self._stream_id}, since this object closed its " + "writer" + ) + 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.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] + ) + return answer.last_offset or 0 + + async def close(self) -> None: + """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: + 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 0000000000..c636710c0d --- /dev/null +++ b/tests/contrib/streams/fake_extension.py @@ -0,0 +1,193 @@ +"""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, StreamIdConflictPolicy +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.service import RPCStatusCode +from temporalio.streams import StreamBackend +from temporalio.streams._backend import _ExtensionCalls +from temporalio.streams.proto import v1 as ext_pb + + +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, request_id: str) -> None: + self.request_id = request_id + self.records: list[stream_v1.StreamRecord] = [] + self.closed = False + self.changed = asyncio.Event() + + def wake(self) -> None: + self.changed.set() + 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.creates: list[stream_pb.CreateStreamRequest] = [] + 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.wake() + return record.offset + + def free_writer(self, writer_id: int) -> None: + self.freed.append(writer_id) + + def _create(self, create: stream_pb.CreateStreamRequest) -> bytes: + self.creates.append(create) + self.namespaces.add(create.namespace) + stream_id = create.stream_ref.id + stream = self.streams.get(stream_id) + created = stream is None + if stream is None: + stream = self.streams[stream_id] = _Stream(create.request_id) + elif stream.closed: + raise _error( + RPCStatusCode.FAILED_PRECONDITION, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + elif create.request_id and create.request_id == stream.request_id: + # A retry of the create that made the Stream. + created = True + elif ( + create.id_conflict_policy + != StreamIdConflictPolicy.STREAM_ID_CONFLICT_POLICY_USE_EXISTING + ): + raise _error( + RPCStatusCode.ALREADY_EXISTS, + StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS, + ) + return stream_pb.CreateStreamResponse( + ref=stream_v1.StreamRef(id=stream_id), created=created + ).SerializeToString() + + async def call(self, method: str, input: bytes) -> bytes: + if method == "CreateStream": + return self._create(stream_pb.CreateStreamRequest.FromString(input)) + if method == "OpenWriter": + open_writer = ext_pb.OpenWriterRequest.FromString(input) + writer_id = len(self.writers) + 1 + self.writers[writer_id] = (open_writer.ref.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": + # A close writes no Record. It wakes the waiting reads, which then end. + close = stream_pb.CloseStreamRequest.FromString(input) + stream = self.streams[close.ref.id] + stream.closed = True + stream.wake() + return b"" + if method == "DescribeStream": + describe = stream_pb.DescribeStreamRequest.FromString(input) + records = self.streams[describe.ref.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.ref.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.ref.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 or stream.closed: + return stream_pb.ReadRecordsResponse( + records=found, + next_offset=found[-1].offset + 1 if found else offset, + done=stream.closed, + ) + 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 0000000000..a3b204742f --- /dev/null +++ b/tests/contrib/streams/test_facade.py @@ -0,0 +1,314 @@ +import asyncio +import gc +import subprocess +import sys +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +import temporalio.api.stream.v1 +import temporalio.client +from temporalio import streams +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import StreamErrorReason, StreamIdConflictPolicy +from temporalio.api.stream.v1 import StreamErrorDetails +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError +from temporalio.service import RPCError, RPCStatusCode +from temporalio.streams import _errors, _stream_client +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, streams.StreamFencedError) + assert err.reason == fenced + # A reason with no class of its own gives the base class, with the reason. + for reason in ( + StreamErrorReason.STREAM_ERROR_REASON_SEQUENCE_TOO_LOW, + StreamErrorReason.STREAM_ERROR_REASON_STORAGE_CORRUPT, + ): + other = _errors.from_call_error( + StreamCallError(int(RPCStatusCode.DATA_LOSS), "other", _details(reason)) + ) + assert type(other) is streams.StreamError and other.reason == reason + assert other.non_retryable + unknown = _errors.from_call_error( + StreamCallError(int(RPCStatusCode.INTERNAL), "odd", b"") + ) + assert type(unknown) is streams.StreamError + assert unknown.reason == StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED + exists = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.ALREADY_EXISTS), + "exists", + _details(StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS), + ) + ) + assert isinstance(exists, streams.StreamAlreadyExistsError) + + +async def test_client_writes_reads_and_closes(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream[Turn]("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() + # A close writes no item. The read ends once it gave the last Record. + items = [item async for item in stream.subscribe()] + assert [item.kind for item in items] == [streams.ItemKind.RECORD] * 2 + turn, token = items + assert isinstance(turn, streams.Record) and turn.data == Turn("model") + assert isinstance(token, streams.Record) and token.data == Turn("user", "hello") + assert turn.producer is not None and turn.producer.epoch is not None + assert isinstance(turn.producer.epoch.kind, streams.ProducerEpochKindClient) + 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] + + +async def test_create_fails_on_an_id_that_exists(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream[str]("taken") + with pytest.raises(streams.StreamAlreadyExistsError): + await client.create_stream[str]("taken") + # The create-or-get gets it, and ignores the options. + got = await client.create_or_get_stream[str]("taken", retention=timedelta(hours=1)) + assert got.ref == stream.ref + policies = [create.id_conflict_policy for create in extension.creates] + assert policies == [ + StreamIdConflictPolicy.STREAM_ID_CONFLICT_POLICY_FAIL, + StreamIdConflictPolicy.STREAM_ID_CONFLICT_POLICY_FAIL, + StreamIdConflictPolicy.STREAM_ID_CONFLICT_POLICY_USE_EXISTING, + ] + # Each create call sends its own request id, so a retry of one isn't another create. + request_ids = [create.request_id for create in extension.creates] + assert all(request_ids) and len(set(request_ids)) == 3 + + +async def test_a_stream_gives_its_ref_and_its_wire_message(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream[Turn]("chat-proto") + assert stream.ref == streams.StreamRef(id="chat-proto", backend="fake") + assert stream.proto_ref == temporalio.api.stream.v1.StreamRef( + id="chat-proto", backend="fake" + ) + + +async def test_read_waits_for_create(client: Client): + client, _ = fake_client(client) + stream = client.get_stream[str]("tokens-7") + with pytest.raises(streams.StreamNotFoundError): + async for _ in stream.subscribe(): + pass + + async def create_later() -> None: + created = await client.create_stream[str](stream.ref.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 [item.data for item in items if isinstance(item, streams.Record)] == ["five"] + + +async def test_a_waiting_read_ends_at_a_close(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream[str]("closes-later") + reader = stream.subscribe() + + async def read() -> list[str]: + return [item.data async for item in reader if isinstance(item, streams.Record)] + + task = asyncio.create_task(read()) + await stream.publish("last") + await asyncio.sleep(0.05) + await stream.close() + assert await asyncio.wait_for(task, 5) == ["last"] + + +def test_stream_errors_fail_the_workflow_not_the_task(): + closed = StreamErrorReason.STREAM_ERROR_REASON_CLOSED + err = streams.StreamClosedError("closed") + assert isinstance(err, ApplicationError) + assert err.non_retryable and err.type == "StreamClosedError" + assert err.reason == closed + # A Store outage is the one error that a retry can fix. + assert not streams.StreamStorageUnavailableError("down").non_retryable + mapped = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "closed", _details(closed) + ) + ) + assert isinstance(mapped, streams.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", StreamErrorDetails(reason=closed), type="StreamClosedError" + ) + ) + assert isinstance(returned, streams.StreamClosedError) + assert returned.reason == closed + + +async def test_streams_of_one_client_share_a_producer_and_free_it(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream[str]("two-writers") + other = client.get_stream[str]("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, streams.Record) and item.producer + } + assert len(producers) == 1 + # A Stream that Python frees without a close frees its writer. So a server that makes a + # Stream 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[str]("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, streams.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[str]("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, streams.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 + + +async def test_client_streams_are_of_the_exported_class(client: Client): + client, _ = fake_client(client) + stream = client.get_stream[str]("exported") + assert isinstance(stream, temporalio.client.Stream) + assert "Stream" in temporalio.client.__all__ + + +def test_import_of_temporalio_client_loads_no_stream_code(): + code = "import sys, temporalio.client; print('temporalio.streams' in sys.modules)" + loaded = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ).stdout + assert loaded.strip() == "False" + + +async def test_bad_create_and_read_arguments_raise_with_no_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(ValueError, match="id is empty"): + await client.create_stream[str]("") + stream = client.get_stream[str]("args") + with pytest.raises(ValueError, match="from 1 to 1000"): + stream.subscribe(last=1001) + assert not extension.creates and not extension.namespaces + + +async def test_a_failed_close_is_tried_again(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream[str]("closes-twice") + await stream.publish("a") + call = extension.call + refused: list[str] = [] + + async def refuses_the_first_close(method: str, input: bytes) -> bytes: + if method == "CloseStream" and not refused: + refused.append(method) + raise StreamCallError( + int(RPCStatusCode.UNAVAILABLE), + "down", + _details(StreamErrorReason.STREAM_ERROR_REASON_STORAGE_UNAVAILABLE), + ) + return await call(method, input) + + setattr(extension, "call", refuses_the_first_close) + with pytest.raises(streams.StreamStorageUnavailableError): + async with stream: + pass + # The writer closed before the Store refused, so this object writes no more. + with pytest.raises(streams.StreamClosedError, match="closed its writer"): + stream.publish("late") + # The Stream is still open, and the next close goes through. + assert not extension.streams["closes-twice"].closed + await stream.close() + assert extension.streams["closes-twice"].closed diff --git a/tests/contrib/streams/test_reader.py b/tests/contrib/streams/test_reader.py new file mode 100644 index 0000000000..c4e88f5606 --- /dev/null +++ b/tests/contrib/streams/test_reader.py @@ -0,0 +1,115 @@ +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import Any, cast + +import pytest +from typing_extensions import override + +import temporalio.common +import temporalio.converter +from temporalio.streams import StreamNotFoundError, _backend, _reader +from temporalio.streams._items import Record, Superseded +from temporalio.streams.system import models + + +class _Pages(_backend.StreamServiceCalls): + """Answers each read with the next page, and keeps the requests.""" + + def __init__(self, *pages: models.ReadRecordsResponse[Any]) -> None: + super().__init__(cast(Any, None)) + self.pages = list(pages) + self.requests: list[models.ReadRecordsRequest] = [] + + @override + async def send_read_records( + self, request: models.ReadRecordsRequest, body_type: type[Any] + ) -> models.ReadRecordsResponse[Any]: + assert body_type is temporalio.common.RawValue + self.requests.append(request) + return self.pages.pop(0) + + +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(records=[], next_offset=8), + models.ReadRecordsResponse(records=[record], next_offset=10), + ) + page = await _reader._read_until_records( + service, + models.ReadRecordsRequest(ref=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) + + +def _activity(id: str, attempt: int) -> models.Producer: + return models.Producer( + id=id, + epoch=models.ProducerEpoch( + kind=models.ProducerEpochKindActivity(models.ActivityEpoch(attempt=attempt)) + ), + ) + + +async def test_superseded_lists_each_withdrawn_offset_once(): + failed = _activity("act:run:1:1", 1) + other = _activity("act:run:2:2", 1) + # A batch has consecutive offsets, and other Producers' Records come between batches. + written = [(10, failed), (11, failed), (12, failed), (13, other), (20, failed)] + facts = models.StreamSuperseded( + producer_id=failed.id, + failed_tasks=[], + epoch=_activity(failed.id, 2).epoch, + ) + + async def items() -> AsyncGenerator[_reader.Item, None]: + for offset, producer in written: + yield Record(data=None, offset=offset, producer=producer) + yield Superseded(offset=30, facts=facts) + yield Superseded(offset=31, facts=facts) + + 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), ()] + + +class _Missing(_backend.StreamServiceCalls): + """Answers each read at once with NOT_FOUND, as a Store that ignores create_wait.""" + + def __init__(self) -> None: + super().__init__(cast(Any, None)) + self.calls = 0 + + @override + async def send_read_records( + self, request: models.ReadRecordsRequest, body_type: type[Any] + ) -> models.ReadRecordsResponse[Any]: + self.calls += 1 + raise StreamNotFoundError("nobody created it") + + +async def test_a_wait_for_the_create_pauses_between_reads(): + service = _Missing() + items = _reader.read( + service, + temporalio.converter.DataConverter.default, + "missing", + position=_reader.Position(0), + payload_type=None, + from_offset=None, + wait_new_records=False, + create_wait=timedelta(milliseconds=500), + ) + with pytest.raises(StreamNotFoundError): + await anext(items) + # The pauses double from 50 ms, so the half second has a few reads, not a hot loop. + assert 2 <= service.calls <= 8 diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py new file mode 100644 index 0000000000..f731c6c9c2 --- /dev/null +++ b/tests/contrib/streams/test_stream_client.py @@ -0,0 +1,155 @@ +import asyncio +import uuid +from dataclasses import dataclass +from typing import Any + +import pytest + +from temporalio import streams, 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 = "" + + +async def test_unknown_backend_raises_before_any_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(streams.StreamConfigError, match="no Stream backend redis"): + await client.create_stream[Turn]("chat-1", backend="redis") + with pytest.raises(streams.StreamConfigError, match="It has fake"): + client.get_stream[Turn]("chat-1", backend="redis") + with pytest.raises(streams.StreamConfigError, match="It has fake"): + client.to_stream(streams.StreamRef(id="chat-1", backend="redis")) + assert not extension.streams and not extension.namespaces + stream = await client.create_stream[Turn]("chat-1") + assert stream.ref.backend == "fake" + + +async def test_client_without_backend_has_no_transport(client: Client): + with pytest.raises(streams.StreamConfigError, match="no Stream backend"): + client.get_stream[Turn]("chat-1") + with pytest.raises(streams.StreamConfigError, match="no Stream backend"): + await client.create_stream[Turn]("chat-1") + + +async def test_to_stream_gives_the_stream_of_a_ref(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream[Turn]("chat-2") + converter = DataConverter.default.payload_converter + ref = converter.from_payload( + converter.to_payload(stream.ref), streams.StreamRef[Turn] + ) + assert ref == stream.ref + typed = client.to_stream(ref, payload_type=Turn) + assert typed.ref is ref and await typed.get_offset() == 0 + await typed.publish(Turn("model", "hi")) + [record] = [r async for r in typed.subscribe(wait_new_records=False)] + assert isinstance(record, streams.Record) and record.data == Turn("model", "hi") + # With no payload type, a read decodes each Record as a JSON value. + untyped = client.to_stream(ref) + [raw] = [r async for r in untyped.subscribe(wait_new_records=False)] + assert isinstance(raw, streams.Record) and raw.data == { + "kind": "model", + "text": "hi", + } + + +async def test_a_call_with_no_type_argument_reads_json_values(client: Client): + client, _ = fake_client(client) + typed = await client.create_stream[Turn]("chat-3") + await typed.publish(Turn("model", "hi")) + untyped: Any = client.get_stream("chat-3") + [raw] = [r async for r in untyped.subscribe(wait_new_records=False)] + assert raw.data == {"kind": "model", "text": "hi"} + + +async def test_publish_starts_the_write_at_the_call(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream[str]("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, streams.Record)] == [ + "0", + "1", + "2", + "3", + ] + assert [await futures[0], await futures[2]] == [items[0].offset, items[2].offset] + with pytest.raises(streams.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[str]("lost") + await client.get_stream[str]("lost").close() + late = stream.publish("late") + await asyncio.wait([late]) + with pytest.raises(streams.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_a_client_cant_close_a_stream_bound_to_a_workflow(client: Client): + client, extension = fake_client(client) + await client.create_stream[str]("bound") + owner = streams.Execution( + type=streams.ExecutionType.ExecutionTypeWorkflow, + business_id="wf", + run_id="run", + ) + ref: streams.StreamRef[str] = streams.StreamRef(id="bound", owner=owner) + bound = client.to_stream(ref) + with pytest.raises(streams.StreamNotOwnerError): + await bound.close() + with pytest.raises(streams.StreamNotOwnerError): + async with bound: + pass + assert not extension.streams["bound"].closed + + +@workflow.defn +class HandsOutStream: + @workflow.run + 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's Stream that a query returns. + return {"id": "handed-out"} + + +async def test_query_and_workflow_results_give_refs(client: Client): + client, _ = fake_client(client) + created = await client.create_stream[str]("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", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + result_type=streams.StreamRef[str], + ) + result = await handle.result() + queried = await handle.query("stream", result_type=streams.StreamRef[str]) + # The ref has no backend, so its Stream uses the default one. + assert queried.id == "handed-out" and queried.backend is None + await client.to_stream(queried, payload_type=str).publish("from the query result") + await client.to_stream(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, streams.Record)] == [ + "from the query result", + "from the Workflow result", + ]