From 4eb6692ff5802d862bf8cc560ec6e9193880fec3 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 00:34:33 -0700 Subject: [PATCH 01/12] Built the Workflow Stream on the generated Stream resource. --- scripts/_wit/stream-service.wit | 48 +++- scripts/gen_stream_service.py | 4 +- temporalio/bridge/Cargo.lock | 2 +- temporalio/bridge/Cargo.toml | 2 +- temporalio/contrib/streams/_access.py | 12 +- temporalio/contrib/streams/redis.py | 2 +- temporalio/contrib/streams/system/__init__.py | 10 +- .../streams/system/_resources/__init__.py | 9 + .../streams/system/_resources/stream.py | 270 ++++++++++++++++++ .../system/_system_nexus_interceptor.py | 8 +- temporalio/contrib/streams/system/client.py | 10 +- temporalio/contrib/streams/system/models.py | 34 +-- .../system/operations/append_records.py | 63 ---- .../streams/system/operations/close_stream.py | 56 ---- .../system/operations/create_stream.py | 17 +- .../system/operations/describe_stream.py | 35 --- .../streams/system/operations/read_records.py | 67 ----- temporalio/contrib/streams/system/services.py | 4 +- temporalio/contrib/streams/workflow.py | 29 +- 19 files changed, 385 insertions(+), 297 deletions(-) create mode 100644 temporalio/contrib/streams/system/_resources/__init__.py create mode 100644 temporalio/contrib/streams/system/_resources/stream.py delete mode 100644 temporalio/contrib/streams/system/operations/append_records.py delete mode 100644 temporalio/contrib/streams/system/operations/close_stream.py delete mode 100644 temporalio/contrib/streams/system/operations/describe_stream.py delete mode 100644 temporalio/contrib/streams/system/operations/read_records.py diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index 0b48fafcf..80b59edb7 100644 --- a/scripts/_wit/stream-service.wit +++ b/scripts/_wit/stream-service.wit @@ -290,8 +290,54 @@ interface stream-service { auto-close: option, } + // A handle on one Stream. 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. nexgen has no doc directive + // on a resource yet, so the docs stay on the operations. + resource %stream { + constructor(%stream: stream-ref); + + /// @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.operation "close-stream" + close: func( + producer: option, + %result: option, + failure: option, + previous-completed-task: option, + failed-tasks: list, + ); + + /// @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.operation "describe-stream" + describe: func() -> describe-stream-response; + } + + // The handle is built from the response. The created and status fields of the response are + // not on the handle. + /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" + type create-stream-result = own<%stream>; + /// @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." diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py index d04af620c..722d1563a 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 = "de0c7edd5e36552391cb4f37c96aeb065270035e" nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index c5ae55a30..3c62e3d89 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=88abba51a857da66ac5b03720a99727e01659078#88abba51a857da66ac5b03720a99727e01659078" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index d2e25c4d3..550e78e01 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 = "88abba51a857da66ac5b03720a99727e01659078" default-features = false optional = true diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 80d1eda29..0ce160ecf 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -15,6 +15,7 @@ from temporalio.contrib.streams._handle import StreamHandle from temporalio.contrib.streams._reader import Item from temporalio.contrib.streams._writer import StreamWriter +from temporalio.contrib.streams.system import models T = TypeVar("T") @@ -85,13 +86,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 diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 0b3e698cc..2a3c52f39 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -555,7 +555,7 @@ async def start_system_nexus_operation( return await super().start_system_nexus_operation(input) operation = input.operation_name if operation == "CreateStream": - result: Any = models.CreateStreamResponse( + result: Any = models.CreateStreamResult( stream=_owned_ref(input.input.stream_id) ) elif operation == "AppendRecords": 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..d24eff30c --- /dev/null +++ b/temporalio/contrib/streams/system/_resources/stream.py @@ -0,0 +1,270 @@ +# Generated by nexgen v0.2.9. DO NOT EDIT! + +from __future__ import annotations + +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, + StreamRef, + WorkflowTaskEpoch, +) + + +@dataclasses.dataclass +class Stream: + 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: + 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: + 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: + 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 describe( + self, + ) -> DescribeStreamResponse: + 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/_system_nexus_interceptor.py b/temporalio/contrib/streams/system/_system_nexus_interceptor.py index 44d8d4a6b..38639a325 100644 --- a/temporalio/contrib/streams/system/_system_nexus_interceptor.py +++ b/temporalio/contrib/streams/system/_system_nexus_interceptor.py @@ -32,7 +32,7 @@ async def _start_system_nexus_operation( and input.operation_name == "CreateStream" ): typed_input = typing.cast( - "StartNexusOperationInput[models.CreateStreamRequest, models.CreateStreamResponse]", + "StartNexusOperationInput[models.CreateStreamRequest, models.CreateStreamResult]", input, ) # The dispatch check above establishes that this operation's response type is _OutputT. @@ -111,7 +111,7 @@ async def start_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResponse]: + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: """Intercept the CreateStream operation.""" return await self._next_system_nexus_interceptor().start_create_stream(request) @@ -151,7 +151,7 @@ async def _intercept_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResponse]: + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput from temporalio.workflow import NexusOperationCancellationType @@ -162,7 +162,7 @@ async def start_create_stream( service="temporal.api.streamservice.v1.StreamService", operation="CreateStream", input=request, - output_type=models.CreateStreamResponse, + output_type=models.CreateStreamResult, schedule_to_close_timeout=None, schedule_to_start_timeout=None, start_to_close_timeout=None, diff --git a/temporalio/contrib/streams/system/client.py b/temporalio/contrib/streams/system/client.py index 2883b3f1c..ba98e3a0e 100644 --- a/temporalio/contrib/streams/system/client.py +++ b/temporalio/contrib/streams/system/client.py @@ -19,7 +19,7 @@ CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, - CreateStreamResponse, + CreateStreamResult, DescribeStreamRequest, DescribeStreamResponse, ReadRecordsRequest, @@ -29,7 +29,7 @@ _CloseStreamRequestTransferTypeConverter, _CloseStreamResponseTransferTypeConverter, _CreateStreamRequestTransferTypeConverter, - _CreateStreamResponseTransferTypeConverter, + _CreateStreamResultTransferTypeConverter, _DescribeStreamRequestTransferTypeConverter, _DescribeStreamResponseTransferTypeConverter, _ReadRecordsRequestTransferTypeConverter, @@ -155,15 +155,15 @@ async def _call( [answer] = await self._payload_codec.decode([answer]) return answer - async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResponse: + async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResult: payload = _proto_payload( _CreateStreamRequestTransferTypeConverter().to_transfer_type(request) ) answer = await self._call("CreateStream", payload) - converter = _CreateStreamResponseTransferTypeConverter() + converter = _CreateStreamResultTransferTypeConverter() return converter.from_transfer_type( _proto_message(answer, converter.transfer_type), - CreateStreamResponse, + CreateStreamResult, ) async def append_records( diff --git a/temporalio/contrib/streams/system/models.py b/temporalio/contrib/streams/system/models.py index 2304fc581..cc6570133 100644 --- a/temporalio/contrib/streams/system/models.py +++ b/temporalio/contrib/streams/system/models.py @@ -86,17 +86,13 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class CreateStreamResponse: +class CreateStreamResult: stream: StreamRef | None = None - created: bool | None = None - status: ( - temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType | None - ) = None -class _CreateStreamResponseTransferTypeConverter( +class _CreateStreamResultTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - CreateStreamResponse, + CreateStreamResult, temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, ] ): @@ -109,26 +105,20 @@ class _CreateStreamResponseTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, - type_hint: type[CreateStreamResponse], - ) -> CreateStreamResponse: - return CreateStreamResponse( + type_hint: type[CreateStreamResult], + ) -> CreateStreamResult: + return CreateStreamResult( stream=_StreamRefTransferTypeConverter().from_transfer_type( value.stream, StreamRef ) if value.HasField("stream") else None, - created=value.created if bool(value.created) else None, - status=temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType( - value.status - ) - if value.status != 0 - else None, ) @typing_extensions.override def to_transfer_type( self, - value: CreateStreamResponse, + value: CreateStreamResult, ) -> temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse: message = ( temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse() @@ -137,18 +127,12 @@ def to_transfer_type( message.stream.CopyFrom( _StreamRefTransferTypeConverter().to_transfer_type(value.stream) ) - if value.created is not None: - message.created = value.created - if value.status is not None: - message.status = temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType( - value.status - ) return message temporalio.converter.transfer_type_convertible( - _CreateStreamResponseTransferTypeConverter -)(CreateStreamResponse) # pyright: ignore[reportUnusedCallResult] + _CreateStreamResultTransferTypeConverter +)(CreateStreamResult) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) 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..7d49d5842 100644 --- a/temporalio/contrib/streams/system/operations/create_stream.py +++ b/temporalio/contrib/streams/system/operations/create_stream.py @@ -6,23 +6,30 @@ import temporalio.workflow +from .._resources import Stream from ..models import ( CreateStreamRequest, - CreateStreamResponse, + CreateStreamResult, ) async def _create_stream( request: CreateStreamRequest, -) -> temporalio.workflow.NexusOperationHandle[CreateStreamResponse,]: +) -> Stream: 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, + output_type=CreateStreamResult, + ) + result = await handle + if result.stream is None: + raise RuntimeError("CreateStream answered without `stream`") + return Stream( + stream=result.stream, ) @@ -32,7 +39,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,]: +) -> Stream: """ 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/system/services.py b/temporalio/contrib/streams/system/services.py index f9e297ce5..e5a0445bb 100644 --- a/temporalio/contrib/streams/system/services.py +++ b/temporalio/contrib/streams/system/services.py @@ -10,7 +10,7 @@ CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, - CreateStreamResponse, + CreateStreamResult, DescribeStreamRequest, DescribeStreamResponse, ReadRecordsRequest, @@ -27,7 +27,7 @@ class StreamService: create_stream: Operation[ CreateStreamRequest, - CreateStreamResponse, + CreateStreamResult, ] = Operation(name="CreateStream") append_records: Operation[ diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 9526c334a..02b81e54d 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -23,14 +23,9 @@ 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.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") @@ -182,6 +177,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 +230,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 +257,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.""" @@ -372,9 +367,7 @@ 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 return stream From d0f8132e59cff437cd9d461c43da97374e78a0be Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 00:49:09 -0700 Subject: [PATCH 02/12] Read through the generated read iterator and moved to n8 aff8ff85. --- scripts/_wit/stream-service.wit | 6 +- scripts/gen_stream_service.py | 2 +- temporalio/bridge/Cargo.lock | 2 +- temporalio/bridge/Cargo.toml | 2 +- temporalio/contrib/streams/_reader.py | 69 +++++++++++-------- temporalio/contrib/streams/redis.py | 2 +- .../streams/system/_resources/stream.py | 65 +++++++++++++++++ .../system/_system_nexus_interceptor.py | 8 +-- temporalio/contrib/streams/system/client.py | 43 ++++++++++-- temporalio/contrib/streams/system/models.py | 20 +++--- .../system/operations/create_stream.py | 4 +- temporalio/contrib/streams/system/services.py | 4 +- 12 files changed, 171 insertions(+), 56 deletions(-) diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index 80b59edb7..dc3f40648 100644 --- a/scripts/_wit/stream-service.wit +++ b/scripts/_wit/stream-service.wit @@ -334,10 +334,10 @@ interface stream-service { // The handle is built from the response. The created and status fields of the response are // not on the handle. /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" - type create-stream-result = own<%stream>; + type stream-result = own<%stream>; /// @nexus.operation name="CreateStream" - create-stream: func(request: create-stream-request) -> create-stream-result; + create-stream: func(request: create-stream-request) -> stream-result; /// @nexus.doc /// "In Workflow code, Brook buffers this call in the run, and the facade drops the answer." @@ -347,6 +347,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 722d1563a..5e27ca61b 100644 --- a/scripts/gen_stream_service.py +++ b/scripts/gen_stream_service.py @@ -33,7 +33,7 @@ NEX_GEN_REPOSITORY = "https://github.com/moetemp/nexgen" # The head of moe/brook/v3/n8-stream-resource. A branch moves, so the pin is the commit. -NEX_GEN_REVISION = "de0c7edd5e36552391cb4f37c96aeb065270035e" +NEX_GEN_REVISION = "aff8ff85450ebfbabe17fccb87b00acb7a97aad5" nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index 3c62e3d89..0e8c8a810 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=88abba51a857da66ac5b03720a99727e01659078#88abba51a857da66ac5b03720a99727e01659078" +source = "git+https://github.com/moetemp/sdk-streams?rev=760a53bf6958580fd9d3c2142deb75a9e286ff9d#760a53bf6958580fd9d3c2142deb75a9e286ff9d" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 550e78e01..369fc5956 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 = "88abba51a857da66ac5b03720a99727e01659078" +rev = "760a53bf6958580fd9d3c2142deb75a9e286ff9d" default-features = false optional = true diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/contrib/streams/_reader.py index defd0ffd5..d5d2c98b1 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 @@ -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( @@ -40,46 +39,62 @@ 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, + 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: + # Only the service's own conversion keeps Record bodies raw. + with _convert.raw_bodies(converter.failure_converter): + 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( diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 2a3c52f39..01c4802d7 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -555,7 +555,7 @@ async def start_system_nexus_operation( return await super().start_system_nexus_operation(input) operation = input.operation_name if operation == "CreateStream": - result: Any = models.CreateStreamResult( + result: Any = models.StreamResult( stream=_owned_ref(input.input.stream_id) ) elif operation == "AppendRecords": diff --git a/temporalio/contrib/streams/system/_resources/stream.py b/temporalio/contrib/streams/system/_resources/stream.py index d24eff30c..09d3884eb 100644 --- a/temporalio/contrib/streams/system/_resources/stream.py +++ b/temporalio/contrib/streams/system/_resources/stream.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio +import collections.abc import dataclasses import datetime import typing @@ -20,10 +22,17 @@ 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: @@ -93,6 +102,62 @@ async def read_records( 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: diff --git a/temporalio/contrib/streams/system/_system_nexus_interceptor.py b/temporalio/contrib/streams/system/_system_nexus_interceptor.py index 38639a325..abd8b2fe9 100644 --- a/temporalio/contrib/streams/system/_system_nexus_interceptor.py +++ b/temporalio/contrib/streams/system/_system_nexus_interceptor.py @@ -32,7 +32,7 @@ async def _start_system_nexus_operation( and input.operation_name == "CreateStream" ): typed_input = typing.cast( - "StartNexusOperationInput[models.CreateStreamRequest, models.CreateStreamResult]", + "StartNexusOperationInput[models.CreateStreamRequest, models.StreamResult]", input, ) # The dispatch check above establishes that this operation's response type is _OutputT. @@ -111,7 +111,7 @@ async def start_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + ) -> temporalio.workflow.NexusOperationHandle[models.StreamResult]: """Intercept the CreateStream operation.""" return await self._next_system_nexus_interceptor().start_create_stream(request) @@ -151,7 +151,7 @@ async def _intercept_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + ) -> temporalio.workflow.NexusOperationHandle[models.StreamResult]: from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput from temporalio.workflow import NexusOperationCancellationType @@ -162,7 +162,7 @@ async def start_create_stream( service="temporal.api.streamservice.v1.StreamService", operation="CreateStream", input=request, - output_type=models.CreateStreamResult, + output_type=models.StreamResult, schedule_to_close_timeout=None, schedule_to_start_timeout=None, start_to_close_timeout=None, diff --git a/temporalio/contrib/streams/system/client.py b/temporalio/contrib/streams/system/client.py index ba98e3a0e..d8b4b74ab 100644 --- a/temporalio/contrib/streams/system/client.py +++ b/temporalio/contrib/streams/system/client.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import collections.abc import dataclasses import datetime import typing @@ -19,21 +20,22 @@ CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, - CreateStreamResult, DescribeStreamRequest, DescribeStreamResponse, ReadRecordsRequest, ReadRecordsResponse, + StreamRecord, + StreamResult, _AppendRecordsRequestTransferTypeConverter, _AppendRecordsResponseTransferTypeConverter, _CloseStreamRequestTransferTypeConverter, _CloseStreamResponseTransferTypeConverter, _CreateStreamRequestTransferTypeConverter, - _CreateStreamResultTransferTypeConverter, _DescribeStreamRequestTransferTypeConverter, _DescribeStreamResponseTransferTypeConverter, _ReadRecordsRequestTransferTypeConverter, _ReadRecordsResponseTransferTypeConverter, + _StreamResultTransferTypeConverter, ) @@ -155,15 +157,15 @@ async def _call( [answer] = await self._payload_codec.decode([answer]) return answer - async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResult: + async def create_stream(self, request: CreateStreamRequest) -> StreamResult: payload = _proto_payload( _CreateStreamRequestTransferTypeConverter().to_transfer_type(request) ) answer = await self._call("CreateStream", payload) - converter = _CreateStreamResultTransferTypeConverter() + converter = _StreamResultTransferTypeConverter() return converter.from_transfer_type( _proto_message(answer, converter.transfer_type), - CreateStreamResult, + StreamResult, ) async def append_records( @@ -255,6 +257,37 @@ 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: diff --git a/temporalio/contrib/streams/system/models.py b/temporalio/contrib/streams/system/models.py index cc6570133..988884c0e 100644 --- a/temporalio/contrib/streams/system/models.py +++ b/temporalio/contrib/streams/system/models.py @@ -86,13 +86,13 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class CreateStreamResult: +class StreamResult: stream: StreamRef | None = None -class _CreateStreamResultTransferTypeConverter( +class _StreamResultTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - CreateStreamResult, + StreamResult, temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, ] ): @@ -105,9 +105,9 @@ class _CreateStreamResultTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, - type_hint: type[CreateStreamResult], - ) -> CreateStreamResult: - return CreateStreamResult( + type_hint: type[StreamResult], + ) -> StreamResult: + return StreamResult( stream=_StreamRefTransferTypeConverter().from_transfer_type( value.stream, StreamRef ) @@ -118,7 +118,7 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: CreateStreamResult, + value: StreamResult, ) -> temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse: message = ( temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse() @@ -130,9 +130,9 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible( - _CreateStreamResultTransferTypeConverter -)(CreateStreamResult) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible(_StreamResultTransferTypeConverter)( + StreamResult +) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) diff --git a/temporalio/contrib/streams/system/operations/create_stream.py b/temporalio/contrib/streams/system/operations/create_stream.py index 7d49d5842..f57f13897 100644 --- a/temporalio/contrib/streams/system/operations/create_stream.py +++ b/temporalio/contrib/streams/system/operations/create_stream.py @@ -9,7 +9,7 @@ from .._resources import Stream from ..models import ( CreateStreamRequest, - CreateStreamResult, + StreamResult, ) @@ -23,7 +23,7 @@ async def _create_stream( handle = await nexus_client.start_operation( operation="CreateStream", input=request, - output_type=CreateStreamResult, + output_type=StreamResult, ) result = await handle if result.stream is None: diff --git a/temporalio/contrib/streams/system/services.py b/temporalio/contrib/streams/system/services.py index e5a0445bb..91e05b70b 100644 --- a/temporalio/contrib/streams/system/services.py +++ b/temporalio/contrib/streams/system/services.py @@ -10,11 +10,11 @@ CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, - CreateStreamResult, DescribeStreamRequest, DescribeStreamResponse, ReadRecordsRequest, ReadRecordsResponse, + StreamResult, ) @@ -27,7 +27,7 @@ class StreamService: create_stream: Operation[ CreateStreamRequest, - CreateStreamResult, + StreamResult, ] = Operation(name="CreateStream") append_records: Operation[ From 4dfe5f018ad10ac558873615de267f25e12a2207 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 00:49:49 -0700 Subject: [PATCH 03/12] Formatted the StreamResult answer in the plugin. --- temporalio/contrib/streams/redis.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 01c4802d7..7810d28aa 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -555,9 +555,7 @@ async def start_system_nexus_operation( return await super().start_system_nexus_operation(input) operation = input.operation_name if operation == "CreateStream": - result: Any = models.StreamResult( - stream=_owned_ref(input.input.stream_id) - ) + result: Any = models.StreamResult(stream=_owned_ref(input.input.stream_id)) elif operation == "AppendRecords": result = models.AppendRecordsResponse() elif operation == "CloseStream": From e6eb7542a370ace002890e7dcf256339092a00c7 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:05:40 -0700 Subject: [PATCH 04/12] Dropped tests of attributes, defaults and getters. --- tests/contrib/streams/test_bindings.py | 22 --------------- tests/contrib/streams/test_facade.py | 17 +---------- tests/contrib/streams/test_handles.py | 28 ------------------- tests/contrib/streams/test_stream_client.py | 6 ---- tests/contrib/streams/test_workflow_module.py | 11 -------- 5 files changed, 1 insertion(+), 83 deletions(-) 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..08038f46d 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -8,7 +8,7 @@ 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.proto import v1 as ext_pb from temporalio.contrib.streams.system import models from temporalio.exceptions import ApplicationError @@ -33,16 +33,6 @@ def fake_client(client: Client) -> tuple[Client, FakeExtension]: 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 +152,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( 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_stream_client.py b/tests/contrib/streams/test_stream_client.py index 6cffb7ec0..088c56534 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -200,7 +200,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: @@ -228,8 +227,3 @@ async def test_default_plugin_is_the_default_backend(client: Client): ] stream = await streams.StreamClient(Client(**config)).create(Events, id="d-1") 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") 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 From 7c465639d750ef8c39556cd1b08883a7e98a9d1d Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:09:26 -0700 Subject: [PATCH 05/12] Took backend= on stream_handle and carried each Stream's backend. --- temporalio/contrib/streams/redis.py | 34 ++++++++++++++++---------- temporalio/contrib/streams/workflow.py | 22 ++++++++++++++--- 2 files changed, 40 insertions(+), 16 deletions(-) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 7810d28aa..40c4a9e04 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -43,9 +43,8 @@ # 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): @@ -375,7 +374,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 +515,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 entry in converter.from_payload(payload, list[dict[str, Any]]): + run_stream = state.stream(_owned_ref(entry["id"])) + run_stream.owned = True + run_stream.backend = entry.get("backend") + run_stream.close_with_run = bool(entry.get("closeWithRun")) class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): @@ -571,16 +573,22 @@ 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, and whether it + # closes with the chain. + owned = [ + { + "id": stream_id, + "backend": run_stream.backend, + "closeWithRun": 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.owned 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) diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 02b81e54d..14674020a 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -46,6 +46,9 @@ 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. + owned: bool = False + backend: str | None = None @dataclass @@ -324,7 +327,10 @@ 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. @@ -370,10 +376,14 @@ def create_stream_nowait( _create_stream(stream_id=stream.id, retention=retention, auto_close=auto_close) ) run_stream.close_with_run = close_with_run + run_stream.owned = True + 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 @@ -382,5 +392,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 From 7af8626f288fcf7ad0347182de5d54a96bc4dd3b Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:10:02 -0700 Subject: [PATCH 06/12] Bound a handle that came as data with get_handle(stream=). --- temporalio/contrib/streams/_stream_client.py | 34 ++++++++++++++++--- .../streams/examples/turn_scoped_stream.py | 4 +-- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index f318e3016..8b59045e3 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -286,22 +286,46 @@ async def create( ) return stream._bind(self._client, stream_backend) + @overload def get_handle( self, definition: type[E], *, id: str, backend: str | None = None + ) -> ClientStream[E]: ... + + @overload + def get_handle( + self, definition: type[E], *, stream: StreamHandle[Any] + ) -> ClientStream[E]: ... + + def get_handle( + self, + definition: type[E], + *, + id: str | None = None, + stream: StreamHandle[Any] | None = None, + backend: 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. Its id, owner and backend stay. + backend: The name of the backend. Defaults to the client's default backend. """ definition_of(definition) + if stream is not None: + id, backend = stream.id, stream.backend + if id is None: + raise TypeError("get_handle takes id= or stream=") backend_name, stream_backend = client_backend(self._client, backend) - stream = ClientStream( - id=id, backend=backend_name or None, definition=definition + handle = ClientStream( + id=id, + workflow_id=stream.workflow_id if stream else None, + first_execution_run_id=stream.first_execution_run_id if stream else None, + backend=backend_name or None, + definition=definition, ) - return stream._bind(self._client, stream_backend) + return handle._bind(self._client, stream_backend) def workflow_stream( self, 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() From dfb967a4e19aedf963dcee47936cf598547d9f02 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:12:00 -0700 Subject: [PATCH 07/12] Looked in each named backend in order for a moving Stream. --- temporalio/contrib/streams/_backend.py | 10 +- temporalio/contrib/streams/_stream_client.py | 109 ++++++++++++++++--- tests/contrib/streams/test_stream_client.py | 14 +++ 3 files changed, 117 insertions(+), 16 deletions(-) diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index a2a6ac994..72a14d185 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -164,10 +164,11 @@ async def call( def client_backend( - client: temporalio.client.Client, name: str | None + 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. + 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: StreamTransportMissingError: The client has no Stream plugin. @@ -188,6 +189,11 @@ def client_backend( raise _errors.StreamTransportMissingError( "The client has no Stream plugin, such as RedisStreamPlugin" ) + if not name and only and len(backends) > 1: + raise _errors.StreamUnsupportedError( + "The handle names no backend, and the client has more than one", + reason="stream_unsupported", + ) if not name: defaults = [entry for entry in backends if entry[2]] or backends return defaults[0][0], defaults[0][1] diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index 8b59045e3..1b60292f5 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -18,6 +18,7 @@ from temporalio.contrib.streams._handle import StreamHandle 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") @@ -45,6 +46,14 @@ class ClientTopic(_AsyncTopic[T]): This API is experimental. """ + 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 super().publish(*values) + class ClientStream(StreamHandle[E]): """A Stream that a client writes, reads and closes. Get one from :py:class:`StreamClient`. @@ -80,18 +89,62 @@ 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]] = [] + self._probing: asyncio.Future[None] | None = None 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 that keeps the Stream. + + Raises: + RuntimeError: The handle looks in more than one backend, and no call has found the + Stream yet. Call :py:meth:`get_offset` first. + """ + if self._probe: + raise RuntimeError( + "The Stream's backend is not known yet. Call get_offset() first." + ) + return self._backend + + async def _ready(self) -> None: + if self._probe: + # Calls that start together share one search. + if self._probing is None: + self._probing = asyncio.ensure_future(self._find_backend()) + await asyncio.shield(self._probing) + + async def _find_backend(self) -> None: + client = self._client + assert client is not None + stream_id = await self._stream_id() + found = self._probe[0] + for name, backend in self._probe: + try: + await backend.service(client.namespace).describe_stream( + models.DescribeStreamRequest(stream=models.StreamRef(id=stream_id)) + ) + except _errors.StreamNotFoundError: + continue + found = (name, backend) + break + # With no backend that has the Stream, the first one waits for its create. + self._backend, self._stream_backend = found + self._probe = [] + @property def id(self) -> str: """The Stream id. @@ -202,7 +255,7 @@ def subscribe( 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 +264,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 +282,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) @@ -288,7 +348,11 @@ async def create( @overload def get_handle( - self, definition: type[E], *, id: str, backend: str | None = None + self, + definition: type[E], + *, + id: str, + backend: str | tuple[str, ...] | None = None, ) -> ClientStream[E]: ... @overload @@ -302,22 +366,30 @@ def get_handle( *, id: str | None = None, stream: StreamHandle[Any] | None = None, - backend: str | None = None, + backend: str | tuple[str, ...] | None = None, ) -> ClientStream[E]: """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. - stream: A handle of the Stream. Its id, owner and backend stay. - backend: The name of the backend. Defaults to the client's default backend. + stream: A handle of the Stream. Its id, owner and backend stay. A handle with no + backend name needs a client with one backend. + 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) if stream is not None: - id, backend = stream.id, stream.backend - if id is None: + id = stream.id + backend_name, stream_backend = client_backend( + self._client, stream.backend, only=True + ) + probe: list[tuple[str, StreamBackend]] = [] + elif id is None: raise TypeError("get_handle takes id= or stream=") - backend_name, stream_backend = client_backend(self._client, backend) + else: + backend_name, stream_backend, probe = self._backends(backend) handle = ClientStream( id=id, workflow_id=stream.workflow_id if stream else None, @@ -325,7 +397,15 @@ def get_handle( backend=backend_name or None, definition=definition, ) - return handle._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, @@ -333,7 +413,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. @@ -344,10 +424,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) @@ -359,7 +440,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() @@ -371,4 +452,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/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 088c56534..09923c47e 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -227,3 +227,17 @@ async def test_default_plugin_is_the_default_backend(client: Client): ] stream = await streams.StreamClient(Client(**config)).create(Events, id="d-1") assert stream.backend == "b" + + +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" + # A handle with no backend name is ambiguous when the client has two backends. + with pytest.raises(streams.StreamUnsupportedError): + stream_client.get_handle(Events, stream=streams.ClientStream(id="moving")) From aef30a2f6fff9998b9bc29a6336b76eaf2e4cded Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:14:16 -0700 Subject: [PATCH 08/12] Raised StreamConfigError for config errors and mapped server details. --- temporalio/contrib/streams/__init__.py | 4 +-- .../contrib/streams/_activity_stream.py | 7 ++--- temporalio/contrib/streams/_backend.py | 16 +++++------ temporalio/contrib/streams/_errors.py | 28 +++++++++++++++++-- temporalio/contrib/streams/_stream_client.py | 4 +-- temporalio/contrib/streams/workflow.py | 14 ++++++---- .../contrib/streams/test_activity_handles.py | 6 ++-- tests/contrib/streams/test_facade.py | 20 ++++++++++++- tests/contrib/streams/test_stream_client.py | 10 +++---- 9 files changed, 75 insertions(+), 34 deletions(-) diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 089aa6a10..3fefc7e53 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -36,13 +36,13 @@ from temporalio.contrib.streams._errors import ( ReadOnlyContextError, StreamClosedError, + StreamConfigError, StreamError, StreamExpiredError, StreamFencedError, StreamNotFoundError, StreamOffsetError, StreamStorageError, - StreamTransportMissingError, StreamUnsupportedError, ) from temporalio.contrib.streams._handle import StreamHandle @@ -66,6 +66,7 @@ "StreamBackend", "StreamClient", "StreamClosedError", + "StreamConfigError", "StreamDefinition", "StreamError", "StreamExpiredError", @@ -74,7 +75,6 @@ "StreamNotFoundError", "StreamOffsetError", "StreamStorageError", - "StreamTransportMissingError", "StreamUnsupportedError", "Superseded", "SyncActivityStream", diff --git a/temporalio/contrib/streams/_activity_stream.py b/temporalio/contrib/streams/_activity_stream.py index 1c258b7ad..c9e623096 100644 --- a/temporalio/contrib/streams/_activity_stream.py +++ b/temporalio/contrib/streams/_activity_stream.py @@ -30,14 +30,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 diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 72a14d185..5ef9d5941 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -171,8 +171,8 @@ def client_backend( backend, since a handle with no backend name may be of any of them. Raises: - StreamTransportMissingError: The client has no Stream plugin. - StreamUnsupportedError: The client has no Stream plugin of that name. + StreamConfigError: The client has no Stream plugin, none of that name, or more than + one where ``only`` needs one. """ backends = [ ( @@ -186,13 +186,12 @@ def client_backend( ) ] if not backends: - raise _errors.StreamTransportMissingError( + raise _errors.StreamConfigError( "The client has no Stream plugin, such as RedisStreamPlugin" ) if not name and only and len(backends) > 1: - raise _errors.StreamUnsupportedError( - "The handle names no backend, and the client has more than one", - reason="stream_unsupported", + 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 @@ -200,8 +199,7 @@ def client_backend( for backend_name, backend, _ in backends: if backend_name == name: return backend_name, backend - raise _errors.StreamUnsupportedError( + raise _errors.StreamConfigError( f"The client has no Stream backend {name}. It has " - f"{', '.join(entry[0] for entry in backends)}.", - reason="stream_unsupported", + f"{', '.join(entry[0] for entry in backends)}." ) diff --git a/temporalio/contrib/streams/_errors.py b/temporalio/contrib/streams/_errors.py index 85d2323f9..4e8cfb8ee 100644 --- a/temporalio/contrib/streams/_errors.py +++ b/temporalio/contrib/streams/_errors.py @@ -77,8 +77,14 @@ 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 the name that a Stream or a call gives. + + It comes from the Worker or client config, not from a Store, so it is not a + :py:class:`StreamError`, and ``except StreamError`` does not catch it. In Workflow code it + fails the Workflow Task, so a Worker with the right config can take the run on. In an + Activity the attempt is retried. + """ class ReadOnlyContextError(StreamError): @@ -96,6 +102,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 +117,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/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index 1b60292f5..d5ed12e19 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -184,7 +184,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." ) @@ -319,7 +319,7 @@ 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``. diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 14674020a..c02a4d223 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -94,11 +94,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}." ) @@ -112,7 +112,9 @@ 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]): @@ -148,8 +150,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) @@ -335,7 +337,7 @@ async def create_stream( 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. 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_facade.py b/tests/contrib/streams/test_facade.py index 08038f46d..c727c0aa8 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -2,8 +2,10 @@ 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 @@ -13,7 +15,7 @@ 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 @@ -162,3 +164,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_stream_client.py b/tests/contrib/streams/test_stream_client.py index 09923c47e..94a69f788 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -43,9 +43,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 +53,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 +65,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() @@ -239,5 +239,5 @@ async def test_handle_looks_in_each_backend_during_a_move(client: Client): assert await moving.get_offset() > 0 assert moving.backend == "b" # A handle with no backend name is ambiguous when the client has two backends. - with pytest.raises(streams.StreamUnsupportedError): + with pytest.raises(streams.StreamConfigError): stream_client.get_handle(Events, stream=streams.ClientStream(id="moving")) From c8f5294f5ca4d0fc663357670570c929031d2afb Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:16:45 -0700 Subject: [PATCH 09/12] Took streams.END as a from_offset for new Records only. --- temporalio/contrib/streams/__init__.py | 3 ++- temporalio/contrib/streams/_access.py | 7 +++--- .../contrib/streams/_activity_stream.py | 10 +++++---- temporalio/contrib/streams/_items.py | 11 ++++++++++ temporalio/contrib/streams/_reader.py | 13 +++++++---- temporalio/contrib/streams/_stream_client.py | 6 +++-- .../streams/examples/resume_from_offset.py | 22 ++++++++++++++++--- .../contrib/streams/examples/test_examples.py | 1 + 8 files changed, 56 insertions(+), 17 deletions(-) diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 3fefc7e53..00e75d229 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -46,7 +46,7 @@ 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._stream_client import ( ClientStream, ClientTopic, @@ -55,6 +55,7 @@ from temporalio.contrib.streams.workflow import WorkflowStream, WorkflowTopic __all__ = [ + "END", "ActivityStream", "ActivityTopic", "BufferedPublisher", diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 0ce160ecf..40dd70868 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -13,6 +13,7 @@ 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._items import End from temporalio.contrib.streams._reader import Item from temporalio.contrib.streams._writer import StreamWriter from temporalio.contrib.streams.system import models @@ -42,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]: @@ -62,7 +63,7 @@ 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]: @@ -165,7 +166,7 @@ async 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, ) -> AsyncGenerator[Item, None]: diff --git a/temporalio/contrib/streams/_activity_stream.py b/temporalio/contrib/streams/_activity_stream.py index c9e623096..def9b1446 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") @@ -104,7 +105,7 @@ def topic( 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, @@ -117,7 +118,8 @@ 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. Defaults to @@ -171,7 +173,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]: @@ -243,7 +245,7 @@ def topic( 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, diff --git a/temporalio/contrib/streams/_items.py b/temporalio/contrib/streams/_items.py index 9cf0f8870..3ffa65aba 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,16 @@ T = TypeVar("T") +class End(enum.Enum): + """The type of :py:data:`END`.""" + + END = "END" + + +END = End.END +"""A ``from_offset`` that starts a read after the newest Record, for new Records only.""" + + @dataclass(frozen=True) class Record(Generic[T]): """A data Record of a Stream. diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/contrib/streams/_reader.py index d5d2c98b1..bf74a22f4 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/contrib/streams/_reader.py @@ -11,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 @@ -29,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]: @@ -48,8 +48,13 @@ def 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, + # 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, diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index d5ed12e19..e9ce90d3e 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -16,6 +16,7 @@ from temporalio.contrib.streams._backend import StreamBackend, client_backend 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._reader import Item from temporalio.contrib.streams._writer import StreamWriter from temporalio.contrib.streams.system import models @@ -234,7 +235,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, @@ -248,7 +249,8 @@ 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 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"], ) From c656d8ebb19e7fc876cb04448c08af0820d2fc93 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:21:15 -0700 Subject: [PATCH 10/12] Simplified the backend lookup and the get_handle overload. --- temporalio/contrib/streams/_access.py | 39 ++------- .../contrib/streams/_activity_stream.py | 29 +------ temporalio/contrib/streams/_backend.py | 5 -- temporalio/contrib/streams/_errors.py | 7 +- temporalio/contrib/streams/_handle.py | 27 ++++++ temporalio/contrib/streams/_items.py | 3 +- temporalio/contrib/streams/_stream_client.py | 82 +++++-------------- temporalio/contrib/streams/redis.py | 21 ++--- temporalio/contrib/streams/workflow.py | 24 +----- tests/contrib/streams/test_stream_client.py | 3 + 10 files changed, 75 insertions(+), 165 deletions(-) diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 40dd70868..09e904f51 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -5,14 +5,14 @@ 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 @@ -139,28 +139,15 @@ 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( @@ -179,21 +166,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 def9b1446..f4891fcc7 100644 --- a/temporalio/contrib/streams/_activity_stream.py +++ b/temporalio/contrib/streams/_activity_stream.py @@ -91,14 +91,7 @@ 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) @@ -110,21 +103,7 @@ def subscribe( 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, 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. 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, @@ -238,7 +217,7 @@ 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) @@ -250,7 +229,7 @@ def subscribe( 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 5ef9d5941..7b32e635b 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -48,11 +48,6 @@ 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) -> service_client.StreamServiceClient: """A generated Stream service client for ``namespace``.""" return service_client.StreamServiceClient(_ExtensionTransport(self, namespace)) diff --git a/temporalio/contrib/streams/_errors.py b/temporalio/contrib/streams/_errors.py index 4e8cfb8ee..de203ebed 100644 --- a/temporalio/contrib/streams/_errors.py +++ b/temporalio/contrib/streams/_errors.py @@ -78,12 +78,9 @@ class StreamUnsupportedError(StreamError): class StreamConfigError(RuntimeError): - """No Stream plugin, or no backend of the name that a Stream or a call gives. + """No Stream plugin, or no backend of a given name. It is config, not a Store error. - It comes from the Worker or client config, not from a Store, so it is not a - :py:class:`StreamError`, and ``except StreamError`` does not catch it. In Workflow code it - fails the Workflow Task, so a Worker with the right config can take the run on. In an - Activity the attempt is retried. + ``except StreamError`` does not catch it. It fails a Workflow Task, and an Activity retries. """ 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 3ffa65aba..cda16568b 100644 --- a/temporalio/contrib/streams/_items.py +++ b/temporalio/contrib/streams/_items.py @@ -13,13 +13,12 @@ class End(enum.Enum): - """The type of :py:data:`END`.""" + """A ``from_offset`` after the newest Record, for new Records only. Use ``END``.""" END = "END" END = End.END -"""A ``from_offset`` that starts a read after the newest Record, for new Records only.""" @dataclass(frozen=True) diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index e9ce90d3e..c08998708 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -47,14 +47,6 @@ class ClientTopic(_AsyncTopic[T]): This API is experimental. """ - 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 super().publish(*values) - class ClientStream(StreamHandle[E]): """A Stream that a client writes, reads and closes. Get one from :py:class:`StreamClient`. @@ -92,7 +84,6 @@ def __init__( 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]] = [] - self._probing: asyncio.Future[None] | None = None def _bind( self, @@ -109,42 +100,27 @@ def _bind( @property def backend(self) -> str | None: - """Name of the backend that keeps the Stream. - - Raises: - RuntimeError: The handle looks in more than one backend, and no call has found the - Stream yet. Call :py:meth:`get_offset` first. - """ - if self._probe: - raise RuntimeError( - "The Stream's backend is not known yet. Call get_offset() first." - ) - return self._backend + """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: - if self._probe: - # Calls that start together share one search. - if self._probing is None: - self._probing = asyncio.ensure_future(self._find_backend()) - await asyncio.shield(self._probing) - - async def _find_backend(self) -> None: - client = self._client - assert client is not None - stream_id = await self._stream_id() - found = self._probe[0] + # 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(client.namespace).describe_stream( - models.DescribeStreamRequest(stream=models.StreamRef(id=stream_id)) - ) + await backend.service(self._client.namespace).describe_stream(request) except _errors.StreamNotFoundError: continue - found = (name, backend) - break - # With no backend that has the Stream, the first one waits for its create. - self._backend, self._stream_backend = found - self._probe = [] + 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: @@ -348,20 +324,6 @@ async def create( ) return stream._bind(self._client, stream_backend) - @overload - def get_handle( - self, - definition: type[E], - *, - id: str, - backend: str | tuple[str, ...] | None = None, - ) -> ClientStream[E]: ... - - @overload - def get_handle( - self, definition: type[E], *, stream: StreamHandle[Any] - ) -> ClientStream[E]: ... - def get_handle( self, definition: type[E], @@ -375,25 +337,25 @@ def get_handle( Args: definition: The Stream definition class. id: The Stream id. - stream: A handle of the Stream. Its id, owner and backend stay. A handle with no - backend name needs a client with one backend. + 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) if stream is not None: - id = stream.id + # 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 None: - raise TypeError("get_handle takes id= or stream=") - else: + 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=id, + 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, diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 40c4a9e04..5898ed55f 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -523,11 +523,11 @@ def _take_owned_streams( return state = streams_workflow._state() converter = temporalio.converter.PayloadConverter.default - for entry in converter.from_payload(payload, list[dict[str, Any]]): - run_stream = state.stream(_owned_ref(entry["id"])) - run_stream.owned = True - run_stream.backend = entry.get("backend") - run_stream.close_with_run = bool(entry.get("closeWithRun")) + 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 +573,11 @@ async def start_system_nexus_operation( return _DoneHandle(result) def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: - # The next run needs each owned Stream's backend for stream_handle, and whether it - # closes with the chain. + # The next run needs each owned Stream's backend for stream_handle. owned = [ - { - "id": stream_id, - "backend": run_stream.backend, - "closeWithRun": run_stream.close_with_run, - } + [stream_id, run_stream.backend, run_stream.close_with_run] for stream_id, run_stream in streams_workflow._state().streams.items() - if run_stream.owned and not run_stream.closed + if run_stream.backend is not None and not run_stream.closed ] if owned: converter = temporalio.converter.PayloadConverter.default diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index c02a4d223..43415b427 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -16,13 +16,13 @@ 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._resources.stream import Stream as _Stream @@ -47,7 +47,6 @@ class _RunStream: close_with_run: bool = False closed: bool = False # Set for a Stream that the run chain created, which continue-as-new passes on. - owned: bool = False backend: str | None = None @@ -117,29 +116,13 @@ def _now(call: Coroutine[Any, Any, _R]) -> _R: ) -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. @@ -378,7 +361,6 @@ def create_stream_nowait( _create_stream(stream_id=stream.id, retention=retention, auto_close=auto_close) ) run_stream.close_with_run = close_with_run - run_stream.owned = True run_stream.backend = stream.backend return stream diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 94a69f788..543a1c34d 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -238,6 +238,9 @@ async def test_handle_looks_in_each_backend_during_a_move(client: Client): 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")) From 9f5f5474ae6e4a361abfa288ac315bac7fe0c666 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:29:13 -0700 Subject: [PATCH 11/12] Closed and created through the generated client with its codec path. --- scripts/_wit/stream-service.wit | 36 ++- scripts/gen_stream_service.py | 2 +- temporalio/bridge/Cargo.lock | 2 +- temporalio/bridge/Cargo.toml | 2 +- temporalio/contrib/streams/_access.py | 22 +- temporalio/contrib/streams/_backend.py | 16 +- temporalio/contrib/streams/_convert.py | 42 +-- temporalio/contrib/streams/_reader.py | 8 +- temporalio/contrib/streams/_stream_client.py | 13 +- temporalio/contrib/streams/redis.py | 4 +- .../streams/system/_resources/stream.py | 6 + .../system/_system_nexus_interceptor.py | 8 +- temporalio/contrib/streams/system/client.py | 294 ++++++++++++++---- temporalio/contrib/streams/system/models.py | 49 ++- .../system/operations/create_stream.py | 17 +- temporalio/contrib/streams/system/services.py | 4 +- tests/contrib/streams/test_redis.py | 17 + 17 files changed, 387 insertions(+), 155 deletions(-) diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index dc3f40648..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,13 +281,15 @@ interface stream-service { auto-close: option, } - // A handle on one Stream. 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. nexgen has no doc directive - // on a resource yet, so the docs stay on the operations. + // 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, @@ -307,6 +300,7 @@ interface stream-service { 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, @@ -316,6 +310,7 @@ interface stream-service { 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, @@ -327,17 +322,24 @@ interface stream-service { 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 response. The created and status fields of the response are - // not on the handle. + // 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" - type stream-result = own<%stream>; + 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) -> stream-result; + 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." diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py index 5e27ca61b..292303a52 100644 --- a/scripts/gen_stream_service.py +++ b/scripts/gen_stream_service.py @@ -33,7 +33,7 @@ NEX_GEN_REPOSITORY = "https://github.com/moetemp/nexgen" # The head of moe/brook/v3/n8-stream-resource. A branch moves, so the pin is the commit. -NEX_GEN_REVISION = "aff8ff85450ebfbabe17fccb87b00acb7a97aad5" +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 0e8c8a810..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=760a53bf6958580fd9d3c2142deb75a9e286ff9d#760a53bf6958580fd9d3c2142deb75a9e286ff9d" +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 369fc5956..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 = "760a53bf6958580fd9d3c2142deb75a9e286ff9d" +rev = "2cc3eee3760a24b2feef30dc9471010cd49750e6" default-features = false optional = true diff --git a/temporalio/contrib/streams/_access.py b/temporalio/contrib/streams/_access.py index 09e904f51..e7601de54 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/contrib/streams/_access.py @@ -7,7 +7,6 @@ from datetime import timedelta 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 @@ -16,6 +15,7 @@ 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") @@ -68,7 +68,7 @@ async def _read( 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, @@ -99,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( diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 7b32e635b..53073f716 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -10,6 +10,7 @@ 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 @@ -48,9 +49,18 @@ def __init__(self, extension: temporalio.worker.WorkerExtension) -> None: """Create a backend over an extension.""" self._extension = extension - def service(self, namespace: str) -> service_client.StreamServiceClient: - """A generated Stream service client for ``namespace``.""" - return service_client.StreamServiceClient(_ExtensionTransport(self, namespace)) + def service( + self, + namespace: str, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> service_client.StreamServiceClient: + """A generated Stream service client for ``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.""" 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/_reader.py b/temporalio/contrib/streams/_reader.py index bf74a22f4..5fbb3f98e 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/contrib/streams/_reader.py @@ -71,9 +71,7 @@ def request() -> models.ReadRecordsRequest: try: while True: try: - # Only the service's own conversion keeps Record bodies raw. - with _convert.raw_bodies(converter.failure_converter): - record = await records.__anext__() + record = await records.__anext__() except StopAsyncIteration: return item = await _item(converter, record, types) @@ -111,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 c08998708..a6fe5f6e3 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -304,15 +304,10 @@ async def create( """ 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( diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 5898ed55f..9d42a9c98 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -557,7 +557,9 @@ async def start_system_nexus_operation( return await super().start_system_nexus_operation(input) operation = input.operation_name if operation == "CreateStream": - result: Any = models.StreamResult(stream=_owned_ref(input.input.stream_id)) + result: Any = models.CreateStreamResponse( + stream=_owned_ref(input.input.stream_id) + ) elif operation == "AppendRecords": result = models.AppendRecordsResponse() elif operation == "CloseStream": diff --git a/temporalio/contrib/streams/system/_resources/stream.py b/temporalio/contrib/streams/system/_resources/stream.py index 09d3884eb..457699b68 100644 --- a/temporalio/contrib/streams/system/_resources/stream.py +++ b/temporalio/contrib/streams/system/_resources/stream.py @@ -36,6 +36,8 @@ @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( @@ -47,6 +49,7 @@ async def append_records( 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, @@ -67,6 +70,7 @@ async def close( 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, @@ -89,6 +93,7 @@ async def read_records( 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, @@ -161,6 +166,7 @@ async def iter_read_records( async def describe( self, ) -> DescribeStreamResponse: + """Describe the Stream: its status, newest offset, Retention and auto-close.""" request = DescribeStreamRequest( stream=self.stream, ) diff --git a/temporalio/contrib/streams/system/_system_nexus_interceptor.py b/temporalio/contrib/streams/system/_system_nexus_interceptor.py index abd8b2fe9..44d8d4a6b 100644 --- a/temporalio/contrib/streams/system/_system_nexus_interceptor.py +++ b/temporalio/contrib/streams/system/_system_nexus_interceptor.py @@ -32,7 +32,7 @@ async def _start_system_nexus_operation( and input.operation_name == "CreateStream" ): typed_input = typing.cast( - "StartNexusOperationInput[models.CreateStreamRequest, models.StreamResult]", + "StartNexusOperationInput[models.CreateStreamRequest, models.CreateStreamResponse]", input, ) # The dispatch check above establishes that this operation's response type is _OutputT. @@ -111,7 +111,7 @@ async def start_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.StreamResult]: + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResponse]: """Intercept the CreateStream operation.""" return await self._next_system_nexus_interceptor().start_create_stream(request) @@ -151,7 +151,7 @@ async def _intercept_system_nexus_operation( async def start_create_stream( self, request: models.CreateStreamRequest - ) -> temporalio.workflow.NexusOperationHandle[models.StreamResult]: + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResponse]: from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput from temporalio.workflow import NexusOperationCancellationType @@ -162,7 +162,7 @@ async def start_create_stream( service="temporal.api.streamservice.v1.StreamService", operation="CreateStream", input=request, - output_type=models.StreamResult, + output_type=models.CreateStreamResponse, schedule_to_close_timeout=None, schedule_to_start_timeout=None, start_to_close_timeout=None, diff --git a/temporalio/contrib/streams/system/client.py b/temporalio/contrib/streams/system/client.py index d8b4b74ab..1dc9685ff 100644 --- a/temporalio/contrib/streams/system/client.py +++ b/temporalio/contrib/streams/system/client.py @@ -4,6 +4,7 @@ import asyncio import collections.abc +import contextlib import dataclasses import datetime import typing @@ -12,30 +13,36 @@ 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, CloseStreamResponse, CreateStreamRequest, + CreateStreamResponse, DescribeStreamRequest, DescribeStreamResponse, + Producer, ReadRecordsRequest, ReadRecordsResponse, StreamRecord, - StreamResult, + StreamRef, + WorkflowTaskEpoch, _AppendRecordsRequestTransferTypeConverter, _AppendRecordsResponseTransferTypeConverter, _CloseStreamRequestTransferTypeConverter, _CloseStreamResponseTransferTypeConverter, _CreateStreamRequestTransferTypeConverter, + _CreateStreamResponseTransferTypeConverter, _DescribeStreamRequestTransferTypeConverter, _DescribeStreamResponseTransferTypeConverter, _ReadRecordsRequestTransferTypeConverter, _ReadRecordsResponseTransferTypeConverter, - _StreamResultTransferTypeConverter, ) @@ -115,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. @@ -125,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) -> StreamResult: - payload = _proto_payload( - _CreateStreamRequestTransferTypeConverter().to_transfer_type(request) - ) + async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResponse: + 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 = _StreamResultTransferTypeConverter() - return converter.from_transfer_type( - _proto_message(answer, converter.transfer_type), - StreamResult, - ) + converter = _CreateStreamResponseTransferTypeConverter() + 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 @@ -174,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, @@ -291,12 +356,135 @@ async def iter_read_records( 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 988884c0e..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. @@ -86,13 +90,17 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class StreamResult: +class CreateStreamResponse: stream: StreamRef | None = None + created: bool | None = None + status: ( + temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType | None + ) = None -class _StreamResultTransferTypeConverter( +class _CreateStreamResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - StreamResult, + CreateStreamResponse, temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, ] ): @@ -105,20 +113,26 @@ class _StreamResultTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse, - type_hint: type[StreamResult], - ) -> StreamResult: - return StreamResult( + type_hint: type[CreateStreamResponse], + ) -> CreateStreamResponse: + return CreateStreamResponse( stream=_StreamRefTransferTypeConverter().from_transfer_type( value.stream, StreamRef ) if value.HasField("stream") else None, + created=value.created if bool(value.created) else None, + status=temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType( + value.status + ) + if value.status != 0 + else None, ) @typing_extensions.override def to_transfer_type( self, - value: StreamResult, + value: CreateStreamResponse, ) -> temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse: message = ( temporalio.api.streamservice.v1.stream_service_pb2.CreateStreamResponse() @@ -127,12 +141,18 @@ def to_transfer_type( message.stream.CopyFrom( _StreamRefTransferTypeConverter().to_transfer_type(value.stream) ) + if value.created is not None: + message.created = value.created + if value.status is not None: + message.status = temporalio.api.streamservice.v1.stream_service_pb2.StreamStatus.ValueType( + value.status + ) return message -temporalio.converter.transfer_type_convertible(_StreamResultTransferTypeConverter)( - StreamResult -) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible( + _CreateStreamResponseTransferTypeConverter +)(CreateStreamResponse) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) @@ -1232,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/create_stream.py b/temporalio/contrib/streams/system/operations/create_stream.py index f57f13897..ad7dd041b 100644 --- a/temporalio/contrib/streams/system/operations/create_stream.py +++ b/temporalio/contrib/streams/system/operations/create_stream.py @@ -9,13 +9,14 @@ from .._resources import Stream from ..models import ( CreateStreamRequest, - StreamResult, + CreateStreamResponse, + CreateStreamResult, ) async def _create_stream( request: CreateStreamRequest, -) -> Stream: +) -> CreateStreamResult: nexus_client = temporalio.workflow.create_nexus_client( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", @@ -23,13 +24,17 @@ async def _create_stream( handle = await nexus_client.start_operation( operation="CreateStream", input=request, - output_type=StreamResult, + output_type=CreateStreamResponse, ) result = await handle if result.stream is None: raise RuntimeError("CreateStream answered without `stream`") - return Stream( - stream=result.stream, + return CreateStreamResult( + stream=Stream( + stream=result.stream, + ), + created=result.created, + status=result.status, ) @@ -39,7 +44,7 @@ async def create_stream( retention: datetime.timedelta | None = None, withdraw_run: str | None = None, auto_close: datetime.timedelta | None = None, -) -> Stream: +) -> CreateStreamResult: """ Args: stream_id: Required from Workflow code. Empty elsewhere asks the service for an diff --git a/temporalio/contrib/streams/system/services.py b/temporalio/contrib/streams/system/services.py index 91e05b70b..f9e297ce5 100644 --- a/temporalio/contrib/streams/system/services.py +++ b/temporalio/contrib/streams/system/services.py @@ -10,11 +10,11 @@ CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, + CreateStreamResponse, DescribeStreamRequest, DescribeStreamResponse, ReadRecordsRequest, ReadRecordsResponse, - StreamResult, ) @@ -27,7 +27,7 @@ class StreamService: create_stream: Operation[ CreateStreamRequest, - StreamResult, + CreateStreamResponse, ] = Operation(name="CreateStream") append_records: Operation[ diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index d132fe86e..fca5dff48 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -400,6 +400,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: From 75890122b28c62a395cc5382394781dfa07671ce Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 01:43:22 -0700 Subject: [PATCH 12/12] Made StreamBackend internal and moved the plugin base into the package. --- temporalio/contrib/streams/__init__.py | 4 +- temporalio/contrib/streams/_backend.py | 88 ++++++++------------ temporalio/contrib/streams/_plugin.py | 86 +++++++++++++++++++ temporalio/contrib/streams/_stream_client.py | 3 +- temporalio/contrib/streams/redis.py | 63 ++++++-------- temporalio/contrib/streams/testing.py | 4 +- tests/contrib/streams/conftest.py | 3 +- tests/contrib/streams/test_facade.py | 10 ++- tests/contrib/streams/test_redis.py | 15 ++-- tests/contrib/streams/test_stream_client.py | 11 ++- 10 files changed, 181 insertions(+), 106 deletions(-) create mode 100644 temporalio/contrib/streams/_plugin.py diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 00e75d229..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, @@ -47,6 +46,7 @@ ) from temporalio.contrib.streams._handle import StreamHandle 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, @@ -64,7 +64,6 @@ "Closed", "ReadOnlyContextError", "Record", - "StreamBackend", "StreamClient", "StreamClosedError", "StreamConfigError", @@ -75,6 +74,7 @@ "StreamHandle", "StreamNotFoundError", "StreamOffsetError", + "StreamPlugin", "StreamStorageError", "StreamUnsupportedError", "Superseded", diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/contrib/streams/_backend.py index 53073f716..76740fb1e 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/contrib/streams/_backend.py @@ -3,13 +3,12 @@ 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 @@ -35,14 +34,43 @@ } -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``.""" + ... - 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 open_writer(self, request: ext_pb.OpenWriterRequest) -> int: + """Open the writer of an Activity or client Producer, and return its id.""" + ... - .. warning:: - This API is experimental. + 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.""" + ... + + 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.""" + ... + + +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: @@ -142,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 @@ -166,45 +194,3 @@ async def call( }, data=response.SerializeToString(), ) - - -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 = [ - ( - 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.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/_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/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index a6fe5f6e3..8ea7b681d 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -13,10 +13,11 @@ 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 diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 9d42a9c98..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__) @@ -47,12 +47,8 @@ _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, @@ -68,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 @@ -126,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 @@ -136,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: @@ -169,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 @@ -344,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 @@ -606,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 @@ -702,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() @@ -741,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/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/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/test_facade.py b/tests/contrib/streams/test_facade.py index c727c0aa8..2b11a5c04 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -11,10 +11,10 @@ from temporalio.common import RawValue from temporalio.contrib import streams 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 RPCError, RPCStatusCode from temporalio.worker import ExtensionError from tests.contrib.streams.fake_extension import FakeExtension @@ -28,8 +28,12 @@ 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 diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index fca5dff48..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 @@ -509,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()}" @@ -587,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 543a1c34d..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