diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 771c454e0..6de435e36 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -128,8 +128,9 @@ class ReadSpec: def read_spec( handle: StreamHandle[Any], + # Any is a Topic of Workflow code, which this module can't import. topics: Sequence[ - str | TopicDefinition[Any] | _AsyncTopic[Any] | SyncTopicBase[Any] + str | TopicDefinition[Any] | _AsyncTopic[Any] | SyncTopicBase[Any] | Any ], *, prefixes: Sequence[str], @@ -141,6 +142,8 @@ def read_spec( ) -> ReadSpec: """The Topic names of a read, the decode type of each known Topic, and its options. + A Topic object of any context, such as a ``WorkflowTopic``, gives its name and type. + Raises: ValueError: ``last`` is not above zero, or it comes with ``from_offset``. """ @@ -149,14 +152,14 @@ def read_spec( types = handle._known_types() names = [] for topic in topics: - if isinstance(topic, (_AsyncTopic, SyncTopicBase)): - name, topic_type = topic.name, topic.type - elif isinstance(topic, TopicDefinition): + if isinstance(topic, TopicDefinition): name, topic_type = handle._topic_spec(topic, None) - else: + elif isinstance(topic, str): # A read names a Topic outside the definition with no type, and its values # decode with the read's default type. name, topic_type = topic, types.get(topic) + else: + name, topic_type = topic.name, topic.type names.append(name) if topic_type is not None: types[name] = topic_type diff --git a/temporalio/contrib/streams/_errors.py b/temporalio/contrib/streams/_errors.py index 55dbb0084..6ae947520 100644 --- a/temporalio/contrib/streams/_errors.py +++ b/temporalio/contrib/streams/_errors.py @@ -142,6 +142,31 @@ def from_call_error(err: Exception) -> StreamError: ) +_ERROR_CLASSES: list[type[StreamError]] = [StreamError, *_BY_REASON.values()] +_BY_NAME: dict[str, type[StreamError]] = {cls.__name__: cls for cls in _ERROR_CLASSES} + + +def from_application_error( + err: temporalio.exceptions.ApplicationError, +) -> StreamError | None: + """The Stream error that an application error of a failed call carries, if any. + + A Stream error's type is its class name, and its first detail is its reason. A server may + give the reason as the type, or at the start of the message, such as ``stream_closed: ...``. + """ + token = err.message.split(":", 1)[0].strip() + reason = err.details[0] if err.details and isinstance(err.details[0], str) else "" + cls = _BY_NAME.get(err.type or "") + if cls is None: + reason = next( + (name for name in (err.type or "", token) if name in _BY_REASON), "" + ) + cls = _BY_REASON.get(reason) + if cls is None: + return None + return cls(err.message, reason=reason) + + def from_reason(reason: str, message: str) -> StreamError: """The Stream error for a reason token.""" return _BY_REASON.get(reason, StreamError)(message, reason=reason) diff --git a/temporalio/contrib/streams/_workflow_reader.py b/temporalio/contrib/streams/_workflow_reader.py new file mode 100644 index 000000000..650b66b1c --- /dev/null +++ b/temporalio/contrib/streams/_workflow_reader.py @@ -0,0 +1,191 @@ +"""Reads of a Stream in Workflow code. + +Each page is a call whose answer is in History, so a replay reads the same items. A Stream +plugin's outbound interceptor reads a page with ``_temporal_stream_read_page``. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, Mapping, Sequence +from datetime import datetime, timedelta +from typing import Any + +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.common +import temporalio.converter +import temporalio.exceptions +import temporalio.nexus.system +import temporalio.workflow +from temporalio.contrib.streams import _convert, _errors +from temporalio.contrib.streams._items import Closed, End, Record, Superseded +from temporalio.contrib.streams.system import models + +Item = Record[Any] | Superseded | Closed + +# The service waits at most this long for a Stream that nobody created yet, per call. +_MAX_CREATE_WAIT = timedelta(seconds=60) +# A backend that can't park a read pauses on a durable timer after a page with no Record, +# from the first pause up to the last. +_FIRST_PAUSE = timedelta(seconds=1) +_LAST_PAUSE = timedelta(seconds=60) + + +def reader_of(backend: str) -> Any: + """The outbound interceptor of the Stream plugin with this backend name.""" + outbound = getattr(temporalio.workflow._Runtime.current(), "_outbound", None) + while outbound is not None: + if ( + getattr(outbound, "_temporal_stream_transport", False) + and getattr(outbound, "_temporal_stream_backend_name", None) == backend + ): + return outbound + outbound = getattr(outbound, "next", None) + raise RuntimeError(f"No Stream plugin of the Worker has the backend {backend}") + + +async def read( + reader: Any, + ref: stream_pb.StreamRef, + *, + topics: Sequence[str] | None, + prefixes: Sequence[str], + types: Mapping[str, Any], + from_offset: int | End | None, + wait_new_records: bool, + create_wait: timedelta | None, +) -> AsyncGenerator[Item, None]: + """Give the items of a Stream in order, until a close, or the newest Record when + ``wait_new_records`` is False. + """ + request = stream_pb.ReadRecordsRequest( + stream=ref, topics=list(topics or []), topic_prefixes=list(prefixes) + ) + if from_offset is End.END: + # END is only the first page's start. Each answer's next offset pins later pages. + request.start = stream_pb.ReadStart.READ_START_END + elif from_offset: + request.from_offset = from_offset + else: + request.start = stream_pb.ReadStart.READ_START_BEGINNING + create_until = temporalio.workflow.now() + create_wait if create_wait else None + pause = _FIRST_PAUSE + while True: + if temporalio.workflow.unsafe.is_read_only(): + raise _errors.ReadOnlyContextError("Can't read a Stream in read-only code") + page_request = stream_pb.ReadRecordsRequest() + page_request.CopyFrom(request) + if wait_new_records: + page_request.wait.FromTimedelta(reader._temporal_stream_read_wait) + left = _left(create_until) + if left: + page_request.create_wait.FromTimedelta(min(left, _MAX_CREATE_WAIT)) + try: + page = await reader._temporal_stream_read_page(page_request) + except _errors.StreamNotFoundError: + if not _left(create_until): + raise + await temporalio.workflow.sleep(pause) + pause = min(pause * 2, _LAST_PAUSE) + continue + for record in page.records: + item = workflow_item(record, types) + if item is None: + continue + yield item + if isinstance(item, Closed): + return + if page.done: + return + if page.next_offset: + request.from_offset = page.next_offset + request.ClearField("start") + if page.records: + pause = _FIRST_PAUSE + continue + if not wait_new_records: + return + if not reader._temporal_stream_parks_reads: + # The server wake comes later. Until then a quiet Stream costs one timer and one + # read per pause. + await temporalio.workflow.sleep(pause) + pause = min(pause * 2, _LAST_PAUSE) + + +def _left(until: datetime | None) -> timedelta | None: + if until is None: + return None + left = until - temporalio.workflow.now() + return left if left > timedelta() else None + + +def workflow_item( + record: stream_pb.StreamRecord, types: Mapping[str, Any] +) -> Item | None: + """One read item, with the body decoded by the Workflow's converter and its Topic's type. + + The codec already ran outside the sandbox, on the page. + """ + converter = temporalio.workflow.payload_converter() + # The generated converter makes the producer and the supersession models, and the body + # stays a Payload that the Topic's type decodes. + with temporalio.nexus.system._user_converter_context( + temporalio.nexus.system._SystemNexusUserConverters( + payload_converter=_convert._RawPayloadConverter(), + failure_converter=_failure_converter(), + ) + ): + model = models._StreamRecordTransferTypeConverter().from_transfer_type( + record, models.StreamRecord + ) + offset = model.offset or 0 + if model.kind == stream_pb.StreamRecordKind.STREAM_RECORD_KIND_DATA: + topic = model.topic or "" + return Record( + topic=topic, + data=_decode(converter, model.body, types.get(topic)), + offset=offset, + producer=model.producer, + ) + if model.kind == stream_pb.StreamRecordKind.STREAM_RECORD_KIND_SUPERSEDED: + if model.supersession is None: + return None + return Superseded(offset=offset, supersession=model.supersession) + if model.kind == stream_pb.StreamRecordKind.STREAM_RECORD_KIND_CLOSE: + return Closed( + offset=offset, + result=_decode(converter, model.close_result, None), + failure=model.close_failure, + ) + return None + + +def _decode( + converter: temporalio.converter.PayloadConverter, body: Any, type_hint: Any +) -> Any: + if not isinstance(body, temporalio.common.RawValue): + return body + return converter.from_payload(body.payload, type_hint) + + +def _failure_converter() -> temporalio.converter.FailureConverter: + runtime = temporalio.workflow._Runtime.current() + return getattr( + runtime, + "_workflow_context_failure_converter", + temporalio.converter.DefaultFailureConverter(), + ) + + +def error_of(err: BaseException) -> BaseException: + """The Stream error that a failed page carries as its cause, or the error itself. + + A page's failure is in History, so a replay raises the same error. + """ + cause: BaseException | None = err + while cause is not None: + if isinstance(cause, temporalio.exceptions.ApplicationError): + stream_error = _errors.from_application_error(cause) + if stream_error is not None: + return stream_error + cause = cause.__cause__ + return err diff --git a/temporalio/contrib/streams/native.py b/temporalio/contrib/streams/native.py index 1e9935f75..25ff8520d 100644 --- a/temporalio/contrib/streams/native.py +++ b/temporalio/contrib/streams/native.py @@ -8,12 +8,26 @@ from datetime import timedelta +import temporalio.api.streamservice.v1 as stream_pb import temporalio.bridge.client import temporalio.bridge.streams +import temporalio.exceptions +import temporalio.nexus.system import temporalio.service +import temporalio.worker +import temporalio.workflow +from temporalio.contrib.streams import _workflow_reader from temporalio.contrib.streams._backend import _ExtensionStreamBackend from temporalio.contrib.streams.proto import v1 as ext_pb -from temporalio.contrib.streams.redis import _ExtensionStreamPlugin, _runtime_of +from temporalio.contrib.streams.redis import ( + _ExtensionStreamPlugin, + _runtime_of, + _WorkflowOutbound, +) + +_SERVICE = "temporal.api.streamservice.v1.StreamService" + +_WAIT_COMPLETED = temporalio.workflow.NexusOperationCancellationType.WAIT_COMPLETED class _ConnectionRuntime: @@ -55,6 +69,37 @@ async def _bridge_client( return await connect() +class _NativeWorkflowOutbound(_WorkflowOutbound): + # The server parks a read at the end of the Stream and completes it on an append or a + # close, so a reader needs no timer. + _temporal_stream_parks_reads = True + _temporal_stream_read_wait = timedelta(seconds=60) + + async def _temporal_stream_read_page( + self, request: stream_pb.ReadRecordsRequest + ) -> stream_pb.ReadRecordsResponse: + # The page is a System Nexus operation, so its answer is in History. + handle = await self.next.start_system_nexus_operation( + temporalio.worker.StartNexusOperationInput( + endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT, + service=_SERVICE, + operation="ReadRecords", + input=request, + output_type=stream_pb.ReadRecordsResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=_WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) + try: + return await handle + except temporalio.exceptions.NexusOperationError as err: + raise _workflow_reader.error_of(err) from None + + class NativeStreamPlugin(_ExtensionStreamPlugin): """Carries Stream calls to the Temporal server, which keeps the Streams. @@ -69,6 +114,7 @@ class NativeStreamPlugin(_ExtensionStreamPlugin): This API is experimental. """ + _outbound_class = _NativeWorkflowOutbound _server_closes_with_run = True def _new_backend( diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 594baa5b6..57cd7d8d3 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -28,7 +28,13 @@ import temporalio.service import temporalio.worker import temporalio.workflow -from temporalio.contrib.streams import _activity, _ids, _writes +from temporalio.contrib.streams import ( + _activity, + _errors, + _ids, + _workflow_reader, + _writes, +) from temporalio.contrib.streams import workflow as streams_workflow from temporalio.contrib.streams._backend import StreamBackend, _ExtensionStreamBackend from temporalio.contrib.streams._errors import StreamConfigError @@ -41,6 +47,9 @@ logger = logging.getLogger(__name__) +# How long a Workflow read's Local Activity waits for new Records. Its timeout adds a margin. +_READ_WAIT = timedelta(seconds=10) +_READ_MARGIN = timedelta(seconds=10) # The plugin puts this key into the headers of each Activity and Local Activity that a Workflow # schedules. It keeps the Producer id the same across retries. _KEY_HEADER = "__temporal_stream_key" @@ -51,6 +60,7 @@ class _ExtensionStreamPlugin(StreamPlugin): # The Redis, memory and native plugins: the Stream calls go to the temporalio-streams # runtime. + _outbound_class: type[_WorkflowOutbound] # Whether the Store closes a Stream with its owner's run chain on the server. _server_closes_with_run = False @@ -83,6 +93,8 @@ def __init__( runtime or temporalio.runtime.Runtime.default(), True ) super().__init__(name, backend_name=backend_name, default=default) + # The read activity's name has the backend name, which the base keeps. + self._read_activity = _read_activity(self) def _runtime_backend( self, runtime: temporalio.runtime.Runtime, memory: bool @@ -161,6 +173,9 @@ def configure_worker( ) config["interceptors"] = [*config.get("interceptors", []), mine] mine.converter = client.data_converter + if not self._outbound_class._temporal_stream_parks_reads: + # Workflow reads run as this Local Activity, so the Worker must have it. + config["activities"] = [*config.get("activities", []), self._read_activity] runner = config.get("workflow_runner") if isinstance(runner, SandboxedWorkflowRunner): # The interceptor and the facade must share the run's Stream state, also when the @@ -237,6 +252,67 @@ def check_worker_backends(config: temporalio.worker.WorkerConfig) -> None: ) +def read_activity_name(backend: str) -> str: + """The name of the Local Activity that reads a page of a backend for Workflow code.""" + return f"__temporal_stream_read:{backend}" + + +def _read_activity(plugin: _ExtensionStreamPlugin) -> Callable[..., Any]: + @temporalio.activity.defn(name=read_activity_name(plugin.backend_name)) + async def read_page( + request: stream_pb.ReadRecordsRequest, + ) -> stream_pb.ReadRecordsResponse: + request.namespace = temporalio.activity.info().namespace + # The Worker's client gave the attempt its Stream runtime. + streams = _attempt_streams(plugin.backend_name) + try: + page = await cast(_ExtensionStreamBackend, streams.backend).call_async( + "ReadRecords", request, stream_pb.ReadRecordsResponse + ) + except _errors.StreamError as err: + # The reason goes in the details, so Workflow code gets the same Stream error. + raise temporalio.exceptions.ApplicationError( + str(err.message), + err.reason, + type=type(err).__name__, + non_retryable=err.non_retryable, + ) from None + # Workflow code can't run the codec, so the page's bodies go to it decoded. The page + # as a whole goes through the codec as the Local Activity's result. + await _decode_page(streams.converter, page) + return page + + return read_page + + +def _attempt_streams(backend: str) -> _activity.ActivityStreams: + attempt = _activity.current.get() + streams = attempt.get(backend) if attempt is not None else None + if streams is None: + raise StreamConfigError( + f"No Stream plugin of the backend {backend} runs this Local Activity" + ) + return streams + + +async def _decode_page( + converter: temporalio.converter.DataConverter, page: stream_pb.ReadRecordsResponse +) -> None: + codec = converter.payload_codec + if codec is None: + return + payloads = [ + payload + for record in page.records + for payload in (record.body, record.close_result) + if payload.ByteSize() + ] + if not payloads: + return + for payload, decoded in zip(payloads, await codec.decode(payloads)): + payload.CopyFrom(decoded) + + def _runtime_of( service: temporalio.service.ServiceClient, ) -> temporalio.runtime.Runtime: @@ -333,6 +409,7 @@ def workflow_interceptor_class( { "_backend_name": self.plugin.backend_name, "_default": self.plugin.default, + "_outbound_class": self.plugin._outbound_class, "_server_closes_with_run": self.plugin._server_closes_with_run, "_interceptor": self if self.backend is not None else None, }, @@ -360,11 +437,13 @@ class _RunClose: class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): _backend_name = "" _default = False + # Each plugin's interceptor sets it on its own class. + _outbound_class: type[_WorkflowOutbound] # pyright: ignore[reportUninitializedInstanceVariable] _server_closes_with_run = False _interceptor: _StreamInterceptor | None = None def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: - super().init(_WorkflowOutbound(outbound, self._backend_name, self._default)) + super().init(self._outbound_class(outbound, self._backend_name, self._default)) async def execute_workflow( self, input: temporalio.worker.ExecuteWorkflowInput @@ -557,6 +636,9 @@ def _take_owned_streams( class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): _temporal_stream_transport = True + # The runtime can't park a read yet, so a reader pauses on a timer after an empty page. + _temporal_stream_parks_reads = False + _temporal_stream_read_wait = _READ_WAIT def __init__( self, @@ -579,6 +661,25 @@ def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: converter = temporalio.converter.PayloadConverter.default return {_KEY_HEADER: converter.to_payload(str(self._activity_calls))} + async def _temporal_stream_read_page( + self, request: stream_pb.ReadRecordsRequest + ) -> stream_pb.ReadRecordsResponse: + # The page is the Local Activity's result, so it is in History. + try: + return await temporalio.workflow.execute_local_activity( + read_activity_name(self._temporal_stream_backend_name), + request, + result_type=stream_pb.ReadRecordsResponse, + start_to_close_timeout=_READ_WAIT + _READ_MARGIN, + ) + except ( + temporalio.exceptions.ActivityError, + temporalio.exceptions.ApplicationError, + ) as err: + # A Local Activity's failure may come as its application error, with no + # ActivityError around it. + raise _workflow_reader.error_of(err) from None + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: # The next run needs each owned Stream's backend for stream_handle. owned = [ @@ -781,3 +882,6 @@ def _activity_streams( previous_completed_task=previous, failed_tasks=failed, ) + + +_ExtensionStreamPlugin._outbound_class = _WorkflowOutbound diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 4b1711921..cea5d81e8 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator, Sequence from dataclasses import dataclass, field from datetime import timedelta from types import TracebackType @@ -27,7 +28,8 @@ import temporalio.converter import temporalio.exceptions import temporalio.workflow -from temporalio.contrib.streams import _errors, _ids +from temporalio.contrib.streams import _errors, _ids, _workflow_reader +from temporalio.contrib.streams._access import read_spec from temporalio.contrib.streams._defn import TopicDefinition, definition_of from temporalio.contrib.streams._handle import ( Creating, @@ -35,7 +37,8 @@ TopicBase, exit_failure, ) -from temporalio.contrib.streams._items import close_status +from temporalio.contrib.streams._items import End, close_status +from temporalio.contrib.streams._workflow_reader import Item from temporalio.contrib.streams.system import models E = TypeVar("E") @@ -342,9 +345,24 @@ def publish_nowait(self, *values: T) -> None: """ _write_nowait(self._stream._append(self._name, values)) + def subscribe( + self, + *, + from_offset: int | End | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + ) -> AsyncIterator[Item]: + """Read this Topic only. See :py:meth:`WorkflowStream.subscribe`.""" + return self._stream.subscribe( + self, + from_offset=from_offset, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + class WorkflowStream(StreamHandle[E]): - """A Stream in Workflow code. Workflow code writes and closes it, and can't read it. + """A Stream in Workflow code. Workflow code writes, reads and closes it. Pass it to an Activity, a child Workflow or the next run of continue-as-new as an argument. Use it in ``async with`` to close it at the end of the block. @@ -394,6 +412,69 @@ def topic( name, topic_type = self._topic_spec(topic, type) return WorkflowTopic(self, name, topic_type) + def subscribe( + self, + *topics: str | TopicDefinition[Any] | WorkflowTopic[Any], + from_offset: int | End | None = None, + wait_new_records: bool = True, + prefixes: Sequence[str] = (), + create_wait: timedelta | None = None, + ) -> AsyncIterator[Item]: + """Read the Stream's items in order, in Workflow code. + + It is an async iterator, as on the client, and it yields the same ``Record``, + ``Superseded`` and ``Closed`` items. Each page is a call whose answer is in History, + so a replay reads the same items. On a native Stream a page is a System Nexus + ``ReadRecords`` operation, and a read at the end of the Stream waits on the server + until a Record comes or the Stream closes. On Redis a page is a Local Activity that + waits a short time for new Records, and a page with none is followed by a durable + timer from 1 s up to 60 s. + + A read sees this run's writes once they went out: at the end of the Workflow Task, + or on Redis before the page's Local Activity. + + Args: + topics: Read these Topics. With no Topics and no prefixes, read every Topic. + from_offset: Start at this offset, inclusive, or at ``streams.END`` for new + Records only. Defaults to the first Record. + wait_new_records: Wait for new Records at the end of the Stream. When False, + the read ends at the newest Record. + prefixes: Also read the Topics whose names start with one of these. + create_wait: How long to wait for a Stream that nobody created yet. Defaults to + no wait. + + Raises: + StreamNotFoundError: The Stream does not exist, or it was not created in time. + ReadOnlyContextError: The code must not change state, such as a query handler. + StreamConfigError: No Stream plugin is on the Worker, or the Stream is of a + backend that the Worker does not have. It fails the Workflow Task, not the run. + """ + read = read_spec( + self, + topics, + prefixes=prefixes, + default_type=None, + from_offset=from_offset, + last=None, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + reader = _workflow_reader.reader_of(_check_backend(self.backend, "read")) + return _workflow_reader.read( + reader, + stream_pb.StreamRef( + id=self.id, + workflow_id=self.workflow_id or "", + first_execution_run_id=self.first_execution_run_id or "", + ), + topics=read.topics, + prefixes=prefixes, + types=read.types, + from_offset=from_offset, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + def _check( self, action: str, *, backend: str | None = None ) -> tuple[_RunStream, str]: diff --git a/tests/contrib/streams/conftest.py b/tests/contrib/streams/conftest.py index c2cde1e5e..04f5d0434 100644 --- a/tests/contrib/streams/conftest.py +++ b/tests/contrib/streams/conftest.py @@ -10,6 +10,7 @@ from temporalio.client import Client from temporalio.contrib.streams import StreamPlugin +from temporalio.contrib.streams.native import NativeStreamPlugin from temporalio.contrib.streams.redis import RedisStreamPlugin from temporalio.contrib.streams.testing import MemoryStreamPlugin @@ -70,3 +71,20 @@ async def stream_client( config = client.config() config["plugins"] = [plugin] yield Client(**config) + + +# A server that serves the Stream service, such as a dev server of moetemp/temporal with the +# native Stream backend. The native lane runs only when it is set. +NATIVE_TARGET = os.environ.get("TEMPORAL_STREAMS_NATIVE_TARGET", "") +NATIVE_NAMESPACE = os.environ.get("TEMPORAL_STREAMS_NATIVE_NAMESPACE", "default") + + +@pytest.fixture +async def native_client() -> Client: + if not NATIVE_TARGET: + pytest.skip( + "TEMPORAL_STREAMS_NATIVE_TARGET names no server with the Stream service" + ) + return await Client.connect( + NATIVE_TARGET, namespace=NATIVE_NAMESPACE, plugins=[NativeStreamPlugin()] + ) diff --git a/tests/contrib/streams/test_native_server.py b/tests/contrib/streams/test_native_server.py new file mode 100644 index 000000000..071825045 --- /dev/null +++ b/tests/contrib/streams/test_native_server.py @@ -0,0 +1,132 @@ +"""The native lane: the same calls against a server that serves the Stream service. + +Set ``TEMPORAL_STREAMS_NATIVE_TARGET`` to its address to run it. +""" + +import concurrent.futures +import uuid +from datetime import timedelta + +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams.native import NativeStreamPlugin +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.worker import Replayer +from tests.contrib.streams.test_workflow_reads import ( + Events, + Reader, + ReadInput, + Writer, +) +from tests.helpers import new_worker + + +@activity.defn +async def write_native(stream: streams.ActivityStream[Events]) -> None: + await stream.topic(Events.tokens).publish("activity") + + +@activity.defn +def write_native_sync(stream: streams.SyncActivityStream[Events]) -> None: + stream.topic(Events.tokens).publish("sync activity") + + +@workflow.defn +class WritesNativeAndRedis: + @workflow.run + async def run(self) -> list[str | None]: + on_redis = await stream_workflow.create_stream(streams.with_name(Events, "old")) + on_native = await stream_workflow.create_stream( + streams.with_name(Events, "new"), backend="native" + ) + await on_redis.topic(Events.tokens).publish("redis") + await on_native.topic(Events.tokens).publish("workflow 1", "workflow 2") + # The Activities write after the run's own writes, as on Redis. + await workflow.execute_activity( + write_native, + streams.ActivityStream.from_handle(on_native), + start_to_close_timeout=timedelta(seconds=10), + ) + await workflow.execute_activity( + write_native_sync, + streams.SyncActivityStream.from_handle(on_native), + start_to_close_timeout=timedelta(seconds=10), + ) + await on_native.topic(Events.tokens).publish("last") + return [on_redis.backend, on_native.backend] + + +async def test_native_client_round_trip(native_client: Client): + stream = await streams.StreamClient(native_client).create( + Events, id=f"native-{uuid.uuid4()}", auto_close=timedelta(minutes=5) + ) + assert await stream.topic(Events.tokens).publish("a", "b") + await stream.close(result="done") + items = [item async for item in stream.subscribe()] + assert [item.data for item in items if isinstance(item, streams.Record)] == [ + "a", + "b", + ] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" + + +async def test_native_workflow_reads_and_replays(native_client: Client): + async with new_worker(native_client, Writer, Reader) as worker: + writer = await native_client.start_workflow( + Writer.run, None, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + stream = streams.StreamClient(native_client).workflow_stream(writer, Events) + reader = await native_client.start_workflow( + Reader.run, + ReadInput(stream_workflow.WorkflowStream.from_handle(stream)), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await reader.result() == [ + "turn user hi", + "a", + "b", + "c", + "closed written", + ] + replayer = Replayer(workflows=[Writer, Reader], plugins=[NativeStreamPlugin()]) + await replayer.replay_workflow(await writer.fetch_history()) + await replayer.replay_workflow(await reader.fetch_history()) + + +async def test_redis_and_native_on_a_native_server( + native_client: Client, redis_url: str +): + config = native_client.config() + config["plugins"] = [ + RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True), + NativeStreamPlugin(), + ] + two = Client(**config) + async with new_worker( + two, + WritesNativeAndRedis, + activities=[write_native, write_native_sync], + activity_executor=concurrent.futures.ThreadPoolExecutor(max_workers=2), + ) as worker: + handle = await two.start_workflow( + WritesNativeAndRedis.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == ["redis", "native"] + new = streams.StreamClient(two).workflow_stream( + handle, streams.with_name(Events, "new"), backend="native" + ) + items = [item async for item in new.subscribe()] + assert [item.data for item in items if isinstance(item, streams.Record)] == [ + "workflow 1", + "workflow 2", + "activity", + "sync activity", + "last", + ] + assert isinstance(items[-1], streams.Closed) + assert items[-1].result == ["redis", "native"] diff --git a/tests/contrib/streams/test_workflow_reads.py b/tests/contrib/streams/test_workflow_reads.py new file mode 100644 index 000000000..f5cf4d207 --- /dev/null +++ b/tests/contrib/streams/test_workflow_reads.py @@ -0,0 +1,253 @@ +import dataclasses +import uuid +from collections.abc import Coroutine +from dataclasses import dataclass +from datetime import timedelta +from typing import Any, cast + +import pytest + +from temporalio import workflow +from temporalio.client import Client, WorkflowFailureError +from temporalio.contrib import streams +from temporalio.contrib.streams import _errors +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.converter import DataConverter +from temporalio.exceptions import ApplicationError +from temporalio.worker import Replayer +from tests.contrib.streams.test_redis import _Rot13Codec +from tests.helpers import new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + + +@workflow.defn +class Writer: + @workflow.run + async def run(self, backend: str | None) -> str: + stream = await stream_workflow.create_stream(Events, backend=backend) + await stream.topic(Events.turns).publish(Turn("user", "hi")) + for token in ("a", "b", "c"): + # Each token goes out in its own Workflow Task, so the reader waits for some. + await workflow.sleep(0.3) + await stream.topic(Events.tokens).publish(token) + return "written" + + +@dataclass +class ReadInput: + stream: stream_workflow.WorkflowStream[Events] + tokens_only: bool = False + wait_new_records: bool = True + + +@workflow.defn +class Reader: + @workflow.run + async def run(self, input: ReadInput) -> list[Any]: + seen: list[Any] = [] + items = ( + input.stream.topic(Events.tokens).subscribe( + wait_new_records=input.wait_new_records, + create_wait=timedelta(seconds=20), + ) + if input.tokens_only + else input.stream.subscribe( + wait_new_records=input.wait_new_records, + create_wait=timedelta(seconds=20), + ) + ) + async for item in items: + # The Topic's type decodes each body in Workflow code. + if isinstance(item, streams.Record) and isinstance(item.data, Turn): + seen.append(f"turn {item.data.kind} {item.data.text}") + elif isinstance(item, streams.Record): + seen.append(item.data) + elif isinstance(item, streams.Closed): + seen.append(f"closed {item.result}") + return seen + + +def plugins_of(client: Client) -> list[Any]: + return [ + plugin + for plugin in client.config()["plugins"] + if isinstance(plugin, streams.StreamPlugin) + ] + + +async def write_and_read( + client: Client, backend: str | None, **read_options: Any +) -> tuple[list[Any], Any]: + async with new_worker(client, Writer, Reader) as worker: + writer = await client.start_workflow( + Writer.run, + backend, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream = streams.StreamClient(client).workflow_stream( + writer, Events, backend=backend + ) + reader = await client.start_workflow( + Reader.run, + ReadInput( + stream_workflow.WorkflowStream.from_handle(stream), **read_options + ), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + result = await reader.result() + assert await writer.result() == "written" + return result, reader + + +async def test_workflow_reads_another_workflows_stream(stream_client: Client): + seen, reader = await write_and_read(stream_client, None) + assert seen == ["turn user hi", "a", "b", "c", "closed written"] + # Each page is in History, so a replay reads the same items. + replayer = Replayer(workflows=[Reader], plugins=plugins_of(stream_client)) + await replayer.replay_workflow(await reader.fetch_history()) + + +async def test_workflow_reads_one_topic(stream_client: Client): + seen, _ = await write_and_read(stream_client, None, tokens_only=True) + assert seen == ["a", "b", "c", "closed written"] + + +@workflow.defn +class ReadsMissing: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = stream_workflow.WorkflowStream[Events](id=stream_id) + try: + async for _ in stream.subscribe(): + pass + except streams.StreamNotFoundError as err: + return f"not found: {err.reason}" + return "read" + + +async def test_a_missing_stream_raises_in_workflow_code(stream_client: Client): + async with new_worker(stream_client, ReadsMissing) as worker: + result = await stream_client.execute_workflow( + ReadsMissing.run, + f"missing-{uuid.uuid4()}", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == "not found: stream_not_found" + + +@workflow.defn +class ReadsOwnWrites: + @workflow.run + async def run(self) -> list[Any]: + stream = await stream_workflow.create_stream(Events, close_with_run=False) + await stream.topic(Events.tokens).publish("mine") + # A read's Local Activity first writes the run's buffer, so the read sees it. + return [ + item.data + async for item in stream.subscribe(wait_new_records=False) + if isinstance(item, streams.Record) + ] + + +async def test_a_redis_read_sees_the_runs_own_writes(stream_client: Client): + async with new_worker(stream_client, ReadsOwnWrites) as worker: + assert await stream_client.execute_workflow( + ReadsOwnWrites.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) == ["mine"] + + +@workflow.defn +class ReadsInQuery: + def __init__(self) -> None: + self._stream: stream_workflow.WorkflowStream[Events] | None = None + + @workflow.run + async def run(self) -> None: + # Set in the first task, since a query may come before a create's Local Activity ends. + self._stream = stream_workflow.create_stream_nowait(Events) + await workflow.wait_condition(lambda: False) + + @workflow.query + def read(self) -> str: + assert self._stream is not None + items = self._stream.subscribe() + try: + # A query can't await, so it takes the read's first step by hand. + cast(Coroutine[Any, Any, Any], items.__anext__()).send(None) + except streams.ReadOnlyContextError as err: + return str(err) + except StopIteration: + pass + return "read" + + +async def test_a_read_in_a_query_raises(stream_client: Client): + async with new_worker(stream_client, ReadsInQuery) as worker: + handle = await stream_client.start_workflow( + ReadsInQuery.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert "read-only" in await handle.query(ReadsInQuery.read) + await handle.terminate() + + +@workflow.defn +class FailsOnStreamError: + @workflow.run + async def run(self, stream_id: str) -> None: + stream = stream_workflow.WorkflowStream[Events](id=stream_id) + async for _ in stream.subscribe(): + pass + + +async def test_an_uncaught_read_error_fails_the_run(stream_client: Client): + # A Stream error is in History, so it fails the run, not the Workflow Task. + async with new_worker(stream_client, FailsOnStreamError) as worker: + with pytest.raises(WorkflowFailureError) as failure: + await stream_client.execute_workflow( + FailsOnStreamError.run, + f"missing-{uuid.uuid4()}", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert isinstance(failure.value.cause, ApplicationError) + assert failure.value.cause.type == "StreamNotFoundError" + + +async def test_workflow_reads_go_through_the_codec(stream_client: Client): + # The Store keeps encoded bodies, and Workflow code can't run the codec, so the read's + # Local Activity decodes them. + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=_Rot13Codec() + ) + coded = Client(**config) + seen, _ = await write_and_read(coded, None) + assert seen == ["turn user hi", "a", "b", "c", "closed written"] + + +def test_a_server_failure_maps_by_its_reason(): + by_type = _errors.from_application_error( + ApplicationError("closed", type="stream_closed") + ) + assert isinstance(by_type, streams.StreamClosedError) + by_message = _errors.from_application_error( + ApplicationError("stream_not_found: Nobody created the Stream") + ) + assert isinstance(by_message, streams.StreamNotFoundError) + assert _errors.from_application_error(ApplicationError("other")) is None