From c0c0a8603c0551cf6d142f2712194b7c15c10067 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 14:51:13 -0700 Subject: [PATCH] Added Stream backends and `stream_backends=` on the client and Replayer. A backend registers a name that Stream handles store, so any process can route a handle that it receives. A Worker uses its client's backends. --- temporalio/client/_client.py | 19 ++ temporalio/common/__init__.py | 3 + temporalio/common/_streams/__init__.py | 2 + temporalio/common/_streams/_backend.py | 347 +++++++++++++++++++++++++ temporalio/common/_streams/_plugin.py | 221 ++++++++++++++++ temporalio/testing/_workflow.py | 16 ++ temporalio/worker/_replayer.py | 15 +- temporalio/worker/_worker.py | 7 + tests/contrib/streams/test_backends.py | 96 +++++++ 9 files changed, 725 insertions(+), 1 deletion(-) create mode 100644 temporalio/common/_streams/_backend.py create mode 100644 temporalio/common/_streams/_plugin.py create mode 100644 tests/contrib/streams/test_backends.py diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index efef5cb528..f171bfa5c6 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -110,6 +110,8 @@ ) if TYPE_CHECKING: + from temporalio.common import StreamBackend + from ._interceptor import Interceptor from ._plugin import Plugin @@ -156,6 +158,7 @@ async def connect( grpc_compression: GrpcCompression = GrpcCompression.GZIP, payload_limits: PayloadLimitsConfig = PayloadLimitsConfig(), header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ) -> Self: """Connect to a Temporal server. @@ -225,6 +228,12 @@ async def connect( payload_limits: Warning thresholds for outbound payload/memo sizes. Over-threshold fields are logged but still sent. Set a threshold to 0 to disable it. header_codec_behavior: Encoding behavior for headers sent by the client. + stream_backends: The Stream backends of this client and its Workers. Names must be + unique. With more than one backend, exactly one must have ``default=True``. See + :py:class:`temporalio.common.StreamBackend`. + + .. warning:: + This API is experimental. """ connect_config = temporalio.service.ConnectConfig( target_host=target_host, @@ -262,6 +271,7 @@ def make_lambda( default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, plugins=plugins, + stream_backends=stream_backends, ) def __init__( @@ -275,6 +285,7 @@ def __init__( default_workflow_query_reject_condition: None | (temporalio.common.QueryRejectCondition) = None, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ): """Create a Temporal client from a service client. @@ -289,11 +300,17 @@ def __init__( interceptors=interceptors, default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, + stream_backends=list(stream_backends), ) self._initial_config = config.copy() for plugin in plugins: config = plugin.configure_client(config) + if config.get("stream_backends"): + # Only a client with Stream backends loads the Streams package. + from temporalio.common._streams import _plugin as _stream_plugin + + config = _stream_plugin.configure_client(config) self._init_from_config(config) @@ -3029,6 +3046,7 @@ class ClientConnectConfig(TypedDict, total=False): grpc_compression: GrpcCompression payload_limits: PayloadLimitsConfig header_codec_behavior: HeaderCodecBehavior + stream_backends: Sequence[StreamBackend] class ClientConfig(TypedDict, total=False): @@ -3043,3 +3061,4 @@ class ClientConfig(TypedDict, total=False): temporalio.common.QueryRejectCondition | None ] header_codec_behavior: Required[HeaderCodecBehavior] + stream_backends: Required[Sequence[StreamBackend]] diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index d4ae79334f..c1104fce3d 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -62,6 +62,7 @@ Record, Superseded, ) + from temporalio.common._streams._plugin import StreamBackend __all__ = [ "ActivityIDConflictPolicy", @@ -102,6 +103,7 @@ "SearchAttributeValueType", "SearchAttributeValues", "SearchAttributes", + "StreamBackend", "StreamClosedError", "StreamConfigError", "StreamDefinition", @@ -1589,6 +1591,7 @@ def wait_sync(self, timeout: float | None = None) -> None: "END", "End", "Record", + "StreamBackend", "StreamClosedError", "StreamConfigError", "StreamDefinition", diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 517ce48883..299ee3cedb 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -24,6 +24,7 @@ Record, Superseded, ) +from temporalio.common._streams._plugin import StreamBackend __all__ = [ "END", @@ -35,6 +36,7 @@ "StreamClosedError", "StreamConfigError", "StreamDefinition", + "StreamBackend", "StreamError", "StreamExpiredError", "StreamFencedError", diff --git a/temporalio/common/_streams/_backend.py b/temporalio/common/_streams/_backend.py new file mode 100644 index 0000000000..2e98eae1df --- /dev/null +++ b/temporalio/common/_streams/_backend.py @@ -0,0 +1,347 @@ +"""Calls to the temporalio-streams runtime in the bridge, for Activities and clients.""" + +from __future__ import annotations + +import asyncio +import contextlib +import contextvars +import functools +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast + +import google.protobuf.message +from typing_extensions import override + +import temporalio.api.common.v1 +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.bridge.streams +import temporalio.converter +import temporalio.nexus.system +import temporalio.service +from temporalio.common._streams import _errors +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.common._streams.system import client as service_client +from temporalio.converter._payload_converter import _get_transfer_type_converter + +if TYPE_CHECKING: + # The writer bindings import this package's protos, so _writers_class() imports them on first + # use. + from temporalio.common._streams.writer_system import client as writer_client + +_M = TypeVar("_M", bound=google.protobuf.message.Message) + +_RESPONSES: dict[str, type[google.protobuf.message.Message]] = { + "CreateStream": stream_pb.CreateStreamResponse, + "AppendRecords": stream_pb.AppendRecordsResponse, + "CloseStream": stream_pb.CloseStreamResponse, + "ReadRecords": stream_pb.ReadRecordsResponse, + "DescribeStream": stream_pb.DescribeStreamResponse, + "OpenWriter": ext_pb.OpenWriterResponse, + "WriteRecords": ext_pb.WriteRecordsResponse, + "CloseWriter": ext_pb.CloseWriterResponse, + "FreeWriter": ext_pb.FreeWriterResponse, +} + + +@dataclass(frozen=True) +class CallOptions: + """The ``rpc_metadata`` and ``rpc_timeout`` of a client's Stream call.""" + + metadata: Mapping[str, str | bytes] + timeout: timedelta | None + + +# The options of the client call in progress. The runtime's Nexus client reads them, so the +# generated clients and the writer don't need a new argument for them. +call_options: contextvars.ContextVar[CallOptions | None] = contextvars.ContextVar( + "__temporal_stream_call_options", default=None +) + + +def options_of( + rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None +) -> CallOptions | None: + """The call options, or None for a call with neither.""" + if not rpc_metadata and rpc_timeout is None: + return None + return CallOptions(metadata=rpc_metadata, timeout=rpc_timeout) + + +@contextlib.contextmanager +def using(options: CallOptions | None) -> Iterator[None]: + """Apply these options to the calls in the block.""" + if options is None: + yield + return + token = call_options.set(options) + try: + yield + finally: + call_options.reset(token) + + +def longest_wait(longest: timedelta) -> timedelta: + """The longest wait of one call, at most ``longest`` and below the call's rpc_timeout.""" + options = call_options.get() + if options is None or options.timeout is None: + return longest + # A long poll returns before the call's deadline, as the server's long polls do, so a + # waiting read starts its next poll and doesn't fail. + timeout = options.timeout + return min(longest, max(timeout - timedelta(seconds=1), timeout / 2)) + + +class BackendCalls(Protocol): + """The calls that a Stream backend gives the handles of one client.""" + + def service( + self, + namespace: str, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> service_client.StreamServiceClientExtension: + """The Stream service for ``namespace``.""" + ... + + def writers( + self, namespace: str, data_converter: temporalio.converter.DataConverter + ) -> writer_client.StreamWriterClientExtension: + """The writers of Activity and client Producers for ``namespace``.""" + ... + + def free_writer(self, writer_id: int) -> None: + """Release one handle of a writer without waiting, such as when Python frees it.""" + ... + + +class StreamRuntime(Protocol): + """The calls of the crate's runtime object, which the bridge builds.""" + + async def call(self, method: str, input: bytes) -> bytes: + """Call a method with encoded input, and return the encoded output.""" + ... + + def free_writer(self, writer_id: int) -> None: + """Release one handle of a writer without waiting.""" + ... + + +class _ExtensionCalls: + """Makes Stream calls to the crate's runtime object, which serves one client. + + The generated clients start the Stream service and writer operations. Their Nexus client is + the runtime. + """ + + def __init__(self, runtime: StreamRuntime) -> None: + """Create a backend over a runtime object.""" + self._runtime = runtime + + def service( + self, + namespace: str, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> service_client.StreamServiceClientExtension: + """The generated Stream service client for ``namespace``. + + With a ``data_converter``, the client encodes user values in the requests with it and + its codec. + """ + return _StreamService(_RuntimeNexusClient(self, namespace, data_converter)) + + def writers( + self, namespace: str, data_converter: temporalio.converter.DataConverter + ) -> writer_client.StreamWriterClientExtension: + """The generated writer client over the runtime.""" + return _writers_class()(_RuntimeNexusClient(self, namespace, data_converter)) + + def free_writer(self, writer_id: int) -> None: + """Release one handle of a writer without waiting.""" + try: + self._runtime.free_writer(writer_id) + except temporalio.bridge.streams.StreamCallError: + # A writer that the runtime already dropped, such as at shutdown, has nothing to + # release. + pass + + async def call_async( + self, method: str, request: google.protobuf.message.Message, response: type[_M] + ) -> _M: + """Call a method of the runtime.""" + try: + data = await self._runtime.call(method, request.SerializeToString()) + except temporalio.bridge.streams.StreamCallError as err: + raise _errors.from_call_error(err) from None + return response.FromString(data) + + async def close_writers(self) -> int: + """Close every open writer of the runtime, and return how many failed.""" + response = await self.call_async( + "CloseWriters", ext_pb.CloseWritersRequest(), ext_pb.CloseWritersResponse + ) + return response.failed + + +class _StreamService(service_client.StreamServiceClientExtension): + def __init__(self, nexus_client: _RuntimeNexusClient) -> None: + self._nexus_client = nexus_client + + @override + def _create_nexus_client( + self, *, service: str, endpoint: str | None + ) -> _RuntimeNexusClient: + del service, endpoint + return self._nexus_client + + +@functools.cache +def _writers_class() -> Callable[ + [_RuntimeNexusClient], writer_client.StreamWriterClientExtension +]: + from temporalio.common._streams.writer_system import client + + class _StreamWriters(client.StreamWriterClientExtension): + def __init__(self, nexus_client: _RuntimeNexusClient) -> None: + self._nexus_client = nexus_client + + @override + def _create_nexus_client( + self, *, service: str, endpoint: str | None + ) -> _RuntimeNexusClient: + del service, endpoint + return self._nexus_client + + return _StreamWriters + + +class _RuntimeNexusClient: + # The generated clients start each operation on a Nexus client. This one starts it on the + # runtime in place of an endpoint, and the runtime answers before the start returns. + def __init__( + self, + backend: _ExtensionCalls, + namespace: str, + data_converter: temporalio.converter.DataConverter | None, + ) -> None: + self._backend = backend + self._namespace = namespace + self._data_converter = ( + data_converter or temporalio.converter.DataConverter.default + ) + + def _converters(self) -> contextlib.AbstractContextManager[None]: + # A model's user values and failures convert with the client's converters, as in + # Workflow code. + 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, + ) + ) + + async def start_operation( + self, + operation: str, + input: Any, + *, + output_type: type[Any] | None = None, + ) -> Awaitable[Any]: + with self._converters(): + request = cast( + google.protobuf.message.Message, + _transfer_type_converter(type(input)).to_transfer_type(input), + ) + # The IDL leaves the namespace out, since the server knows it. + if "namespace" in request.DESCRIPTOR.fields_by_name: + setattr(request, "namespace", self._namespace) + codec = self._data_converter.payload_codec + await _encode_payloads(request, codec) + response = await self._call(operation, request) + await _decode_payloads(response, codec) + result: Any = None + if output_type is not None: + with self._converters(): + result = _transfer_type_converter(output_type).from_transfer_type( + response, output_type + ) + done: asyncio.Future[Any] = asyncio.get_running_loop().create_future() + done.set_result(result) + return done + + async def _call( + self, operation: str, request: google.protobuf.message.Message + ) -> google.protobuf.message.Message: + # The runtime makes no RPC, so it has no use for rpc_metadata. The rpc_timeout is the + # call's deadline here too. + options = call_options.get() + timeout = options.timeout if options is not None else None + call = self._backend.call_async(operation, request, _RESPONSES[operation]) + if timeout is None: + return await call + try: + return await asyncio.wait_for(call, timeout.total_seconds()) + except asyncio.TimeoutError: + raise temporalio.service.RPCError( + f"The Stream call {operation} passed its rpc_timeout of {timeout}", + temporalio.service.RPCStatusCode.DEADLINE_EXCEEDED, + b"", + ) from None + + +def _transfer_type_converter( + model: type[Any], +) -> temporalio.converter.TransferTypeConverter[Any, Any]: + converter = _get_transfer_type_converter(model) + if converter is None: + raise TypeError(f"{model.__name__} has no proto transfer type") + return converter + + +# The codec runs on the user payloads inside a request or a response, failure details +# included, and never on the message itself, which the runtime 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 = cast( + "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 + # One codec call for each message, since a remote codec costs a round trip for each call. + payloads = _payloads_in(message) + if payloads: + encoded = await codec.encode(payloads) + for payload, new in zip(payloads, encoded, strict=True): + payload.CopyFrom(new) + + +async def _decode_payloads( + message: google.protobuf.message.Message, + codec: temporalio.converter.PayloadCodec | None, +) -> None: + if codec is None: + return + payloads = _payloads_in(message) + if payloads: + decoded = await codec.decode(payloads) + for payload, new in zip(payloads, decoded, strict=True): + payload.CopyFrom(new) diff --git a/temporalio/common/_streams/_plugin.py b/temporalio/common/_streams/_plugin.py new file mode 100644 index 0000000000..1763ea9b81 --- /dev/null +++ b/temporalio/common/_streams/_plugin.py @@ -0,0 +1,221 @@ +"""The base of the Stream backends, and the rules of their registration.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +import temporalio.client +import temporalio.service +import temporalio.worker +from temporalio.common._streams import _errors +from temporalio.common._streams._backend import BackendCalls +from temporalio.plugin import SimplePlugin + +# The default backend when no registered backend has default=True. This SDK has no native +# backend, so a call that needs the default then raises StreamConfigError. +NATIVE = "native" + + +class StreamBackend(SimplePlugin): + """Base class of the Stream backends. Each instance registers one backend under its name. + + Register backends on the client with ``Client.connect(..., stream_backends=[...])``. A + Worker uses its client's backends. A ``Replayer`` takes ``stream_backends=`` too. Names + must be unique. With more than one registered backend, exactly one must have + ``default=True``. With no backend marked, the default is ``native``. This SDK has no + ``native`` backend, so a create with no backend raises ``StreamConfigError``. Mark a + backend with ``default=True``, or name one with ``backend=``. + + A backend is not a plugin, so passing it in ``plugins=`` raises ``StreamConfigError``. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + *, + backend_name: str, + default: bool, + **options: Any, + ) -> None: + """Create the backend. ``options`` go to :py:class:`temporalio.plugin.SimplePlugin`.""" + self._backend_name = backend_name + self._default = default + super().__init__(name, **options) + + @property + def backend_name(self) -> str: + """The backend name that Stream handles store.""" + return self._backend_name + + @property + def default(self) -> bool: + """Whether a create with no backend name uses this backend.""" + return self._default + + def _backend_for( + self, + client: temporalio.client.Client, # pyright: ignore[reportUnusedParameter] + ) -> BackendCalls: + # Each backend subclass returns its calls for the client and its Workers. + raise NotImplementedError(f"{type(self).__name__} gives no calls") + + async def connect_service_client( + self, + config: temporalio.service.ConnectConfig, + next: Callable[ + [temporalio.service.ConnectConfig], + Awaitable[temporalio.service.ServiceClient], + ], + ) -> temporalio.service.ServiceClient: + """Refuse a backend in ``plugins=``, before the client connects.""" + raise _not_a_plugin(self) + + def configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + def configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + def configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + # Stream registration calls these, so the public hooks above only run for a backend that + # was passed in plugins= by mistake. + + def _configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + return super().configure_client(config) + + def _configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + return super().configure_worker(config) + + def _configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + return super().configure_replayer(config) + + +def _not_a_plugin(backend: StreamBackend) -> _errors.StreamConfigError: + return _errors.StreamConfigError( + f"{type(backend).__name__} is a Stream backend, not a plugin. Pass it in " + "stream_backends= of Client.connect, or of the Replayer." + ) + + +def check_backends(backends: Sequence[object], owner: str) -> list[StreamBackend]: + """Check the registered backends, and return them. + + A replayed create uses the default backend, so the default must not depend on the order + of the list. + + Raises: + TypeError: An entry is not a ``StreamBackend``. + StreamConfigError: Two backends have the same name, or there is more than one backend + and not exactly one has ``default=True``. + """ + checked: list[StreamBackend] = [] + for backend in backends: + if not isinstance(backend, StreamBackend): + raise TypeError(f"{backend!r} is not a StreamBackend") + if any(other.backend_name == backend.backend_name for other in checked): + raise _errors.StreamConfigError( + f"The {owner} has two Stream backends named {backend.backend_name}" + ) + checked.append(backend) + if len(checked) < 2: + return checked + defaults = [b.backend_name for b in checked if b.default] + if len(defaults) != 1: + names = ", ".join(b.backend_name for b in checked) + marked = f"{', '.join(defaults)} have" if defaults else "none has" + raise _errors.StreamConfigError( + f"The {owner} has the Stream backends {names}, and {marked} default=True. " + "Mark exactly one with default=True." + ) + return checked + + +def default_name(backends: Sequence[StreamBackend]) -> str: + """The name of the backend with ``default=True``, or ``native`` if none has it.""" + return next((b.backend_name for b in backends if b.default), NATIVE) + + +def configure_client( + config: temporalio.client.ClientConfig, +) -> temporalio.client.ClientConfig: + """Check the client's backends, and apply the plugin configuration of each one.""" + backends = check_backends(config.get("stream_backends", ()), "client") + config["stream_backends"] = backends + for backend in backends: + config = backend._configure_client(config) + return config + + +def configure_worker( + backends: Sequence[StreamBackend], config: temporalio.worker.WorkerConfig +) -> temporalio.worker.WorkerConfig: + """Apply the plugin configuration of the client's backends to a Worker.""" + for backend in backends: + config = backend._configure_worker(config) + return config + + +def configure_replayer( + config: temporalio.worker.ReplayerConfig, +) -> temporalio.worker.ReplayerConfig: + """Check the Replayer's backends, and apply the plugin configuration of each one.""" + backends = check_backends(config.get("stream_backends", ()), "Replayer") + config["stream_backends"] = backends + for backend in backends: + config = backend._configure_replayer(config) + return config + + +def client_backends(client: temporalio.client.Client) -> list[StreamBackend]: + """The backends registered on a client.""" + return list(client.config(active_config=True).get("stream_backends", ())) + + +def client_backend( + client: temporalio.client.Client, name: str | None, *, hint: str = "backend=" +) -> tuple[str, BackendCalls]: + """The name and the calls of a client's backend. A ``name`` of None picks the default. + + Raises: + StreamConfigError: The client has no backend of that name, or none has + ``default=True`` for a call with no name. + """ + backends = client_backends(client) + wanted = name or default_name(backends) + for backend in backends: + if backend.backend_name == wanted: + return wanted, backend._backend_for(client) + if not backends: + raise _errors.StreamConfigError( + "The client has no Stream backend. Register one with " + "Client.connect(..., stream_backends=[...])." + ) + if not name: + raise _errors.no_default_backend( + "client", [b.backend_name for b in backends], hint + ) + raise _errors.StreamConfigError( + f"The client has no Stream backend {name}. It has " + f"{', '.join(b.backend_name for b in backends)}." + ) diff --git a/temporalio/testing/_workflow.py b/temporalio/testing/_workflow.py index 6c8354162c..b4af035279 100644 --- a/temporalio/testing/_workflow.py +++ b/temporalio/testing/_workflow.py @@ -8,6 +8,7 @@ from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timedelta, timezone from typing import ( + TYPE_CHECKING, Any, cast, ) @@ -31,6 +32,9 @@ TimeSkippingConfig, ) +if TYPE_CHECKING: + from temporalio.common import StreamBackend + logger = logging.getLogger(__name__) @@ -102,6 +106,7 @@ async def start_local( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> WorkflowEnvironment: """Start a full Temporal server locally, downloading if necessary. @@ -165,6 +170,8 @@ async def start_local( become ready after spawning. Defaults to 5 seconds and does not include download time. ui_port: UI port to use if UI is enabled. + stream_backends: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started CLI dev server workflow environment. @@ -194,6 +201,7 @@ async def start_local( dev_server_download_ttl=dev_server_download_ttl, dev_server_startup_timeout=dev_server_startup_timeout, ui_port=ui_port, + stream_backends=stream_backends, ) @classmethod @@ -216,6 +224,7 @@ async def start_time_skipping( test_server_extra_args: Sequence[str] = [], test_server_download_ttl: timedelta | None = None, test_server_startup_timeout: timedelta | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> WorkflowEnvironment: """Start a V1 time skipping workflow environment. @@ -276,6 +285,8 @@ async def start_time_skipping( test_server_startup_timeout: Maximum time to wait for the test server to become ready after spawning. Defaults to 5 seconds and does not include download time. + stream_backends: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started workflow environment with time skipping. @@ -296,6 +307,7 @@ async def start_time_skipping( test_server_extra_args=test_server_extra_args, test_server_download_ttl=test_server_download_ttl, test_server_startup_timeout=test_server_startup_timeout, + stream_backends=stream_backends, ) @classmethod @@ -582,6 +594,7 @@ async def _bootstrap_dev_server( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> tuple[ temporalio.bridge.testing.EphemeralServer, temporalio.client.Client, @@ -649,6 +662,7 @@ async def _bootstrap_dev_server( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_backends=stream_backends, ) except: try: @@ -719,6 +733,7 @@ async def _create( test_server_extra_args: Sequence[str], test_server_download_ttl: timedelta | None, test_server_startup_timeout: timedelta | None, + stream_backends: Sequence[StreamBackend], ) -> _V1WorkflowEnvironment: runtime = runtime or temporalio.runtime.Runtime.default() download_ttl_ms = None @@ -753,6 +768,7 @@ async def _create( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_backends=stream_backends, ) except: try: diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index b3eb1a4d12..80dcf8c0da 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass +from typing import TYPE_CHECKING from typing_extensions import TypedDict @@ -32,6 +33,9 @@ ) from .workflow_sandbox import SandboxedWorkflowRunner +if TYPE_CHECKING: + from temporalio.common import StreamBackend + logger = logging.getLogger(__name__) @@ -56,13 +60,16 @@ def __init__( runtime: temporalio.runtime.Runtime | None = None, disable_safe_workflow_eviction: bool = False, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ) -> None: """Create a replayer to replay workflows from history. See :py:meth:`temporalio.worker.Worker.__init__` for a description of most of the arguments. Most of the same arguments need to be passed to the replayer that were passed to the worker when the workflow originally - ran. + ran. Pass ``stream_backends`` with the same backends as the Worker's + client. A replay makes no Stream call, so the Store doesn't need to be + reachable. Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that @@ -86,6 +93,7 @@ def __init__( runtime=runtime, disable_safe_workflow_eviction=disable_safe_workflow_eviction, header_codec_behavior=header_codec_behavior, + stream_backends=list(stream_backends), ) self._initial_config = self._config.copy() self._default_workflow_logic_flags = set(_DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS) @@ -94,6 +102,10 @@ def __init__( self.plugins = plugins for plugin in plugins: self._config = plugin.configure_replayer(self._config) + if self._config.get("stream_backends"): + from temporalio.common._streams import _plugin as _stream_plugin + + self._config = _stream_plugin.configure_replayer(self._config) # Validate workflows after plugin configuration if not self._config.get("workflows"): @@ -433,6 +445,7 @@ class ReplayerConfig(TypedDict, total=False): runtime: temporalio.runtime.Runtime | None disable_safe_workflow_eviction: bool header_codec_behavior: HeaderCodecBehavior + stream_backends: Sequence[StreamBackend] @dataclass(frozen=True) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 3c4c479962..1a8bba21eb 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -418,6 +418,13 @@ def __init__( self._plugins = plugins for plugin in plugins: config = plugin.configure_worker(config) + # Configure Streams after the plugins, so a plugin that replaces the Workflow runner + # doesn't lose the Streams sandbox passthrough. + stream_backends = client.config(active_config=True).get("stream_backends") + if stream_backends: + from temporalio.common._streams import _plugin as _stream_plugin + + config = _stream_plugin.configure_worker(stream_backends, config) self._init_from_config(client, config) diff --git a/tests/contrib/streams/test_backends.py b/tests/contrib/streams/test_backends.py new file mode 100644 index 0000000000..878012fafb --- /dev/null +++ b/tests/contrib/streams/test_backends.py @@ -0,0 +1,96 @@ +from typing import Any, cast + +import pytest + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.plugin import SimplePlugin +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Replayer, WorkerConfig +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner +from tests.helpers import new_worker + + +class _Backend(common.StreamBackend): + # Registration makes no Stream call, so this backend needs no Store. + def __init__(self, name: str = "test", *, default: bool = False) -> None: + super().__init__(f"test-streams-{name}", backend_name=name, default=default) + + +class _KeepsTheRunner(_Backend): + runner: object = None + + def _configure_worker(self, config: WorkerConfig) -> WorkerConfig: + self.runner = config.get("workflow_runner") + return super()._configure_worker(config) + + +@workflow.defn +class Idle: + @workflow.run + async def run(self) -> None: + pass + + +async def test_client_checks_its_backends(client: Client): + config = client.config() + config["stream_backends"] = [_Backend("a"), _Backend("b", default=True)] + two = Client(**config) + names = [b.backend_name for b in two.config(active_config=True)["stream_backends"]] + assert names == ["a", "b"] + # The default is part of Workflow code, so it must not depend on the order of the list. + for backends in ( + [_Backend("a"), _Backend("b")], + [_Backend("a", default=True), _Backend("b", default=True)], + ): + config["stream_backends"] = backends + with pytest.raises(common.StreamConfigError, match="exactly one"): + Client(**config) + config["stream_backends"] = [_Backend(), _Backend()] + with pytest.raises(common.StreamConfigError, match="two Stream backends"): + Client(**config) + config["stream_backends"] = cast(Any, [object()]) + with pytest.raises(TypeError, match="not a StreamBackend"): + Client(**config) + # A generator of backends works too, and a copy of the client keeps them. + config["stream_backends"] = cast(Any, (b for b in [_Backend()])) + copied = Client(**Client(**config).config()) + assert len(copied.config(active_config=True)["stream_backends"]) == 1 + + +async def test_backend_in_plugins_raises(client: Client): + # The connect refuses it before it connects, so the unreachable target doesn't matter. + with pytest.raises(common.StreamConfigError, match="not a plugin"): + await Client.connect("localhost:1", plugins=[_Backend()]) + config = client.config() + config["plugins"] = [_Backend()] + with pytest.raises(common.StreamConfigError, match="not a plugin"): + Client(**config) + with pytest.raises(common.StreamConfigError, match="not a plugin"): + new_worker(client, Idle, plugins=[_Backend()]) + with pytest.raises(common.StreamConfigError, match="not a plugin"): + Replayer(workflows=[Idle], plugins=[_Backend()]) + + +def test_replayer_checks_its_backends(): + with pytest.raises(common.StreamConfigError, match="two Stream backends"): + Replayer(workflows=[Idle], stream_backends=[_Backend(), _Backend()]) + replayer = Replayer(workflows=[Idle], stream_backends=[_Backend()]) + assert len(replayer.config(active_config=True).get("stream_backends", ())) == 1 + + +async def test_worker_configures_backends_after_its_plugins(client: Client): + # A plugin can replace the Workflow runner, and the backend must see the runner in use. + backend = _KeepsTheRunner(default=True) + config = client.config() + config["stream_backends"] = [backend] + runner = SandboxedWorkflowRunner() + plugin = SimplePlugin("replaces-runner", workflow_runner=runner) + new_worker(Client(**config), Idle, plugins=[plugin]) + assert backend.runner is runner + + +async def test_workflow_environment_gives_backends_to_its_client(): + backend = _Backend() + async with await WorkflowEnvironment.start_local(stream_backends=[backend]) as env: + assert list(env.client.config()["stream_backends"]) == [backend]