diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index 0b48fafcf..63bdb7fa7 100644 --- a/scripts/_wit/stream-service.wit +++ b/scripts/_wit/stream-service.wit @@ -174,15 +174,6 @@ interface stream-service { stream-status-expired, } - /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" - record create-stream-response { - %stream: option, - /// @nexus.doc "False when the Stream existed with the same Retention." - created: option, - /// @nexus.doc "Closed when an earlier create and close used the id." - status: option, - } - // The meta key of each Stream has a format version v. Lua refuses a write to a Stream with a // higher v with FAILED_PRECONDITION and the detail stream_format_too_new. /// @nexus.doc @@ -290,8 +281,65 @@ interface stream-service { auto-close: option, } + // Its stream-ref is its serialized form, so a handle passes through Workflow input, Activity + // input and Signals as the ref. Each method fills in the ref of its request and takes the rest. + // The wire operations stay the same. A handle never changes: offsets are values that the + // methods return. + /// @nexus.doc "A handle on one Stream. Build it from a stream-ref, or get it from a create." + resource %stream { + constructor(%stream: stream-ref); + + /// @nexus.doc "Append Records. The answer has the first and the last offset." + /// @nexus.operation "append-records" + append-records: func( + producer: option, + sequence: option, + records: list, + previous-completed-task: option, + failed-tasks: list, + withdraw-run: option, + ) -> append-records-response; + + /// @nexus.doc "Close the Stream. Readers drain what is left, then end." + /// @nexus.operation "close-stream" + close: func( + producer: option, + %result: option, + failure: option, + previous-completed-task: option, + failed-tasks: list, + ); + + /// @nexus.doc "Read one page of Records. The iterator of the same name reads them all." + /// @nexus.operation "read-records" + read-records: func( + from-offset: option, + start: option, + topics: list, + topic-prefixes: list, + max-records: option, + wait: option, + create-wait: option, + ) -> read-records-response; + + /// @nexus.doc "Describe the Stream: its status, newest offset, Retention and auto-close." + /// @nexus.operation "describe-stream" + describe: func() -> describe-stream-response; + } + + // The handle is built from the stream field of the response. The other fields stay beside + // it, since they describe this create, not the Stream. + /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" + record create-stream-result { + %stream: own<%stream>, + /// @nexus.doc "False when the Stream existed with the same Retention." + created: option, + /// @nexus.doc "Closed when an earlier create and close used the id." + status: option, + } + /// @nexus.operation name="CreateStream" - create-stream: func(request: create-stream-request) -> create-stream-response; + create-stream: func(request: create-stream-request) -> create-stream-result; /// @nexus.doc /// "In Workflow code, Brook buffers this call in the run, and the facade drops the answer." @@ -301,6 +349,8 @@ interface stream-service { /// @nexus.operation name="CloseStream" close-stream: func(request: close-stream-request) -> close-stream-response; + // The long poll moves from-offset, so nexgen also reads its pages as one async iterator. A WIT + // stream can't cross a Nexus call, so the read stays request and response on the wire. /// @nexus.doc "Not for Workflow code in Brook. A long poll. NOT_FOUND for an uncreated Stream." /// @nexus.operation name="ReadRecords" /// @nexus.long-poll diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py index d04af620c..292303a52 100644 --- a/scripts/gen_stream_service.py +++ b/scripts/gen_stream_service.py @@ -32,8 +32,8 @@ output_dir = base_dir / "temporalio" / "contrib" / "streams" / "system" NEX_GEN_REPOSITORY = "https://github.com/moetemp/nexgen" -# The head of moe/brook/v3/n7-long-poll-loop. A branch moves, so the pin is the commit. -NEX_GEN_REVISION = "25b9c1be8f51245618297415aa51df6928c58504" +# The head of moe/brook/v3/n8-stream-resource. A branch moves, so the pin is the commit. +NEX_GEN_REVISION = "10fc35b40fcf48651c25e654a260bc71585f1429" nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index c5ae55a30..9304de8ac 100644 --- a/temporalio/bridge/Cargo.lock +++ b/temporalio/bridge/Cargo.lock @@ -2879,7 +2879,7 @@ dependencies = [ [[package]] name = "temporalio-streams" version = "0.1.0" -source = "git+https://github.com/moetemp/sdk-streams?rev=daeb5fc2a13e1f3457a461f109e7f5b8eb31951c#daeb5fc2a13e1f3457a461f109e7f5b8eb31951c" +source = "git+https://github.com/moetemp/sdk-streams?rev=2cc3eee3760a24b2feef30dc9471010cd49750e6#2cc3eee3760a24b2feef30dc9471010cd49750e6" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index d2e25c4d3..93f011cfb 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -94,7 +94,7 @@ url = "2.2" # crates, so the bridge sees one type of each. [dependencies.temporalio-streams] git = "https://github.com/moetemp/sdk-streams" -rev = "daeb5fc2a13e1f3457a461f109e7f5b8eb31951c" +rev = "2cc3eee3760a24b2feef30dc9471010cd49750e6" default-features = false optional = true diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 089aa6a10..289bf22da 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -25,7 +25,6 @@ SyncActivityStream, SyncActivityTopic, ) -from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._buffered import BufferedPublisher from temporalio.contrib.streams._defn import ( StreamDefinition, @@ -36,17 +35,18 @@ from temporalio.contrib.streams._errors import ( ReadOnlyContextError, StreamClosedError, + StreamConfigError, StreamError, StreamExpiredError, StreamFencedError, StreamNotFoundError, StreamOffsetError, StreamStorageError, - StreamTransportMissingError, StreamUnsupportedError, ) from temporalio.contrib.streams._handle import StreamHandle -from temporalio.contrib.streams._items import Closed, Record, Superseded +from temporalio.contrib.streams._items import END, Closed, Record, Superseded +from temporalio.contrib.streams._plugin import StreamPlugin from temporalio.contrib.streams._stream_client import ( ClientStream, ClientTopic, @@ -55,6 +55,7 @@ from temporalio.contrib.streams.workflow import WorkflowStream, WorkflowTopic __all__ = [ + "END", "ActivityStream", "ActivityTopic", "BufferedPublisher", @@ -63,9 +64,9 @@ "Closed", "ReadOnlyContextError", "Record", - "StreamBackend", "StreamClient", "StreamClosedError", + "StreamConfigError", "StreamDefinition", "StreamError", "StreamExpiredError", @@ -73,8 +74,8 @@ "StreamHandle", "StreamNotFoundError", "StreamOffsetError", + "StreamPlugin", "StreamStorageError", - "StreamTransportMissingError", "StreamUnsupportedError", "Superseded", "SyncActivityStream", diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 80d1eda29..e7601de54 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -5,16 +5,18 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import timedelta -from typing import Any, Generic, TypeVar +from typing import Any, TypeVar -import temporalio.api.streamservice.v1 as stream_pb import temporalio.converter from temporalio.contrib.streams import _convert, _reader from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._defn import TopicDefinition -from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._handle import StreamHandle, TopicBase +from temporalio.contrib.streams._items import End from temporalio.contrib.streams._reader import Item from temporalio.contrib.streams._writer import StreamWriter +from temporalio.contrib.streams.system import client as service_client +from temporalio.contrib.streams.system import models T = TypeVar("T") @@ -41,7 +43,7 @@ def read( topics: list[str] | None, prefixes: Sequence[str], types: dict[str, Any], - from_offset: int | None, + from_offset: int | End | None, wait_new_records: bool, create_wait: timedelta | None, ) -> AsyncGenerator[Item, None]: @@ -61,12 +63,12 @@ async def _read( topics: list[str] | None, prefixes: Sequence[str], types: dict[str, Any], - from_offset: int | None, + from_offset: int | End | None, wait_new_records: bool, create_wait: timedelta | None, ) -> AsyncGenerator[Item, None]: items = _reader.read( - self.backend.service(self.namespace), + self.backend.service(self.namespace, _convert.raw(self.converter)), self.converter, await self.stream_id(), topics=topics, @@ -85,13 +87,10 @@ async def get_offset(self) -> int: It is 0 for a Stream with no Record, and ``from_offset=0`` reads from the start. Offsets are sparse and above zero, so 0 is never a Record's offset. """ - response = await self.backend.call_async( - "DescribeStream", - stream_pb.DescribeStreamRequest( - namespace=self.namespace, - stream=stream_pb.StreamRef(id=await self.stream_id()), - ), - stream_pb.DescribeStreamResponse, + response = await self.backend.service(self.namespace).describe_stream( + models.DescribeStreamRequest( + stream=models.StreamRef(id=await self.stream_id()) + ) ) # Offsets start above zero, so zero means the Stream has no Record. return response.last_offset + 1 if response.last_offset else 0 @@ -100,20 +99,14 @@ async def close(self, result: Any | None, failure: BaseException | None) -> None """Write this Producer's Records, then close the Stream.""" writer = self.writer() await writer.close() - request = stream_pb.CloseStreamRequest( - namespace=self.namespace, - stream=stream_pb.StreamRef(id=await self.stream_id()), - producer=writer.producer, + stream = service_client.Stream( + self.backend.service(self.namespace, self.converter), + models.StreamRef(id=await self.stream_id()), ) - if result is not None: - request.result.CopyFrom( - (await _convert.encode(self.converter, result)).payload - ) - if failure is not None: - await self.converter.encode_failure(failure, request.failure) - await self.backend.call_async( - "CloseStream", request, stream_pb.CloseStreamResponse + producer = models._ProducerTransferTypeConverter().from_transfer_type( + writer.producer, models.Producer ) + await stream.close(producer=producer, result=result, failure=failure) def read_spec( @@ -140,34 +133,21 @@ def read_spec( return names or None, types -class _AsyncTopic(Generic[T]): +class _AsyncTopic(TopicBase[T]): # A Topic of an Activity or client handle. Each publish writes before it returns. - def __init__(self, stream: Any, name: str, type: Any) -> None: - self._stream = stream - self._name = name - self._type = type - - @property - def name(self) -> str: - """Name of the Topic.""" - return self._name - - @property - def type(self) -> Any: - """Type that a read decodes the Topic's values with.""" - return self._type async def publish(self, *values: T) -> int: """Write one Record for each value, in order, and return the last Record's offset. A list is one value. """ + await self._stream._ready() return await self._stream._access().publish(self._name, values) def subscribe( self, *, - from_offset: int | None = None, + from_offset: int | End | None = None, wait_new_records: bool = True, create_wait: timedelta | None = None, ) -> AsyncGenerator[Item, None]: @@ -180,21 +160,5 @@ def subscribe( ) -class SyncTopicBase(Generic[T]): +class SyncTopicBase(TopicBase[T]): """Base of a Topic whose calls block, for sync Activities.""" - - def __init__(self, stream: Any, name: str, type: Any) -> None: - """Create a Topic. Use the Stream's ``topic`` instead.""" - self._stream = stream - self._name = name - self._type = type - - @property - def name(self) -> str: - """Name of the Topic.""" - return self._name - - @property - def type(self) -> Any: - """Type that a read decodes the Topic's values with.""" - return self._type diff --git a/temporalio/contrib/streams/_activity_stream.py b/temporalio/contrib/streams/_activity_stream.py index 1c258b7ad..f4891fcc7 100644 --- a/temporalio/contrib/streams/_activity_stream.py +++ b/temporalio/contrib/streams/_activity_stream.py @@ -17,6 +17,7 @@ ) from temporalio.contrib.streams._defn import TopicDefinition from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._items import End from temporalio.contrib.streams._reader import Item E = TypeVar("E") @@ -30,14 +31,13 @@ def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: streams = _activity.current.get() if streams is None: - raise _errors.StreamTransportMissingError( + raise _errors.StreamConfigError( "No Stream plugin runs this Activity, such as RedisStreamPlugin" ) if handle.backend and handle.backend != streams.backend_name: - raise _errors.StreamUnsupportedError( + raise _errors.StreamConfigError( f"The Stream {handle.id} is of the backend {handle.backend}. The Worker has only " - f"{streams.backend_name}.", - reason="stream_unsupported", + f"{streams.backend_name}." ) return streams @@ -91,39 +91,19 @@ def topic(self, topic: str, *, type: type[T]) -> ActivityTopic[T]: ... def topic( self, topic: str | TopicDefinition[Any], *, type: Any = None ) -> ActivityTopic[Any]: - """A Topic of this Stream. - - Args: - topic: A Topic of the Stream's definition, such as ``Events.turns``, or a name. - A name of the definition takes its type. - type: The type of a Topic outside the definition, or of a Stream with no - definition. - """ + """A Topic of this Stream. See :py:meth:`ClientStream.topic`.""" name, topic_type = self._topic_spec(topic, type) return ActivityTopic(self, name, topic_type) def subscribe( self, *topics: str | TopicDefinition[Any] | ActivityTopic[Any], - from_offset: int | None = None, + from_offset: int | End | None = None, wait_new_records: bool = True, prefixes: Sequence[str] = (), create_wait: timedelta | None = None, ) -> AsyncGenerator[Item, None]: - """Read the Stream's items in order. - - The read ends after a ``Closed`` item, or at the newest Record when - ``wait_new_records`` is False. A ``Superseded`` item withdraws earlier Records of a - failed attempt, and it passes every Topic filter. - - Args: - topics: Read these Topics. With no Topics and no prefixes, read every Topic. - from_offset: Start at this offset, inclusive. Defaults to the first Record. - wait_new_records: Wait for new Records at the end of the Stream. - 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. - """ + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" names, types = read_spec(self, topics) return self._access().read( topics=names, @@ -172,7 +152,7 @@ def publish(self, *values: T) -> int: def subscribe( self, *, - from_offset: int | None = None, + from_offset: int | End | None = None, wait_new_records: bool = True, create_wait: timedelta | None = None, ) -> Iterator[Item]: @@ -237,19 +217,19 @@ def topic(self, topic: str, *, type: type[T]) -> SyncActivityTopic[T]: ... def topic( self, topic: str | TopicDefinition[Any], *, type: Any = None ) -> SyncActivityTopic[Any]: - """A Topic of this Stream. See :py:meth:`ActivityStream.topic`.""" + """A Topic of this Stream. See :py:meth:`ClientStream.topic`.""" name, topic_type = self._topic_spec(topic, type) return SyncActivityTopic(self, name, topic_type) def subscribe( self, *topics: str | TopicDefinition[Any] | SyncActivityTopic[Any], - from_offset: int | None = None, + from_offset: int | End | None = None, wait_new_records: bool = True, prefixes: Sequence[str] = (), create_wait: timedelta | None = None, ) -> Iterator[Item]: - """Read the Stream's items in order. See :py:meth:`ActivityStream.subscribe`.""" + """Read the Stream's items in order. See :py:meth:`ClientStream.subscribe`.""" names, types = read_spec(self, topics) items = self._access().read( topics=names, diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index a2a6ac994..76740fb1e 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -3,13 +3,13 @@ from __future__ import annotations from datetime import timedelta -from typing import TypeVar +from typing import Protocol, TypeVar import google.protobuf.message import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb -import temporalio.client +import temporalio.converter import temporalio.worker from temporalio.contrib.streams import _errors from temporalio.contrib.streams.proto import v1 as ext_pb @@ -34,28 +34,61 @@ } -class StreamBackend: - """Carries Stream calls to a worker extension with the temporalio-streams methods. +class StreamBackend(Protocol): + """What a Stream plugin gives the handles. Not public. Each backend has its own class.""" + + def service( + self, + namespace: str, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> service_client.StreamServiceClient: + """The Stream service for ``namespace``.""" + ... + + async def open_writer(self, request: ext_pb.OpenWriterRequest) -> int: + """Open the writer of an Activity or client Producer, and return its id.""" + ... + + async def write_records( + self, writer_id: int, records: list[stream_pb.AppendRecord] + ) -> int: + """Write Records for a writer, and return the last Record's offset.""" + ... + + async def close_writer(self, writer_id: int) -> int: + """Drop a writer. The Stream stays open.""" + ... - A Stream plugin gives one to the facade. The five Stream service operations go through the - generated client, and the writer methods are the crate's own. + async def flush_run_buffer( + self, namespace: str, run_id: str, workflow_task: stream_pb.WorkflowTaskEpoch + ) -> int: + """Write a run's buffered calls before a Local Activity of the run starts.""" + ... - .. warning:: - This API is experimental. + +class _ExtensionStreamBackend: + """Carries Stream calls to a worker extension with the temporalio-streams methods. + + The five Stream service operations go through the generated client, and the writer + methods are the crate's own. """ def __init__(self, extension: temporalio.worker.WorkerExtension) -> None: """Create a backend over an extension.""" self._extension = extension - @property - def extension(self) -> temporalio.worker.WorkerExtension: - """The extension that the backend calls.""" - return self._extension + def service( + self, + namespace: str, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> service_client.StreamServiceClient: + """A generated Stream service client for ``namespace``. - def service(self, namespace: str) -> service_client.StreamServiceClient: - """A generated Stream service client for ``namespace``.""" - return service_client.StreamServiceClient(_ExtensionTransport(self, namespace)) + With a ``data_converter``, user values in the requests go through it and its codec. + """ + return service_client.StreamServiceClient( + _ExtensionTransport(self, namespace), data_converter=data_converter + ) def call(self, method: str, request: google.protobuf.message.Message) -> bytes: """Call a synchronous extension method.""" @@ -137,7 +170,7 @@ async def drain_closes(self, timeout: timedelta) -> int: class _ExtensionTransport: # The generated client takes a transport, so the five operations reach the extension as # they would reach a server later. - def __init__(self, backend: StreamBackend, namespace: str) -> None: + def __init__(self, backend: _ExtensionStreamBackend, namespace: str) -> None: self._backend = backend self._namespace = namespace @@ -161,41 +194,3 @@ async def call( }, data=response.SerializeToString(), ) - - -def client_backend( - client: temporalio.client.Client, name: str | None -) -> tuple[str, StreamBackend]: - """The name and the backend of a client's Stream plugin. None picks the plugin marked - as the default, or else the first one. - - Raises: - StreamTransportMissingError: The client has no Stream plugin. - StreamUnsupportedError: The client has no Stream plugin of that name. - """ - backends = [ - ( - getattr(plugin, "backend_name", ""), - backend, - getattr(plugin, "default", False), - ) - for plugin in client.config()["plugins"] - if isinstance( - backend := getattr(plugin, "_temporal_stream_backend", None), StreamBackend - ) - ] - if not backends: - raise _errors.StreamTransportMissingError( - "The client has no Stream plugin, such as RedisStreamPlugin" - ) - if not name: - defaults = [entry for entry in backends if entry[2]] or backends - return defaults[0][0], defaults[0][1] - for backend_name, backend, _ in backends: - if backend_name == name: - return backend_name, backend - raise _errors.StreamUnsupportedError( - f"The client has no Stream backend {name}. It has " - f"{', '.join(entry[0] for entry in backends)}.", - reason="stream_unsupported", - ) diff --git a/temporalio/contrib/streams/_convert.py b/temporalio/contrib/streams/_convert.py index e0980c2e4..ba24d6490 100644 --- a/temporalio/contrib/streams/_convert.py +++ b/temporalio/contrib/streams/_convert.py @@ -2,19 +2,18 @@ from __future__ import annotations -import contextlib -from collections.abc import Iterator, Sequence +import dataclasses +from collections.abc import Sequence from typing import Any import temporalio.api.common.v1 import temporalio.common import temporalio.converter -import temporalio.nexus.system class _RawPayloadConverter(temporalio.converter.PayloadConverter): - # The generated models carry Record bodies as RawValue, so the facade runs the user's - # converter and codec on the bodies only, never on the outer request. + # The generated client runs the codec on Record bodies and keeps them as RawValue, so the + # facade converts each body with its Topic's type. def to_payloads( self, values: Sequence[Any] ) -> list[temporalio.api.common.v1.Payload]: @@ -33,33 +32,20 @@ def from_payloads( return [temporalio.common.RawValue(payload) for payload in payloads] -@contextlib.contextmanager -def raw_bodies( - failure_converter: temporalio.converter.FailureConverter, -) -> Iterator[None]: - """Convert generated models with Record bodies kept as RawValue.""" - with temporalio.nexus.system._user_converter_context( - temporalio.nexus.system._SystemNexusUserConverters( - payload_converter=_RawPayloadConverter(), - failure_converter=failure_converter, - ) - ): - yield +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) -async def encode( - converter: temporalio.converter.DataConverter, value: Any -) -> temporalio.common.RawValue: - """Convert and encode one value for a Record body.""" - [payload] = await converter.encode([value]) - return temporalio.common.RawValue(payload) - - -async def decode( +def decode( converter: temporalio.converter.DataConverter, body: Any, type_hint: Any ) -> Any: - """Decode a Record body that a read gave as RawValue.""" + """Convert a Record body that a read gave as RawValue, after the codec.""" if not isinstance(body, temporalio.common.RawValue): return body - [value] = await converter.decode([body.payload], [type_hint] if type_hint else None) + [value] = converter.payload_converter.from_payloads( + [body.payload], [type_hint] if type_hint else None + ) return value diff --git a/temporalio/contrib/streams/_errors.py b/temporalio/contrib/streams/_errors.py index 85d2323f9..de203ebed 100644 --- a/temporalio/contrib/streams/_errors.py +++ b/temporalio/contrib/streams/_errors.py @@ -77,8 +77,11 @@ class StreamUnsupportedError(StreamError): """The Store lacks a capability that the call needs. The message ends with its name.""" -class StreamTransportMissingError(StreamError): - """No plugin carries Stream calls, such as ``RedisStreamPlugin``.""" +class StreamConfigError(RuntimeError): + """No Stream plugin, or no backend of a given name. It is config, not a Store error. + + ``except StreamError`` does not catch it. It fails a Workflow Task, and an Activity retries. + """ class ReadOnlyContextError(StreamError): @@ -96,6 +99,13 @@ class ReadOnlyContextError(StreamError): } +# The details move to temporal.api.streamservice.v1. Both names carry the same message. +_DETAILS_TYPES = ( + "temporal.api.streamservice.v1.StreamErrorDetails", + StreamErrorDetails.DESCRIPTOR.full_name, +) + + def from_extension_error(err: temporalio.worker.ExtensionError) -> StreamError: """Turn a failed extension call into the Stream error for its reason.""" reason = "" @@ -104,3 +114,14 @@ def from_extension_error(err: temporalio.worker.ExtensionError) -> StreamError: return _BY_REASON.get(reason, StreamError)( err.message, reason=reason, code=err.code ) + + +def from_rpc_error(err: temporalio.service.RPCError) -> Exception: + """The Stream error for a server failure with Stream error details, else ``err``.""" + for detail in err.grpc_status.details: + if detail.type_url.rsplit("/", 1)[-1] in _DETAILS_TYPES: + reason = StreamErrorDetails.FromString(detail.value).reason + return _BY_REASON.get(reason, StreamError)( + err.message, reason=reason, code=err.status + ) + return err diff --git a/temporalio/contrib/streams/_handle.py b/temporalio/contrib/streams/_handle.py index 2c8219b31..65c04083f 100644 --- a/temporalio/contrib/streams/_handle.py +++ b/temporalio/contrib/streams/_handle.py @@ -12,6 +12,29 @@ from temporalio.contrib.streams._defn import TopicDefinition, definition_of E = TypeVar("E") +T = TypeVar("T") + + +class TopicBase(Generic[T]): + """Base of the Topic of each handle class. All Topics of a Stream share one order.""" + + def __init__(self, stream: Any, name: str, type: Any) -> None: + """Create a Topic. Use the Stream's ``topic`` instead.""" + self._stream = stream + self._name = name + self._type = type + + @property + def name(self) -> str: + """Name of the Topic.""" + return self._name + + @property + def type(self) -> Any: + """Type that a read decodes the Topic's values with.""" + return self._type + + H = TypeVar("H", bound="StreamHandle[Any]") @@ -185,6 +208,10 @@ def _known_types(self) -> dict[str, Any]: topics = definition_of(self._definition).topics return {name: topic.type for name, topic in topics.items()} + async def _ready(self) -> None: + # A handle that must look up its backend first does it here. + pass + def _fields(self) -> dict[str, str]: # The fields that name the Stream. They need no call, so a handle whose id is not known # yet still compares, hashes and prints. diff --git a/temporalio/contrib/streams/_items.py b/temporalio/contrib/streams/_items.py index 9cf0f8870..cda16568b 100644 --- a/temporalio/contrib/streams/_items.py +++ b/temporalio/contrib/streams/_items.py @@ -2,6 +2,7 @@ from __future__ import annotations +import enum from dataclasses import dataclass from typing import Any, Generic, TypeVar @@ -11,6 +12,15 @@ T = TypeVar("T") +class End(enum.Enum): + """A ``from_offset`` after the newest Record, for new Records only. Use ``END``.""" + + END = "END" + + +END = End.END + + @dataclass(frozen=True) class Record(Generic[T]): """A data Record of a Stream. diff --git a/temporalio/contrib/streams/_plugin.py b/temporalio/contrib/streams/_plugin.py new file mode 100644 index 000000000..56a92b6f9 --- /dev/null +++ b/temporalio/contrib/streams/_plugin.py @@ -0,0 +1,86 @@ +"""The base of the Stream plugins.""" + +from __future__ import annotations + +from typing import Any + +import temporalio.client +from temporalio.contrib.streams import _errors +from temporalio.contrib.streams._backend import StreamBackend +from temporalio.plugin import SimplePlugin + + +class StreamPlugin(SimplePlugin): + """Base of the Stream plugins. Each registers one backend under its name. + + Put a Stream plugin on the client. A Worker takes its client's plugins. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + *, + backend_name: str, + default: bool, + stream_backend: StreamBackend, + **options: Any, + ) -> None: + """Create the plugin. ``options`` go to :py:class:`temporalio.plugin.SimplePlugin`.""" + self._backend_name = backend_name + self._default = default + self._stream_backend = stream_backend + super().__init__(name, **options) + + @property + def backend_name(self) -> str: + """The name that Stream handles carry for this backend.""" + return self._backend_name + + @property + def default(self) -> bool: + """Whether a client with more than one Stream plugin uses this one by default.""" + return self._default + + @property + def stream_backend(self) -> StreamBackend: + """The backend that carries the calls.""" + return self._stream_backend + + +def client_backend( + client: temporalio.client.Client, name: str | None, *, only: bool = False +) -> tuple[str, StreamBackend]: + """The name and the backend of a client's Stream plugin. None picks the plugin marked + as the default, or else the first one. With ``only``, None needs a client with one + backend, since a handle with no backend name may be of any of them. + + Raises: + StreamConfigError: The client has no Stream plugin, none of that name, or more than + one where ``only`` needs one. + """ + backends = [ + (plugin.backend_name, plugin.stream_backend, plugin.default) + for plugin in client.config()["plugins"] + if isinstance(plugin, StreamPlugin) + ] + if not backends: + raise _errors.StreamConfigError( + "The client has no Stream plugin, such as RedisStreamPlugin" + ) + if not name and only and len(backends) > 1: + raise _errors.StreamConfigError( + "The handle names no backend, and the client has more than one" + ) + if not name: + defaults = [entry for entry in backends if entry[2]] or backends + return defaults[0][0], defaults[0][1] + for backend_name, backend, _ in backends: + if backend_name == name: + return backend_name, backend + raise _errors.StreamConfigError( + f"The client has no Stream backend {name}. It has " + f"{', '.join(entry[0] for entry in backends)}." + ) diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/contrib/streams/_reader.py index defd0ffd5..5fbb3f98e 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/contrib/streams/_reader.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import dataclasses from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import timedelta from typing import Any @@ -10,7 +11,7 @@ import temporalio.converter from temporalio.api.streamservice.v1 import ReadStart, StreamRecordKind from temporalio.contrib.streams import _convert, _errors -from temporalio.contrib.streams._items import Closed, Record, Superseded +from temporalio.contrib.streams._items import Closed, End, Record, Superseded from temporalio.contrib.streams.system import client as service_client from temporalio.contrib.streams.system import models @@ -18,8 +19,6 @@ # The service waits at most 60 s per read, so a longer wait is a loop of reads. _MAX_WAIT = timedelta(seconds=60) -# How long one loop of the generated long poll may run before the reader looks again. -_POLL_DEADLINE = timedelta(minutes=5) async def read( @@ -30,7 +29,7 @@ async def read( topics: Sequence[str] | None, topic_prefixes: Sequence[str] = (), types: Mapping[str, Any], - from_offset: int | None, + from_offset: int | End | None, wait_new_records: bool, create_wait: timedelta | None, ) -> AsyncGenerator[Item, None]: @@ -40,46 +39,65 @@ async def read( loop = asyncio.get_running_loop() create_until = loop.time() + create_wait.total_seconds() if create_wait else None offset = from_offset - names = list(topics or []) - while True: + + def request() -> models.ReadRecordsRequest: create_left = ( timedelta(seconds=max(create_until - loop.time(), 0)) if create_until is not None else None ) - request = models.ReadRecordsRequest( + return models.ReadRecordsRequest( stream=models.StreamRef(id=stream_id), - from_offset=offset, - start=None if offset is not None else ReadStart.READ_START_BEGINNING, - topics=names, + # END is only the first page's start. Each answer's next offset pins later pages. + from_offset=offset if isinstance(offset, int) else None, + start=ReadStart.READ_START_END + if offset is End.END + else ReadStart.READ_START_BEGINNING + if offset is None + else None, + topics=list(topics or []), topic_prefixes=list(topic_prefixes), create_wait=min(create_left, _MAX_WAIT) if create_left else None, ) + + while True: + # The generated iterator pages the long poll. A read with no wait reads pages until + # one has no Record. + records = ( + service.iter_read_records(request()) + if wait_new_records + else _pages_until_empty(service, request()) + ) try: - with _convert.raw_bodies(converter.failure_converter): - if wait_new_records: - answer = await service.read_records_until_records( - request, deadline=_POLL_DEADLINE - ) - else: - answer = await service.read_records(request) + while True: + try: + record = await records.__anext__() + except StopAsyncIteration: + return + item = await _item(converter, record, types) + if item is None: + continue + yield item + if isinstance(item, Closed): + return except _errors.StreamNotFoundError: - if create_until is not None and loop.time() < create_until: - continue - raise + if create_until is None or loop.time() >= create_until: + raise + finally: + if isinstance(records, AsyncGenerator): + await records.aclose() + + +async def _pages_until_empty( + service: service_client.StreamServiceClient, request: models.ReadRecordsRequest +) -> AsyncGenerator[models.StreamRecord, None]: + while True: + answer = await service.read_records(request) for record in answer.records or []: - item = await _item(converter, record, types) - if item is None: - continue - yield item - if isinstance(item, Closed): - return - if answer.next_offset is not None: - offset = answer.next_offset - if answer.records: - continue - if answer.done or not wait_new_records: + yield record + if answer.done or not answer.records or answer.next_offset is None: return + request = dataclasses.replace(request, from_offset=answer.next_offset) async def _item( @@ -91,13 +109,13 @@ async def _item( if record.kind == StreamRecordKind.STREAM_RECORD_KIND_DATA: topic = record.topic or "" type_hint = types.get(topic) - data = await _convert.decode(converter, record.body, type_hint) + data = _convert.decode(converter, record.body, type_hint) return Record(topic=topic, data=data, offset=offset, producer=record.producer) if record.kind == StreamRecordKind.STREAM_RECORD_KIND_SUPERSEDED: if record.supersession is None: return None return Superseded(offset=offset, supersession=record.supersession) if record.kind == StreamRecordKind.STREAM_RECORD_KIND_CLOSE: - result = await _convert.decode(converter, record.close_result, None) + result = _convert.decode(converter, record.close_result, None) return Closed(offset=offset, result=result, failure=record.close_failure) return None diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index f318e3016..8ea7b681d 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -13,11 +13,14 @@ import temporalio.client from temporalio.contrib.streams import _errors, _ids from temporalio.contrib.streams._access import Access, _AsyncTopic, read_spec -from temporalio.contrib.streams._backend import StreamBackend, client_backend +from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._defn import TopicDefinition, definition_of from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._items import End +from temporalio.contrib.streams._plugin import client_backend from temporalio.contrib.streams._reader import Item from temporalio.contrib.streams._writer import StreamWriter +from temporalio.contrib.streams.system import models E = TypeVar("E") T = TypeVar("T") @@ -80,18 +83,46 @@ def __init__( # The Stream name of a Workflow's Stream whose first run id is not known yet. self._stream_name = "" self._finding: asyncio.Future[str] | None = None + # The backends to look in, in order, until one has the Stream. + self._probe: list[tuple[str, StreamBackend]] = [] def _bind( self, client: temporalio.client.Client, backend: StreamBackend, find_first_run: Callable[[], Awaitable[str]] | None = None, + probe: list[tuple[str, StreamBackend]] | None = None, ) -> ClientStream[E]: self._client = client self._stream_backend = backend self._find_first_run = find_first_run + self._probe = probe or [] return self + @property + def backend(self) -> str | None: + """Name of the backend. A handle that looks in many backends knows it after a call.""" + return None if self._probe else self._backend + + async def _ready(self) -> None: + # The first backend that has the Stream keeps it. Two first calls may both look. + for name, backend in self._probe: + assert self._client is not None + request = models.DescribeStreamRequest( + stream=models.StreamRef(id=await self._stream_id()) + ) + try: + await backend.service(self._client.namespace).describe_stream(request) + except _errors.StreamNotFoundError: + continue + self._backend, self._stream_backend, self._probe = name, backend, [] + return + if self._probe: + raise _errors.StreamNotFoundError( + f"No backend of {[name for name, _ in self._probe]} has the Stream {self._id}", + reason="stream_not_found", + ) + @property def id(self) -> str: """The Stream id. @@ -131,7 +162,7 @@ def _fields(self) -> dict[str, str]: def _access(self) -> Access: client, backend = self._client, self._stream_backend if client is None or backend is None: - raise _errors.StreamTransportMissingError( + raise _errors.StreamConfigError( "This ClientStream came from a payload, so it has no client. " "Use StreamClient.get_handle." ) @@ -181,7 +212,7 @@ def topic( def subscribe( self, *topics: str | TopicDefinition[Any] | ClientTopic[Any], - from_offset: int | None = None, + from_offset: int | End | None = None, wait_new_records: bool = True, prefixes: Sequence[str] = (), create_wait: timedelta | None = None, @@ -195,14 +226,15 @@ def subscribe( Args: topics: Read these Topics. With no Topics and no prefixes, read every Topic. - from_offset: Start at this offset, inclusive. Defaults to the first Record. + 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. 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, such as the Stream of a Workflow that just started. Defaults to no wait. """ names, types = read_spec(self, topics) - return self._access().read( + return self._read( topics=names, prefixes=prefixes, types=types, @@ -211,8 +243,14 @@ def subscribe( create_wait=create_wait, ) + async def _read(self, **options: Any) -> AsyncGenerator[Item, None]: + await self._ready() + async for item in self._access().read(**options): + yield item + async def get_offset(self) -> int: """The next offset of the Stream, the newest offset plus one.""" + await self._ready() return await self._access().get_offset() async def close( @@ -223,6 +261,7 @@ async def close( Readers get the Records before it, then a ``Closed`` item with ``result`` or ``failure``. """ + await self._ready() await self._access().close(result, failure) @@ -259,22 +298,17 @@ async def create( backend: The name of the backend. Defaults to the client's first Stream plugin. Raises: - StreamUnsupportedError: The client has no backend of that name. No call is made. + StreamConfigError: The client has no backend of that name. No call is made. StreamClosedError: A Stream of that id was created and closed before. StreamError: The Stream exists with another Retention or ``auto_close``, with the reason ``stream_already_exists``. """ definition_of(definition) backend_name, stream_backend = client_backend(self._client, backend) - request = stream_pb.CreateStreamRequest( - namespace=self._client.namespace, stream_id=id - ) - if retention is not None: - request.retention.FromTimedelta(retention) - if auto_close is not None: - request.auto_close.FromTimedelta(auto_close) - response = await stream_backend.call_async( - "CreateStream", request, stream_pb.CreateStreamResponse + response = await stream_backend.service(self._client.namespace).create_stream( + models.CreateStreamRequest( + stream_id=id, retention=retention, auto_close=auto_close + ) ) if response.status == stream_pb.StreamStatus.STREAM_STATUS_CLOSED: raise _errors.StreamClosedError( @@ -287,21 +321,51 @@ async def create( return stream._bind(self._client, stream_backend) def get_handle( - self, definition: type[E], *, id: str, backend: str | None = None + self, + definition: type[E], + *, + id: str | None = None, + stream: StreamHandle[Any] | None = None, + backend: str | tuple[str, ...] | None = None, ) -> ClientStream[E]: - """A Stream by its id, such as one that this or another client created. No call. + """A Stream by its id, or a handle that came as data, such as in a Record. No call. Args: definition: The Stream definition class. id: The Stream id. - backend: The name of the backend. Defaults to the client's first Stream plugin. + stream: A handle of the Stream, such as one in a Record. Its id, owner and + backend stay. + backend: The name of the backend. Defaults to the client's default backend. A + tuple, such as ``("native", "redis")`` during a move, looks in each in order + on first use and keeps the first that has the Stream. """ definition_of(definition) - backend_name, stream_backend = client_backend(self._client, backend) - stream = ClientStream( - id=id, backend=backend_name or None, definition=definition + if stream is not None: + # An empty name is of any backend, so it needs a client with one. + backend_name, stream_backend = client_backend( + self._client, stream.backend, only=True + ) + probe: list[tuple[str, StreamBackend]] = [] + elif id is not None: + backend_name, stream_backend, probe = self._backends(backend) + else: + raise TypeError("get_handle takes id= or stream=") + handle = ClientStream( + id=stream.id if stream else id or "", + workflow_id=stream.workflow_id if stream else None, + first_execution_run_id=stream.first_execution_run_id if stream else None, + backend=backend_name or None, + definition=definition, ) - return stream._bind(self._client, stream_backend) + return handle._bind(self._client, stream_backend, probe=probe) + + def _backends( + self, backend: str | tuple[str, ...] | None + ) -> tuple[str, StreamBackend, list[tuple[str, StreamBackend]]]: + if not isinstance(backend, tuple): + return (*client_backend(self._client, backend), []) + probe = [client_backend(self._client, name) for name in backend] + return (*probe[0], probe if len(probe) > 1 else []) def workflow_stream( self, @@ -309,7 +373,7 @@ def workflow_stream( definition: type[E], *, name: str | None = None, - backend: str | None = None, + backend: str | tuple[str, ...] | None = None, ) -> ClientStream[E]: """A Stream that a Workflow created, by the Workflow handle and the definition. @@ -320,10 +384,11 @@ def workflow_stream( handle: A handle of the Workflow, best the one that ``start_workflow`` returned. definition: The Stream definition class. name: The Stream name that the Workflow gave. Defaults to the definition's name. - backend: The name of the backend. Defaults to the client's first Stream plugin. + backend: The name of the backend, or a tuple of names to look in, as for + :py:meth:`get_handle`. Defaults to the client's default backend. """ stream_name = name or definition_of(definition).name - backend_name, stream_backend = client_backend(self._client, backend) + backend_name, stream_backend, probe = self._backends(backend) first_run = handle.first_execution_run_id stream = ClientStream( id=_ids.scoped_stream_id(handle.id, first_run, stream_name) @@ -335,7 +400,7 @@ def workflow_stream( definition=definition, ) if first_run: - return stream._bind(self._client, stream_backend) + return stream._bind(self._client, stream_backend, probe=probe) async def find_first_run() -> str: description = await handle.describe() @@ -347,4 +412,4 @@ async def find_first_run() -> str: return first stream._stream_name = stream_name - return stream._bind(self._client, stream_backend, find_first_run) + return stream._bind(self._client, stream_backend, find_first_run, probe) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 0b3e698cc..58398cbe7 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -26,11 +26,11 @@ import temporalio.workflow from temporalio.contrib.streams import _activity, _errors, _ids from temporalio.contrib.streams import workflow as streams_workflow -from temporalio.contrib.streams._backend import StreamBackend +from temporalio.contrib.streams._backend import _ExtensionStreamBackend from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._plugin import StreamPlugin from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.system import models -from temporalio.plugin import SimplePlugin from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner logger = logging.getLogger(__name__) @@ -43,17 +43,12 @@ # 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" -# Continue-as-new passes the ids of the open Streams that close with the run chain to the next -# run, in this header. -_CLOSE_WITH_RUN_HEADER = "__temporal_stream_close_with_run" +# Continue-as-new passes the run chain's open Streams to the next run in this header. +_OWNED_STREAMS_HEADER = "__temporal_owned_streams" -class StreamPlugin(SimplePlugin): - """Base of the plugins that carry Stream calls to a temporalio-streams extension. - - .. warning:: - This API is experimental. - """ +class _ExtensionStreamPlugin(StreamPlugin): + # The Redis and memory plugins: the Stream calls go to a temporalio-streams extension. def __init__( self, @@ -69,25 +64,19 @@ def __init__( ``backend_name`` is the name that Stream handles carry for this backend. ``default`` makes it the client's default backend when the client has more than one. """ - self._backend_name = backend_name - self._default = default config.backend = backend_name self._extension = temporalio.worker.WorkerExtension( extension_name, config.SerializeToString(), runtime=runtime ) - self._temporal_stream_backend = StreamBackend(self._extension) + self._backend = _ExtensionStreamBackend(self._extension) self._converter = temporalio.converter.DataConverter.default - super().__init__(name, interceptors=[_StreamInterceptor(self)]) - - @property - def backend_name(self) -> str: - """The name that Stream handles carry for this backend.""" - return self._backend_name - - @property - def default(self) -> bool: - """Whether a client with more than one Stream plugin uses this one by default.""" - return self._default + super().__init__( + name, + backend_name=backend_name, + default=default, + stream_backend=self._backend, + interceptors=[_StreamInterceptor(self)], + ) def configure_worker( self, config: temporalio.worker.WorkerConfig @@ -127,7 +116,7 @@ async def run_worker( try: await super().run_worker(worker, next) finally: - pending = await self._temporal_stream_backend.drain_closes(_DRAIN_TIMEOUT) + pending = await self._backend.drain_closes(_DRAIN_TIMEOUT) if pending: logger.warning( "Worker shutdown left %d Stream closes in flight", pending @@ -137,7 +126,7 @@ async def shutdown(self) -> None: """Flush and drop every open writer, including the client's, and wait for the closes in flight. Call it once the process no longer writes Streams. """ - backend = self._temporal_stream_backend + backend = self._backend failed = await backend.close_writers() pending = await backend.drain_closes(_DRAIN_TIMEOUT) if failed or pending: @@ -170,10 +159,10 @@ async def on_workflow_activation_complete( ) # The call only adds to the run's buffer under a short lock, so it doesn't block the # event loop. - self._temporal_stream_backend.call("BufferWorkflowCalls", request) + self._backend.call("BufferWorkflowCalls", request) -class RedisStreamPlugin(StreamPlugin): +class RedisStreamPlugin(_ExtensionStreamPlugin): """Carries Stream calls to Redis, through the temporalio-streams runtime in the bridge. Put it on the client. A Worker takes its client's plugins, so one plugin serves Workflow @@ -345,20 +334,19 @@ def operation_token(self) -> str | None: class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): - def __init__(self, plugin: StreamPlugin) -> None: + def __init__(self, plugin: _ExtensionStreamPlugin) -> None: super().__init__() self._plugin = plugin - # Workflow code learns the Worker's backend name from the interceptor class. - self._inbound_class = type( - "_WorkflowInbound", - (_WorkflowInbound,), - {"_backend_name": plugin.backend_name}, - ) def workflow_interceptor_class( self, input: temporalio.worker.WorkflowInterceptorClassInput ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: - return self._inbound_class + # Workflow code learns the Worker's backend name from the interceptor class. + return type( + "_WorkflowInbound", + (_WorkflowInbound,), + {"_backend_name": self._plugin.backend_name}, + ) def intercept_activity( self, next: temporalio.worker.ActivityInboundInterceptor @@ -375,7 +363,7 @@ def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: async def execute_workflow( self, input: temporalio.worker.ExecuteWorkflowInput ) -> Any: - _take_close_with_run(input.headers) + _take_owned_streams(input.headers) try: result = await super().execute_workflow(input) except temporalio.workflow.ContinueAsNewError: @@ -516,16 +504,19 @@ def _owned_ref(stream_id: str) -> models.StreamRef: ) -def _take_close_with_run( +def _take_owned_streams( headers: Mapping[str, temporalio.api.common.v1.Payload], ) -> None: - payload = headers.get(_CLOSE_WITH_RUN_HEADER) + payload = headers.get(_OWNED_STREAMS_HEADER) if payload is None: return state = streams_workflow._state() converter = temporalio.converter.PayloadConverter.default - for stream_id in converter.from_payload(payload, list[str]): - state.stream(_owned_ref(stream_id)).close_with_run = True + for stream_id, backend, close_with_run in converter.from_payload( + payload, list[Any] + ): + run_stream = state.stream(_owned_ref(stream_id)) + run_stream.backend, run_stream.close_with_run = backend, close_with_run class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): @@ -573,16 +564,17 @@ async def start_system_nexus_operation( return _DoneHandle(result) def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: - open_ids = [ - stream_id + # The next run needs each owned Stream's backend for stream_handle. + owned = [ + [stream_id, run_stream.backend, run_stream.close_with_run] for stream_id, run_stream in streams_workflow._state().streams.items() - if run_stream.close_with_run and not run_stream.closed + if run_stream.backend is not None and not run_stream.closed ] - if open_ids: + if owned: converter = temporalio.converter.PayloadConverter.default input.headers = { **input.headers, - _CLOSE_WITH_RUN_HEADER: converter.to_payload(open_ids), + _OWNED_STREAMS_HEADER: converter.to_payload(owned), } super().continue_as_new(input) @@ -603,7 +595,7 @@ class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): def __init__( self, next: temporalio.worker.ActivityInboundInterceptor, - plugin: StreamPlugin, + plugin: _ExtensionStreamPlugin, ) -> None: super().__init__(next) self._plugin = plugin @@ -699,7 +691,7 @@ async def _close_quietly(streams: _activity.ActivityStreams) -> None: def _activity_streams( - input: temporalio.worker.ExecuteActivityInput, plugin: StreamPlugin + input: temporalio.worker.ExecuteActivityInput, plugin: _ExtensionStreamPlugin ) -> _activity.ActivityStreams: # A sync Activity has no client, so the plugin gives the backend and the converter. info = temporalio.activity.info() @@ -738,7 +730,7 @@ def _activity_streams( ) producer.epoch.CopyFrom(epoch) return _activity.ActivityStreams( - backend=plugin._temporal_stream_backend, + backend=plugin._backend, converter=plugin._converter, namespace=info.namespace, producer=producer, diff --git a/temporalio/contrib/streams/system/__init__.py b/temporalio/contrib/streams/system/__init__.py index 5919cad90..7113dd269 100644 --- a/temporalio/contrib/streams/system/__init__.py +++ b/temporalio/contrib/streams/system/__init__.py @@ -10,11 +10,13 @@ import temporalio.converter from . import services as _services -from .operations.append_records import append_records -from .operations.close_stream import close_stream +from ._resources.stream import ( + append_records, + close_stream, + describe_stream, + read_records, +) from .operations.create_stream import create_stream -from .operations.describe_stream import describe_stream -from .operations.read_records import read_records __all__ = [ "append_records", diff --git a/temporalio/contrib/streams/system/_resources/__init__.py b/temporalio/contrib/streams/system/_resources/__init__.py new file mode 100644 index 000000000..66de98b6e --- /dev/null +++ b/temporalio/contrib/streams/system/_resources/__init__.py @@ -0,0 +1,9 @@ +# Generated by nexgen v0.2.9. DO NOT EDIT! + +from __future__ import annotations + +from .stream import Stream + +__all__ = [ + "Stream", +] diff --git a/temporalio/contrib/streams/system/_resources/stream.py b/temporalio/contrib/streams/system/_resources/stream.py new file mode 100644 index 000000000..457699b68 --- /dev/null +++ b/temporalio/contrib/streams/system/_resources/stream.py @@ -0,0 +1,341 @@ +# Generated by nexgen v0.2.9. DO NOT EDIT! + +from __future__ import annotations + +import asyncio +import collections.abc +import dataclasses +import datetime +import typing + +import temporalio.api.streamservice.v1.stream_service_pb2 +import temporalio.workflow + +from ..models import ( + AppendRecord, + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + DescribeStreamRequest, + DescribeStreamResponse, + Producer, + ReadRecordsRequest, + ReadRecordsResponse, + StreamRecord, + StreamRef, + WorkflowTaskEpoch, +) + +# An answer with nothing that comes back in less than half of its wait is +# followed by a pause, so a service that ignores the wait is not asked again +# at once. The pause doubles up to the last one. +_LONG_POLL_FIRST_PAUSE = 0.05 +_LONG_POLL_LAST_PAUSE = 2.0 + + +@dataclasses.dataclass +class Stream: + """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" + + stream: StreamRef + + async def append_records( + self, + producer: Producer | None = None, + sequence: int | None = None, + records: list[AppendRecord] | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + withdraw_run: str | None = None, + ) -> AppendRecordsResponse: + """Append Records. The answer has the first and the last offset.""" + request = AppendRecordsRequest( + stream=self.stream, + producer=producer, + sequence=sequence, + records=records or [], + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + withdraw_run=withdraw_run, + ) + handle = await _append_records(request) + return await handle + + async def close( + self, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + ) -> None: + """Close the Stream. Readers drain what is left, then end.""" + request = CloseStreamRequest( + stream=self.stream, + producer=producer, + result=result, + failure=failure, + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + ) + handle = await _close_stream(request) + await handle + + async def read_records( + self, + from_offset: int | None = None, + start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType + | None = None, + topics: list[str] | None = None, + topic_prefixes: list[str] | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> ReadRecordsResponse: + """Read one page of Records. The iterator of the same name reads them all.""" + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + topics=topics or [], + max_records=max_records, + wait=wait, + create_wait=create_wait, + topic_prefixes=topic_prefixes or [], + ) + handle = await _read_records(request) + return await handle + + async def iter_read_records( + self, + from_offset: int | None = None, + start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType + | None = None, + topics: list[str] | None = None, + topic_prefixes: list[str] | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + *, + deadline: datetime.timedelta | None = None, + ) -> collections.abc.AsyncIterator[StreamRecord]: + """Yield each item of ``records``, page after page, as one sequence. + + Each page is one ``ReadRecords`` operation call, which sets ``wait`` to what is + left of the deadline. The next page starts at the answer's ``next_offset``. + The iterator ends after an answer with ``done`` set, or when the deadline + passes. With no deadline, it ends only on ``done``. + """ + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + topics=topics or [], + max_records=max_records, + wait=wait, + create_wait=create_wait, + topic_prefixes=topic_prefixes or [], + ) + clock = asyncio.get_running_loop() + end = None if deadline is None else clock.time() + deadline.total_seconds() + pause = _LONG_POLL_FIRST_PAUSE + while True: + page_wait = 60.0 if end is None else max(end - clock.time(), 0.001) + page_wait = min(page_wait, 60.0) + request = dataclasses.replace( + request, wait=datetime.timedelta(seconds=page_wait) + ) + asked = clock.time() + handle = await _read_records(request) + answer = await handle + for item in answer.records or []: + yield item + if answer.done: + return + if answer.next_offset is not None: + request = dataclasses.replace(request, from_offset=answer.next_offset) + if end is not None and clock.time() >= end: + return + if not answer.records and clock.time() - asked < page_wait / 2: + await asyncio.sleep(pause) + pause = min(pause * 2, _LONG_POLL_LAST_PAUSE) + else: + pause = _LONG_POLL_FIRST_PAUSE + + async def describe( + self, + ) -> DescribeStreamResponse: + """Describe the Stream: its status, newest offset, Retention and auto-close.""" + request = DescribeStreamRequest( + stream=self.stream, + ) + handle = await _describe_stream(request) + return await handle + + +async def _append_records( + request: AppendRecordsRequest, +) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="AppendRecords", + input=request, + output_type=AppendRecordsResponse, + ) + + +async def append_records( + *, + stream: StreamRef | None = None, + producer: Producer | None = None, + sequence: int | None = None, + records: list[AppendRecord] | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + withdraw_run: str | None = None, +) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: + """In Workflow code, Brook buffers this call in the run, and the facade drops the + answer. + + Args: + producer: Required outside a Workflow. In a Workflow the transport sets it. + sequence: The batch number in this epoch, 1 or more. It is part of the dedupe + key. + previous_completed_task: Workflow only. The previous completed Workflow Task of + the run. + failed_tasks: Workflow only. Failed or timed out tasks since this run's last + flush here. + withdraw_run: Retry run only. The previous run, whose Records this withdraws. + """ + request = AppendRecordsRequest( + stream=stream, + producer=producer, + sequence=sequence, + records=records or [], + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + withdraw_run=withdraw_run, + ) + return await _append_records(request) + + +async def _close_stream( + request: CloseStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="CloseStream", + input=request, + output_type=CloseStreamResponse, + ) + + +async def close_stream( + *, + stream: StreamRef | None = None, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, +) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: + """ + Args: + previous_completed_task: Workflow only. The previous completed Workflow Task of + the run. + failed_tasks: Workflow only. Failed or timed out tasks since this run's last + flush here. + """ + request = CloseStreamRequest( + stream=stream, + producer=producer, + result=result, + failure=failure, + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + ) + return await _close_stream(request) + + +async def _read_records( + request: ReadRecordsRequest, +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="ReadRecords", + input=request, + output_type=ReadRecordsResponse, + ) + + +async def read_records( + *, + stream: StreamRef | None = None, + from_offset: int | None = None, + start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType + | None = None, + topics: list[str] | None = None, + topic_prefixes: list[str] | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: + """Not for Workflow code in Brook. A long poll. NOT_FOUND for an uncreated Stream. + + Args: + from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the beginning. + topics: Topics to read, by exact match. With topic-prefixes, empty reads every + Topic. + topic_prefixes: Also read the Topics that start with one of these. + max_records: At most this many Records. 1 to 1000. + wait: How long to wait for a new Record before an empty answer. At most 60 s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. At + most 60 s. + """ + request = ReadRecordsRequest( + stream=stream, + from_offset=from_offset, + start=start, + topics=topics or [], + topic_prefixes=topic_prefixes or [], + max_records=max_records, + wait=wait, + create_wait=create_wait, + ) + return await _read_records(request) + + +async def _describe_stream( + request: DescribeStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="DescribeStream", + input=request, + output_type=DescribeStreamResponse, + ) + + +async def describe_stream( + *, + stream: StreamRef | None = None, +) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: + request = DescribeStreamRequest( + stream=stream, + ) + return await _describe_stream(request) diff --git a/temporalio/contrib/streams/system/client.py b/temporalio/contrib/streams/system/client.py index 2883b3f1c..1dc9685ff 100644 --- a/temporalio/contrib/streams/system/client.py +++ b/temporalio/contrib/streams/system/client.py @@ -3,6 +3,8 @@ from __future__ import annotations import asyncio +import collections.abc +import contextlib import dataclasses import datetime import typing @@ -11,9 +13,12 @@ import google.protobuf.message import temporalio.api.common.v1 +import temporalio.api.streamservice.v1.stream_service_pb2 import temporalio.converter +import temporalio.nexus.system from .models import ( + AppendRecord, AppendRecordsRequest, AppendRecordsResponse, CloseStreamRequest, @@ -22,8 +27,12 @@ CreateStreamResponse, DescribeStreamRequest, DescribeStreamResponse, + Producer, ReadRecordsRequest, ReadRecordsResponse, + StreamRecord, + StreamRef, + WorkflowTaskEpoch, _AppendRecordsRequestTransferTypeConverter, _AppendRecordsResponseTransferTypeConverter, _CloseStreamRequestTransferTypeConverter, @@ -113,6 +122,50 @@ def _proto_message( return message +# The codec runs on the user payloads inside a proto-backed model, Record bodies and +# failure details included, and never on the model's own payload, which the service reads. +def _payloads_in( + message: google.protobuf.message.Message, +) -> list[temporalio.api.common.v1.Payload]: + if isinstance(message, temporalio.api.common.v1.Payload): + return [message] + found: list[temporalio.api.common.v1.Payload] = [] + for _, value in message.ListFields(): + if isinstance(value, google.protobuf.message.Message): + found.extend(_payloads_in(value)) + elif not isinstance(value, (str, bytes, int, float, bool)): + items = typing.cast( + "collections.abc.Iterable[object]", + value.values() if hasattr(value, "values") else value, + ) + for item in items: + if isinstance(item, google.protobuf.message.Message): + found.extend(_payloads_in(item)) + return found + + +async def _encode_payloads( + message: google.protobuf.message.Message, + codec: temporalio.converter.PayloadCodec | None, +) -> None: + if codec is None: + return + for payload in _payloads_in(message): + [encoded] = await codec.encode([payload]) + payload.CopyFrom(encoded) + + +async def _decode_payloads( + message: google.protobuf.message.Message, + codec: temporalio.converter.PayloadCodec | None, +) -> None: + if codec is None: + return + for payload in _payloads_in(message): + [decoded] = await codec.decode([payload]) + payload.CopyFrom(decoded) + + class StreamServiceClient: """Calls the `temporal.api.streamservice.v1.StreamService` operations over a transport. @@ -123,48 +176,53 @@ def __init__( self, transport: Transport, *, - payload_converter: temporalio.converter.PayloadConverter | None = None, - payload_codec: temporalio.converter.PayloadCodec | None = None, + data_converter: temporalio.converter.DataConverter | None = None, ) -> None: """Create a caller. Args: transport: Carries each call to the service. - payload_converter: Converts the models that are not proto-backed. - Defaults to the SDK's default payload converter. - payload_codec: Encodes each request payload and decodes each answer. + data_converter: Converts and encodes the user values. A proto-backed + model travels as is, and each payload and failure inside it goes + through the converter and the codec, as in Workflow code. A model + that is not proto-backed is one user value. Defaults to the SDK's + default data converter. """ self._transport = transport - self._payload_converter = ( - payload_converter - or temporalio.converter.DataConverter.default.payload_converter + self._data_converter = ( + data_converter or temporalio.converter.DataConverter.default + ) + + def _converters(self) -> contextlib.AbstractContextManager[None]: + return temporalio.nexus.system._user_converter_context( + temporalio.nexus.system._SystemNexusUserConverters( + payload_converter=self._data_converter.payload_converter, + failure_converter=self._data_converter.failure_converter, + ) ) - self._payload_codec = payload_codec async def _call( self, operation: str, payload: temporalio.api.common.v1.Payload ) -> temporalio.api.common.v1.Payload: - if self._payload_codec is not None: - [payload] = await self._payload_codec.encode([payload]) - answer = await self._transport.call( + return await self._transport.call( service="temporal.api.streamservice.v1.StreamService", operation=operation, input=payload, ) - if self._payload_codec is not None: - [answer] = await self._payload_codec.decode([answer]) - return answer async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResponse: - payload = _proto_payload( - _CreateStreamRequestTransferTypeConverter().to_transfer_type(request) - ) + with self._converters(): + message = _CreateStreamRequestTransferTypeConverter().to_transfer_type( + request + ) + await _encode_payloads(message, self._data_converter.payload_codec) + payload = _proto_payload(message) answer = await self._call("CreateStream", payload) converter = _CreateStreamResponseTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - CreateStreamResponse, - ) + message = _proto_message(answer, converter.transfer_type) + await _decode_payloads(message, self._data_converter.payload_codec) + with self._converters(): + return converter.from_transfer_type(message, CreateStreamResponse) async def append_records( self, request: AppendRecordsRequest @@ -172,38 +230,47 @@ async def append_records( """In Workflow code, Brook buffers this call in the run, and the facade drops the answer. """ - payload = _proto_payload( - _AppendRecordsRequestTransferTypeConverter().to_transfer_type(request) - ) + with self._converters(): + message = _AppendRecordsRequestTransferTypeConverter().to_transfer_type( + request + ) + await _encode_payloads(message, self._data_converter.payload_codec) + payload = _proto_payload(message) answer = await self._call("AppendRecords", payload) converter = _AppendRecordsResponseTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - AppendRecordsResponse, - ) + message = _proto_message(answer, converter.transfer_type) + await _decode_payloads(message, self._data_converter.payload_codec) + with self._converters(): + return converter.from_transfer_type(message, AppendRecordsResponse) async def close_stream(self, request: CloseStreamRequest) -> CloseStreamResponse: - payload = _proto_payload( - _CloseStreamRequestTransferTypeConverter().to_transfer_type(request) - ) + with self._converters(): + message = _CloseStreamRequestTransferTypeConverter().to_transfer_type( + request + ) + await _encode_payloads(message, self._data_converter.payload_codec) + payload = _proto_payload(message) answer = await self._call("CloseStream", payload) converter = _CloseStreamResponseTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - CloseStreamResponse, - ) + message = _proto_message(answer, converter.transfer_type) + await _decode_payloads(message, self._data_converter.payload_codec) + with self._converters(): + return converter.from_transfer_type(message, CloseStreamResponse) async def read_records(self, request: ReadRecordsRequest) -> ReadRecordsResponse: """Not for Workflow code in Brook. A long poll. NOT_FOUND for an uncreated Stream.""" - payload = _proto_payload( - _ReadRecordsRequestTransferTypeConverter().to_transfer_type(request) - ) + with self._converters(): + message = _ReadRecordsRequestTransferTypeConverter().to_transfer_type( + request + ) + await _encode_payloads(message, self._data_converter.payload_codec) + payload = _proto_payload(message) answer = await self._call("ReadRecords", payload) converter = _ReadRecordsResponseTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - ReadRecordsResponse, - ) + message = _proto_message(answer, converter.transfer_type) + await _decode_payloads(message, self._data_converter.payload_codec) + with self._converters(): + return converter.from_transfer_type(message, ReadRecordsResponse) async def read_records_until_records( self, @@ -255,15 +322,169 @@ async def read_records_until_records( else: pause = _LONG_POLL_FIRST_PAUSE + async def iter_read_records( + self, + request: ReadRecordsRequest, + *, + deadline: datetime.timedelta | None = None, + ) -> collections.abc.AsyncIterator[StreamRecord]: + """Yield each item of ``records``, page after page, as one sequence. + + Each page is one ``read_records_until_records`` call, which sets ``wait`` to + what is left of the deadline. The next page starts at the answer's + ``next_offset``. + The iterator ends after an answer with ``done`` set, or when the deadline + passes. With no deadline, it ends only on ``done``. + """ + clock = asyncio.get_running_loop() + end = None if deadline is None else clock.time() + deadline.total_seconds() + attempt = request + while True: + left = 60.0 if end is None else end - clock.time() + if left <= 0: + return + answer = await self.read_records_until_records( + attempt, deadline=datetime.timedelta(seconds=left) + ) + for item in answer.records or []: + yield item + if answer.done: + return + if answer.next_offset is not None: + attempt = dataclasses.replace(attempt, from_offset=answer.next_offset) + async def describe_stream( self, request: DescribeStreamRequest ) -> DescribeStreamResponse: - payload = _proto_payload( - _DescribeStreamRequestTransferTypeConverter().to_transfer_type(request) - ) + with self._converters(): + message = _DescribeStreamRequestTransferTypeConverter().to_transfer_type( + request + ) + await _encode_payloads(message, self._data_converter.payload_codec) + payload = _proto_payload(message) answer = await self._call("DescribeStream", payload) converter = _DescribeStreamResponseTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - DescribeStreamResponse, + message = _proto_message(answer, converter.transfer_type) + await _decode_payloads(message, self._data_converter.payload_codec) + with self._converters(): + return converter.from_transfer_type(message, DescribeStreamResponse) + + +class Stream: + """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" + + def __init__( + self, + client: StreamServiceClient, + stream: StreamRef, + ) -> None: + """Bind the handle to a caller.""" + self._client = client + self.stream = stream + + async def append_records( + self, + producer: Producer | None = None, + sequence: int | None = None, + records: list[AppendRecord] | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + withdraw_run: str | None = None, + ) -> AppendRecordsResponse: + """Append Records. The answer has the first and the last offset.""" + request = AppendRecordsRequest( + stream=self.stream, + producer=producer, + sequence=sequence, + records=records or [], + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + withdraw_run=withdraw_run, + ) + return await self._client.append_records(request) + + async def close( + self, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + ) -> None: + """Close the Stream. Readers drain what is left, then end.""" + request = CloseStreamRequest( + stream=self.stream, + producer=producer, + result=result, + failure=failure, + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + ) + await self._client.close_stream(request) + + async def read_records( + self, + from_offset: int | None = None, + start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType + | None = None, + topics: list[str] | None = None, + topic_prefixes: list[str] | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> ReadRecordsResponse: + """Read one page of Records. The iterator of the same name reads them all.""" + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + topics=topics or [], + max_records=max_records, + wait=wait, + create_wait=create_wait, + topic_prefixes=topic_prefixes or [], + ) + return await self._client.read_records(request) + + async def iter_read_records( + self, + from_offset: int | None = None, + start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType + | None = None, + topics: list[str] | None = None, + topic_prefixes: list[str] | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + *, + deadline: datetime.timedelta | None = None, + ) -> collections.abc.AsyncIterator[StreamRecord]: + """Yield each item of ``records``, page after page, as one sequence. + + Each page is one ``read_records_until_records`` call, which sets ``wait`` to + what is left of the deadline. The next page starts at the answer's + ``next_offset``. + The iterator ends after an answer with ``done`` set, or when the deadline + passes. With no deadline, it ends only on ``done``. + """ + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + topics=topics or [], + max_records=max_records, + wait=wait, + create_wait=create_wait, + topic_prefixes=topic_prefixes or [], + ) + async for item in self._client.iter_read_records(request, deadline=deadline): + yield item + + async def describe( + self, + ) -> DescribeStreamResponse: + """Describe the Stream: its status, newest offset, Retention and auto-close.""" + request = DescribeStreamRequest( + stream=self.stream, ) + return await self._client.describe_stream(request) diff --git a/temporalio/contrib/streams/system/models.py b/temporalio/contrib/streams/system/models.py index 2304fc581..8d19543f7 100644 --- a/temporalio/contrib/streams/system/models.py +++ b/temporalio/contrib/streams/system/models.py @@ -13,6 +13,10 @@ import temporalio.converter import temporalio.nexus.system._support as _support +if typing.TYPE_CHECKING: + from ._resources import Stream + + StreamOffset = typing.NewType("StreamOffset", int) """A position in an ordered log, as the service gives it. @@ -1248,3 +1252,12 @@ def to_transfer_type( temporalio.converter.transfer_type_convertible( _DescribeStreamResponseTransferTypeConverter )(DescribeStreamResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CreateStreamResult: + stream: Stream + created: bool | None = None + status: ( + temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType | None + ) = None diff --git a/temporalio/contrib/streams/system/operations/append_records.py b/temporalio/contrib/streams/system/operations/append_records.py deleted file mode 100644 index 43324b0ee..000000000 --- a/temporalio/contrib/streams/system/operations/append_records.py +++ /dev/null @@ -1,63 +0,0 @@ -# Generated by nexgen v0.2.9. DO NOT EDIT! - -from __future__ import annotations - -import temporalio.workflow - -from ..models import ( - AppendRecord, - AppendRecordsRequest, - AppendRecordsResponse, - Producer, - StreamRef, - WorkflowTaskEpoch, -) - - -async def _append_records( - request: AppendRecordsRequest, -) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: - nexus_client = temporalio.workflow.create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="AppendRecords", - input=request, - output_type=AppendRecordsResponse, - ) - - -async def append_records( - *, - stream: StreamRef | None = None, - producer: Producer | None = None, - sequence: int | None = None, - records: list[AppendRecord] | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, - withdraw_run: str | None = None, -) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: - """In Workflow code, Brook buffers this call in the run, and the facade drops the - answer. - - Args: - producer: Required outside a Workflow. In a Workflow the transport sets it. - sequence: The batch number in this epoch, 1 or more. It is part of the dedupe - key. - previous_completed_task: Workflow only. The previous completed Workflow Task of - the run. - failed_tasks: Workflow only. Failed or timed out tasks since this run's last - flush here. - withdraw_run: Retry run only. The previous run, whose Records this withdraws. - """ - request = AppendRecordsRequest( - stream=stream, - producer=producer, - sequence=sequence, - records=records or [], - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - withdraw_run=withdraw_run, - ) - return await _append_records(request) diff --git a/temporalio/contrib/streams/system/operations/close_stream.py b/temporalio/contrib/streams/system/operations/close_stream.py deleted file mode 100644 index 80b8b5d09..000000000 --- a/temporalio/contrib/streams/system/operations/close_stream.py +++ /dev/null @@ -1,56 +0,0 @@ -# Generated by nexgen v0.2.9. DO NOT EDIT! - -from __future__ import annotations - -import typing - -import temporalio.workflow - -from ..models import ( - CloseStreamRequest, - CloseStreamResponse, - Producer, - StreamRef, - WorkflowTaskEpoch, -) - - -async def _close_stream( - request: CloseStreamRequest, -) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: - nexus_client = temporalio.workflow.create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="CloseStream", - input=request, - output_type=CloseStreamResponse, - ) - - -async def close_stream( - *, - stream: StreamRef | None = None, - producer: Producer | None = None, - result: typing.Any | None = None, - failure: BaseException | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, -) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: - """ - Args: - previous_completed_task: Workflow only. The previous completed Workflow Task of - the run. - failed_tasks: Workflow only. Failed or timed out tasks since this run's last - flush here. - """ - request = CloseStreamRequest( - stream=stream, - producer=producer, - result=result, - failure=failure, - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - ) - return await _close_stream(request) diff --git a/temporalio/contrib/streams/system/operations/create_stream.py b/temporalio/contrib/streams/system/operations/create_stream.py index 2135918aa..ad7dd041b 100644 --- a/temporalio/contrib/streams/system/operations/create_stream.py +++ b/temporalio/contrib/streams/system/operations/create_stream.py @@ -6,24 +6,36 @@ import temporalio.workflow +from .._resources import Stream from ..models import ( CreateStreamRequest, CreateStreamResponse, + CreateStreamResult, ) async def _create_stream( request: CreateStreamRequest, -) -> temporalio.workflow.NexusOperationHandle[CreateStreamResponse,]: +) -> CreateStreamResult: nexus_client = temporalio.workflow.create_nexus_client( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - return await nexus_client.start_operation( + handle = await nexus_client.start_operation( operation="CreateStream", input=request, output_type=CreateStreamResponse, ) + result = await handle + if result.stream is None: + raise RuntimeError("CreateStream answered without `stream`") + return CreateStreamResult( + stream=Stream( + stream=result.stream, + ), + created=result.created, + status=result.status, + ) async def create_stream( @@ -32,7 +44,7 @@ async def create_stream( retention: datetime.timedelta | None = None, withdraw_run: str | None = None, auto_close: datetime.timedelta | None = None, -) -> temporalio.workflow.NexusOperationHandle[CreateStreamResponse,]: +) -> CreateStreamResult: """ Args: stream_id: Required from Workflow code. Empty elsewhere asks the service for an diff --git a/temporalio/contrib/streams/system/operations/describe_stream.py b/temporalio/contrib/streams/system/operations/describe_stream.py deleted file mode 100644 index 3f42e3898..000000000 --- a/temporalio/contrib/streams/system/operations/describe_stream.py +++ /dev/null @@ -1,35 +0,0 @@ -# Generated by nexgen v0.2.9. DO NOT EDIT! - -from __future__ import annotations - -import temporalio.workflow - -from ..models import ( - DescribeStreamRequest, - DescribeStreamResponse, - StreamRef, -) - - -async def _describe_stream( - request: DescribeStreamRequest, -) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: - nexus_client = temporalio.workflow.create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="DescribeStream", - input=request, - output_type=DescribeStreamResponse, - ) - - -async def describe_stream( - *, - stream: StreamRef | None = None, -) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: - request = DescribeStreamRequest( - stream=stream, - ) - return await _describe_stream(request) diff --git a/temporalio/contrib/streams/system/operations/read_records.py b/temporalio/contrib/streams/system/operations/read_records.py deleted file mode 100644 index 2bfce568d..000000000 --- a/temporalio/contrib/streams/system/operations/read_records.py +++ /dev/null @@ -1,67 +0,0 @@ -# Generated by nexgen v0.2.9. DO NOT EDIT! - -from __future__ import annotations - -import datetime - -import temporalio.api.streamservice.v1.stream_service_pb2 -import temporalio.workflow - -from ..models import ( - ReadRecordsRequest, - ReadRecordsResponse, - StreamRef, -) - - -async def _read_records( - request: ReadRecordsRequest, -) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: - nexus_client = temporalio.workflow.create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="ReadRecords", - input=request, - output_type=ReadRecordsResponse, - ) - - -async def read_records( - *, - stream: StreamRef | None = None, - from_offset: int | None = None, - start: temporalio.api.streamservice.v1.stream_service_pb2.ReadStart.ValueType - | None = None, - topics: list[str] | None = None, - topic_prefixes: list[str] | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, -) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: - """Not for Workflow code in Brook. A long poll. NOT_FOUND for an uncreated Stream. - - Args: - from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest - offset + 1. - start: Where to start when from-offset is empty. Unspecified is the beginning. - topics: Topics to read, by exact match. With topic-prefixes, empty reads every - Topic. - topic_prefixes: Also read the Topics that start with one of these. - max_records: At most this many Records. 1 to 1000. - wait: How long to wait for a new Record before an empty answer. At most 60 s. - create_wait: How long to wait for an uncreated Stream. Empty fails at once. At - most 60 s. - """ - request = ReadRecordsRequest( - stream=stream, - from_offset=from_offset, - start=start, - topics=topics or [], - topic_prefixes=topic_prefixes or [], - max_records=max_records, - wait=wait, - create_wait=create_wait, - ) - return await _read_records(request) diff --git a/temporalio/contrib/streams/testing.py b/temporalio/contrib/streams/testing.py index c933bc16a..2e25392be 100644 --- a/temporalio/contrib/streams/testing.py +++ b/temporalio/contrib/streams/testing.py @@ -10,12 +10,12 @@ import temporalio.runtime from temporalio.contrib.streams.proto import v1 as ext_pb -from temporalio.contrib.streams.redis import StreamPlugin +from temporalio.contrib.streams.redis import _ExtensionStreamPlugin MEMORY_EXTENSION_NAME = "temporalio-streams-memory" -class MemoryStreamPlugin(StreamPlugin): +class MemoryStreamPlugin(_ExtensionStreamPlugin): """Keeps Streams in this process, with the same rules as on Redis, for tests. Put it on the client, as with ``RedisStreamPlugin``. Every Worker and client that shares one diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 9526c334a..43415b427 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -16,21 +16,16 @@ from dataclasses import dataclass, field from datetime import timedelta from types import TracebackType -from typing import Any, Generic, TypeVar, overload +from typing import Any, TypeVar, overload import temporalio.exceptions import temporalio.workflow from temporalio.contrib.streams import _errors, _ids from temporalio.contrib.streams._defn import TopicDefinition, definition_of -from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams._handle import StreamHandle, TopicBase +from temporalio.contrib.streams.system import create_stream as _create_stream from temporalio.contrib.streams.system import models -from temporalio.contrib.streams.system.operations import ( - append_records as _append_records, -) -from temporalio.contrib.streams.system.operations import close_stream as _close_stream -from temporalio.contrib.streams.system.operations import ( - create_stream as _create_stream, -) +from temporalio.contrib.streams.system._resources.stream import Stream as _Stream E = TypeVar("E") T = TypeVar("T") @@ -51,6 +46,8 @@ class _RunStream: ref: models.StreamRef close_with_run: bool = False closed: bool = False + # Set for a Stream that the run chain created, which continue-as-new passes on. + backend: str | None = None @dataclass @@ -96,11 +93,11 @@ def _check_backend(backend: str | None, action: str) -> str: # and a Worker with the right config retries it. An ApplicationError would fail the run. worker_backend = _backend_name() if worker_backend is None: - raise RuntimeError( + raise _errors.StreamConfigError( "No Stream plugin is on the Worker, such as RedisStreamPlugin" ) if backend and backend != worker_backend: - raise RuntimeError( + raise _errors.StreamConfigError( f"Can't {action} a Stream of the backend {backend}. The Worker has only " f"{worker_backend}." ) @@ -114,32 +111,18 @@ def _now(call: Coroutine[Any, Any, _R]) -> _R: except StopIteration as done: return done.value call.close() - raise RuntimeError("No Stream plugin took the call, such as RedisStreamPlugin") + raise _errors.StreamConfigError( + "No Stream plugin took the call, such as RedisStreamPlugin" + ) -class WorkflowTopic(Generic[T]): +class WorkflowTopic(TopicBase[T]): """One Topic of a :py:class:`WorkflowStream`. All Topics of a Stream share one order. .. warning:: This API is experimental. """ - def __init__(self, stream: WorkflowStream[Any], name: str, type: Any) -> None: - """Create a Topic. Use :py:meth:`WorkflowStream.topic` instead.""" - self._stream = stream - self._name = name - self._type = type - - @property - def name(self) -> str: - """Name of the Topic.""" - return self._name - - @property - def type(self) -> Any: - """Type of the Topic's values.""" - return self._type - async def publish(self, *values: T) -> None: """Add one Record for each value, in order. @@ -150,8 +133,8 @@ async def publish(self, *values: T) -> None: Raises: StreamClosedError: This run closed the Stream. ReadOnlyContextError: The code must not change state, such as a query handler. - RuntimeError: 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. + 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. """ self.publish_nowait(*values) @@ -182,6 +165,11 @@ def _model_ref(self) -> models.StreamRef: first_execution_run_id=self.first_execution_run_id, ) + def _resource(self) -> _Stream: + # The generated Stream resource makes the System Nexus calls. This class adds the + # definition's typed Topics and the run's replay rules. + return _Stream(stream=self._model_ref()) + @overload def topic(self, topic: TopicDefinition[T]) -> WorkflowTopic[T]: ... @@ -230,11 +218,10 @@ def _append(self, topic: str, values: tuple[Any, ...]) -> None: self._check("write to") # The backend's answer has no offsets that replay could give again, so it is dropped. _now( - _append_records.append_records( - stream=self._model_ref(), + self._resource().append_records( records=[ models.AppendRecord(topic=topic, body=value) for value in values - ], + ] ) ) @@ -258,11 +245,7 @@ def close_nowait( """ run_stream = self._check("close") run_stream.closed = True - _now( - _close_stream.close_stream( - stream=self._model_ref(), result=result, failure=failure - ) - ) + _now(self._resource().close(result=result, failure=failure)) async def __aenter__(self) -> WorkflowStream[E]: """Use the Stream in ``async with``, which closes it at the end.""" @@ -329,12 +312,15 @@ async def create_stream( long, so a Stream that nobody closes does not stay open. Defaults to never. close_with_run: Close the Stream when the run chain closes, with the result or the failure of its last run. A run that continues as new passes it on. - backend: The name of the backend. Defaults to the Worker's. + backend: The name of the backend. Defaults to the Worker's. The default is part of + the Workflow's code, since a replay must create on the same backend. To move a + Workflow to another backend, gate it with a patch, for example + ``backend="native" if workflow.patched("streams-native") else "redis"``. Raises: StreamClosedError: This run closed the Stream. ReadOnlyContextError: The code must not change state, such as a query handler. - RuntimeError: No Stream plugin is on the Worker, or it has no backend of that name. + StreamConfigError: No Stream plugin is on the Worker, or it has no backend of that name. It fails the Workflow Task, not the run. ValueError: ``close_with_run=False`` for a Stream that already closes with the run chain. @@ -372,15 +358,16 @@ def create_stream_nowait( "close_with_run=False can't turn that off" ) _now( - _create_stream.create_stream( - stream_id=stream.id, retention=retention, auto_close=auto_close - ) + _create_stream(stream_id=stream.id, retention=retention, auto_close=auto_close) ) run_stream.close_with_run = close_with_run + run_stream.backend = stream.backend return stream -def stream_handle(definition: type[E], *, name: str | None = None) -> WorkflowStream[E]: +def stream_handle( + definition: type[E], *, name: str | None = None, backend: str | None = None +) -> WorkflowStream[E]: """A Stream that this run chain created, by its definition and name. It makes no call. A later run of the chain, such as one after continue-as-new, uses it to write the Stream @@ -389,5 +376,11 @@ def stream_handle(definition: type[E], *, name: str | None = None) -> WorkflowSt Args: definition: The Stream definition class. name: The Stream name. Defaults to the definition's name. + backend: The backend of the Stream. Defaults to the one that the run chain created it + on, and else to the Worker's. """ - return _own_handle(definition, name, _backend_name() or None) + stream = _own_handle(definition, name, backend) + if backend is None: + known = _state().streams.get(stream.id) + stream._backend = (known.backend if known else None) or _backend_name() or None + return stream diff --git a/tests/contrib/streams/conftest.py b/tests/contrib/streams/conftest.py index 7a29e0b64..a2a518f4d 100644 --- a/tests/contrib/streams/conftest.py +++ b/tests/contrib/streams/conftest.py @@ -8,7 +8,8 @@ import pytest from temporalio.client import Client -from temporalio.contrib.streams.redis import RedisStreamPlugin, StreamPlugin +from temporalio.contrib.streams import StreamPlugin +from temporalio.contrib.streams.redis import RedisStreamPlugin from temporalio.contrib.streams.testing import MemoryStreamPlugin # Ports that other local services own. diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py index e0c3f81a0..955d7d7fd 100644 --- a/tests/contrib/streams/examples/resume_from_offset.py +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -1,7 +1,8 @@ """Brook API proposal, "A client resumes a scoped Stream from an offset". The Workflow writes one Stream per turn with a ``name=`` override. The Reader saves the offset -after each Record itself, stops, and later resumes from it with ``from_offset``. +after each Record itself, stops, and later resumes from it with ``from_offset``. A view that +wants only new Records starts at ``streams.END``. """ from __future__ import annotations @@ -74,14 +75,29 @@ async def read_some( return seen -async def main(client: Client, task_queue: str) -> tuple[list[str], list[str]]: +async def read_new(client: Client, handle: WorkflowHandle[Turns, None]) -> list[str]: + # A view that shows only what comes next starts after the newest Record. + turn3 = streams.StreamClient(client).workflow_stream(handle, Events, name="turn-3") + return [ + item.data + async for item in turn3.topic(Events.tokens).subscribe( + from_offset=streams.END, wait_new_records=False + ) + if isinstance(item, streams.Record) + ] + + +async def main( + client: Client, task_queue: str +) -> tuple[list[str], list[str], list[str]]: async with Worker(client, task_queue=task_queue, workflows=[Turns]): handle = await client.start_workflow( Turns.run, ["a", "b"], id=f"turns-{uuid.uuid4()}", task_queue=task_queue ) store = OffsetStore() before = await read_some(client, handle, store, 2) + new = await read_new(client, handle) await handle.signal(Turns.add, ["c", "d"]) after = await read_some(client, handle, store, 2) await handle.result() - return before, after + return before, new, after diff --git a/tests/contrib/streams/examples/test_examples.py b/tests/contrib/streams/examples/test_examples.py index ec4939196..858fe706c 100644 --- a/tests/contrib/streams/examples/test_examples.py +++ b/tests/contrib/streams/examples/test_examples.py @@ -33,6 +33,7 @@ async def test_ui_backend_reads(stream_client: Client): async def test_resume_from_offset(stream_client: Client): assert await resume_from_offset.main(stream_client, task_queue()) == ( ["a", "b"], + [], ["c", "d"], ) diff --git a/tests/contrib/streams/examples/turn_scoped_stream.py b/tests/contrib/streams/examples/turn_scoped_stream.py index d615bcc63..e1640e46c 100644 --- a/tests/contrib/streams/examples/turn_scoped_stream.py +++ b/tests/contrib/streams/examples/turn_scoped_stream.py @@ -104,8 +104,8 @@ async def main(client: Client, task_queue: str) -> dict[int, list[str]]: async for item in index.subscribe(): if not isinstance(item, streams.Record): continue - # A Stream from a payload has no client, so the reader gets it again by id. - turn = stream_client.get_handle(Events, id=item.data.stream.id) + # A Stream from a payload has no client, so the reader binds it to this one. + turn = stream_client.get_handle(Events, stream=item.data.stream) turns[item.data.n] = [ token.data async for token in turn.subscribe() diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index 8328dac03..fd34e21f5 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -186,8 +186,8 @@ async def test_retry_withdraws_a_received_stream(stream_client: Client): async def wrong_backend(stream: streams.ActivityStream[Events]) -> str: try: await stream.topic(Events.tokens).publish("x") - except streams.StreamUnsupportedError as err: - return err.reason + except streams.StreamConfigError: + return "config" return "wrote" @@ -222,7 +222,7 @@ async def test_activity_closes_and_checks_the_backend(stream_client: Client): handle = await stream_client.start_workflow( Closing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - assert await handle.result() == "stream_unsupported" + assert await handle.result() == "config" items = await read(stream_client, handle) assert [r.data for r in records(items)] == ["last"] assert isinstance(items[-1], streams.Closed) diff --git a/tests/contrib/streams/test_bindings.py b/tests/contrib/streams/test_bindings.py index 0bfc603b4..2780c8123 100644 --- a/tests/contrib/streams/test_bindings.py +++ b/tests/contrib/streams/test_bindings.py @@ -1,10 +1,7 @@ from datetime import timedelta -import pytest - import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb -import temporalio.contrib.streams.proto.v1 as extension_pb import temporalio.converter import temporalio.nexus.system as nexus_system from temporalio.contrib.streams.system import client, models @@ -98,22 +95,3 @@ def test_models_round_trip_through_the_proto(): assert message.records[0].body.metadata["encoding"] == b"json/plain" assert back.stream == request.stream assert back.records and back.records[0].body == {"n": 1} - - -def test_extension_protos_reference_the_service_protos(): - call = extension_pb.WorkflowCall( - append=stream_pb.AppendRecordsRequest(stream=stream_pb.StreamRef(id="wf:run")), - when_run_closes=False, - ) - request = extension_pb.BufferWorkflowCallsRequest(run_id="run", calls=[call]) - parsed = extension_pb.BufferWorkflowCallsRequest.FromString( - request.SerializeToString() - ) - assert parsed.calls[0].append.stream.id == "wf:run" - - -async def test_append_stub_runs_only_in_a_workflow(): - from temporalio.contrib.streams.system import append_records - - with pytest.raises(Exception): - await append_records(stream=models.StreamRef(id="wf:run")) diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index b3c178e1c..2b11a5c04 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -2,18 +2,20 @@ from dataclasses import dataclass from datetime import timedelta +import google.protobuf.any_pb2 import pytest +import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb from temporalio.client import Client from temporalio.common import RawValue from temporalio.contrib import streams -from temporalio.contrib.streams import _errors, _ids +from temporalio.contrib.streams import _errors +from temporalio.contrib.streams._backend import _ExtensionStreamBackend from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.system import models from temporalio.exceptions import ApplicationError -from temporalio.plugin import SimplePlugin -from temporalio.service import RPCStatusCode +from temporalio.service import RPCError, RPCStatusCode from temporalio.worker import ExtensionError from tests.contrib.streams.fake_extension import FakeExtension @@ -26,23 +28,17 @@ class Turn: def fake_client(client: Client) -> tuple[Client, FakeExtension]: extension = FakeExtension() - plugin = SimplePlugin("fake-streams") - setattr(plugin, "_temporal_stream_backend", streams.StreamBackend(extension)) + plugin = streams.StreamPlugin( + "fake-streams", + backend_name="fake", + default=False, + stream_backend=_ExtensionStreamBackend(extension), + ) config = client.config() config["plugins"] = [plugin] return Client(**config), extension -def test_ids_match_the_crate(): - assert ( - _ids.scoped_stream_id("a:b/c", "r%1", "é name") - == "a%3Ab%2Fc:r%251:%C3%A9%20name" - ) - assert _ids.workflow_producer_id("run:1") == "wf:run%3A1" - assert _ids.activity_producer_id("run:1", "a/1", "s-1") == "act:run%3A1:a%2F1:s-1" - assert _ids.client_producer_id("1@host", "p-1") == "client:1%40host:p-1" - - def test_extension_errors_map_to_stream_errors(): details = ext_pb.StreamErrorDetails(reason="stream_fenced").SerializeToString() err = _errors.from_extension_error( @@ -162,11 +158,6 @@ async def test_each_client_handle_is_its_own_producer(client: Client): assert len(producers) == 2 -def test_workflow_stream_has_no_reads(): - assert not hasattr(streams.WorkflowStream, "subscribe") - assert not hasattr(streams.WorkflowStream, "get_offset") - - def test_unsupported_capability_maps_to_its_error(): details = ext_pb.StreamErrorDetails(reason="stream_unsupported").SerializeToString() err = _errors.from_extension_error( @@ -177,3 +168,19 @@ def test_unsupported_capability_maps_to_its_error(): assert isinstance(err, streams.StreamUnsupportedError) assert err.non_retryable and err.code == RPCStatusCode.UNIMPLEMENTED assert err.message.endswith("replica_wait") + + +def test_server_failure_maps_to_its_stream_error(): + detail = google.protobuf.any_pb2.Any() + detail.Pack(ext_pb.StreamErrorDetails(reason="stream_closed")) + status = temporalio.api.common.v1.GrpcStatus( + code=9, message="closed", details=[detail] + ) + err = _errors.from_rpc_error( + RPCError( + "closed", RPCStatusCode.FAILED_PRECONDITION, status.SerializeToString() + ) + ) + assert isinstance(err, streams.StreamClosedError) + plain = RPCError("down", RPCStatusCode.UNAVAILABLE, b"") + assert _errors.from_rpc_error(plain) is plain diff --git a/tests/contrib/streams/test_handles.py b/tests/contrib/streams/test_handles.py index 303526fe2..2d8d6267a 100644 --- a/tests/contrib/streams/test_handles.py +++ b/tests/contrib/streams/test_handles.py @@ -44,15 +44,6 @@ class PlainStream(streams.StreamHandle[E]): ) -def test_definition_takes_the_name_and_the_topics(): - events = definition_of(Events) - assert events.name == "events" - assert list(events.topics) == ["turns", "tokens", "raw-values"] - assert Events.turns.name == "turns" and Events.turns.type is Turn - assert Events.raw.name == "raw-values" - assert definition_of(Index).name == "Index" - - def test_definition_rejects_two_topics_of_one_name(): with pytest.raises(ValueError, match="two Topics named t"): @@ -93,12 +84,6 @@ def test_topic_of_the_definition_keeps_its_type(): OWNED._topic_spec(Events.turns, int) -def test_stream_without_definition_takes_any_topic(): - plain: PlainStream[Any] = PlainStream(id="s") - assert plain._topic_spec("t", None) == ("t", None) - assert plain._topic_spec(Events.turns, None) == ("turns", Turn) - - def test_handle_is_the_proto_json_of_the_ref(): converter = DataConverter.default.payload_converter payload = converter.to_payload(OWNED) @@ -160,19 +145,6 @@ def test_handle_in_a_pydantic_field(): assert decoded.stream == OWNED -class OtherStream(streams.StreamHandle[E]): - pass - - -def test_from_handle_keeps_the_stream(): - other: OtherStream[Events] = OtherStream.from_handle(OWNED) - assert type(other) is OtherStream - assert other.definition is Events - assert DataConverter.default.payload_converter.to_payload( - other - ) == DataConverter.default.payload_converter.to_payload(OWNED) - - def test_any_type_argument_has_no_definition(): converter = DataConverter.default.payload_converter decoded = converter.from_payload(converter.to_payload(OWNED), PlainStream[Any]) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index d132fe86e..12ae75fa7 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -15,11 +15,12 @@ from temporalio.client import Client, WorkflowFailureError, WorkflowHandle from temporalio.common import RawValue, RetryPolicy from temporalio.contrib import streams +from temporalio.contrib.streams import StreamPlugin from temporalio.contrib.streams import workflow as stream_workflow from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.redis import ( RedisStreamPlugin, - StreamPlugin, + _ExtensionStreamPlugin, _WorkflowInbound, ) from temporalio.contrib.streams.testing import MemoryStreamPlugin @@ -400,6 +401,23 @@ async def test_records_go_through_the_payload_codec(stream_client: Client): assert first.data.payload.metadata["encoding"] == b"json/plain" +async def test_client_close_goes_through_the_payload_codec(stream_client: Client): + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=_Rot13Codec() + ) + stream_client_ = streams.StreamClient(Client(**config)) + done = await stream_client_.create(Events, id=f"done-{uuid.uuid4()}") + await done.close(result=Turn("done")) + failed = await stream_client_.create(Events, id=f"failed-{uuid.uuid4()}") + await failed.close(failure=ApplicationError("bad turn")) + [closed] = [item async for item in done.subscribe()] + assert isinstance(closed, streams.Closed) + assert closed.result == {"kind": "done", "text": ""} + [closed] = [item async for item in failed.subscribe()] + assert isinstance(closed, streams.Closed) and str(closed.failure) == "bad turn" + + @workflow.defn class WaitsForSignal: def __init__(self) -> None: @@ -492,7 +510,9 @@ async def test_worker_shutdown_drains_closes(stream_client: Client): async def test_plugin_shutdown_after_client_writes(stream_client: Client): [plugin] = [ - p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) + p + for p in stream_client.config()["plugins"] + if isinstance(p, _ExtensionStreamPlugin) ] stream = await streams.StreamClient(stream_client).create( Events, id=f"shutdown-{uuid.uuid4()}" @@ -570,17 +590,19 @@ async def test_eviction_prints_nothing( async def test_workflow_calls_carry_the_owner(stream_client: Client): [plugin] = [ - p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) + p + for p in stream_client.config()["plugins"] + if isinstance(p, _ExtensionStreamPlugin) ] sent: list[bytes] = [] - call = plugin._temporal_stream_backend.call + call = plugin._backend.call def record(method: str, request: Any) -> bytes: if method == "BufferWorkflowCalls": sent.append(request.SerializeToString()) return call(method, request) - plugin._temporal_stream_backend.call = record # type: ignore[method-assign] + plugin._backend.call = record # type: ignore[method-assign] async with new_worker(stream_client, QuickWriter) as worker: handle = await stream_client.start_workflow( QuickWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 6cffb7ec0..9c499497e 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -11,9 +11,9 @@ from temporalio.client import Client from temporalio.contrib import streams from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams._backend import _ExtensionStreamBackend from temporalio.contrib.streams.testing import MemoryStreamPlugin from temporalio.converter import DataConverter -from temporalio.plugin import SimplePlugin from tests.contrib.streams.fake_extension import FakeExtension from tests.helpers import new_worker @@ -32,9 +32,12 @@ class Events: def fake_client(client: Client) -> tuple[Client, FakeExtension]: extension = FakeExtension() - plugin = SimplePlugin("fake-streams") - setattr(plugin, "_temporal_stream_backend", streams.StreamBackend(extension)) - setattr(plugin, "backend_name", "fake") + plugin = streams.StreamPlugin( + "fake-streams", + backend_name="fake", + default=False, + stream_backend=_ExtensionStreamBackend(extension), + ) config = client.config() config["plugins"] = [plugin] return Client(**config), extension @@ -43,9 +46,9 @@ def fake_client(client: Client) -> tuple[Client, FakeExtension]: async def test_unknown_backend_raises_before_any_call(client: Client): client, extension = fake_client(client) stream_client = streams.StreamClient(client) - with pytest.raises(streams.StreamUnsupportedError, match="It has fake"): + with pytest.raises(streams.StreamConfigError, match="It has fake"): await stream_client.create(Events, id="chat-1", backend="redis") - with pytest.raises(streams.StreamUnsupportedError): + with pytest.raises(streams.StreamConfigError): stream_client.get_handle(Events, id="chat-1", backend="redis") assert not extension.streams and not extension.namespaces stream = await stream_client.create(Events, id="chat-1") @@ -53,7 +56,7 @@ async def test_unknown_backend_raises_before_any_call(client: Client): async def test_client_without_plugin_has_no_transport(client: Client): - with pytest.raises(streams.StreamTransportMissingError): + with pytest.raises(streams.StreamConfigError): streams.StreamClient(client).get_handle(Events, id="chat-1") @@ -65,7 +68,7 @@ async def test_decoded_client_stream_has_no_client(client: Client): converter.to_payload(stream), streams.ClientStream[Events] ) assert decoded == stream - with pytest.raises(streams.StreamTransportMissingError, match="get_handle"): + with pytest.raises(streams.StreamConfigError, match="get_handle"): await decoded.get_offset() @@ -200,7 +203,6 @@ async def test_unresolved_workflow_stream_compares_and_prints(client: Client): other = stream_client.workflow_stream(handle, Events, name="other") key = hash(first) assert first == same and first != other - assert "agent-2" in repr(first) and "events" in repr(first) describes = [] async def describe() -> Any: @@ -230,6 +232,18 @@ async def test_default_plugin_is_the_default_backend(client: Client): assert stream.backend == "b" -def test_workflow_stream_has_no_get_prefix(): - assert hasattr(streams.StreamClient, "workflow_stream") - assert not hasattr(streams.StreamClient, "get_workflow_stream") +async def test_handle_looks_in_each_backend_during_a_move(client: Client): + config = client.config() + config["plugins"] = [MemoryStreamPlugin(name="a"), MemoryStreamPlugin(name="b")] + stream_client = streams.StreamClient(Client(**config)) + created = await stream_client.create(Events, id="moving", backend="b") + await created.topic(Events.tokens).publish("on b") + moving = stream_client.get_handle(Events, id="moving", backend=("a", "b")) + assert await moving.get_offset() > 0 + assert moving.backend == "b" + nowhere = stream_client.get_handle(Events, id="nowhere", backend=("a", "b")) + with pytest.raises(streams.StreamNotFoundError): + await nowhere.get_offset() + # A handle with no backend name is ambiguous when the client has two backends. + with pytest.raises(streams.StreamConfigError): + stream_client.get_handle(Events, stream=streams.ClientStream(id="moving")) diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index 62aadd51d..49323dfee 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -15,8 +15,6 @@ ) from temporalio.contrib import streams from temporalio.contrib.streams import workflow as stream_workflow -from temporalio.contrib.streams.redis import RedisStreamPlugin -from temporalio.contrib.streams.testing import MemoryStreamPlugin from temporalio.exceptions import ApplicationError from tests.helpers import assert_task_fail_eventually, new_worker @@ -284,15 +282,6 @@ async def test_cron_runs_share_an_open_stream(stream_client: Client): assert not any(isinstance(item, streams.Closed) for item in items) -def test_plugins_take_their_backend_name(): - redis = RedisStreamPlugin( - "redis://127.0.0.1:1", response_timeout=timedelta(seconds=1), default=True - ) - assert (redis.backend_name, redis.default) == ("redis", True) - assert RedisStreamPlugin("redis://127.0.0.1:1", name="old").backend_name == "old" - assert MemoryStreamPlugin().backend_name == "memory" - - @workflow.defn class SyncHandlers: @workflow.init