diff --git a/changelog/added/babbling-brook.md b/changelog/added/babbling-brook.md new file mode 100644 index 0000000000..d038ff2b7a --- /dev/null +++ b/changelog/added/babbling-brook.md @@ -0,0 +1,3 @@ +Add experimental Streams to `temporalio.common`. A Stream is an ordered log of Records that Workflows, Activities and clients write, and that clients and Activities read. Register a backend with `stream_backends=` on `Client.connect`, `Replayer` or `WorkflowEnvironment`, then use `workflow.create_stream`, `Client.stream` or `WorkflowHandle.stream`. +Add a Redis Stream backend in `temporalio.contrib.stream`, and an in-memory one for tests in `temporalio.contrib.stream.testing`. The Redis backend needs a build with the bridge's `streams` feature, which the official builds have. +Add `activity.Info.workflow_task`, the Workflow Task attempt in which a Worker dispatched a Local Activity. diff --git a/changelog/changed/babbling-brook.md b/changelog/changed/babbling-brook.md new file mode 100644 index 0000000000..f22c0d8a31 --- /dev/null +++ b/changelog/changed/babbling-brook.md @@ -0,0 +1 @@ +Convert values of a `transfer_type_convertible` class inside JSON values too, such as list items and dataclass fields, not only at the top level of a payload. A dataclass with the decorator is still written as its fields inside a JSON value. diff --git a/pyproject.toml b/pyproject.toml index ebbdc02bcf..9f5333e595 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -186,9 +186,9 @@ exclude = [ 'temporalio/api', 'temporalio/bridge/proto', 'temporalio/nexus/system/workflow_service', - 'temporalio/contrib/streams/proto', - 'temporalio/contrib/streams/system', - 'temporalio/contrib/streams/writer_system', + 'temporalio/common/_streams/proto', + 'temporalio/common/_streams/system', + 'temporalio/common/_streams/writer_system', ] [[tool.mypy.overrides]] @@ -197,9 +197,9 @@ ignore_errors = true [[tool.mypy.overrides]] module = [ - "temporalio.contrib.streams.proto.*", - "temporalio.contrib.streams.system.*", - "temporalio.contrib.streams.writer_system.*", + "temporalio.common._streams.proto.*", + "temporalio.common._streams.system.*", + "temporalio.common._streams.writer_system.*", ] ignore_errors = true @@ -239,6 +239,8 @@ privacy = [ "HIDDEN:temporalio.worker.workflow_instance", "HIDDEN:temporalio.worker.workflow_sandbox.importer", "HIDDEN:temporalio.worker.workflow_sandbox.in_sandbox", + "HIDDEN:temporalio.common._streams.system.client", + "HIDDEN:temporalio.common._streams.writer_system.client", "HIDDEN:**.*_pb2*", ] project-name = "Temporal Python" @@ -272,7 +274,7 @@ exclude = [ "temporalio/api", "temporalio/bridge/proto", "temporalio/bridge/_visitor.py", - "temporalio/contrib/streams/proto", + "temporalio/common/_streams/proto", "tests/worker/workflow_sandbox/testmodules/proto", ] diff --git a/scripts/_proto/temporal/api/streamservice/v1/stream_service.proto b/scripts/_proto/temporal/api/streamservice/v1/stream_service.proto index 4a9e1468ca..9f047f7270 100644 --- a/scripts/_proto/temporal/api/streamservice/v1/stream_service.proto +++ b/scripts/_proto/temporal/api/streamservice/v1/stream_service.proto @@ -131,8 +131,7 @@ message StreamRecord { CloseStatus close_status = 9; } -// Idempotent for one id, one Retention and one auto_close. ALREADY_EXISTS for another -// Retention or another auto_close. +// Idempotent for one id and one Retention. ALREADY_EXISTS for another Retention. message CreateStreamRequest { string namespace = 1; // Required from Workflow code. @@ -141,11 +140,8 @@ message CreateStreamRequest { google.protobuf.Duration retention = 3; // Workflow only. A retry run's first create: the previous run, whose Records it withdraws. string withdraw_run = 4; - // Close the Stream with no result when no Record was appended for this long, counted from - // the create. Empty never closes on its own. It is separate from Retention and the safety - // TTL. The Store applies it on the next access, so no call sees the Stream open after the - // period. UNIMPLEMENTED (detail stream_unsupported) on a Store that can't. - google.protobuf.Duration auto_close = 5; + reserved 5; + reserved "auto_close"; // The Workflow run that owns the Stream. Workflow kind, with its run_id and epoch. Producer owner = 6; // Close the Stream with the result or the failure of the owner's run chain when the chain @@ -282,8 +278,8 @@ message DescribeStreamResponse { // Zero when the Stream has no Record. int64 last_offset = 2; google.protobuf.Duration retention = 3; - // The auto_close of the create. Empty when the Stream has none. - google.protobuf.Duration auto_close = 4; + reserved 4; + reserved "auto_close"; } // The details of a failed Stream call, in the gRPC status details and in the details of the diff --git a/scripts/_proto/temporal/sdk/streams/v1/extension.proto b/scripts/_proto/temporal/sdk/streams/v1/extension.proto index c61590419e..88d0a7cf90 100644 --- a/scripts/_proto/temporal/sdk/streams/v1/extension.proto +++ b/scripts/_proto/temporal/sdk/streams/v1/extension.proto @@ -8,7 +8,7 @@ package temporal.sdk.streams.v1; import "google/protobuf/duration.proto"; import "temporal/api/streamservice/v1/stream_service.proto"; -// The settings of the Redis Stream extension. Lang builds it from the plugin options. +// The settings of the Redis Stream extension. Lang builds it from the backend options. message RedisStreamConfig { // One URL, or the seed URLs of a cluster. Credentials and TLS come from the URL. repeated string urls = 1; @@ -96,7 +96,7 @@ message FreeWriterRequest { message FreeWriterResponse { } -// Method `CloseWriters`, asynchronous. Drops every open writer, for example when the plugin's +// Method `CloseWriters`, asynchronous. Drops every open writer, for example when the backend's // client or Worker shuts down. message CloseWritersRequest { } diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index 188860c9ae..0222790f7c 100644 --- a/scripts/_wit/stream-service.wit +++ b/scripts/_wit/stream-service.wit @@ -169,17 +169,15 @@ interface stream-service { close-status: option, } - /// @nexus.doc "Create a Stream idempotently. ALREADY_EXISTS for another Retention or auto-close." + /// @nexus.doc "Create a Stream idempotently. ALREADY_EXISTS for another Retention." /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamRequest" record create-stream-request { - /// @nexus.doc "Required from Workflow code. Empty elsewhere asks the service for an id." + /// @nexus.doc "Required on Redis and memory. On native, empty asks the server for an id." stream-id: string, /// @nexus.doc "The age after which a Record is dropped. Empty is the Store default." retention: option, /// @nexus.doc "Retry run only. The previous run, whose Records this withdraws." withdraw-run: option, - /// @nexus.doc "Close with no result after this long with no append. Empty never auto-closes." - auto-close: option, /// @nexus.doc "The Workflow run that owns the Stream. The Workflow transport sets it." owner: option, /// @nexus.doc "Close with the owner run chain's result when it closes. Native only." @@ -330,8 +328,6 @@ interface stream-service { /// @nexus.offset "stream-offset" last-offset: option, retention: option, - /// @nexus.doc "The auto-close of the create. Empty when the Stream has none." - auto-close: option, } // Its stream-ref is its serialized form, so a handle passes through Workflow input, Activity @@ -375,7 +371,7 @@ interface stream-service { create-wait: option, ) -> read-records-response; - /// @nexus.doc "Describe the Stream: its status, newest offset, Retention and auto-close." + /// @nexus.doc "Describe the Stream: its status, newest offset and Retention." /// @nexus.operation "describe-stream" describe: func() -> describe-stream-response; } diff --git a/scripts/gen_protos.py b/scripts/gen_protos.py index 1ecfd31c6d..4f6c2ad862 100644 --- a/scripts/gen_protos.py +++ b/scripts/gen_protos.py @@ -33,7 +33,7 @@ api_out_dir = base_dir / "temporalio" / "api" sdk_out_dir = base_dir / "temporalio" / "bridge" / "proto" -streams_out_dir = base_dir / "temporalio" / "contrib" / "streams" / "proto" +streams_out_dir = base_dir / "temporalio" / "common" / "_streams" / "proto" py_fixes = [ partial(re.compile(r"from temporal\.api\.").sub, r"from temporalio.api."), @@ -53,11 +53,11 @@ ), partial( re.compile(r"from temporal\.sdk\.streams\.").sub, - r"from temporalio.contrib.streams.proto.", + r"from temporalio.common._streams.proto.", ), partial( re.compile(r"'__module__' : 'temporal\.sdk\.streams\.").sub, - r"'__module__' : 'temporalio.contrib.streams.proto.", + r"'__module__' : 'temporalio.common._streams.proto.", ), partial( re.compile(r"'__module__' : 'temporal\.api\.").sub, @@ -82,7 +82,7 @@ partial(re.compile(r"temporal\.sdk\.core\.").sub, r"temporalio.bridge.proto."), partial( re.compile(r"temporal\.sdk\.streams\.").sub, - r"temporalio.contrib.streams.proto.", + r"temporalio.common._streams.proto.", ), ] @@ -231,7 +231,7 @@ def generate_protos(output_dir: Path): for p in (output_dir / "temporal" / "sdk" / "core").iterdir(): shutil.rmtree(sdk_out_dir / p.name, ignore_errors=True) p.replace(sdk_out_dir / p.name) - # The stream extension's own protos go with the contrib package that uses them + # Only the Streams package imports the stream extension protos streams_out_dir.mkdir(parents=True, exist_ok=True) for p in (output_dir / "temporal" / "sdk" / "streams").iterdir(): shutil.rmtree(streams_out_dir / p.name, ignore_errors=True) diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py index 473825f756..c5d5bbc306 100644 --- a/scripts/gen_stream_service.py +++ b/scripts/gen_stream_service.py @@ -1,10 +1,9 @@ -"""Generates the Stream service bindings of temporalio.contrib.streams. +"""Generates the Stream service bindings of the Streams API in temporalio.common. The WIT and the protos are copies of the temporalio-streams crate's, in scripts/_wit and scripts/_proto. Run scripts/gen_protos.py first for the proto modules. No nexgen release has the -offset, long-poll and client features yet, so the script installs nexgen at -NEX_GEN_REVISION from moetemp/nexgen. NEX_GEN_STREAMS_BIN may name a build of that revision -instead. +offset, long-poll and client features, so the script installs nexgen at NEX_GEN_REVISION from +moetemp/nexgen. Set NEX_GEN_STREAMS_BIN to use a local build of that revision instead. """ import os @@ -30,12 +29,12 @@ / "v1" / "stream_service.proto" ) -output_dir = base_dir / "temporalio" / "contrib" / "streams" / "system" -# The writer is client-only, so it has a caller and its resource class, and no Workflow code. -writer_output_dir = base_dir / "temporalio" / "contrib" / "streams" / "writer_system" +output_dir = base_dir / "temporalio" / "common" / "_streams" / "system" +# Only clients call the writer service, so its bindings have no Workflow code. +writer_output_dir = base_dir / "temporalio" / "common" / "_streams" / "writer_system" 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. +# A commit on the moe/brook/v3/n8-stream-resource branch. NEX_GEN_REVISION = "204b566ee996939b165ccd1d18c8c432bfdcc0dd" nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION @@ -81,14 +80,14 @@ def generate_stream_service() -> None: raise RuntimeError(f"missing generator input: {path}") with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir: descriptor_path = Path(temp_dir) / "stream_service.bin" - # The extension messages live in temporalio.contrib.streams.proto, so the descriptor - # names their file by that module path. + # The generated extension messages are in temporalio.common._streams.proto, so the + # descriptor must name their file by that module path. extension_root = Path(temp_dir) / "extension" extension_path = ( extension_root / "temporalio" - / "contrib" - / "streams" + / "common" + / "_streams" / "proto" / "v1" / "extension.proto" diff --git a/temporalio/activity.py b/temporalio/activity.py index 434f707bdf..712d0a30ee 100644 --- a/temporalio/activity.py +++ b/temporalio/activity.py @@ -112,7 +112,7 @@ class WorkflowTaskAttempt: """ID of the ``WorkflowTaskStarted`` event of the run's last completed task. 0 if none.""" failed_tasks: Sequence[int] = () - """``WorkflowTaskStarted`` event IDs of the run's failed and timed out tasks, the last 100.""" + """``WorkflowTaskStarted`` event IDs of the run's last 100 failed and timed out tasks.""" @dataclass(frozen=True) @@ -167,7 +167,7 @@ class Info: workflow_task: WorkflowTaskAttempt | None = None """For a local activity, the Workflow Task in which the Worker dispatched this execution. - Retries that the Worker runs without a new task keep it. None for other activities. + Retries that the Worker runs in the same task have the same value. None for other activities. .. warning:: This API is experimental. diff --git a/temporalio/api/streamservice/v1/stream_service_pb2.py b/temporalio/api/streamservice/v1/stream_service_pb2.py index 2c3894ef0c..19295a86d4 100644 --- a/temporalio/api/streamservice/v1/stream_service_pb2.py +++ b/temporalio/api/streamservice/v1/stream_service_pb2.py @@ -26,7 +26,7 @@ ) DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n2temporal/api/streamservice/v1/stream_service.proto\x12\x1dtemporal.api.streamservice.v1\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/api/failure/v1/message.proto"\x8b\x01\n\tStreamRef\x12\n\n\x02id\x18\x01 \x01(\t\x12\x13\n\x0bworkflow_id\x18\x02 \x01(\t\x12\x1e\n\x16\x66irst_execution_run_id\x18\x03 \x01(\t\x12\x0f\n\x07\x62\x61\x63kend\x18\x04 \x01(\t\x12\x13\n\x0b\x61\x63tivity_id\x18\x05 \x01(\t\x12\x17\n\x0f\x61\x63tivity_run_id\x18\x06 \x01(\t"\xc8\x01\n\rProducerEpoch\x12\x1d\n\x15task_started_event_id\x18\x01 \x01(\x03\x12\x14\n\x0ctask_attempt\x18\x02 \x01(\x05\x12\x35\n\x11task_started_time\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12:\n\x16\x61ttempt_scheduled_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x04 \x01(\x05"\x9e\x01\n\x08Producer\x12\x39\n\x04kind\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.ProducerKind\x12\n\n\x02id\x18\x02 \x01(\t\x12;\n\x05\x65poch\x18\x03 \x01(\x0b\x32,.temporal.api.streamservice.v1.ProducerEpoch\x12\x0e\n\x06run_id\x18\x04 \x01(\t"p\n\x11WorkflowTaskEpoch\x12\x18\n\x10started_event_id\x18\x01 \x01(\x03\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x05\x12\x30\n\x0cstarted_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"L\n\x0c\x41ppendRecord\x12\r\n\x05topic\x18\x01 \x01(\t\x12-\n\x04\x62ody\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload"\xb3\x02\n\x0cSupersession\x12\x42\n\rproducer_kind\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.ProducerKind\x12\x13\n\x0bproducer_id\x18\x02 \x01(\t\x12;\n\x05\x65poch\x18\x03 \x01(\x0b\x32,.temporal.api.streamservice.v1.ProducerEpoch\x12Q\n\x17previous_completed_task\x18\x04 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x05 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x06 \x01(\t\x12\x0e\n\x06run_id\x18\x07 \x01(\t"\xcb\x03\n\x0cStreamRecord\x12=\n\x04kind\x18\x01 \x01(\x0e\x32/.temporal.api.streamservice.v1.StreamRecordKind\x12\x0e\n\x06offset\x18\x02 \x01(\x03\x12\r\n\x05topic\x18\x03 \x01(\t\x12\x39\n\x08producer\x18\x04 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12-\n\x04\x62ody\x18\x05 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x41\n\x0csupersession\x18\x06 \x01(\x0b\x32+.temporal.api.streamservice.v1.Supersession\x12\x35\n\x0c\x63lose_result\x18\x07 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x37\n\rclose_failure\x18\x08 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12@\n\x0c\x63lose_status\x18\t \x01(\x0e\x32*.temporal.api.streamservice.v1.CloseStatus"\xc9\x02\n\x13\x43reateStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\tretention\x18\x03 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x14\n\x0cwithdraw_run\x18\x04 \x01(\t\x12-\n\nauto_close\x18\x05 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x36\n\x05owner\x18\x06 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12\x16\n\x0e\x63lose_with_run\x18\x07 \x01(\x08\x12\x35\n\x03ref\x18\x08 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x12\n\ntask_token\x18\t \x01(\x0c"\x9e\x01\n\x14\x43reateStreamResponse\x12\x38\n\x06stream\x18\x01 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12;\n\x06status\x18\x03 \x01(\x0e\x32+.temporal.api.streamservice.v1.StreamStatus"\x9b\x03\n\x14\x41ppendRecordsRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x39\n\x08producer\x18\x03 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12\x10\n\x08sequence\x18\x04 \x01(\x03\x12<\n\x07records\x18\x05 \x03(\x0b\x32+.temporal.api.streamservice.v1.AppendRecord\x12Q\n\x17previous_completed_task\x18\x06 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x07 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x08 \x01(\t\x12\x18\n\x10\x64\x65\x66\x65r_withdrawal\x18\t \x01(\x08\x12\x12\n\ntask_token\x18\n \x01(\x0c"B\n\x15\x41ppendRecordsResponse\x12\x14\n\x0c\x66irst_offset\x18\x01 \x01(\x03\x12\x13\n\x0blast_offset\x18\x02 \x01(\x03"\xcf\x03\n\x12\x43loseStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x39\n\x08producer\x18\x03 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12/\n\x06result\x18\x04 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x31\n\x07\x66\x61ilure\x18\x05 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12Q\n\x17previous_completed_task\x18\x06 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x07 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x08 \x01(\t\x12:\n\x06status\x18\t \x01(\x0e\x32*.temporal.api.streamservice.v1.CloseStatus\x12\x12\n\ntask_token\x18\n \x01(\x0c"\x15\n\x13\x43loseStreamResponse"\xb4\x03\n\x12ReadRecordsRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x13\n\x0b\x66rom_offset\x18\x03 \x01(\x03\x12\x37\n\x05start\x18\x04 \x01(\x0e\x32(.temporal.api.streamservice.v1.ReadStart\x12\x0e\n\x06topics\x18\x05 \x03(\t\x12\x13\n\x0bmax_records\x18\x06 \x01(\x05\x12\'\n\x04wait\x18\x07 \x01(\x0b\x32\x19.google.protobuf.Duration\x12.\n\x0b\x63reate_wait\x18\x08 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x16\n\x0etopic_prefixes\x18\n \x03(\t\x12\x11\n\tto_offset\x18\x0b \x01(\x03\x12\x11\n\twake_only\x18\x0c \x01(\x08\x12\x0c\n\x04last\x18\r \x01(\x05\x12\x11\n\tmax_bytes\x18\x0e \x01(\x03\x12\x12\n\nfloor_from\x18\x0f \x01(\x03J\x04\x08\t\x10\nR\x0cidle_timeout"\x99\x01\n\x13ReadRecordsResponse\x12<\n\x07records\x18\x01 \x03(\x0b\x32+.temporal.api.streamservice.v1.StreamRecord\x12\x13\n\x0bnext_offset\x18\x02 \x01(\x03\x12\x0c\n\x04\x64one\x18\x03 \x01(\x08\x12!\n\x19\x63ontinue_as_new_suggested\x18\x04 \x01(\x08"d\n\x15\x44\x65scribeStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef"\xc7\x01\n\x16\x44\x65scribeStreamResponse\x12;\n\x06status\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.StreamStatus\x12\x13\n\x0blast_offset\x18\x02 \x01(\x03\x12,\n\tretention\x18\x03 \x01(\x0b\x32\x19.google.protobuf.Duration\x12-\n\nauto_close\x18\x04 \x01(\x0b\x32\x19.google.protobuf.Duration"$\n\x12StreamErrorDetails\x12\x0e\n\x06reason\x18\x01 \x01(\t*\xc8\x01\n\x0cProducerKind\x12\x1d\n\x19PRODUCER_KIND_UNSPECIFIED\x10\x00\x12\x1a\n\x16PRODUCER_KIND_WORKFLOW\x10\x01\x12 \n\x1cPRODUCER_KIND_LOCAL_ACTIVITY\x10\x02\x12\x1a\n\x16PRODUCER_KIND_ACTIVITY\x10\x03\x12\x18\n\x14PRODUCER_KIND_CLIENT\x10\x04\x12%\n!PRODUCER_KIND_STANDALONE_ACTIVITY\x10\x05*\xb3\x01\n\x10StreamRecordKind\x12"\n\x1eSTREAM_RECORD_KIND_UNSPECIFIED\x10\x00\x12\x1b\n\x17STREAM_RECORD_KIND_DATA\x10\x01\x12!\n\x1dSTREAM_RECORD_KIND_SUPERSEDED\x10\x02\x12\x1c\n\x18STREAM_RECORD_KIND_CLOSE\x10\x03"\x04\x08\x04\x10\x04*\x17STREAM_RECORD_KIND_IDLE*\xb4\x01\n\x0b\x43loseStatus\x12\x1c\n\x18\x43LOSE_STATUS_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43LOSE_STATUS_COMPLETED\x10\x01\x12\x17\n\x13\x43LOSE_STATUS_FAILED\x10\x02\x12\x19\n\x15\x43LOSE_STATUS_CANCELED\x10\x03\x12\x1b\n\x17\x43LOSE_STATUS_TERMINATED\x10\x04\x12\x1a\n\x16\x43LOSE_STATUS_TIMED_OUT\x10\x05*z\n\x0cStreamStatus\x12\x1d\n\x19STREAM_STATUS_UNSPECIFIED\x10\x00\x12\x16\n\x12STREAM_STATUS_OPEN\x10\x01\x12\x18\n\x14STREAM_STATUS_CLOSED\x10\x02\x12\x19\n\x15STREAM_STATUS_EXPIRED\x10\x03*U\n\tReadStart\x12\x1a\n\x16READ_START_UNSPECIFIED\x10\x00\x12\x18\n\x14READ_START_BEGINNING\x10\x01\x12\x12\n\x0eREAD_START_END\x10\x02\x42\xb2\x01\n io.temporal.api.streamservice.v1B\x12StreamServiceProtoP\x01Z1go.temporal.io/api/streamservice/v1;streamservice\xaa\x02\x1fTemporalio.Api.StreamService.V1\xea\x02"Temporalio::Api::StreamService::V1b\x06proto3' + b'\n2temporal/api/streamservice/v1/stream_service.proto\x12\x1dtemporal.api.streamservice.v1\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/api/failure/v1/message.proto"\x8b\x01\n\tStreamRef\x12\n\n\x02id\x18\x01 \x01(\t\x12\x13\n\x0bworkflow_id\x18\x02 \x01(\t\x12\x1e\n\x16\x66irst_execution_run_id\x18\x03 \x01(\t\x12\x0f\n\x07\x62\x61\x63kend\x18\x04 \x01(\t\x12\x13\n\x0b\x61\x63tivity_id\x18\x05 \x01(\t\x12\x17\n\x0f\x61\x63tivity_run_id\x18\x06 \x01(\t"\xc8\x01\n\rProducerEpoch\x12\x1d\n\x15task_started_event_id\x18\x01 \x01(\x03\x12\x14\n\x0ctask_attempt\x18\x02 \x01(\x05\x12\x35\n\x11task_started_time\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12:\n\x16\x61ttempt_scheduled_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x04 \x01(\x05"\x9e\x01\n\x08Producer\x12\x39\n\x04kind\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.ProducerKind\x12\n\n\x02id\x18\x02 \x01(\t\x12;\n\x05\x65poch\x18\x03 \x01(\x0b\x32,.temporal.api.streamservice.v1.ProducerEpoch\x12\x0e\n\x06run_id\x18\x04 \x01(\t"p\n\x11WorkflowTaskEpoch\x12\x18\n\x10started_event_id\x18\x01 \x01(\x03\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x05\x12\x30\n\x0cstarted_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"L\n\x0c\x41ppendRecord\x12\r\n\x05topic\x18\x01 \x01(\t\x12-\n\x04\x62ody\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload"\xb3\x02\n\x0cSupersession\x12\x42\n\rproducer_kind\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.ProducerKind\x12\x13\n\x0bproducer_id\x18\x02 \x01(\t\x12;\n\x05\x65poch\x18\x03 \x01(\x0b\x32,.temporal.api.streamservice.v1.ProducerEpoch\x12Q\n\x17previous_completed_task\x18\x04 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x05 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x06 \x01(\t\x12\x0e\n\x06run_id\x18\x07 \x01(\t"\xcb\x03\n\x0cStreamRecord\x12=\n\x04kind\x18\x01 \x01(\x0e\x32/.temporal.api.streamservice.v1.StreamRecordKind\x12\x0e\n\x06offset\x18\x02 \x01(\x03\x12\r\n\x05topic\x18\x03 \x01(\t\x12\x39\n\x08producer\x18\x04 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12-\n\x04\x62ody\x18\x05 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x41\n\x0csupersession\x18\x06 \x01(\x0b\x32+.temporal.api.streamservice.v1.Supersession\x12\x35\n\x0c\x63lose_result\x18\x07 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x37\n\rclose_failure\x18\x08 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12@\n\x0c\x63lose_status\x18\t \x01(\x0e\x32*.temporal.api.streamservice.v1.CloseStatus"\xac\x02\n\x13\x43reateStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\tretention\x18\x03 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x14\n\x0cwithdraw_run\x18\x04 \x01(\t\x12\x36\n\x05owner\x18\x06 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12\x16\n\x0e\x63lose_with_run\x18\x07 \x01(\x08\x12\x35\n\x03ref\x18\x08 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x12\n\ntask_token\x18\t \x01(\x0cJ\x04\x08\x05\x10\x06R\nauto_close"\x9e\x01\n\x14\x43reateStreamResponse\x12\x38\n\x06stream\x18\x01 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12;\n\x06status\x18\x03 \x01(\x0e\x32+.temporal.api.streamservice.v1.StreamStatus"\x9b\x03\n\x14\x41ppendRecordsRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x39\n\x08producer\x18\x03 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12\x10\n\x08sequence\x18\x04 \x01(\x03\x12<\n\x07records\x18\x05 \x03(\x0b\x32+.temporal.api.streamservice.v1.AppendRecord\x12Q\n\x17previous_completed_task\x18\x06 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x07 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x08 \x01(\t\x12\x18\n\x10\x64\x65\x66\x65r_withdrawal\x18\t \x01(\x08\x12\x12\n\ntask_token\x18\n \x01(\x0c"B\n\x15\x41ppendRecordsResponse\x12\x14\n\x0c\x66irst_offset\x18\x01 \x01(\x03\x12\x13\n\x0blast_offset\x18\x02 \x01(\x03"\xcf\x03\n\x12\x43loseStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x39\n\x08producer\x18\x03 \x01(\x0b\x32\'.temporal.api.streamservice.v1.Producer\x12/\n\x06result\x18\x04 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x31\n\x07\x66\x61ilure\x18\x05 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12Q\n\x17previous_completed_task\x18\x06 \x01(\x0b\x32\x30.temporal.api.streamservice.v1.WorkflowTaskEpoch\x12\x14\n\x0c\x66\x61iled_tasks\x18\x07 \x03(\x03\x12\x14\n\x0cwithdraw_run\x18\x08 \x01(\t\x12:\n\x06status\x18\t \x01(\x0e\x32*.temporal.api.streamservice.v1.CloseStatus\x12\x12\n\ntask_token\x18\n \x01(\x0c"\x15\n\x13\x43loseStreamResponse"\xb4\x03\n\x12ReadRecordsRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef\x12\x13\n\x0b\x66rom_offset\x18\x03 \x01(\x03\x12\x37\n\x05start\x18\x04 \x01(\x0e\x32(.temporal.api.streamservice.v1.ReadStart\x12\x0e\n\x06topics\x18\x05 \x03(\t\x12\x13\n\x0bmax_records\x18\x06 \x01(\x05\x12\'\n\x04wait\x18\x07 \x01(\x0b\x32\x19.google.protobuf.Duration\x12.\n\x0b\x63reate_wait\x18\x08 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x16\n\x0etopic_prefixes\x18\n \x03(\t\x12\x11\n\tto_offset\x18\x0b \x01(\x03\x12\x11\n\twake_only\x18\x0c \x01(\x08\x12\x0c\n\x04last\x18\r \x01(\x05\x12\x11\n\tmax_bytes\x18\x0e \x01(\x03\x12\x12\n\nfloor_from\x18\x0f \x01(\x03J\x04\x08\t\x10\nR\x0cidle_timeout"\x99\x01\n\x13ReadRecordsResponse\x12<\n\x07records\x18\x01 \x03(\x0b\x32+.temporal.api.streamservice.v1.StreamRecord\x12\x13\n\x0bnext_offset\x18\x02 \x01(\x03\x12\x0c\n\x04\x64one\x18\x03 \x01(\x08\x12!\n\x19\x63ontinue_as_new_suggested\x18\x04 \x01(\x08"d\n\x15\x44\x65scribeStreamRequest\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x38\n\x06stream\x18\x02 \x01(\x0b\x32(.temporal.api.streamservice.v1.StreamRef"\xaa\x01\n\x16\x44\x65scribeStreamResponse\x12;\n\x06status\x18\x01 \x01(\x0e\x32+.temporal.api.streamservice.v1.StreamStatus\x12\x13\n\x0blast_offset\x18\x02 \x01(\x03\x12,\n\tretention\x18\x03 \x01(\x0b\x32\x19.google.protobuf.DurationJ\x04\x08\x04\x10\x05R\nauto_close"$\n\x12StreamErrorDetails\x12\x0e\n\x06reason\x18\x01 \x01(\t*\xc8\x01\n\x0cProducerKind\x12\x1d\n\x19PRODUCER_KIND_UNSPECIFIED\x10\x00\x12\x1a\n\x16PRODUCER_KIND_WORKFLOW\x10\x01\x12 \n\x1cPRODUCER_KIND_LOCAL_ACTIVITY\x10\x02\x12\x1a\n\x16PRODUCER_KIND_ACTIVITY\x10\x03\x12\x18\n\x14PRODUCER_KIND_CLIENT\x10\x04\x12%\n!PRODUCER_KIND_STANDALONE_ACTIVITY\x10\x05*\xb3\x01\n\x10StreamRecordKind\x12"\n\x1eSTREAM_RECORD_KIND_UNSPECIFIED\x10\x00\x12\x1b\n\x17STREAM_RECORD_KIND_DATA\x10\x01\x12!\n\x1dSTREAM_RECORD_KIND_SUPERSEDED\x10\x02\x12\x1c\n\x18STREAM_RECORD_KIND_CLOSE\x10\x03"\x04\x08\x04\x10\x04*\x17STREAM_RECORD_KIND_IDLE*\xb4\x01\n\x0b\x43loseStatus\x12\x1c\n\x18\x43LOSE_STATUS_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43LOSE_STATUS_COMPLETED\x10\x01\x12\x17\n\x13\x43LOSE_STATUS_FAILED\x10\x02\x12\x19\n\x15\x43LOSE_STATUS_CANCELED\x10\x03\x12\x1b\n\x17\x43LOSE_STATUS_TERMINATED\x10\x04\x12\x1a\n\x16\x43LOSE_STATUS_TIMED_OUT\x10\x05*z\n\x0cStreamStatus\x12\x1d\n\x19STREAM_STATUS_UNSPECIFIED\x10\x00\x12\x16\n\x12STREAM_STATUS_OPEN\x10\x01\x12\x18\n\x14STREAM_STATUS_CLOSED\x10\x02\x12\x19\n\x15STREAM_STATUS_EXPIRED\x10\x03*U\n\tReadStart\x12\x1a\n\x16READ_START_UNSPECIFIED\x10\x00\x12\x18\n\x14READ_START_BEGINNING\x10\x01\x12\x12\n\x0eREAD_START_END\x10\x02\x42\xb2\x01\n io.temporal.api.streamservice.v1B\x12StreamServiceProtoP\x01Z1go.temporal.io/api/streamservice/v1;streamservice\xaa\x02\x1fTemporalio.Api.StreamService.V1\xea\x02"Temporalio::Api::StreamService::V1b\x06proto3' ) _PRODUCERKIND = DESCRIPTOR.enum_types_by_name["ProducerKind"] @@ -283,16 +283,16 @@ if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None DESCRIPTOR._serialized_options = b'\n io.temporal.api.streamservice.v1B\022StreamServiceProtoP\001Z1go.temporal.io/api/streamservice/v1;streamservice\252\002\037Temporalio.Api.StreamService.V1\352\002"Temporalio::Api::StreamService::V1' - _PRODUCERKIND._serialized_start = 4099 - _PRODUCERKIND._serialized_end = 4299 - _STREAMRECORDKIND._serialized_start = 4302 - _STREAMRECORDKIND._serialized_end = 4481 - _CLOSESTATUS._serialized_start = 4484 - _CLOSESTATUS._serialized_end = 4664 - _STREAMSTATUS._serialized_start = 4666 - _STREAMSTATUS._serialized_end = 4788 - _READSTART._serialized_start = 4790 - _READSTART._serialized_end = 4875 + _PRODUCERKIND._serialized_start = 4041 + _PRODUCERKIND._serialized_end = 4241 + _STREAMRECORDKIND._serialized_start = 4244 + _STREAMRECORDKIND._serialized_end = 4423 + _CLOSESTATUS._serialized_start = 4426 + _CLOSESTATUS._serialized_end = 4606 + _STREAMSTATUS._serialized_start = 4608 + _STREAMSTATUS._serialized_end = 4730 + _READSTART._serialized_start = 4732 + _READSTART._serialized_end = 4817 _STREAMREF._serialized_start = 228 _STREAMREF._serialized_end = 367 _PRODUCEREPOCH._serialized_start = 370 @@ -308,25 +308,25 @@ _STREAMRECORD._serialized_start = 1236 _STREAMRECORD._serialized_end = 1695 _CREATESTREAMREQUEST._serialized_start = 1698 - _CREATESTREAMREQUEST._serialized_end = 2027 - _CREATESTREAMRESPONSE._serialized_start = 2030 - _CREATESTREAMRESPONSE._serialized_end = 2188 - _APPENDRECORDSREQUEST._serialized_start = 2191 - _APPENDRECORDSREQUEST._serialized_end = 2602 - _APPENDRECORDSRESPONSE._serialized_start = 2604 - _APPENDRECORDSRESPONSE._serialized_end = 2670 - _CLOSESTREAMREQUEST._serialized_start = 2673 - _CLOSESTREAMREQUEST._serialized_end = 3136 - _CLOSESTREAMRESPONSE._serialized_start = 3138 - _CLOSESTREAMRESPONSE._serialized_end = 3159 - _READRECORDSREQUEST._serialized_start = 3162 - _READRECORDSREQUEST._serialized_end = 3598 - _READRECORDSRESPONSE._serialized_start = 3601 - _READRECORDSRESPONSE._serialized_end = 3754 - _DESCRIBESTREAMREQUEST._serialized_start = 3756 - _DESCRIBESTREAMREQUEST._serialized_end = 3856 - _DESCRIBESTREAMRESPONSE._serialized_start = 3859 - _DESCRIBESTREAMRESPONSE._serialized_end = 4058 - _STREAMERRORDETAILS._serialized_start = 4060 - _STREAMERRORDETAILS._serialized_end = 4096 + _CREATESTREAMREQUEST._serialized_end = 1998 + _CREATESTREAMRESPONSE._serialized_start = 2001 + _CREATESTREAMRESPONSE._serialized_end = 2159 + _APPENDRECORDSREQUEST._serialized_start = 2162 + _APPENDRECORDSREQUEST._serialized_end = 2573 + _APPENDRECORDSRESPONSE._serialized_start = 2575 + _APPENDRECORDSRESPONSE._serialized_end = 2641 + _CLOSESTREAMREQUEST._serialized_start = 2644 + _CLOSESTREAMREQUEST._serialized_end = 3107 + _CLOSESTREAMRESPONSE._serialized_start = 3109 + _CLOSESTREAMRESPONSE._serialized_end = 3130 + _READRECORDSREQUEST._serialized_start = 3133 + _READRECORDSREQUEST._serialized_end = 3569 + _READRECORDSRESPONSE._serialized_start = 3572 + _READRECORDSRESPONSE._serialized_end = 3725 + _DESCRIBESTREAMREQUEST._serialized_start = 3727 + _DESCRIBESTREAMREQUEST._serialized_end = 3827 + _DESCRIBESTREAMRESPONSE._serialized_start = 3830 + _DESCRIBESTREAMRESPONSE._serialized_end = 4000 + _STREAMERRORDETAILS._serialized_start = 4002 + _STREAMERRORDETAILS._serialized_end = 4038 # @@protoc_insertion_point(module_scope) diff --git a/temporalio/api/streamservice/v1/stream_service_pb2.pyi b/temporalio/api/streamservice/v1/stream_service_pb2.pyi index d664aba13f..5066333fe4 100644 --- a/temporalio/api/streamservice/v1/stream_service_pb2.pyi +++ b/temporalio/api/streamservice/v1/stream_service_pb2.pyi @@ -517,9 +517,7 @@ class StreamRecord(google.protobuf.message.Message): global___StreamRecord = StreamRecord class CreateStreamRequest(google.protobuf.message.Message): - """Idempotent for one id, one Retention and one auto_close. ALREADY_EXISTS for another - Retention or another auto_close. - """ + """Idempotent for one id and one Retention. ALREADY_EXISTS for another Retention.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @@ -527,7 +525,6 @@ class CreateStreamRequest(google.protobuf.message.Message): STREAM_ID_FIELD_NUMBER: builtins.int RETENTION_FIELD_NUMBER: builtins.int WITHDRAW_RUN_FIELD_NUMBER: builtins.int - AUTO_CLOSE_FIELD_NUMBER: builtins.int OWNER_FIELD_NUMBER: builtins.int CLOSE_WITH_RUN_FIELD_NUMBER: builtins.int REF_FIELD_NUMBER: builtins.int @@ -541,13 +538,6 @@ class CreateStreamRequest(google.protobuf.message.Message): withdraw_run: builtins.str """Workflow only. A retry run's first create: the previous run, whose Records it withdraws.""" @property - def auto_close(self) -> google.protobuf.duration_pb2.Duration: - """Close the Stream with no result when no Record was appended for this long, counted from - the create. Empty never closes on its own. It is separate from Retention and the safety - TTL. The Store applies it on the next access, so no call sees the Stream open after the - period. UNIMPLEMENTED (detail stream_unsupported) on a Store that can't. - """ - @property def owner(self) -> global___Producer: """The Workflow run that owns the Stream. Workflow kind, with its run_id and epoch.""" close_with_run: builtins.bool @@ -570,7 +560,6 @@ class CreateStreamRequest(google.protobuf.message.Message): stream_id: builtins.str = ..., retention: google.protobuf.duration_pb2.Duration | None = ..., withdraw_run: builtins.str = ..., - auto_close: google.protobuf.duration_pb2.Duration | None = ..., owner: global___Producer | None = ..., close_with_run: builtins.bool = ..., ref: global___StreamRef | None = ..., @@ -579,21 +568,12 @@ class CreateStreamRequest(google.protobuf.message.Message): def HasField( self, field_name: typing_extensions.Literal[ - "auto_close", - b"auto_close", - "owner", - b"owner", - "ref", - b"ref", - "retention", - b"retention", + "owner", b"owner", "ref", b"ref", "retention", b"retention" ], ) -> builtins.bool: ... def ClearField( self, field_name: typing_extensions.Literal[ - "auto_close", - b"auto_close", "close_with_run", b"close_with_run", "namespace", @@ -1072,34 +1052,24 @@ class DescribeStreamResponse(google.protobuf.message.Message): STATUS_FIELD_NUMBER: builtins.int LAST_OFFSET_FIELD_NUMBER: builtins.int RETENTION_FIELD_NUMBER: builtins.int - AUTO_CLOSE_FIELD_NUMBER: builtins.int status: global___StreamStatus.ValueType last_offset: builtins.int """Zero when the Stream has no Record.""" @property def retention(self) -> google.protobuf.duration_pb2.Duration: ... - @property - def auto_close(self) -> google.protobuf.duration_pb2.Duration: - """The auto_close of the create. Empty when the Stream has none.""" def __init__( self, *, status: global___StreamStatus.ValueType = ..., last_offset: builtins.int = ..., retention: google.protobuf.duration_pb2.Duration | None = ..., - auto_close: google.protobuf.duration_pb2.Duration | None = ..., ) -> None: ... def HasField( - self, - field_name: typing_extensions.Literal[ - "auto_close", b"auto_close", "retention", b"retention" - ], + self, field_name: typing_extensions.Literal["retention", b"retention"] ) -> builtins.bool: ... def ClearField( self, field_name: typing_extensions.Literal[ - "auto_close", - b"auto_close", "last_offset", b"last_offset", "retention", diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index 67939bd6cd..6591afb5a3 100644 --- a/temporalio/bridge/Cargo.lock +++ b/temporalio/bridge/Cargo.lock @@ -2938,7 +2938,7 @@ dependencies = [ [[package]] name = "temporalio-streams" version = "0.1.0" -source = "git+https://github.com/moetemp/sdk-streams?rev=879563a2f47212c10052423dbdfb15f44bcbeef8#879563a2f47212c10052423dbdfb15f44bcbeef8" +source = "git+https://github.com/moetemp/sdk-streams?rev=ee6d1e51d91c96f0c661fb59ad0957a97a579cbe#ee6d1e51d91c96f0c661fb59ad0957a97a579cbe" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 22a5ce272c..e103624418 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -39,9 +39,8 @@ tls-aws-lc = [ "tokio-rustls/aws_lc_rs", ] fips = ["tls-aws-lc", "dep:aws-lc-rs", "streams", "temporalio-streams?/fips"] -# Links the Stream runtime, which a Stream plugin builds for each client. No stream code runs -# unless a plugin builds it. The official builds turn it on. It is off by default, since its git -# dependencies would break offline builds of the sdist. +# Off by default, since the git dependencies of temporalio-streams break offline sdist builds. +# The official builds turn it on. No stream code runs unless a Stream plugin builds a runtime. streams = ["dep:temporalio-streams"] [dependencies] @@ -90,11 +89,11 @@ tonic = "0.14" tracing = "0.1" url = "2.2" -# The rev moves with the sdk-core pin, and the patch below makes the crate use the submodule's -# crates, so the bridge sees one type of each. +# Update this rev with the sdk-core pin. The patch below makes this crate use the submodule's +# crates, so the bridge sees one copy of each type. [dependencies.temporalio-streams] git = "https://github.com/moetemp/sdk-streams" -rev = "879563a2f47212c10052423dbdfb15f44bcbeef8" +rev = "ee6d1e51d91c96f0c661fb59ad0957a97a579cbe" default-features = false optional = true diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index b565c43354..f782636659 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit b565c433547496f2c14a911549c288a1076e72b6 +Subproject commit f782636659f25f2765260c9d061c9165afa34b3a diff --git a/temporalio/bridge/src/streams.rs b/temporalio/bridge/src/streams.rs index 789b093f45..c2dde7255c 100644 --- a/temporalio/bridge/src/streams.rs +++ b/temporalio/bridge/src/streams.rs @@ -1,5 +1,5 @@ -//! The Stream runtime of the temporalio-streams crate, when this build links it. A Stream plugin -//! builds one for each client from its config, and Python calls it by method name. +//! Bindings to the Stream runtime of the temporalio-streams crate. A Stream plugin builds one +//! runtime for each client from its config. Python calls it by method name. use pyo3::exceptions::PyException; #[cfg(feature = "streams")] @@ -8,10 +8,10 @@ use pyo3::prelude::*; #[cfg(feature = "streams")] use crate::runtime; -// Its args are the gRPC code, the message and the encoded details, as with `RPCError`. +// The exception args are the gRPC code, the message and the encoded details, as in `RPCError`. pyo3::create_exception!(temporal_sdk_bridge, StreamCallError, PyException); -/// The crate's runtime object. Python keeps it on the plugin for the client that it serves. +/// The crate's runtime object. Python keeps it with the client that it serves. #[cfg(feature = "streams")] #[pyclass(frozen)] pub struct StreamRuntimeRef { @@ -41,8 +41,8 @@ impl StreamRuntimeRef { }) } - /// Releases one handle of a writer with no wait. It needs no runtime, so a finalizer of a - /// handle that Python drops can call it. + /// Releases one handle of a writer without waiting. It needs no Tokio runtime, so a Python + /// finalizer can call it. fn free_writer(&self, writer_id: u64) -> PyResult<()> { self.streams.free_writer(writer_id).map_err(call_error) } @@ -64,7 +64,7 @@ pub fn new_stream_runtime( let runtime = runtime_ref.runtime.clone(); runtime.assert_same_process("create a Stream runtime")?; - // The runtime may start tasks on the Tokio runtime that it lives on. + // Building the Stream runtime can spawn Tokio tasks, which need an entered Tokio runtime. let _guard = runtime.core.tokio_handle().enter(); let meter = runtime.core.telemetry().get_temporal_metric_meter(); let streams = if memory { diff --git a/temporalio/bridge/streams.py b/temporalio/bridge/streams.py index 23097bc503..4486ea32ff 100644 --- a/temporalio/bridge/streams.py +++ b/temporalio/bridge/streams.py @@ -31,7 +31,7 @@ def __init__( @property def name(self) -> str: - """Name the runtime gives itself.""" + """Name that the runtime reports.""" return self._ref.name async def call(self, method: str, input: bytes) -> bytes: @@ -43,5 +43,5 @@ async def call(self, method: str, input: bytes) -> bytes: return await self._ref.call(method, input) def free_writer(self, writer_id: int) -> None: - """Release one handle of a writer with no wait, such as from a finalizer.""" + """Release one handle of a writer without waiting, such as from a finalizer.""" self._ref.free_writer(writer_id) diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index efef5cb528..73dd8c600c 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -13,6 +13,7 @@ TYPE_CHECKING, Any, Concatenate, + TypeVar, cast, overload, ) @@ -110,9 +111,13 @@ ) if TYPE_CHECKING: + from temporalio.common import ClientStream, StreamBackend, StreamHandle + from ._interceptor import Interceptor from ._plugin import Plugin +_StreamDefinitionT = TypeVar("_StreamDefinitionT") + class Client: """Client for accessing Temporal. @@ -156,6 +161,7 @@ async def connect( grpc_compression: GrpcCompression = GrpcCompression.GZIP, payload_limits: PayloadLimitsConfig = PayloadLimitsConfig(), header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ) -> Self: """Connect to a Temporal server. @@ -225,6 +231,13 @@ async def connect( payload_limits: Warning thresholds for outbound payload/memo sizes. Over-threshold fields are logged but still sent. Set a threshold to 0 to disable it. header_codec_behavior: Encoding behavior for headers sent by the client. + stream_backends: The Stream backends of this client and its Workers, such as + :py:class:`temporalio.contrib.stream.RedisStreamBackend`. Names must be unique. + With more than one backend, exactly one must have ``default=True``. See + :py:class:`temporalio.common.StreamBackend`. + + .. warning:: + This API is experimental. """ connect_config = temporalio.service.ConnectConfig( target_host=target_host, @@ -262,6 +275,7 @@ def make_lambda( default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, plugins=plugins, + stream_backends=stream_backends, ) def __init__( @@ -275,6 +289,7 @@ def __init__( default_workflow_query_reject_condition: None | (temporalio.common.QueryRejectCondition) = None, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ): """Create a Temporal client from a service client. @@ -289,11 +304,17 @@ def __init__( interceptors=interceptors, default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, + stream_backends=list(stream_backends), ) self._initial_config = config.copy() for plugin in plugins: config = plugin.configure_client(config) + if config.get("stream_backends"): + # Only a client with Stream backends loads the Streams package. + from temporalio.common._streams import _plugin as _stream_plugin + + config = _stream_plugin.configure_client(config) self._init_from_config(config) @@ -894,6 +915,42 @@ def get_workflow_handle( result_type=result_type, ) + def stream( + self, + definition: type[_StreamDefinitionT], + *, + id: str | None = None, + stream: StreamHandle[Any] | None = None, + backend: str | None = None, + ) -> ClientStream[_StreamDefinitionT]: + """Get a Stream by its id, or from a handle that came in data, such as in a Record. + + This method makes no call. To create a Stream that the client owns, use a backend's + ``create`` function, such as :py:func:`temporalio.contrib.stream.create`. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. + id: The Stream id. + stream: A handle of the Stream, such as one in a Record. The result keeps its id, + owner and backend. + backend: The name of the backend. Defaults to the handle's backend, or else to the + registered backend with ``default=True``. + + Raises: + StreamConfigError: The client has no backend with that name, or no backend has + ``default=True`` and the call gives no backend. + TypeError: Neither or both of ``id`` and ``stream`` are given. + ValueError: The id is empty, or ``backend`` is not the handle's backend. + """ + from temporalio.common._streams import _stream_client + + return _stream_client.client_stream( + self, definition, id=id, stream=stream, backend=backend + ) + def get_workflow_handle_for( self, workflow: ( @@ -3029,6 +3086,7 @@ class ClientConnectConfig(TypedDict, total=False): grpc_compression: GrpcCompression payload_limits: PayloadLimitsConfig header_codec_behavior: HeaderCodecBehavior + stream_backends: Sequence[StreamBackend] class ClientConfig(TypedDict, total=False): @@ -3043,3 +3101,4 @@ class ClientConfig(TypedDict, total=False): temporalio.common.QueryRejectCondition | None ] header_codec_behavior: Required[HeaderCodecBehavior] + stream_backends: Required[Sequence[StreamBackend]] diff --git a/temporalio/client/_workflow.py b/temporalio/client/_workflow.py index 94bc2bb23d..400204ea96 100644 --- a/temporalio/client/_workflow.py +++ b/temporalio/client/_workflow.py @@ -21,6 +21,7 @@ Any, Concatenate, Generic, + TypeVar, cast, overload, ) @@ -78,9 +79,13 @@ ) if TYPE_CHECKING: + from temporalio.common import ClientStream + from ._client import Client from ._interceptor import ListWorkflowsInput +_StreamDefinitionT = TypeVar("_StreamDefinitionT") + class WorkflowHistoryEventFilterType(IntEnum): """Type of history events to get for a workflow. @@ -187,6 +192,32 @@ def first_execution_run_id(self) -> str | None: """ return self._first_execution_run_id + def stream( + self, definition: type[_StreamDefinitionT], *, backend: str | None = None + ) -> ClientStream[_StreamDefinitionT]: + """Get a Stream that this Workflow created, by its definition. + + This method builds the Stream id from the Workflow id and the first run id, so it makes + no call. A handle with no first run id makes one Describe call when one of its Streams + is first used, and caches the result. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class, with the same name as in the Workflow's + create. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. + + Raises: + StreamConfigError: The client has no backend with that name, or no backend has + ``default=True`` and the call gives no backend. + """ + from temporalio.common._streams import _stream_client + + return _stream_client.workflow_stream(self, definition, backend=backend) + async def result( self, *, diff --git a/temporalio/common.py b/temporalio/common/__init__.py similarity index 89% rename from temporalio/common.py rename to temporalio/common/__init__.py index 578d714c6a..2e6ce3b1af 100644 --- a/temporalio/common.py +++ b/temporalio/common/__init__.py @@ -1,4 +1,24 @@ -"""Common code used in the Temporal SDK.""" +"""Common code used in the Temporal SDK. + +This module also has the Streams API. A Stream is an ordered log of Records that Workflows, +Activities and clients write and read. A Stream definition is a class with ``@common.defn``. +It gives the Stream name and the type of each Topic. Workflow code creates a Stream with +:py:func:`temporalio.workflow.create_stream`. A client gets a Workflow's Stream with +:py:meth:`temporalio.client.WorkflowHandle.stream`, and any Stream with +:py:meth:`temporalio.client.Client.stream`. A Stream is data. A Workflow can pass it as an +argument to an Activity, a child Workflow or the next run. The parameter's type sets the handle +class for that context, such as :py:class:`ActivityStream`. + +A :py:class:`StreamBackend` makes the Stream calls and registers a backend name. Examples are +:py:class:`temporalio.contrib.stream.RedisStreamBackend` (``redis``) and +:py:class:`temporalio.contrib.stream.testing.MemoryStreamBackend` (``memory``). Register +backends with ``Client.connect(..., stream_backends=[...])``. A handle stores the name of its +backend. With more than one registered backend, exactly one must have ``default=True``. A +create with no ``backend=`` uses that backend. A Worker serves only one backend. + +.. warning:: + The Streams API is experimental. +""" from __future__ import annotations @@ -13,6 +33,7 @@ from datetime import datetime, timedelta from enum import IntEnum from typing import ( + TYPE_CHECKING, Any, ClassVar, Generic, @@ -32,6 +53,130 @@ import temporalio.api.workflow.v1 import temporalio.types +if TYPE_CHECKING: + from temporalio.common._streams._activity_stream import ( + ActivityStream, + ActivityTopic, + SyncActivityStream, + SyncActivityTopic, + SyncSubscription, + ) + from temporalio.common._streams._buffered import BufferedPublisher + from temporalio.common._streams._defn import ( + StreamDefinition, + TopicDefinition, + defn, + topic, + with_name, + ) + from temporalio.common._streams._errors import ( + ReadOnlyContextError, + StreamClosedError, + StreamConfigError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetError, + StreamSeqDivergentError, + StreamStorageError, + StreamTooLargeError, + StreamUnsupportedError, + ) + from temporalio.common._streams._handle import Creating, StreamHandle + from temporalio.common._streams._items import ( + END, + Closed, + CloseStatus, + End, + Record, + Superseded, + ) + from temporalio.common._streams._plugin import StreamBackend + from temporalio.common._streams._reader import Subscription + from temporalio.common._streams._stream_client import ClientStream, ClientTopic + from temporalio.common._streams._workflow import WorkflowStream, WorkflowTopic + +__all__ = [ + "ActivityIDConflictPolicy", + "ActivityIDReusePolicy", + "ActivityStream", + "ActivityTopic", + "AutoUpgradeVersioningOverride", + "BufferedPublisher", + "ClientStream", + "ClientTopic", + "CloseStatus", + "Closed", + "Creating", + "END", + "End", + "HeaderCodecBehavior", + "MetricAttributes", + "MetricCommon", + "MetricCounter", + "MetricGauge", + "MetricGaugeFloat", + "MetricHistogram", + "MetricHistogramFloat", + "MetricHistogramTimedelta", + "MetricMeter", + "NexusOperationCancellationState", + "NexusOperationExecutionStatus", + "NexusOperationIDConflictPolicy", + "NexusOperationIDReusePolicy", + "OneTimeVersioningOverride", + "PendingNexusOperationExecutionState", + "PinnedVersioningOverride", + "Priority", + "QueryRejectCondition", + "RawValue", + "ReadOnlyContextError", + "Record", + "RetryPolicy", + "SearchAttributeIndexedValueType", + "SearchAttributeKey", + "SearchAttributePair", + "SearchAttributeUpdate", + "SearchAttributeValue", + "SearchAttributeValueType", + "SearchAttributeValues", + "SearchAttributes", + "StreamBackend", + "StreamClosedError", + "StreamConfigError", + "StreamDefinition", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Subscription", + "Superseded", + "SyncActivityStream", + "SyncActivityTopic", + "SyncSubscription", + "TopicDefinition", + "TypedSearchAttributes", + "VersioningBehavior", + "VersioningOverride", + "WorkerDeploymentVersion", + "WorkflowIDConflictPolicy", + "WorkflowIDReusePolicy", + "WorkflowStream", + "WorkflowTopic", + "defn", + "topic", + "with_name", +] + @dataclass class RetryPolicy: @@ -1482,3 +1627,57 @@ def wait_sync(self, timeout: float | None = None) -> None: if not self.thread_event: raise RuntimeError("Missing event") self.thread_event.wait(timeout) + + +# The Stream modules import temporalio.client and temporalio.workflow, which import this module, +# so this module imports the Stream names on first use. +_STREAM_NAMES = frozenset( + [ + "ActivityStream", + "ActivityTopic", + "BufferedPublisher", + "ClientStream", + "ClientTopic", + "CloseStatus", + "Closed", + "Creating", + "END", + "End", + "ReadOnlyContextError", + "Record", + "StreamBackend", + "StreamClosedError", + "StreamConfigError", + "StreamDefinition", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Subscription", + "Superseded", + "SyncActivityStream", + "SyncActivityTopic", + "SyncSubscription", + "TopicDefinition", + "WorkflowStream", + "WorkflowTopic", + "defn", + "topic", + "with_name", + ] +) + + +def __getattr__(name: str) -> Any: + if name in _STREAM_NAMES: + import temporalio.common._streams + + return getattr(temporalio.common._streams, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py new file mode 100644 index 0000000000..a4001dbd31 --- /dev/null +++ b/temporalio/common/_streams/__init__.py @@ -0,0 +1,86 @@ +"""The Streams API, which :py:mod:`temporalio.common` exports.""" + +from temporalio.common._streams._activity_stream import ( + ActivityStream, + ActivityTopic, + SyncActivityStream, + SyncActivityTopic, + SyncSubscription, +) +from temporalio.common._streams._buffered import BufferedPublisher +from temporalio.common._streams._defn import ( + StreamDefinition, + TopicDefinition, + defn, + topic, + with_name, +) +from temporalio.common._streams._errors import ( + ReadOnlyContextError, + StreamClosedError, + StreamConfigError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetError, + StreamSeqDivergentError, + StreamStorageError, + StreamTooLargeError, + StreamUnsupportedError, +) +from temporalio.common._streams._handle import Creating, StreamHandle +from temporalio.common._streams._items import ( + END, + Closed, + CloseStatus, + End, + Record, + Superseded, +) +from temporalio.common._streams._plugin import StreamBackend +from temporalio.common._streams._reader import Subscription +from temporalio.common._streams._stream_client import ClientStream, ClientTopic +from temporalio.common._streams._workflow import WorkflowStream, WorkflowTopic + +__all__ = [ + "END", + "ActivityStream", + "ActivityTopic", + "BufferedPublisher", + "ClientStream", + "ClientTopic", + "CloseStatus", + "Closed", + "Creating", + "End", + "ReadOnlyContextError", + "Record", + "StreamClosedError", + "StreamConfigError", + "StreamDefinition", + "StreamBackend", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Subscription", + "Superseded", + "SyncActivityStream", + "SyncActivityTopic", + "SyncSubscription", + "TopicDefinition", + "WorkflowStream", + "WorkflowTopic", + "defn", + "topic", + "with_name", +] diff --git a/temporalio/contrib/streams/_access.py b/temporalio/common/_streams/_access.py similarity index 68% rename from temporalio/contrib/streams/_access.py rename to temporalio/common/_streams/_access.py index 814bc5c17e..3dba3ff860 100644 --- a/temporalio/contrib/streams/_access.py +++ b/temporalio/common/_streams/_access.py @@ -5,17 +5,17 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import timedelta -from typing import Any, TypeVar, cast +from typing import Any, TypeVar import temporalio.converter -from temporalio.contrib.streams import _convert, _reader -from temporalio.contrib.streams._backend import CallOptions, StreamBackend -from temporalio.contrib.streams._defn import TopicDefinition -from temporalio.contrib.streams._handle import StreamHandle, TopicBase -from temporalio.contrib.streams._items import End, close_status -from temporalio.contrib.streams._reader import Item, Position, Subscription -from temporalio.contrib.streams._writer import StreamWriter -from temporalio.contrib.streams.system import models +from temporalio.common._streams import _convert, _reader +from temporalio.common._streams._backend import BackendCalls, CallOptions +from temporalio.common._streams._defn import TopicDefinition +from temporalio.common._streams._handle import StreamHandle, TopicBase +from temporalio.common._streams._items import End, close_status +from temporalio.common._streams._reader import Item, Position, Subscription +from temporalio.common._streams._writer import StreamWriter +from temporalio.common._streams.system import models T = TypeVar("T") @@ -24,12 +24,12 @@ class Access: """The calls of one Stream handle, bound to a backend.""" - backend: StreamBackend + backend: BackendCalls converter: temporalio.converter.DataConverter namespace: str stream_id: Callable[[], Awaitable[str]] writer: Callable[[], StreamWriter] - # The run chain that a close names, which the Store checks against the Stream's owner. + # A close sends these, and the Store checks them against the Stream's owner. workflow_id: str | None = None first_execution_run_id: str | None = None @@ -55,10 +55,10 @@ async def items( await self.stream_id(), position=position, topics=read.topics, - topic_prefixes=read.prefixes, + topic_prefixes=read.topic_prefixes, types=read.types, default_type=read.default_type, - from_offset=read.from_offset, + from_offset=read.offset, last=read.last, wait_new_records=read.wait_new_records, create_wait=read.create_wait, @@ -70,17 +70,15 @@ async def items( await items.aclose() async def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one. + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. - 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. + Offsets are sparse and above zero, so ``offset=0`` reads from the start. """ 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 async def close(self, result: Any | None, failure: BaseException | None) -> None: @@ -110,10 +108,10 @@ class ReadSpec: """What one read asks for.""" topics: list[str] | None - prefixes: Sequence[str] + topic_prefixes: Sequence[str] types: dict[str, Any] default_type: Any - from_offset: int | End | None + offset: int | End | None last: int | None wait_new_records: bool create_wait: timedelta | None @@ -125,9 +123,9 @@ def read_spec( str | TopicDefinition[Any] | _AsyncTopic[Any] | SyncTopicBase[Any] ], *, - prefixes: Sequence[str], + topic_prefixes: Sequence[str], default_type: Any, - from_offset: int | End | None, + offset: int | End | None, last: int | None, wait_new_records: bool, create_wait: timedelta | None, @@ -135,10 +133,10 @@ def read_spec( """The Topic names of a read, the decode type of each known Topic, and its options. Raises: - ValueError: ``last`` is not above zero, or it comes with ``from_offset``. + ValueError: ``last`` is not above zero, or the call also gives ``offset``. """ - if last is not None and (last < 1 or from_offset is not None): - raise ValueError("last takes a count above zero and no from_offset") + if last is not None and (last < 1 or offset is not None): + raise ValueError("last takes a count above zero and no offset") types = handle._known_types() names = [] for topic in topics: @@ -147,7 +145,7 @@ def read_spec( elif isinstance(topic, TopicDefinition): name, topic_type = handle._topic_spec(topic, None) else: - # A read names a Topic outside the definition with no type, and its values + # A Topic name that the definition doesn't have gets no type here, so its values # decode with the read's default type. name, topic_type = topic, types.get(topic) names.append(name) @@ -155,10 +153,10 @@ def read_spec( types[name] = topic_type return ReadSpec( topics=names or None, - prefixes=prefixes, + topic_prefixes=topic_prefixes, types=types, default_type=default_type, - from_offset=from_offset, + offset=offset, last=last, wait_new_records=wait_new_records, create_wait=create_wait, @@ -171,12 +169,11 @@ def subscription( options: CallOptions | None = None, ) -> Subscription[Any]: """A read over the items that ``items`` gives.""" - # Before its first page, a read from the start resumes at 0, and a read from END or of the - # last Records has no offset yet. - if read.last is not None or read.from_offset is End.END: + # A read from END or of the last Records has no resume offset until its first page. + if read.last is not None or read.offset is End.END: position = Position(None) else: - position = Position(read.from_offset or 0) + position = Position(read.offset or 0) return Subscription(items(position), position, options) @@ -188,31 +185,8 @@ async def publish(self, *values: T) -> int: A list is one value. """ - await self._stream._ready() return await self._stream._access().publish(self._name, values) - def subscribe( - self, - *, - from_offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - ) -> Subscription[T]: - """Read this Topic only. Its Records have the Topic's type. See the Stream's - ``subscribe``. - """ - return cast( - Subscription[T], - self._stream.subscribe( - self, - from_offset=from_offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ), - ) - class SyncTopicBase(TopicBase[T]): - """Base of a Topic whose calls block, for sync Activities.""" + """Base class of a Topic whose calls block, for sync Activities.""" diff --git a/temporalio/contrib/streams/_activity.py b/temporalio/common/_streams/_activity.py similarity index 77% rename from temporalio/contrib/streams/_activity.py rename to temporalio/common/_streams/_activity.py index 09b2eafc4a..09ff5360b4 100644 --- a/temporalio/contrib/streams/_activity.py +++ b/temporalio/common/_streams/_activity.py @@ -9,23 +9,25 @@ import temporalio.api.streamservice.v1 as stream_pb import temporalio.converter -from temporalio.contrib.streams._backend import StreamBackend -from temporalio.contrib.streams._writer import StreamWriter +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._writer import StreamWriter @dataclass class ActivityStreams: - """What a Stream plugin gives the Activity handles for one Activity attempt.""" + """What a Stream backend gives the Activity handles for one Activity attempt.""" - backend: StreamBackend + backend: BackendCalls converter: temporalio.converter.DataConverter namespace: str producer: stream_pb.Producer loop: asyncio.AbstractEventLoop backend_name: str = "" - # The Workflow of the attempt, so a close checks that the Stream is its run chain's. + # Whether the Worker's backend is the default, which a handle with no backend name uses. + default: bool = False + # The attempt's Workflow id, which a close sends so the Store can check the Stream's owner. workflow_id: str | None = None - # A Local Activity's run history, which lets its writes withdraw the run's orphaned Records. + # A Local Activity's writes use these to withdraw the run's orphaned Records. previous_completed_task: stream_pb.WorkflowTaskEpoch | None = None failed_tasks: tuple[int, ...] = () writers: dict[str, StreamWriter] = field(default_factory=dict) @@ -52,7 +54,7 @@ async def known_id() -> str: return writer async def close(self) -> None: - """Flush the attempt's publishers, then drop every writer of the attempt.""" + """Flush the attempt's publishers, then close every writer of the attempt.""" errors = [] for publisher in self.publishers: try: diff --git a/temporalio/contrib/streams/_activity_stream.py b/temporalio/common/_streams/_activity_stream.py similarity index 72% rename from temporalio/contrib/streams/_activity_stream.py rename to temporalio/common/_streams/_activity_stream.py index 9898332d24..9212f3d72d 100644 --- a/temporalio/contrib/streams/_activity_stream.py +++ b/temporalio/common/_streams/_activity_stream.py @@ -6,21 +6,21 @@ import concurrent.futures from collections.abc import Coroutine, Generator, Sequence from datetime import timedelta -from typing import Any, Generic, TypeVar, cast, overload +from typing import Any, Generic, TypeVar, overload import temporalio.api.streamservice.v1 as stream_pb -from temporalio.contrib.streams import _activity, _errors -from temporalio.contrib.streams._access import ( +from temporalio.common._streams import _activity, _errors +from temporalio.common._streams._access import ( Access, SyncTopicBase, _AsyncTopic, read_spec, subscription, ) -from temporalio.contrib.streams._defn import TopicDefinition -from temporalio.contrib.streams._handle import StreamHandle -from temporalio.contrib.streams._items import Closed, End, Record, Superseded -from temporalio.contrib.streams._reader import Subscription +from temporalio.common._streams._defn import TopicDefinition +from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._items import Closed, End, Record, Superseded +from temporalio.common._streams._reader import Subscription E = TypeVar("E") T = TypeVar("T") @@ -34,26 +34,30 @@ def _bound(handle: StreamHandle[Any]) -> _activity.ActivityStreams: streams = _activity.current.get() if streams is None: raise _errors.StreamConfigError( - "No Stream plugin runs this Activity, such as RedisStreamPlugin" + "No Stream backend runs this Activity. Register one on the Worker's client " + "with stream_backends=." ) if handle.backend and handle.backend != streams.backend_name: raise _errors.StreamConfigError( f"The Stream {handle.id} is of the backend {handle.backend}. The Worker has only " f"{streams.backend_name}." ) + if not handle.backend and not streams.default: + raise _errors.no_default_backend( + "Worker", [streams.backend_name], "a handle that has its backend name" + ) return streams def _check_close(handle: StreamHandle[Any]) -> None: - # A Workflow's Stream closes only from its run chain, and a client's only from a client. So - # only a Local Activity of the chain may close here. The Store refuses the others too, and - # this error comes with no call. + # Only Workflow code can close a Workflow's Stream, and only a client can close a client's + # Stream. This check refuses some closes without a call, and the Store refuses the rest. streams = _bound(handle) local = streams.producer.kind == stream_pb.ProducerKind.PRODUCER_KIND_LOCAL_ACTIVITY - if not local or (handle.workflow_id and handle.workflow_id != streams.workflow_id): + if not local or handle.workflow_id: raise _errors.StreamNotOwnerError( f"Can't close the Stream {handle.id} from this Activity. A Workflow's Stream " - "closes only from its run chain, and a client's only from a client.", + "closes only from Workflow code, and a client's only from a client.", reason="stream_not_owner", ) @@ -71,8 +75,8 @@ async def known_id() -> str: namespace=streams.namespace, stream_id=known_id, writer=lambda: streams.writer(stream_id), - # A handle with only an id names no chain, so the close names this Activity's own, and - # the Store refuses it when another chain owns the Stream. + # A handle with only an id has no run chain, so the close sends the Activity's Workflow + # id. The Store refuses that close for a Workflow's Stream. workflow_id=handle.workflow_id or streams.workflow_id, first_execution_run_id=handle.first_execution_run_id, ) @@ -87,7 +91,7 @@ class ActivityTopic(_AsyncTopic[T]): class ActivityStream(StreamHandle[E]): - """A Stream in an async Activity. The Activity writes, reads and closes it. + """A Stream in an async Activity. The Activity can write and read it. An Activity gets a Stream only as an argument, such as a field of its input. Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. @@ -118,10 +122,10 @@ def topic( def subscribe( self, *topics: str | TopicDefinition[Any] | ActivityTopic[Any], - from_offset: int | End | None = None, + offset: int | End | None = None, last: int | None = None, wait_new_records: bool = True, - prefixes: Sequence[str] = (), + topic_prefixes: Sequence[str] = (), create_wait: timedelta | None = None, type: Any = None, ) -> Subscription[Any]: @@ -129,9 +133,9 @@ def subscribe( read = read_spec( self, topics, - prefixes=prefixes, + topic_prefixes=topic_prefixes, default_type=type, - from_offset=from_offset, + offset=offset, last=last, wait_new_records=wait_new_records, create_wait=create_wait, @@ -140,10 +144,9 @@ def subscribe( return subscription(lambda position: access.items(read, position), read) async def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one. + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. - 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. + Offsets are sparse and above zero, so ``offset=0`` reads from the start. """ return await self._access().get_offset() @@ -152,12 +155,15 @@ async def close( ) -> None: """Write this attempt's Records, then close the Stream. - Readers get the Records before it, then a ``Closed`` item with ``result`` or - ``failure``. A ``CancelledError`` failure closes it as canceled. + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. - Only a Local Activity of the run chain that created the Stream may close it. Any other - Activity gets :py:class:`StreamNotOwnerError`, so it can't end the Stream for every - reader. + An Activity can't close a Workflow's Stream or a client's Stream, and this call raises + :py:class:`StreamNotOwnerError`. The rule includes a Local Activity of the run chain + that created the Stream. A Local Activity runs before its Workflow Task commits. If the + task then failed, the close would stay, and the task's retry couldn't create the Stream + again. Close a Workflow's Stream from Workflow code. """ _check_close(self) await self._access().close(result, failure) @@ -179,35 +185,11 @@ def publish(self, *values: T) -> int: assert isinstance(stream, SyncActivityStream) return stream._run(stream._access().publish(self._name, values)) - def subscribe( - self, - *, - from_offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - ) -> SyncSubscription[T]: - """Read this Topic only. Its Records have the Topic's type. See the Stream's - ``subscribe``. - """ - stream = self._stream - assert isinstance(stream, SyncActivityStream) - return cast( - SyncSubscription[T], - stream.subscribe( - self, - from_offset=from_offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ), - ) - class SyncSubscription(Generic[T]): - """The items of one read of a Stream in a sync Activity, in order. Use it in ``for``. + """The items of one read of a Stream in a sync Activity, in order. Iterate it with ``for``. - Each step blocks the Activity's thread. See :py:class:`Subscription`. + Each iteration blocks the Activity's thread. See :py:class:`Subscription`. .. warning:: This API is experimental. @@ -264,15 +246,14 @@ def _run(self, call: Coroutine[Any, Any, _R]) -> _R: future = asyncio.run_coroutine_threadsafe(call, loop) try: while True: - # Short waits let an exception that the Worker raises in this thread, such as - # a cancel, come in. + # Short waits let the Worker raise an exception in this thread, such as a cancel. try: return future.result(timeout=_WAIT_STEP) except concurrent.futures.TimeoutError: continue except BaseException: - # A cancel of the Activity thread, such as from a heartbeat, stops the call too, - # so a long poll does not run on in the Worker's loop. + # Cancel the call too, so a long poll doesn't keep running in the Worker's loop after + # the Activity thread is canceled, such as from a heartbeat. future.cancel() raise @@ -295,10 +276,10 @@ def topic( def subscribe( self, *topics: str | TopicDefinition[Any] | SyncActivityTopic[Any], - from_offset: int | End | None = None, + offset: int | End | None = None, last: int | None = None, wait_new_records: bool = True, - prefixes: Sequence[str] = (), + topic_prefixes: Sequence[str] = (), create_wait: timedelta | None = None, type: Any = None, ) -> SyncSubscription[Any]: @@ -306,9 +287,9 @@ def subscribe( read = read_spec( self, topics, - prefixes=prefixes, + topic_prefixes=topic_prefixes, default_type=type, - from_offset=from_offset, + offset=offset, last=last, wait_new_records=wait_new_records, create_wait=create_wait, @@ -332,18 +313,18 @@ def _items( self._run(items.aclose()) def get_offset(self) -> int: - """The next offset of the Stream, the newest offset plus one. + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. - 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. + Offsets are sparse and above zero, so ``offset=0`` reads from the start. """ return self._run(self._access().get_offset()) def close( self, *, result: Any | None = None, failure: BaseException | None = None ) -> None: - """Write this attempt's Records, then close the Stream. See - :py:meth:`ActivityStream.close`. + """Write this attempt's Records, then close the Stream. + + See :py:meth:`ActivityStream.close`. """ _check_close(self) self._run(self._access().close(result, failure)) diff --git a/temporalio/contrib/streams/_backend.py b/temporalio/common/_streams/_backend.py similarity index 79% rename from temporalio/contrib/streams/_backend.py rename to temporalio/common/_streams/_backend.py index 197ee7d5fd..9d77298f5c 100644 --- a/temporalio/contrib/streams/_backend.py +++ b/temporalio/common/_streams/_backend.py @@ -17,13 +17,13 @@ import temporalio.bridge.streams import temporalio.converter import temporalio.service -from temporalio.contrib.streams import _errors -from temporalio.contrib.streams.proto import v1 as ext_pb -from temporalio.contrib.streams.system import client as service_client +from temporalio.common._streams import _errors +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.common._streams.system import client as service_client if TYPE_CHECKING: - # The writer bindings import this package's protos, so they load after it. - from temporalio.contrib.streams.writer_system import client as writer_client + # The writer bindings import this package's protos, so writers() imports them on first use. + from temporalio.common._streams.writer_system import client as writer_client _M = TypeVar("_M", bound=google.protobuf.message.Message) @@ -60,8 +60,8 @@ class CallOptions: timeout: timedelta | None -# The options of the client call in progress. Each transport reads them, so they reach the -# generated clients' calls and the writer's with no new argument there. +# The options of the client call in progress. Each transport reads them, so the generated +# clients and the writer don't need a new argument for them. call_options: contextvars.ContextVar[CallOptions | None] = contextvars.ContextVar( "__temporal_stream_call_options", default=None ) @@ -78,7 +78,7 @@ def options_of( @contextlib.contextmanager def using(options: CallOptions | None) -> Iterator[None]: - """Give the calls in the block these options.""" + """Apply these options to the calls in the block.""" if options is None: yield return @@ -90,20 +90,18 @@ def using(options: CallOptions | None) -> Iterator[None]: def longest_wait(longest: timedelta) -> timedelta: - """The longest wait of one call. It is less than ``longest`` to answer before the call's - rpc_timeout. - """ + """The longest wait of one call, at most ``longest`` and below the call's rpc_timeout.""" options = call_options.get() if options is None or options.timeout is None: return longest - # A long poll answers before the call's deadline, as the server does for its own long - # polls, so a read that waits goes on with its next poll and does not fail. + # A long poll returns before the call's deadline, as the server's long polls do, so a + # waiting read starts its next poll and doesn't fail. timeout = options.timeout return min(longest, max(timeout - timedelta(seconds=1), timeout / 2)) -class StreamBackend(Protocol): - """What a Stream plugin gives the handles. Not public. Each backend has its own class.""" +class BackendCalls(Protocol): + """The calls that a Stream backend gives the handles of one client.""" def service( self, @@ -120,7 +118,7 @@ def writers( ... def free_writer(self, writer_id: int) -> None: - """Release one handle of a writer with no wait, for a writer that Python drops.""" + """Release one handle of a writer without waiting, such as when Python frees it.""" ... @@ -128,19 +126,19 @@ class StreamRuntime(Protocol): """The calls of the crate's runtime object, which the bridge builds.""" async def call(self, method: str, input: bytes) -> bytes: - """Call a method with encoded input, and give the encoded output.""" + """Call a method with encoded input, and return the encoded output.""" ... def free_writer(self, writer_id: int) -> None: - """Release one handle of a writer with no wait.""" + """Release one handle of a writer without waiting.""" ... -class _ExtensionStreamBackend: - """Carries Stream calls to the crate's runtime object, which serves one client. +class _ExtensionCalls: + """Makes Stream calls to the crate's runtime object, which serves one client. - The five Stream service operations go through the generated client, and the writer - methods are the crate's own. + The five Stream service operations use the generated client. The writer methods are the + crate's own. """ def __init__(self, runtime: StreamRuntime) -> None: @@ -154,7 +152,8 @@ def service( ) -> 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. + With a ``data_converter``, the client encodes user values in the requests with it and + its codec. """ return service_client.StreamServiceClient( _ExtensionTransport(self, namespace), data_converter=data_converter @@ -164,14 +163,14 @@ def writers( self, namespace: str, data_converter: temporalio.converter.DataConverter ) -> writer_client.StreamWriterClient: """The generated writer client over the runtime.""" - from temporalio.contrib.streams.writer_system import client + from temporalio.common._streams.writer_system import client return client.StreamWriterClient( _ExtensionTransport(self, namespace), data_converter=data_converter ) def free_writer(self, writer_id: int) -> None: - """Release one handle of a writer with no wait.""" + """Release one handle of a writer without waiting.""" try: self._runtime.free_writer(writer_id) except temporalio.bridge.streams.StreamCallError: @@ -190,7 +189,7 @@ async def call_async( return response.FromString(data) async def close_writers(self) -> int: - """Drop every open writer of the runtime, and return how many failed.""" + """Close every open writer of the runtime, and return how many failed.""" response = await self.call_async( "CloseWriters", ext_pb.CloseWritersRequest(), ext_pb.CloseWritersResponse ) @@ -198,9 +197,8 @@ async def close_writers(self) -> int: class _ExtensionTransport: - # The generated client takes a transport, so the five operations reach the runtime as they - # would reach a server later. - def __init__(self, backend: _ExtensionStreamBackend, namespace: str) -> None: + # The generated client calls a transport. This one calls the runtime in place of a server. + def __init__(self, backend: _ExtensionCalls, namespace: str) -> None: self._backend = backend self._namespace = namespace diff --git a/temporalio/contrib/streams/_buffered.py b/temporalio/common/_streams/_buffered.py similarity index 78% rename from temporalio/contrib/streams/_buffered.py rename to temporalio/common/_streams/_buffered.py index 3156336461..678681d2ef 100644 --- a/temporalio/contrib/streams/_buffered.py +++ b/temporalio/common/_streams/_buffered.py @@ -7,16 +7,16 @@ from types import TracebackType from typing import Any, Generic, TypeVar, overload -from temporalio.contrib.streams import _activity -from temporalio.contrib.streams._activity_stream import ActivityStream, ActivityTopic -from temporalio.contrib.streams._defn import TopicDefinition -from temporalio.contrib.streams._handle import TopicBase -from temporalio.contrib.streams._stream_client import ClientStream, ClientTopic +from temporalio.common._streams import _activity +from temporalio.common._streams._activity_stream import ActivityStream, ActivityTopic +from temporalio.common._streams._defn import TopicDefinition +from temporalio.common._streams._handle import TopicBase +from temporalio.common._streams._stream_client import ClientStream, ClientTopic T = TypeVar("T") U = TypeVar("U") -# The Records of one batch of the crate's writer. A buffer that holds a batch writes at once. +# Records in one batch of the crate's writer. A full batch is written at once. _BATCH = 1000 @@ -26,23 +26,24 @@ class _AnyTopic: class BufferedPublisher(Generic[T]): - """Batches the publishes of one Topic, or of every Topic of one Stream, for an Activity or - a client that writes often. + """Batches the publishes of one Topic, or of every Topic of one Stream. - :py:meth:`publish` returns at once. The publisher writes the buffered values as one batch - at most ``max_delay`` after the first of them, or at once when they fill a batch. One write - is in flight at a time, and values published during it wait for the next one. Use it in - ``async with``, which writes the rest at the end of the block. In an Activity, the end of - the attempt also writes the rest. + Use it in an Activity or a client that writes often. :py:meth:`publish` returns at once. + The publisher writes the buffered values as one batch at most ``max_delay`` after the first + of them, or at once when they fill a batch. One write is in flight at a time, and values + published during it wait for the next one. Use it in ``async with``, which writes the + remaining values at the end of the block. In an Activity, the end of the attempt also + writes the remaining values. ``await flush()`` is the backpressure point. A caller that publishes faster than the Store - takes Records awaits it now and then, so the buffer does not grow without bound. + accepts Records must await it from time to time, so the buffer doesn't grow without bound. - Over a Stream, ``publish(topic, *values)`` takes the Topic of each call first, and all the - Topics share one buffer, so the Records keep the publish order across Topics. + A publisher over a Stream takes the Topic as the first argument of + ``publish(topic, *values)``. All the Topics share one buffer, so the Records keep the + publish order across Topics. - After a batch fails, the publisher writes nothing more, so a reader sees no gap, and each - later call raises the error. + After a batch fails, the publisher writes nothing more, so a reader sees no gap. Each later + call raises the error. .. warning:: This API is experimental. @@ -75,7 +76,7 @@ def __init__( Args: target: A Topic of an ``ActivityStream`` or a ``ClientStream``, or the Stream itself for a publisher of all its Topics. - max_delay: How long a value waits in the buffer at most. + max_delay: The longest time that a value waits in the buffer. """ self._topic: ActivityTopic[Any] | ClientTopic[Any] | None = None self._stream: ActivityStream[Any] | ClientStream[Any] | None = None @@ -85,12 +86,12 @@ def __init__( self._stream = target self._max_delay = max_delay self._buffer: list[tuple[str, Any]] = [] - # When the oldest buffered value came, so the next batch goes at most max_delay later. + # When the oldest buffered value arrived, so the next batch is at most max_delay later. self._first_at = 0.0 self._timer: asyncio.TimerHandle | None = None self._writing: asyncio.Task[None] | None = None self._published = 0 - # Values that a write took, written or not. + # Count of the values that a write took, written or not. self._settled = 0 self._error: BaseException | None = None self._last_offset = 0 @@ -119,8 +120,8 @@ def publish(self, *values: T) -> None: ... def publish(self, *args: Any) -> None: """Add one Record for each value to the buffer, and return at once. - A publisher over a Stream takes the Topic first. It is a Topic of the definition, such - as ``Events.tokens``, a name of the definition, or a Topic of the Stream. + A publisher over a Stream takes the Topic first. The Topic can be a Topic of the + definition, such as ``Events.tokens``, a Topic name, or a Topic of the Stream. Raises: StreamError: A batch failed to write, so the publisher takes no more values. @@ -151,10 +152,9 @@ def _topic_name(self, topic: Any) -> str: return name async def flush(self) -> int: - """Write every value published before the call, and return the offset of the last - Record written. + """Write every value published before the call, and return the last written offset. - Returns 0 when the publisher wrote nothing yet. + Returns 0 if the publisher hasn't written anything. """ target = self._published while self._settled < target: @@ -171,8 +171,8 @@ def _raise_error(self) -> None: raise self._error async def _flush_at_attempt_end(self) -> None: - # The attempt fails for values that never went out, unless the Activity already saw - # the error and went on. + # Fail the attempt for unwritten values, unless the Activity already got the error and + # continued. reported = self._reported try: await self.flush() @@ -212,14 +212,14 @@ async def _write(self, batch: list[tuple[str, Any]]) -> None: finally: self._settled += len(batch) self._writing = None - # Values published during the write go next, when their delay or a batch is up. + # Schedule the values published during the write, for when their delay ends or they + # fill a batch. self._schedule() async def _send(self, batch: list[tuple[str, Any]]) -> int: if self._topic is not None: return await self._topic.publish(*(value for _, value in batch)) assert self._stream is not None - await self._stream._ready() return await self._stream._access().publish_records(batch) @property diff --git a/temporalio/contrib/streams/_convert.py b/temporalio/common/_streams/_convert.py similarity index 76% rename from temporalio/contrib/streams/_convert.py rename to temporalio/common/_streams/_convert.py index ba24d64902..5b769b82e9 100644 --- a/temporalio/contrib/streams/_convert.py +++ b/temporalio/common/_streams/_convert.py @@ -9,11 +9,12 @@ import temporalio.api.common.v1 import temporalio.common import temporalio.converter +import temporalio.exceptions class _RawPayloadConverter(temporalio.converter.PayloadConverter): # 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. + # handle converts each body with its Topic's type. def to_payloads( self, values: Sequence[Any] ) -> list[temporalio.api.common.v1.Payload]: @@ -39,6 +40,20 @@ def raw( return dataclasses.replace(converter, payload_converter_class=_RawPayloadConverter) +def failure( + converter: temporalio.converter.DataConverter, failure: BaseException +) -> BaseException: + """Convert the RawValue payloads of a read's close failure with the user's converter.""" + if ( + not isinstance(failure, temporalio.exceptions.FailureError) + or not failure.failure + ): + return failure + return converter.failure_converter.from_failure( + failure.failure, converter.payload_converter + ) + + def decode( converter: temporalio.converter.DataConverter, body: Any, type_hint: Any ) -> Any: diff --git a/temporalio/contrib/streams/_defn.py b/temporalio/common/_streams/_defn.py similarity index 87% rename from temporalio/contrib/streams/_defn.py rename to temporalio/common/_streams/_defn.py index 2708a5fe1f..76a37089fa 100644 --- a/temporalio/contrib/streams/_defn.py +++ b/temporalio/common/_streams/_defn.py @@ -1,4 +1,4 @@ -"""Stream definitions: the Stream name and the type of each Topic.""" +"""Stream definitions, which give the Stream name and the type of each Topic.""" from __future__ import annotations @@ -26,7 +26,7 @@ def __init__(self, type: Any, name: str | None = None) -> None: self._owner: type | None = None def __set_name__(self, owner: type, name: str) -> None: - """Take the attribute name as the Topic name, unless the Topic has a name.""" + """Use the attribute name as the Topic name, unless the Topic has one.""" if self._name is None: self._name = name self._owner = owner @@ -61,7 +61,7 @@ def topic(type: Any, *, name: str | None = None) -> TopicDefinition[Any]: ... def topic(type: Any, *, name: str | None = None) -> TopicDefinition[Any]: """Declare a Topic in a Stream definition class. - ``RawValue`` keeps each Record's Payload. + Use ``RawValue`` as the type to read each Record's Payload with no conversion. Args: type: The type of the Topic's values. @@ -101,7 +101,7 @@ def defn( ) -> ClassT | Callable[[ClassT], ClassT]: """Mark a class as a Stream definition. Its :py:func:`topic` attributes are the Topics. - A definition has no backend in it, so Workflows, Activities and clients share it. + A definition doesn't name a backend, so Workflows, Activities and clients can share it. Args: cls: The class. @@ -140,9 +140,9 @@ def decorate(cls: ClassT) -> ClassT: def with_name(definition: ClassT, name: str) -> ClassT: - """The definition with the same Topics and another Stream name, such as one per turn. + """Return a definition with the same Topics and another Stream name, such as one per turn. - It is a subclass of ``definition``, which stays unchanged. Its Topics are the same + The result is a subclass of ``definition``, which doesn't change. Its Topics are the same attributes, so ``Events.tokens`` is a Topic of ``with_name(Events, "turn-3")`` too. Args: @@ -171,9 +171,11 @@ def definition_of(cls: type) -> StreamDefinition: """The definition that :py:func:`defn` gave a class. Raises: - TypeError: The class has no ``@streams.defn``. + TypeError: The class has no ``@common.defn``. """ definition = cls.__dict__.get(_DEFINITION_ATTR) if isinstance(cls, type) else None if not isinstance(definition, StreamDefinition): - raise TypeError(f"{cls!r} is not a Stream definition. Add @streams.defn.") + raise TypeError( + f"{cls!r} is not a Stream definition. Add @temporalio.common.defn." + ) return definition diff --git a/temporalio/contrib/streams/_errors.py b/temporalio/common/_streams/_errors.py similarity index 63% rename from temporalio/contrib/streams/_errors.py rename to temporalio/common/_streams/_errors.py index 4b40bceb91..c8f480b7e1 100644 --- a/temporalio/contrib/streams/_errors.py +++ b/temporalio/common/_streams/_errors.py @@ -1,19 +1,21 @@ -"""Errors of temporalio.contrib.streams.""" +"""Errors of the Streams API in temporalio.common.""" from __future__ import annotations +from collections.abc import Sequence + import temporalio.exceptions import temporalio.service -from temporalio.contrib.streams.proto.v1 import StreamErrorDetails +from temporalio.common._streams.proto.v1 import StreamErrorDetails class StreamError(temporalio.exceptions.ApplicationError): - """Base of every Stream error. + """Base class of every Stream error. - It is an application error that is not retryable, except - :py:class:`StreamStorageError`. So in Workflow code, where only - :py:class:`StreamClosedError` and the plugin errors are raised, an uncaught one fails the - Workflow and does not fail the Workflow Task again and again. + Each Stream error is a non-retryable application error, except + :py:class:`StreamStorageError`. So an uncaught Stream error in Workflow code fails the + Workflow, and doesn't fail the Workflow Task again and again. In Workflow code, only + :py:class:`StreamClosedError` and the backend errors are raised. .. warning:: This API is experimental. @@ -59,7 +61,7 @@ class StreamClosedError(StreamError): class StreamExpiredError(StreamError): - """Retention removed the Stream.""" + """Retention removed the Stream, or the Records at the offset that a read asked for.""" class StreamFencedError(StreamError): @@ -73,37 +75,55 @@ class StreamStorageError(StreamError): class StreamOffsetError(StreamError): - """The offset is past the newest one plus one, or Retention removed it.""" + """The offset is below zero or past the newest one plus one.""" + + +class StreamTooLargeError(StreamError): + """A Record or a batch is over the Store's limit.""" class StreamUnsupportedError(StreamError): - """The Store lacks a capability that the call needs. The message ends with its name.""" + """The Store lacks a capability that the call needs. + + The message ends with the name of the capability. + """ class StreamNotOwnerError(StreamError): - """The caller may not close the Stream or write as that Producer, since the Stream's owner - doesn't allow it. + """The caller isn't allowed to close the Stream or write as that Producer. - A Workflow's Stream closes only from its own run chain, and a client's Stream only from a - client. + Only the Workflow's own run chain can close a Workflow's Stream. Only a client can close a + client's Stream. """ class StreamSeqDivergentError(StreamError): - """A resend has the sequence of the Producer's last batch and other content. + """A resent batch has the sequence number of the Producer's last batch, but other content. - A writer that starts again with the same Producer id and epoch sends its first sequence + A writer that restarts with the same Producer id and epoch sends its first sequence number again. Give it a higher epoch instead. """ class StreamConfigError(RuntimeError): - """No Stream plugin, or no backend of a given name. It is config, not a Store error. + """The Stream backend config is wrong, such as when no backend has a given name. - ``except StreamError`` does not catch it. It fails a Workflow Task, and an Activity retries. + It isn't a :py:class:`StreamError`, so ``except StreamError`` doesn't catch it. In Workflow + code it fails the Workflow Task. An Activity that raises it retries. """ +def no_default_backend( + owner: str, names: Sequence[str], hint: str +) -> StreamConfigError: + """The error of a call with no backend name when no backend has ``default=True``.""" + return StreamConfigError( + f"No Stream backend of the {owner} ({', '.join(names)}) has default=True, so the " + "default is native, which this SDK does not have. Mark a backend with " + f"default=True, or name one, such as with {hint}." + ) + + class ReadOnlyContextError(StreamError): """Workflow code that must not change state, such as a query handler, wrote to a Stream.""" @@ -115,23 +135,17 @@ class ReadOnlyContextError(StreamError): "stream_fenced": StreamFencedError, "stream_storage": StreamStorageError, "stream_offset": StreamOffsetError, + "stream_too_large": StreamTooLargeError, "stream_unsupported": StreamUnsupportedError, "stream_not_owner": StreamNotOwnerError, "seq_divergent": StreamSeqDivergentError, } -# 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_call_error(err: Exception) -> StreamError: """Turn a failed call of the crate's runtime into the Stream error for its reason. - Its args are the gRPC code, the message and the encoded ``StreamErrorDetails``. + ``err.args`` are the gRPC code, the message and the encoded ``StreamErrorDetails``. """ code, message, details = err.args reason = StreamErrorDetails.FromString(details).reason if details else "" @@ -161,14 +175,3 @@ def from_activity_error(err: Exception) -> Exception: reason = str(cause.details[0]) if cause.details else "" return _BY_TYPE[cause.type](cause.message, reason=reason) return err - - -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/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py new file mode 100644 index 0000000000..b27dd1e940 --- /dev/null +++ b/temporalio/common/_streams/_extension.py @@ -0,0 +1,550 @@ +"""The Stream backends that call the temporalio-streams runtime in the bridge.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import logging +import weakref +from collections.abc import Mapping, Sequence +from datetime import timedelta +from typing import Any, NoReturn + +import temporalio.activity +import temporalio.api.common.v1 +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.bridge.streams +import temporalio.client +import temporalio.common +import temporalio.converter +import temporalio.runtime +import temporalio.service +import temporalio.worker +import temporalio.workflow +from temporalio.common._streams import _activity, _ids, _writes +from temporalio.common._streams import _workflow as streams_workflow +from temporalio.common._streams._backend import BackendCalls, _ExtensionCalls +from temporalio.common._streams._errors import StreamConfigError +from temporalio.common._streams._handle import StreamHandle +from temporalio.common._streams._plugin import StreamBackend +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.common._streams.system import models +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner + +logger = logging.getLogger(__name__) + +# The backend puts this key into the headers of each Local Activity that a Workflow schedules, +# and of each Activity of a run that uses Streams. It keeps the Producer id the same on retries. +_KEY_HEADER = "__temporal_stream_key" +# Continue-as-new passes the run chain's open Streams to the next run in this header. +_OWNED_STREAMS_HEADER = "__temporal_owned_streams" +_CLOSE_TIMEOUT = timedelta(seconds=30) +# Only a Store error is retryable. Each Stream error sets its own non_retryable flag. +_CLOSE_RETRY = temporalio.common.RetryPolicy( + initial_interval=timedelta(seconds=1), maximum_interval=timedelta(seconds=30) +) + + +class ExtensionStreamBackend(StreamBackend): + """Base class of the Redis and memory backends, which call the temporalio-streams runtime. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + memory: bool, + config: ext_pb.RedisStreamConfig, + runtime: temporalio.runtime.Runtime | None, + backend_name: str, + default: bool, + ) -> None: + """Create the backend over the bridge's Stream runtime, on Redis or in memory. + + ``backend_name`` is the name that Stream handles store for this backend. ``default`` + makes it the backend of a create with no backend name. + """ + config.backend = backend_name + self._config = config.SerializeToString() + self._backends: weakref.WeakSet[_ExtensionCalls] = weakref.WeakSet() + # Each client gets its own runtime object on Redis. The memory Store has one for the + # backend, so every client and Worker that shares the backend shares its Streams. + self._memory: _ExtensionCalls | None = None + if memory: + self._memory = self._new_backend( + runtime or temporalio.runtime.Runtime.default(), True + ) + super().__init__( + name, + backend_name=backend_name, + default=default, + activities=_writes.ACTIVITIES, + ) + + def _new_backend( + self, runtime: temporalio.runtime.Runtime, memory: bool = False + ) -> _ExtensionCalls: + if self._memory is not None and not memory: + return self._memory + backend = _ExtensionCalls( + temporalio.bridge.streams.StreamRuntime( + runtime._core_runtime, memory, self._config + ) + ) + self._backends.add(backend) + return backend + + def _configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + config = super()._configure_client(config) + backend = self._new_backend(_runtime_of(config["service_client"])) + config["interceptors"] = [ + *config.get("interceptors", []), + _StreamInterceptor(self, backend, config["data_converter"]), + ] + return config + + def _backend_for(self, client: temporalio.client.Client) -> BackendCalls: + # The Stream runtime that this backend gave the client. + interceptor = self._interceptor_of( + client.config(active_config=True).get("interceptors", []) + ) + if interceptor is None or interceptor.backend is None: + raise StreamConfigError( + f"The client has no {type(self).__name__}. Register it with " + "stream_backends=." + ) + return interceptor.backend + + def _interceptor_of(self, interceptors: Sequence[Any]) -> _StreamInterceptor | None: + for interceptor in interceptors: + if ( + isinstance(interceptor, _StreamInterceptor) + and interceptor.plugin is self + ): + return interceptor + return None + + def _configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + # The Worker gets the Stream interceptor from its client. + config = super()._configure_worker(config) + client = config["client"] # type: ignore[reportTypedDictNotRequiredAccess] + interceptors = [ + *client.config(active_config=True).get("interceptors", []), + *config.get("interceptors", []), + ] + if any( + isinstance(other, _StreamInterceptor) and other.plugin is not self + for other in interceptors + ): + raise StreamConfigError( + "The Worker's client has another Stream backend. A Worker serves only one " + "backend." + ) + mine = self._interceptor_of(interceptors) + if mine is None: + raise StreamConfigError( + f"The Worker's client has no {type(self).__name__}. Register it with " + "stream_backends=." + ) + if config.get("no_remote_activities"): + # Another Worker on the task queue can run the close Activity, so it isn't an error. + logger.warning( + "The Worker on the task queue %s polls no Activities, so the close of a " + "Workflow's Stream waits for another Worker on that queue with the backend", + config.get("task_queue"), + ) + # Records get no serialization context, so Workflows, Activities and clients encode + # them alike. + mine.converter = client.data_converter + runner = config.get("workflow_runner") + if isinstance(runner, SandboxedWorkflowRunner): + # The interceptor and temporalio.workflow must share the run's Stream state, even + # when the restrictions don't pass temporalio through. + config["workflow_runner"] = dataclasses.replace( + runner, + restrictions=runner.restrictions.with_passthrough_modules( + "temporalio.common._streams" + ), + ) + return config + + def _configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + # A replay runs no Activity, so it needs the Workflow interceptor and no runtime. + config = super()._configure_replayer(config) + if any( + isinstance(other, _StreamInterceptor) and other.plugin is not self + for other in config.get("interceptors", []) + ): + raise StreamConfigError( + "The Replayer has another Stream backend. A Replayer serves only one " + "backend, as a Worker does." + ) + config["interceptors"] = [ + *config.get("interceptors", []), + _StreamInterceptor(self, None, None), + ] + return config + + async def _close_in_workflow(self, request: stream_pb.CloseStreamRequest) -> int: + # An Activity runs only after the Workflow Task commits, so a failed task leaves the + # Stream open for its retry. The Activity closes the Stream as the run chain that owns it. + return await temporalio.workflow.execute_activity( + streams_workflow.CLOSE_ACTIVITY, + request, + result_type=int, + start_to_close_timeout=_CLOSE_TIMEOUT, + retry_policy=_CLOSE_RETRY, + ) + + async def shutdown(self) -> None: + """Close every open writer, including the client's. + + Call it when the process stops writing to Streams. + """ + failed = 0 + for backend in list(self._backends): + failed += await backend.close_writers() + if failed: + logger.warning("Stream shutdown left %d failed writers", failed) + + +def _runtime_of( + service: temporalio.service.ServiceClient, +) -> temporalio.runtime.Runtime: + # Find the runtime of a client's connection the same way the Worker does. + service = getattr(service, "worker_service_client", service) + runtime = getattr(service, "_runtime", None) or getattr( + getattr(service, "config", None), "runtime", None + ) + return runtime or temporalio.runtime.Runtime.default() + + +class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): + # There is one for each client of the backend, with that client's Stream runtime. The + # client's Workers get it from the client. + def __init__( + self, + plugin: ExtensionStreamBackend, + backend: _ExtensionCalls | None, + converter: temporalio.converter.DataConverter | None, + ) -> None: + super().__init__() + self.plugin = plugin + self.backend = backend + self.converter = converter or temporalio.converter.DataConverter.default + + def workflow_interceptor_class( + self, input: temporalio.worker.WorkflowInterceptorClassInput + ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: + # Workflow code reads the Worker's backend from the interceptor class. + return type( + "_WorkflowInbound", + (_WorkflowInbound,), + { + "_backend": self.plugin, + "_backend_name": self.plugin.backend_name, + "_default": self.plugin.default, + "_kind": self.plugin._kind, + }, + ) + + def intercept_activity( + self, next: temporalio.worker.ActivityInboundInterceptor + ) -> temporalio.worker.ActivityInboundInterceptor: + return _ActivityInbound(next, self) + + +class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): + _backend: ExtensionStreamBackend | None = None + _backend_name = "" + _default = False + _kind = "" + + def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: + assert self._backend is not None + super().init( + _WorkflowOutbound( + outbound, self._backend, self._backend_name, self._default, self._kind + ) + ) + + async def execute_workflow( + self, input: temporalio.worker.ExecuteWorkflowInput + ) -> Any: + _take_owned_streams(input.headers) + _receive(input.args) + try: + result = await super().execute_workflow(input) + except BaseException as err: + if _ends_run(err): + streams_workflow.close_at_run_end(err) + await streams_workflow.drain() + raise + streams_workflow.close_at_run_end(None) + await streams_workflow.drain() + return result + + async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None: + _receive(input.args) + await super().handle_signal(input) + + async def handle_update_handler( + self, input: temporalio.worker.HandleUpdateInput + ) -> Any: + _receive(input.args) + return await super().handle_update_handler(input) + + +def _receive(args: Sequence[Any]) -> None: + # A run that receives a Stream uses Streams. The inputs, Signals and Updates are in + # History, so a replay sets the same flag. + if _stream_arguments(args): + streams_workflow._state().uses_streams = True + + +def _ends_run(err: BaseException) -> bool: + # True if err ends the run, which must send its unawaited writes first. Any other exception + # fails the Workflow Task, and an eviction closes the coroutine outside the Workflow's loop. + if isinstance(err, temporalio.workflow.ContinueAsNewError): + return True + if isinstance(err, asyncio.CancelledError): + return temporalio.workflow.cancellation_reason() is not None + return isinstance( + err, Exception + ) and temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) + + +def _owned_ref(stream_id: str) -> models.StreamRef: + # The run chain owns each Stream that it creates. + info = temporalio.workflow.info() + return models.StreamRef( + id=stream_id, + workflow_id=info.workflow_id, + first_execution_run_id=info.first_execution_run_id, + ) + + +def _take_owned_streams( + headers: Mapping[str, temporalio.api.common.v1.Payload], +) -> None: + payload = headers.get(_OWNED_STREAMS_HEADER) + if payload is None: + return + state = streams_workflow._state() + state.uses_streams = True + converter = temporalio.converter.PayloadConverter.default + for stream_id, backend, *_ in converter.from_payload(payload, list[Any]): + state.stream(_owned_ref(stream_id)).backend = backend + + +class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): + _temporal_stream_transport = True + + def __init__( + self, + next: temporalio.worker.WorkflowOutboundInterceptor, + backend: ExtensionStreamBackend, + backend_name: str, + default: bool, + kind: str, + ) -> None: + super().__init__(next) + self._temporal_stream_backend = backend + self._temporal_stream_backend_name = backend_name + self._temporal_stream_default = default + self._temporal_stream_kind = kind + # Counts the run's Activity calls, so a key is the same on replay and on a retry. It + # doesn't call workflow.uuid4(), so the user's random values don't change. + self._activity_calls = 0 + + def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: + self._activity_calls += 1 + # Both ends use the default converter, so a custom user converter can't break the + # header. + converter = temporalio.converter.PayloadConverter.default + return {_KEY_HEADER: converter.to_payload(str(self._activity_calls))} + + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: + # The next run needs each owned Stream's backend for stream_handle. + owned = [ + [stream_id, run_stream.backend] + for stream_id, run_stream in streams_workflow._state().streams.items() + if run_stream.backend is not None and not run_stream.closed + ] + if owned: + converter = temporalio.converter.PayloadConverter.default + input.headers = { + **input.headers, + _OWNED_STREAMS_HEADER: converter.to_payload(owned), + } + super().continue_as_new(input) + + def start_activity( + self, input: temporalio.worker.StartActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + # The key is in History, so add it only for a run that uses Streams or passes one to + # this Activity. Count every call, so the keys are the same on replay. + headers = self._activity_headers() + state = streams_workflow._state() + if not state.uses_streams and _stream_arguments(input.args): + state.uses_streams = True + if state.uses_streams: + input.headers = {**input.headers, **headers} + return super().start_activity(input) + + def start_local_activity( + self, input: temporalio.worker.StartLocalActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + # A Local Activity's headers aren't in History, so every Local Activity gets the key. + input.headers = {**input.headers, **self._activity_headers()} + return super().start_local_activity(input) + + +class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): + def __init__( + self, + next: temporalio.worker.ActivityInboundInterceptor, + interceptor: _StreamInterceptor, + ) -> None: + super().__init__(next) + self._interceptor = interceptor + + async def execute_activity( + self, input: temporalio.worker.ExecuteActivityInput + ) -> Any: + streams = _activity_streams(input, self._interceptor) + token = _activity.current.set(streams) + try: + if temporalio.activity.info().attempt > 1: + # Withdraw the Records of earlier attempts, even if this attempt publishes + # nothing. This creates no Stream. + stream_ids = [handle.id for handle in _stream_arguments(input.args)] + for stream_id in dict.fromkeys(stream_ids): + await streams.writer(stream_id).withdraw_earlier_attempts() + result = await super().execute_activity(input) + except BaseException: + await _close_quietly(streams) + raise + else: + await streams.close() + return result + finally: + _activity.current.reset(token) + + +def _stream_arguments(args: Sequence[Any]) -> list[StreamHandle[Any]]: + # Find the Stream handles in call arguments. A handle can be in a field, a slot, a list or + # a dict. + handles: list[StreamHandle[Any]] = [] + seen: set[int] = set() + + def walk(value: Any, depth: int) -> None: + if depth > 4 or id(value) in seen: + return + seen.add(id(value)) + if isinstance(value, StreamHandle): + handles.append(value) + elif isinstance(value, (str, bytes, int, float, bool)) or value is None: + return + elif isinstance(value, dict): + for item in value.values(): + walk(item, depth + 1) + elif isinstance(value, (list, tuple, set, frozenset)): + for item in value: + walk(item, depth + 1) + else: + for item in _fields_of(value): + walk(item, depth + 1) + + for arg in args: + walk(arg, 0) + return handles + + +def _fields_of(value: Any) -> list[Any]: + fields = list(getattr(value, "__dict__", {}).values()) + for klass in type(value).__mro__: + slots = getattr(klass, "__slots__", ()) + for slot in (slots,) if isinstance(slots, str) else slots: + if hasattr(value, slot): + fields.append(getattr(value, slot)) + return fields + + +async def _close_quietly(streams: _activity.ActivityStreams) -> None: + try: + await asyncio.shield(streams.close()) + except Exception: + temporalio.activity.logger.warning( + "Failed flushing Stream Records of a failed Activity", exc_info=True + ) + + +def _activity_streams( + input: temporalio.worker.ExecuteActivityInput, interceptor: _StreamInterceptor +) -> _activity.ActivityStreams: + # A sync Activity has no client, so the interceptor gives the backend and the converter. + assert interceptor.backend is not None + info = temporalio.activity.info() + converter = temporalio.converter.PayloadConverter.default + key = input.headers.get(_KEY_HEADER) + key_text = converter.from_payload(key, str) if key is not None else "" + epoch = stream_pb.ProducerEpoch(attempt=info.attempt) + previous: stream_pb.WorkflowTaskEpoch | None = None + failed: tuple[int, ...] = () + if info.is_local: + task = info.workflow_task + if task is not None: + # Always set it, with 0 if no task completed. If it's unset, the Store doesn't + # withdraw the run's orphaned Local Activity Records. + previous = stream_pb.WorkflowTaskEpoch( + started_event_id=task.previous_started_event_id + ) + failed = tuple(task.failed_tasks) + epoch.task_started_event_id = task.started_event_id + epoch.task_attempt = task.attempt + if task.started_time is not None: + epoch.task_started_time.FromDatetime(task.started_time) + run_id = info.workflow_run_id or "" + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_LOCAL_ACTIVITY, + id=_ids.local_activity_producer_id(run_id, key_text), + run_id=run_id, + ) + else: + epoch.attempt_scheduled_time.FromDatetime(info.current_attempt_scheduled_time) + if info.in_workflow: + run_id = info.workflow_run_id or "" + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_ACTIVITY, + id=_ids.activity_producer_id(run_id, info.activity_id, key_text), + run_id=run_id, + ) + else: + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_STANDALONE_ACTIVITY, + id=_ids.standalone_activity_producer_id( + info.activity_id, info.activity_run_id or "" + ), + ) + producer.epoch.CopyFrom(epoch) + return _activity.ActivityStreams( + backend=interceptor.backend, + converter=interceptor.converter, + namespace=info.namespace, + producer=producer, + loop=asyncio.get_running_loop(), + backend_name=interceptor.plugin.backend_name, + default=interceptor.plugin.default, + workflow_id=info.workflow_id, + previous_completed_task=previous, + failed_tasks=failed, + ) diff --git a/temporalio/contrib/streams/_handle.py b/temporalio/common/_streams/_handle.py similarity index 85% rename from temporalio/contrib/streams/_handle.py rename to temporalio/common/_streams/_handle.py index f0f03e1d6e..f657dcb6d1 100644 --- a/temporalio/contrib/streams/_handle.py +++ b/temporalio/common/_streams/_handle.py @@ -12,8 +12,8 @@ import temporalio.converter import temporalio.exceptions -from temporalio.contrib.streams import _errors -from temporalio.contrib.streams._defn import TopicDefinition, definition_of +from temporalio.common._streams import _errors +from temporalio.common._streams._defn import TopicDefinition, definition_of E = TypeVar("E") T = TypeVar("T") @@ -23,7 +23,7 @@ def exit_failure(exc: BaseException | None) -> tuple[bool, BaseException | None] """Whether the end of an ``async with`` block closes the Stream, and with which failure. A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as - an eviction, leaves it open, since the block did not end in the caller's code. + ``KeyboardInterrupt``, leaves it open, since the caller's code didn't end the block. """ if isinstance(exc, asyncio.CancelledError): return True, temporalio.exceptions.CancelledError(str(exc) or "Cancelled") @@ -45,26 +45,26 @@ async def __aexit__( class Creating(Generic[S]): - """A Stream create. Await it for the handle, or use it in ``async with``, which closes the - Stream at the end of the block. + """A pending Stream create. - The create starts when it is awaited or entered, not before. + Await it to get the handle, or use it in ``async with``, which closes the Stream at the end + of the block. The create starts when it's awaited or entered, not before. .. warning:: This API is experimental. """ def __init__(self, create: Callable[[], Coroutine[Any, Any, S]]) -> None: - """Create it. Use a context's create call instead.""" + """Create the object. Use a context's create function instead.""" self._create = create self._stream: S | None = None def __await__(self) -> Generator[Any, None, S]: - """Create the Stream and give its handle.""" + """Create the Stream and return its handle.""" return self._create().__await__() async def __aenter__(self) -> S: - """Create the Stream and give its handle for the block.""" + """Create the Stream and return its handle for the block.""" self._stream = await self._create() return self._stream @@ -80,7 +80,7 @@ async def __aexit__( class TopicBase(Generic[T]): - """Base of the Topic of each handle class. All Topics of a Stream share one order.""" + """Base class 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.""" @@ -147,8 +147,8 @@ def _from_json(value: Any, type_hint: Any) -> StreamHandle[Any]: class _StreamHandleJson( temporalio.converter.TransferTypeConverter["StreamHandle[Any]", dict[str, str]] ): - # One form for every handle type, and the parameter's type hint picks the class, so a - # Workflow passes a Stream to an Activity as data. + # Every handle type has the same JSON form, and the parameter's type hint picks the class. + # So a Workflow can pass a Stream to an Activity as data. transfer_type: type[dict[str, str]] | None = dict def to_transfer_type(self, value: StreamHandle[Any]) -> dict[str, str]: @@ -161,11 +161,11 @@ def from_transfer_type( class StreamHandle(Generic[E]): - """Base of the Stream handle of each context. + """Base class of the Stream handle of each context. - A handle is data: it passes as a Workflow or Activity argument, as the proto JSON of the - ``StreamRef`` message. The parameter's type, such as ``ActivityStream[Events]``, gives the - handle class and the definition. + A handle is data. It's passed as a Workflow or Activity argument in the proto JSON form of + the ``StreamRef`` message. The parameter's type, such as ``ActivityStream[Events]``, sets + the handle class and the definition. .. warning:: This API is experimental. @@ -180,7 +180,7 @@ def __init__( backend: str | None = None, definition: type[E] | None = None, ) -> None: - """Create a handle. Use the create and get calls of a context instead.""" + """Create a handle. Use a context's create and get functions instead.""" self._id = id self._workflow_id = workflow_id self._first_execution_run_id = first_execution_run_id @@ -191,11 +191,10 @@ def __init__( @classmethod def from_handle(cls: type[H], handle: StreamHandle[Any]) -> H: - """The same Stream as a handle of this class, such as an ``ActivityStream`` from a - ``WorkflowStream``. + """Return the same Stream as a handle of this class. - A handle passes as data, so a Workflow uses it to give an Activity's input the type - that the Activity declares. + For example, a Workflow converts a ``WorkflowStream`` to an ``ActivityStream`` for an + Activity's input, so the input has the type that the Activity declares. """ return cls( id=handle.id, @@ -222,7 +221,7 @@ def first_execution_run_id(self) -> str | None: @property def backend(self) -> str | None: - """Name of the backend that keeps the Stream. None means the default one.""" + """Name of the backend that stores the Stream. None means the default one.""" return self._backend @property @@ -234,7 +233,7 @@ def _topic_spec( self, topic: str | TopicDefinition[Any], type: Any ) -> tuple[str, Any]: # The name and the decode type of a Topic. A Topic of a definition keeps its type, - # and a name outside the definition needs a type, so a typo does not pass. + # and a name outside the definition needs a type, so a typo raises an error. definition = definition_of(self._definition) if self._definition else None if isinstance(topic, TopicDefinition): if type is not None: @@ -272,10 +271,6 @@ 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. @@ -296,7 +291,7 @@ def __eq__(self, other: object) -> bool: return self._fields() == other._fields() def __hash__(self) -> int: - """Hash of the fields that never change, so it stays when a client learns the id.""" + """Hash of the fields that never change, so it's the same after a client learns the id.""" return hash((type(self), self._backend, self._workflow_id or self._id)) def __repr__(self) -> str: @@ -309,7 +304,7 @@ def __copy__(self) -> Self: return self def __deepcopy__(self, memo: dict[int, Any]) -> Self: - """A deep copy is the same handle, since dataclasses.asdict copies its fields.""" + """Return the same handle, since dataclasses.asdict deep-copies a handle field.""" return self @classmethod diff --git a/temporalio/contrib/streams/_ids.py b/temporalio/common/_streams/_ids.py similarity index 100% rename from temporalio/contrib/streams/_ids.py rename to temporalio/common/_streams/_ids.py diff --git a/temporalio/contrib/streams/_items.py b/temporalio/common/_streams/_items.py similarity index 79% rename from temporalio/contrib/streams/_items.py rename to temporalio/common/_streams/_items.py index 9148e8157e..94b9388e34 100644 --- a/temporalio/contrib/streams/_items.py +++ b/temporalio/common/_streams/_items.py @@ -2,21 +2,20 @@ from __future__ import annotations -import asyncio import enum from dataclasses import dataclass from typing import Any, Generic, TypeVar from temporalio.api.streamservice.v1 import CloseStatus as _CloseStatus from temporalio.api.streamservice.v1 import ProducerKind -from temporalio.contrib.streams.system import models -from temporalio.exceptions import CancelledError +from temporalio.common._streams.system import models +from temporalio.exceptions import is_cancelled_exception T = TypeVar("T") class End(enum.Enum): - """A ``from_offset`` after the newest Record, for new Records only. Use ``END``.""" + """An ``offset`` past the newest Record, to read only new Records. Use ``END``.""" END = "END" @@ -39,7 +38,7 @@ class Record(Generic[T]): """Value of the Record, decoded with the Topic's type.""" offset: int - """Place of the Record in the Stream. Ordered and unique, with gaps.""" + """Position of the Record in the Stream. Offsets are ordered and unique, with gaps.""" producer: models.Producer | None """Producer that wrote the Record.""" @@ -47,11 +46,11 @@ class Record(Generic[T]): @dataclass(frozen=True) class Superseded: - """A failed attempt's Records are withdrawn. + """An item that withdraws the Records of a failed attempt. - It comes at the same point for every Reader. A live view takes back the Records in - :py:attr:`withdrawn_offsets`. A view that keeps Records from an earlier read, such as one - before a reconnect, uses :py:meth:`withdraws` on them. + Every reader gets it at the same point in the Stream. A live view removes the Records in + :py:attr:`withdrawn_offsets`. A view that keeps Records from an earlier read, such as from + before a reconnect, checks them with :py:meth:`withdraws`. The next write of a run to a Stream withdraws the Records that the run's Local Activities wrote there in Workflow Tasks that never completed. A Stream that the run never writes @@ -62,7 +61,7 @@ class Superseded: """ offset: int - """Place of the item in the Stream.""" + """Position of the item in the Stream.""" supersession: models.Supersession """Which earlier Records the item withdraws.""" @@ -70,7 +69,7 @@ class Superseded: withdrawn_offsets: tuple[int, ...] = () """The offsets of the Records that this read gave and this item withdraws, in order. - A read gives each offset once. Records from before the read's start are not in it. + A read gives each offset once. It doesn't include Records from before the read's start. """ def withdraws(self, record: Record[Any]) -> bool: @@ -92,7 +91,7 @@ class CloseStatus(enum.IntEnum): """Closed with a failure.""" CANCELED = int(_CloseStatus.CLOSE_STATUS_CANCELED) - """Closed with a ``CancelledError`` failure, such as by a cancelled owner run.""" + """Closed with a cancel failure, such as by a cancelled owner run.""" TERMINATED = int(_CloseStatus.CLOSE_STATUS_TERMINATED) """The server closed it, since the owner run was terminated.""" @@ -103,7 +102,7 @@ class CloseStatus(enum.IntEnum): def close_status(failure: BaseException | None) -> _CloseStatus.ValueType | None: """The status a caller's close asks for. None lets the backend pick completed or failed.""" - if isinstance(failure, (asyncio.CancelledError, CancelledError)): + if failure is not None and is_cancelled_exception(failure): return _CloseStatus.CLOSE_STATUS_CANCELED return None @@ -117,17 +116,19 @@ class Closed: """ offset: int - """Place of the close in the Stream.""" + """Position of the close in the Stream.""" result: Any | None - """Result that the close carries, decoded with no type. None if it has none.""" + """Result of the close, decoded with no type hint. None if it has none.""" failure: BaseException | None - """Failure that the close carries. None if it has none.""" + """Failure of the close. None if it has none.""" status: CloseStatus = CloseStatus.COMPLETED - """How the Stream closed, so a reader tells a cancelled or terminated owner from a failed - one with no check of the failure's type. + """How the Stream closed. + + A reader can tell a cancelled or terminated owner from a failed one without checking the + failure's type. """ @@ -160,9 +161,7 @@ def _completed(task: models.WorkflowTaskEpoch) -> tuple[int, int, tuple[int, int def producer_key(producer: models.Producer) -> tuple[Any, ...]: - """The fields of a Producer that :py:func:`withdraws` looks at, so equal keys are withdrawn - together. - """ + """The Producer fields that :py:func:`withdraws` reads, so equal keys are withdrawn together.""" epoch = _epoch(producer.epoch) if producer.epoch is not None else None return (producer.kind or 0, producer.id, producer.run_id or "", epoch) @@ -180,7 +179,7 @@ def withdraws( ) -> bool: """Whether a Superseded item withdraws a Record of ``producer``. - It is the same rule as the crate's, so every SDK withdraws the same Records. + The rule matches the crate's, so every SDK withdraws the same Records. """ if producer is None: return False @@ -215,6 +214,6 @@ def withdraws( return False if (completed.attempt or 0) > 0: return old < _completed(completed) - # Core gives the completed task with attempt 0 when it no longer has the attempt. The - # attempts of one task start in time order, so the started time still tells the lower ones. + # Core sends the completed task with attempt 0 when it no longer knows the attempt. A task's + # attempts start in time order, so the started time still orders them. return old[2] < _time(completed.started_time) diff --git a/temporalio/common/_streams/_plugin.py b/temporalio/common/_streams/_plugin.py new file mode 100644 index 0000000000..d3c9fcc17f --- /dev/null +++ b/temporalio/common/_streams/_plugin.py @@ -0,0 +1,226 @@ +"""The base of the Stream backends, and the rules of their registration.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +import temporalio.client +import temporalio.service +import temporalio.worker +from temporalio.common._streams import _errors +from temporalio.common._streams._backend import BackendCalls +from temporalio.plugin import SimplePlugin + +# The default backend when no registered backend has default=True. This SDK has no native +# backend, so a call that needs the default then raises StreamConfigError. +NATIVE = "native" + + +class StreamBackend(SimplePlugin): + """Base class of the Stream backends. Each instance registers one backend under its name. + + Register backends on the client with ``Client.connect(..., stream_backends=[...])``. A + Worker uses its client's backends. A ``Replayer`` takes ``stream_backends=`` too. Names + must be unique. With more than one registered backend, exactly one must have + ``default=True``. With no backend marked, the default is ``native``. This SDK has no + ``native`` backend, so a create with no backend raises ``StreamConfigError``. Mark a + backend with ``default=True``, or use the functions of its package, such as + ``temporalio.contrib.stream.workflow.create_stream``. + + A backend is not a plugin, so passing it in ``plugins=`` raises ``StreamConfigError``. + + .. warning:: + This API is experimental. + """ + + # The kind of Store, such as "redis". The functions of a backend's package use it to find + # the registered backend of their kind. + _kind = "" + + def __init__( + self, + name: str, + *, + backend_name: str, + default: bool, + **options: Any, + ) -> None: + """Create the backend. ``options`` go to :py:class:`temporalio.plugin.SimplePlugin`.""" + self._backend_name = backend_name + self._default = default + super().__init__(name, **options) + + @property + def backend_name(self) -> str: + """The backend name that Stream handles store.""" + return self._backend_name + + @property + def default(self) -> bool: + """Whether a create with no backend name uses this backend.""" + return self._default + + def _backend_for( + self, + client: temporalio.client.Client, # pyright: ignore[reportUnusedParameter] + ) -> BackendCalls: + # Each backend subclass returns its calls for the client and its Workers. + raise NotImplementedError(f"{type(self).__name__} gives no calls") + + async def connect_service_client( + self, + config: temporalio.service.ConnectConfig, + next: Callable[ + [temporalio.service.ConnectConfig], + Awaitable[temporalio.service.ServiceClient], + ], + ) -> temporalio.service.ServiceClient: + """Refuse a backend in ``plugins=``, before the client connects.""" + raise _not_a_plugin(self) + + def configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + def configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + def configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + """Refuse a backend in ``plugins=``.""" + raise _not_a_plugin(self) + + # Stream registration calls these, so the public hooks above only run for a backend that + # was passed in plugins= by mistake. + + def _configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + return super().configure_client(config) + + def _configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + return super().configure_worker(config) + + def _configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + return super().configure_replayer(config) + + +def _not_a_plugin(backend: StreamBackend) -> _errors.StreamConfigError: + return _errors.StreamConfigError( + f"{type(backend).__name__} is a Stream backend, not a plugin. Pass it in " + "stream_backends= of Client.connect, or of the Replayer." + ) + + +def check_backends(backends: Sequence[object], owner: str) -> list[StreamBackend]: + """Check the registered backends, and return them. + + A replayed create uses the default backend, so the default must not depend on the order + of the list. + + Raises: + TypeError: An entry is not a ``StreamBackend``. + StreamConfigError: Two backends have the same name, or there is more than one backend + and not exactly one has ``default=True``. + """ + checked: list[StreamBackend] = [] + for backend in backends: + if not isinstance(backend, StreamBackend): + raise TypeError(f"{backend!r} is not a StreamBackend") + if any(other.backend_name == backend.backend_name for other in checked): + raise _errors.StreamConfigError( + f"The {owner} has two Stream backends named {backend.backend_name}" + ) + checked.append(backend) + if len(checked) < 2: + return checked + defaults = [b.backend_name for b in checked if b.default] + if len(defaults) != 1: + names = ", ".join(b.backend_name for b in checked) + marked = f"{', '.join(defaults)} have" if defaults else "none has" + raise _errors.StreamConfigError( + f"The {owner} has the Stream backends {names}, and {marked} default=True. " + "Mark exactly one with default=True." + ) + return checked + + +def default_name(backends: Sequence[StreamBackend]) -> str: + """The name of the backend with ``default=True``, or ``native`` if none has it.""" + return next((b.backend_name for b in backends if b.default), NATIVE) + + +def configure_client( + config: temporalio.client.ClientConfig, +) -> temporalio.client.ClientConfig: + """Check the client's backends, and apply the plugin configuration of each one.""" + backends = check_backends(config.get("stream_backends", ()), "client") + config["stream_backends"] = backends + for backend in backends: + config = backend._configure_client(config) + return config + + +def configure_worker( + backends: Sequence[StreamBackend], config: temporalio.worker.WorkerConfig +) -> temporalio.worker.WorkerConfig: + """Apply the plugin configuration of the client's backends to a Worker.""" + for backend in backends: + config = backend._configure_worker(config) + return config + + +def configure_replayer( + config: temporalio.worker.ReplayerConfig, +) -> temporalio.worker.ReplayerConfig: + """Check the Replayer's backends, and apply the plugin configuration of each one.""" + backends = check_backends(config.get("stream_backends", ()), "Replayer") + config["stream_backends"] = backends + for backend in backends: + config = backend._configure_replayer(config) + return config + + +def client_backends(client: temporalio.client.Client) -> list[StreamBackend]: + """The backends registered on a client.""" + return list(client.config(active_config=True).get("stream_backends", ())) + + +def client_backend( + client: temporalio.client.Client, name: str | None, *, hint: str = "backend=" +) -> tuple[str, BackendCalls]: + """The name and the calls of a client's backend. A ``name`` of None picks the default. + + Raises: + StreamConfigError: The client has no backend of that name, or none has + ``default=True`` for a call with no name. + """ + backends = client_backends(client) + wanted = name or default_name(backends) + for backend in backends: + if backend.backend_name == wanted: + return wanted, backend._backend_for(client) + if not backends: + raise _errors.StreamConfigError( + "The client has no Stream backend. Register one with " + "Client.connect(..., stream_backends=[...])." + ) + if not name: + raise _errors.no_default_backend( + "client", [b.backend_name for b in backends], hint + ) + raise _errors.StreamConfigError( + f"The client has no Stream backend {name}. It has " + f"{', '.join(b.backend_name for b in backends)}." + ) diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/common/_streams/_reader.py similarity index 65% rename from temporalio/contrib/streams/_reader.py rename to temporalio/common/_streams/_reader.py index a2b6f48353..7713031108 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/common/_streams/_reader.py @@ -9,9 +9,9 @@ from typing import Any, Generic, TypeVar import temporalio.converter -from temporalio.api.streamservice.v1 import ReadStart, StreamRecordKind -from temporalio.contrib.streams import _backend, _convert, _errors -from temporalio.contrib.streams._items import ( +from temporalio.api.streamservice.v1 import ProducerKind, ReadStart, StreamRecordKind +from temporalio.common._streams import _backend, _convert, _errors +from temporalio.common._streams._items import ( Closed, CloseStatus, End, @@ -20,8 +20,8 @@ producer_key, withdraws, ) -from temporalio.contrib.streams.system import client as service_client -from temporalio.contrib.streams.system import models +from temporalio.common._streams.system import client as service_client +from temporalio.common._streams.system import models T = TypeVar("T") @@ -33,15 +33,15 @@ @dataclasses.dataclass class Position: - """Where a read goes on from. The read moves it past each page it finished.""" + """Where a read resumes from. The read moves it past each page that it finished.""" next_offset: int | None class Subscription(Generic[T]): - """The items of one read of a Stream, in order. Use it in ``async for``. + """The items of one read of a Stream, in order. Iterate it with ``async for``. - It gives :py:class:`Record` items, :py:class:`Superseded` items that withdraw earlier + It yields :py:class:`Record` items, :py:class:`Superseded` items that withdraw earlier Records, and a :py:class:`Closed` item last when the Stream is closed. .. warning:: @@ -59,18 +59,19 @@ def __init__( self._position = position # A client read's rpc_metadata and rpc_timeout, for each call of each step. self._options = options - # The Producer and the offsets of each Record this read gave, with no values, so a - # Superseded item can give the offsets it withdraws. + # The Producer and the offsets of each Record that this read gave, without values, so a + # Superseded item can list the offsets that it withdraws. self._shown: dict[tuple[Any, ...], tuple[models.Producer, list[int]]] = {} @property def next_offset(self) -> int | None: - """The offset to resume from, after the last item given or the last page read. + """The offset to resume from, after the last item or the last page that the read got. - After a ``Superseded`` item it is that item's own offset, since the first Record of its - batch has the same offset and comes next. A read from there gives the ``Superseded`` - item again, which withdraws nothing new. Pass it as ``from_offset`` to a later read. It - is None only before the first page of a read from ``END``. + Pass it as ``offset`` to a later read. After a ``Superseded`` item, it's that item's own + offset, since the first Record of its batch has the same offset and comes next. A read + from there gives the ``Superseded`` item again, which withdraws nothing new. After a + Record that failed to decode, it's past that Record. It's None before the first page of + a read from ``END`` or of a read with ``last``. """ return self._position.next_offset @@ -79,17 +80,27 @@ def __aiter__(self) -> Subscription[T]: return self async def __anext__(self) -> Record[T] | Superseded | Closed: - """The next item.""" + """The next item. + + Raises: + StreamError: A Record's value failed to decode with its Topic's type. The read + ends, and :py:attr:`next_offset` is past that Record. + """ with _backend.using(self._options): item = await self._items.__anext__() if isinstance(item, Superseded): - # The Record of the same entry shares the offset and is not given yet. + # The Record of the same entry has the same offset and comes next. self._position.next_offset = item.offset else: self._position.next_offset = max( self._position.next_offset or 0, item.offset + 1 ) - if isinstance(item, Record) and item.producer is not None: + if ( + isinstance(item, Record) + and item.producer is not None + # Nothing withdraws a client's Records, so a long read of them keeps no offsets. + and item.producer.kind != ProducerKind.PRODUCER_KIND_CLIENT + ): key = producer_key(item.producer) self._shown.setdefault(key, (item.producer, []))[1].append(item.offset) elif isinstance(item, Superseded): @@ -108,7 +119,11 @@ def _withdraw(self, supersession: models.Supersession) -> tuple[int, ...]: return tuple(sorted(offsets)) async def aclose(self) -> None: - """End the read early. It makes no more calls.""" + """End the read early, from the task that reads it. It makes no more calls. + + If another task is waiting in the read for new Records, this raises ``RuntimeError``. + Cancel that task instead. + """ with _backend.using(self._options): await self._items.aclose() @@ -128,8 +143,10 @@ async def read( wait_new_records: bool, create_wait: timedelta | None, ) -> AsyncGenerator[Item, None]: - """Give the items of a Stream in order, until a close, or the newest Record when - ``wait_new_records`` is False. With ``last``, start with the newest ``last`` Records. + """Yield the items of a Stream in order, until a close. + + With ``wait_new_records`` False, stop at the newest Record. With ``last``, start with the + newest ``last`` Records. """ loop = asyncio.get_running_loop() create_until = loop.time() + create_wait.total_seconds() if create_wait else None @@ -142,7 +159,7 @@ def request() -> models.ReadRecordsRequest: ) return models.ReadRecordsRequest( stream=models.StreamRef(id=stream_id), - # END is only the first page's start. Each answer's next offset pins later pages. + # END applies only to the first page. Later pages start at the last page's next offset. from_offset=from_offset if isinstance(from_offset, int) else None, start=ReadStart.READ_START_END if from_offset is End.END @@ -162,7 +179,19 @@ def request() -> models.ReadRecordsRequest: try: async for page in pages: for record in page.records or []: - item = await _item(converter, record, types, default_type) + try: + item = await _item(converter, record, types, default_type) + except Exception as err: + # Move next_offset past the Record, so a resume doesn't fail on it again. + # The error gives the offset, for a read of that Record. + offset = record.offset or 0 + position.next_offset = max( + position.next_offset or 0, offset + 1 + ) + raise _errors.StreamError( + f"Can't decode the Record at the offset {offset} of the Topic " + f"{record.topic or ''}: {err}" + ) from err if item is None: continue yield item @@ -186,10 +215,10 @@ async def _pages( request: models.ReadRecordsRequest, wait_new_records: bool, ) -> AsyncGenerator[models.ReadRecordsResponse, None]: - # A read that waits long-polls each page. A read with no wait reads pages until one has no - # Record. + # A read that waits long-polls each page. A read with no wait reads pages until a page has + # no Record and doesn't advance. if request.last: - # The newest Records come in one page that never waits. A read that follows goes on + # The newest Records come in one page that never waits. A waiting read then continues # after them. page = await service.read_records(request) yield page @@ -206,7 +235,12 @@ async def _pages( yield page if page.done: return - if not wait_new_records and (not page.records or page.next_offset is None): + if not wait_new_records and ( + page.next_offset is None + # A Store scans a limited number of entries per call. So an empty page that advanced + # stopped in Records that a Topic filter skipped, not at the newest Record. + or (not page.records and page.next_offset <= (request.from_offset or 0)) + ): return if page.next_offset is not None: request = dataclasses.replace(request, from_offset=page.next_offset) @@ -231,11 +265,16 @@ async def _item( if record.kind == StreamRecordKind.STREAM_RECORD_KIND_CLOSE: result = _convert.decode(converter, record.close_result, None) failure = record.close_failure - # A close Record from before close statuses has none, and its failure tells. + # An older close Record has no status, so its failure decides the status. status = record.close_status or ( CloseStatus.FAILED if failure is not None else CloseStatus.COMPLETED ) return Closed( - offset=offset, result=result, failure=failure, status=CloseStatus(status) + offset=offset, + result=result, + failure=_convert.failure(converter, failure) + if failure is not None + else None, + status=CloseStatus(status), ) return None diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py new file mode 100644 index 0000000000..49a808053c --- /dev/null +++ b/temporalio/common/_streams/_stream_client.py @@ -0,0 +1,550 @@ +"""Streams in client code.""" + +from __future__ import annotations + +import asyncio +import functools +import itertools +import uuid +import weakref +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from datetime import timedelta +from types import TracebackType +from typing import Any, TypeVar, overload + +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.client +from temporalio.common._streams import _backend, _errors, _ids +from temporalio.common._streams._access import ( + Access, + ReadSpec, + _AsyncTopic, + read_spec, + subscription, +) +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._defn import TopicDefinition, definition_of +from temporalio.common._streams._handle import Creating, StreamHandle, exit_failure +from temporalio.common._streams._items import End +from temporalio.common._streams._plugin import client_backend +from temporalio.common._streams._reader import Item, Position, Subscription +from temporalio.common._streams._writer import StreamWriter +from temporalio.common._streams.system import models + +E = TypeVar("E") +T = TypeVar("T") + +# The handles of one client and Stream write as one Producer. Each client gets its own key, +# since two clients can have two runtime objects, whose writers would send clashing sequences. +_PROCESS_KEY = str(uuid.uuid4()) +_CLIENTS = itertools.count(1) +_CLIENT_KEYS: weakref.WeakKeyDictionary[temporalio.client.Client, str] = ( + weakref.WeakKeyDictionary() +) + + +def _client_producer( + client: temporalio.client.Client, named: tuple[str, int] | None +) -> stream_pb.Producer: + if named is not None: + name, epoch = named + return stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, + id=_ids.named_client_producer_id(name), + epoch=stream_pb.ProducerEpoch(attempt=epoch), + ) + key = _CLIENT_KEYS.get(client) + if key is None: + key = f"{_PROCESS_KEY}-{next(_CLIENTS)}" + _CLIENT_KEYS[client] = key + return stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, + id=_ids.client_producer_id(client.identity, key), + epoch=stream_pb.ProducerEpoch(attempt=1), + ) + + +class ClientTopic(_AsyncTopic[T]): + """One Topic of a :py:class:`ClientStream`. All Topics of a Stream share one order. + + .. warning:: + This API is experimental. + """ + + async def publish( + self, + *values: T, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> int: + """Write one Record for each value, in order, and return the last Record's offset. + + A list is one value. ``rpc_metadata`` and ``rpc_timeout`` are as for + :py:meth:`ClientStream.get_offset`. + """ + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await super().publish(*values) + + +class ClientStream(StreamHandle[E]): + """A Stream in client code. The client can write and read it, and close a client's Stream. + + Get one from :py:meth:`temporalio.client.Client.stream`, + :py:meth:`temporalio.client.WorkflowHandle.stream` or a backend's ``create`` function, such + as :py:func:`temporalio.contrib.stream.create`. + + Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. All + handles of one client and Stream write as one Producer. :py:meth:`producer` returns a + handle that writes as a named Producer. Use a handle in ``async with`` to close the Stream + at the end of the block. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + id: str, + workflow_id: str | None = None, + first_execution_run_id: str | None = None, + backend: str | None = None, + definition: type[E] | None = None, + ) -> None: + """Create a handle. Use :py:meth:`temporalio.client.Client.stream` instead.""" + super().__init__( + id=id, + workflow_id=workflow_id, + first_execution_run_id=first_execution_run_id, + backend=backend, + definition=definition, + ) + self._client: temporalio.client.Client | None = None + self._stream_backend: BackendCalls | None = None + self._find_first_run: Callable[[], Awaitable[str]] | None = None + self._writer: StreamWriter | None = None + # 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 + self._closed = False + # The name and the epoch of a named Producer. + self._named: tuple[str, int] | None = None + + def _bind( + self, + client: temporalio.client.Client, + backend: BackendCalls, + find_first_run: Callable[[], Awaitable[str]] | None = None, + ) -> ClientStream[E]: + self._client = client + self._stream_backend = backend + self._find_first_run = find_first_run + return self + + @property + def id(self) -> str: + """The Stream id. + + Raises: + RuntimeError: The Workflow handle had no first run id, and no call has found it + yet. Call :py:meth:`get_offset` first. + """ + if not self._id: + raise RuntimeError( + "The Stream id is not known yet. Call get_offset() first." + ) + return self._id + + async def _stream_id(self) -> str: + if not self._id: + # Concurrent calls share one Describe call. A failed Describe is discarded, so a + # later call tries again. + if self._finding is None or _failed(self._finding): + assert self._find_first_run is not None + self._finding = asyncio.ensure_future(self._find_first_run()) + first_run = await asyncio.shield(self._finding) + if not self._id: + assert self._workflow_id is not None + self._first_execution_run_id = first_run + self._id = _ids.scoped_stream_id( + self._workflow_id, first_run, self._stream_name + ) + return self._id + + def _fields(self) -> dict[str, str]: + fields = super()._fields() + if not self._id: + # The Stream name tells apart the Streams of one Workflow before the id is known. + fields["name"] = self._stream_name + return fields + + def _access(self) -> Access: + client, backend = self._client, self._stream_backend + if client is None or backend is None: + raise _errors.StreamConfigError( + "This ClientStream came from a payload, so it has no client. " + "Use Client.stream(definition, stream=handle)." + ) + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + writers=backend.writers(client.namespace, client.data_converter), + stream_id=self._stream_id, + producer=_client_producer(client, self._named), + free=backend.free_writer, + ) + return self._writer + + return Access( + backend=backend, + converter=client.data_converter, + namespace=client.namespace, + stream_id=self._stream_id, + writer=writer, + ) + + def producer(self, name: str, epoch: int) -> ClientStream[E]: + """Return this Stream as a handle that writes as the Producer ``name`` at ``epoch``. + + Use it for a writer that must fence its earlier instances, such as a server that + restarts. Keep ``epoch`` in durable storage, and raise it at each start. After a write + at a higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The + Records of the earlier epochs aren't withdrawn, since they're real messages. Run only + one live writer for each name. + + Args: + name: The Producer name, the same across restarts. + epoch: The writer's epoch, 1 or more. + """ + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + handle = ClientStream( + id=self._id, + workflow_id=self._workflow_id, + first_execution_run_id=self._first_execution_run_id, + backend=self._backend, + definition=self._definition, + ) + handle._stream_name = self._stream_name + handle._named = (name, epoch) + if self._client is not None and self._stream_backend is not None: + handle._bind(self._client, self._stream_backend, self._find_first_run) + return handle + + @overload + def topic(self, topic: TopicDefinition[T]) -> ClientTopic[T]: ... + + @overload + def topic(self, topic: str) -> ClientTopic[Any]: ... + + @overload + def topic(self, topic: str, *, type: type[T]) -> ClientTopic[T]: ... + + def topic( + self, topic: str | TopicDefinition[Any], *, type: Any = None + ) -> ClientTopic[Any]: + """A Topic of this Stream. + + Args: + topic: A Topic of the Stream's definition, such as ``Events.turns``, or a name. + A name in the definition gets the definition's type. + type: The type of a Topic outside the definition, or of a Stream with no + definition. + """ + name, topic_type = self._topic_spec(topic, type) + return ClientTopic(self, name, topic_type) + + def subscribe( + self, + *topics: str | TopicDefinition[Any] | ClientTopic[Any], + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + topic_prefixes: Sequence[str] = (), + create_wait: timedelta | None = None, + type: Any = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Subscription[Any]: + """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. To stop after a quiet period, use your own timer, and + resume later with ``offset=`` set to the read's ``next_offset``. A ``Superseded`` item + withdraws earlier Records of a failed attempt. Topic filters never skip it. + + Args: + topics: Read these Topics, by Topic of the definition, name or Topic object. + With no Topics and no ``topic_prefixes``, read every Topic. + offset: Start at this offset, inclusive, or at ``common.END`` for new + Records only. Defaults to the first Record. + last: Start with the newest ``last`` Records that pass the filter, oldest first, + such as for the first view of a chat. Can't be used with ``offset``. A Record + that a later ``Superseded`` item withdraws doesn't count. At most 1000. + wait_new_records: Wait for new Records at the end of the Stream. + topic_prefixes: Also read the Topics whose names start with one of these. + create_wait: How long to wait for a Stream that doesn't exist yet, such as the + Stream of a Workflow that just started. Defaults to no wait. + type: The type of the values of Topics outside the definition, such as Topics that + ``topic_prefixes`` match. Topics of the definition keep their own types. + Defaults to no type, so the converter decodes the values with no type hint. + rpc_metadata: Headers of each call of the read. See :py:meth:`get_offset`. + rpc_timeout: The deadline of each call of the read, not of the whole read. A call + that waits for new Records returns before the deadline. + """ + read = read_spec( + self, + topics, + topic_prefixes=topic_prefixes, + default_type=type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + return subscription( + lambda position: self._items(read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) + + async def _items( + self, read: ReadSpec, position: Position + ) -> AsyncGenerator[Item, None]: + items = self._access().items(read, position) + try: + async for item in items: + yield item + finally: + await items.aclose() + + async def get_offset( + self, + *, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + + Args: + rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so + they send none. + rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with + ``DEADLINE_EXCEEDED``. + """ + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await self._access().get_offset() + + def _check_close(self) -> None: + if self._workflow_id: + # Only the run chain can close a Workflow's Stream. The Store refuses a client's + # close too, but this check makes no call. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {self._id or self._stream_name} from a client, since " + "a Workflow owns it", + reason="stream_not_owner", + ) + + async def close( + self, + *, + result: Any | None = None, + failure: BaseException | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> None: + """Write this handle's Records, then close the Stream. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`get_offset`. + + Raises: + StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. + """ + self._check_close() + self._closed = True + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + await self._access().close(result, failure) + + async def __aenter__(self) -> ClientStream[E]: + """Use the Stream in ``async with``, which closes it at the end. + + Raises: + StreamNotOwnerError: A Workflow owns the Stream, so only its run chain can close it. + """ + # Fail early, since the close at the end would raise and hide the block's own error. + self._check_close() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this + does nothing. + """ + if self._closed: + return + close, failure = exit_failure(exc) + if close: + await self.close(failure=failure) + + +def create( + client: temporalio.client.Client, + definition: type[E], + *, + id: str, + retention: timedelta | None, + backend: str, +) -> Creating[ClientStream[E]]: + """Create a Stream that the client owns, on a registered backend. + + See :py:func:`temporalio.contrib.stream.create`. + """ + + async def create() -> ClientStream[E]: + definition_of(definition) + backend_name, calls = client_backend(client, backend) + response = await calls.service(client.namespace).create_stream( + models.CreateStreamRequest(stream_id=id, retention=retention) + ) + if response.status == stream_pb.StreamStatus.STREAM_STATUS_CLOSED: + raise _errors.StreamClosedError( + f"The Stream {id} was created and closed before", + reason="stream_closed", + ) + if response.status == stream_pb.StreamStatus.STREAM_STATUS_EXPIRED: + raise _errors.StreamExpiredError( + f"The Stream {id} expired", reason="stream_expired" + ) + stream = ClientStream(id=id, backend=backend_name, definition=definition) + return stream._bind(client, calls) + + return Creating(create) + + +def client_stream( + client: temporalio.client.Client, + definition: type[E], + *, + id: str | None, + stream: StreamHandle[Any] | None, + backend: str | None, +) -> ClientStream[E]: + """Get a Stream by its id, or from a handle that came in data. + + See :py:meth:`temporalio.client.Client.stream`. + """ + definition_of(definition) + if (id is None) == (stream is None): + raise TypeError("stream takes one of id= and stream=") + if stream is not None: + if backend and stream.backend and backend != stream.backend: + raise ValueError( + f"The handle is of the backend {stream.backend}, not {backend}" + ) + backend = backend or stream.backend + stream_id = stream.id if stream else id + if not stream_id: + # An empty id would fail only at the first call, with an unclear error. + raise ValueError("The Stream id is empty") + backend_name, calls = client_backend(client, backend) + handle = ClientStream( + id=stream_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, + definition=definition, + ) + return handle._bind(client, calls) + + +# The first run id of a Workflow handle, or the Describe call that fetches it, so the Streams +# of one handle share the call. +_FIRST_RUNS: weakref.WeakKeyDictionary[ + temporalio.client.WorkflowHandle[Any, Any], asyncio.Future[str] | str +] = weakref.WeakKeyDictionary() + + +def _failed(future: asyncio.Future[Any]) -> bool: + return future.done() and (future.cancelled() or future.exception() is not None) + + +async def _describe_first_run( + handle: temporalio.client.WorkflowHandle[Any, Any], +) -> str: + # The Describe uses the call options of the first call that needs the id. + options = _backend.call_options.get() + description = await handle.describe( + rpc_metadata=options.metadata if options else {}, + rpc_timeout=options.timeout if options else None, + ) + first = description.raw_description.workflow_execution_info.first_run_id + if not first: + raise ValueError("The server gave no first run id, so the Stream id is unknown") + return first + + +async def _first_run(handle: temporalio.client.WorkflowHandle[Any, Any]) -> str: + known = _FIRST_RUNS.get(handle) + if isinstance(known, str): + return known + if known is None: + known = asyncio.ensure_future(_describe_first_run(handle)) + _FIRST_RUNS[handle] = known + known.add_done_callback(functools.partial(_settle, weakref.ref(handle))) + return await asyncio.shield(known) + + +def _settle( + handle_ref: weakref.ref[temporalio.client.WorkflowHandle[Any, Any]], + finding: asyncio.Future[str], +) -> None: + # Store a found id as text. Remove a failed call, so a later call tries again and the + # failure's traceback doesn't keep the handle alive. + handle = handle_ref() + if handle is None or _FIRST_RUNS.get(handle) is not finding: + return + if _failed(finding): + del _FIRST_RUNS[handle] + else: + _FIRST_RUNS[handle] = finding.result() + + +def workflow_stream( + handle: temporalio.client.WorkflowHandle[Any, Any], + definition: type[E], + *, + backend: str | None, +) -> ClientStream[E]: + """Get a Stream that a Workflow created. + + See :py:meth:`temporalio.client.WorkflowHandle.stream`. + """ + stream_name = definition_of(definition).name + client = handle._client + backend_name, calls = client_backend(client, backend) + first_run = handle.first_execution_run_id + stream = ClientStream( + id=_ids.scoped_stream_id(handle.id, first_run, stream_name) + if first_run + else "", + workflow_id=handle.id, + first_execution_run_id=first_run, + backend=backend_name, + definition=definition, + ) + if first_run: + return stream._bind(client, calls) + stream._stream_name = stream_name + return stream._bind(client, calls, lambda: _first_run(handle)) diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/common/_streams/_workflow.py similarity index 50% rename from temporalio/contrib/streams/workflow.py rename to temporalio/common/_streams/_workflow.py index b22671361a..bd3955e0a4 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/common/_streams/_workflow.py @@ -1,16 +1,16 @@ """Streams in Workflow code. -Workflow code creates a Stream with :py:func:`create_stream` and gets one again with -:py:func:`stream_handle`. A Stream that the Workflow receives as an argument, such as a -:py:class:`WorkflowStream` that a run passes on through continue-as-new, works the same. - -Each create, publish and close is a Local Activity, so it has a place in the run's commands, -as an Activity call has. To add, remove or move one in a Workflow that has open runs, gate the -change with ``workflow.patched``, else the replay of those runs fails as nondeterministic. -Publishes that wait behind a write to the same Stream go as one Local Activity. - -.. warning:: - This module is experimental. +Workflow code creates a Stream with :py:func:`temporalio.workflow.create_stream` and gets it +again with :py:func:`temporalio.workflow.stream_handle`. A Stream that the Workflow receives as +an argument works the same, such as a :py:class:`WorkflowStream` that a run passes on through +continue-as-new. + +Each create and publish is a Local Activity. Each close is a command that the backend runs after +the Workflow Task commits, and on Redis it's a regular Activity. So each one has a place in the +run's commands. To add, remove or move one in a Workflow that has open runs, gate the change +with ``workflow.patched``, or the replay of those runs fails as nondeterministic. Publishes that +wait behind a write to the same Stream are sent as one Local Activity. If a task fails or the +Worker crashes before the close commits, the Stream stays open for the task's retry. """ from __future__ import annotations @@ -19,7 +19,7 @@ from dataclasses import dataclass, field from datetime import timedelta from types import TracebackType -from typing import Any, TypeVar, overload +from typing import Any, NamedTuple, TypeVar, overload import google.protobuf.message @@ -27,46 +27,36 @@ import temporalio.converter 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 ( +from temporalio.common._streams import _errors, _ids +from temporalio.common._streams._defn import TopicDefinition, definition_of +from temporalio.common._streams._handle import ( Creating, StreamHandle, TopicBase, exit_failure, ) -from temporalio.contrib.streams._items import close_status -from temporalio.contrib.streams.system import models +from temporalio.common._streams._items import close_status +from temporalio.common._streams.system import models E = TypeVar("E") T = TypeVar("T") -# The plugin's Local Activities that write for Workflow code. The marker keeps the answer only, -# since a Local Activity's input is not in History. +# The create and the append are Local Activities, so History has only their results. The Redis +# close is a regular Activity on the Workflow's task queue. CREATE_ACTIVITY = "__temporal_streams_create" APPEND_ACTIVITY = "__temporal_streams_append" CLOSE_ACTIVITY = "__temporal_streams_close" _WRITE_TIMEOUT = timedelta(seconds=30) -# The most that one merged append takes before the codec. The Local Activity checks the size -# after the codec, and sends each call as its own batch when the codec grew the merge past the -# Store's 4 MiB batch limit. Half of it leaves room for most codecs. +# Limits of one merged append before the codec. Half of the Store's 4 MiB batch limit leaves room +# for most codecs. The Local Activity splits a merge by call if the codec grows it past the limit. _MERGE_RECORDS = 1000 _MERGE_BYTES = 2 * 1024 * 1024 -__all__ = [ - "WorkflowStream", - "WorkflowTopic", - "create_stream", - "create_stream_nowait", - "stream_handle", -] - @dataclass class _RunStream: - # What this run knows about one Stream. Many handles may name it, so it is kept by id. + # What this run knows about one Stream. Many handles can name it, so the state is keyed by id. ref: models.StreamRef - close_with_run: bool = False closed: bool = False # Set for a Stream that the run chain created, which continue-as-new passes on. backend: str | None = None @@ -74,8 +64,8 @@ class _RunStream: @dataclass class _Write: - # One Local Activity that writes to a Stream, and the calls it answers. Each call gets the - # Records up to its own last one, and its future. + # One Local Activity that writes to a Stream, and the calls that it answers. Each call has the + # Record count up to its own last Record, and its future. activity: str request: google.protobuf.message.Message calls: list[tuple[int, asyncio.Future[int]]] @@ -83,8 +73,8 @@ class _Write: started: bool = False def take(self, request: google.protobuf.message.Message) -> bool: - # An append that waits for the write before it takes the appends that come after it, - # so a burst of awaits costs two markers and not one for each await. + # An append that waits behind the previous write merges the appends that come after it. + # So a burst of awaits costs two markers, not one for each await. if self.started or not isinstance(request, stream_pb.AppendRecordsRequest): return False mine = self.request @@ -107,10 +97,20 @@ class _RunState: streams: dict[str, _RunStream] = field(default_factory=dict) # The last write of each Stream. Each write waits for it, so writes keep their order. tails: dict[str, asyncio.Task[None]] = field(default_factory=dict) - # The last write of each Stream that has not started, which later appends may join. + # The last write of each Stream that hasn't started, which later appends can join. waiting: dict[str, _Write] = field(default_factory=dict) - # The Streams that a retry run's first write has withdrawn the previous run on. + # The Streams on which a retry run's first write withdrew the previous run's Records. withdrawn: set[str] = field(default_factory=set) + # Whether the run created, bound or received a Stream. Only then do its Activities get a + # Producer key, so a run with no Stream adds nothing to History. + uses_streams: bool = False + # The async with blocks whose close waits for the run's outcome, with each block's failure. + deferred_closes: list[tuple[WorkflowStream[Any], BaseException | None]] = field( + default_factory=list + ) + # The publishes that failed before a caller got the error. A caller can still get it, so + # the run's end checks them again. + failed_publishes: list[asyncio.Future[int]] = field(default_factory=list) def stream(self, ref: models.StreamRef) -> _RunStream: known = self.streams.get(ref.id) @@ -121,8 +121,8 @@ def stream(self, ref: models.StreamRef) -> _RunStream: def _state() -> _RunState: - # The state is an attribute of the run, not a weak map entry, since its tasks refer to the - # run's loop and would keep a map entry alive. + # Keep the state on the run, since its tasks refer to the run's loop and would keep a weak + # map entry alive. runtime = temporalio.workflow._Runtime.current() state = getattr(runtime, "_temporal_streams_state", None) if state is None: @@ -131,31 +131,72 @@ def _state() -> _RunState: return state -def _backend_name() -> str | None: - # A Stream plugin marks its outbound interceptor with its backend name. With none, the - # call would go to the server as a System Nexus command, which only native Streams serve. +class WorkerBackend(NamedTuple): + """The Stream backend of the Worker that runs the Workflow.""" + + name: str + default: bool + kind: str + + +def _transport() -> Any: + # The run's outbound interceptor that the Worker's backend marks, or None. outbound = getattr(temporalio.workflow._Runtime.current(), "_outbound", None) while outbound is not None: if getattr(outbound, "_temporal_stream_transport", False): - return getattr(outbound, "_temporal_stream_backend_name", "") + return outbound outbound = getattr(outbound, "next", None) return None +def worker_backend() -> WorkerBackend | None: + """The Worker's backend, which marks the run's outbound interceptor, or None.""" + outbound = _transport() + if outbound is None: + return None + return WorkerBackend( + getattr(outbound, "_temporal_stream_backend_name", ""), + getattr(outbound, "_temporal_stream_default", False), + getattr(outbound, "_temporal_stream_kind", ""), + ) + + +async def _close_in_workflow(request: stream_pb.CloseStreamRequest) -> int: + # The Worker's backend runs the close after the Workflow Task commits. + outbound = _transport() + if outbound is None: + raise no_worker_backend() + return await outbound._temporal_stream_backend._close_in_workflow(request) + + +def no_worker_backend() -> _errors.StreamConfigError: + """The error of Workflow code on a Worker with no Stream backend.""" + return _errors.StreamConfigError( + "No Stream backend is on the Worker. Register one on its client with " + "stream_backends=." + ) + + def _check_backend(backend: str | None, action: str) -> str: - # The Worker's config gives these errors, not the history, so they fail the Workflow Task - # 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 _errors.StreamConfigError( - "No Stream plugin is on the Worker, such as RedisStreamPlugin" - ) - if backend and backend != worker_backend: + # These errors come from the Worker's config, not from History. So they fail the Workflow + # Task, and a Worker with the right config retries it. + worker = worker_backend() + if worker is None: + raise no_worker_backend() + if not backend: + if not worker.default: + raise _errors.no_default_backend( + "Worker", + [worker.name], + "temporalio.contrib.stream.workflow.create_stream", + ) + return worker.name + if backend != worker.name: raise _errors.StreamConfigError( f"Can't {action} a Stream of the backend {backend}. The Worker has only " - f"{worker_backend}." + f"{worker.name}." ) - return worker_backend + return worker.name def _records(request: google.protobuf.message.Message) -> int: @@ -202,7 +243,7 @@ async def run() -> None: del state.waiting[stream_id] try: if isinstance(write.request, stream_pb.AppendRecordsRequest): - # The Local Activity answers each merged call with its last offset or its error. + # The Local Activity returns each merged call's last offset or error. counts = [ count - before for (before, _), (count, _) in zip( @@ -215,6 +256,9 @@ async def run() -> None: result_type=list[Any], start_to_close_timeout=_WRITE_TIMEOUT, ) + elif activity == CLOSE_ACTIVITY: + assert isinstance(write.request, stream_pb.CloseStreamRequest) + answers = [await _close_in_workflow(write.request)] else: answers = [ await temporalio.workflow.execute_local_activity( @@ -256,7 +300,8 @@ def _write_nowait(future: asyncio.Future[int]) -> None: runtime = temporalio.workflow._Runtime.current() def done(future: asyncio.Future[int]) -> None: - # No caller awaits the write, so its error fails the Workflow Task. + # No caller awaits the write, so the run raises its error after this callback, as an + # error that the run doesn't catch. if not future.cancelled() and future.exception() is not None: setattr(runtime, "_current_activation_error", future.exception()) @@ -264,17 +309,31 @@ def done(future: asyncio.Future[int]) -> None: async def drain() -> None: - """Wait for the run's writes that nobody awaits, before the run ends.""" - for task in list(_state().tails.values()): + """Wait for the run's writes and closes that nobody awaits, before the run ends. + + A publish error that no caller got raises here, as an error that the run doesn't catch. + """ + state = _state() + for task in list(state.tails.values()): try: await task except BaseException: # A cancelled write already cancelled its callers. pass + for future in state.failed_publishes: + error = future.exception() if _unretrieved(future) else None + if error is not None: + raise error + + +def _unretrieved(future: asyncio.Future[int]) -> bool: + # asyncio sets this flag when the future fails and clears it when a caller gets the error, + # with await, result() or exception(). It is what asyncio logs as never retrieved. + return getattr(future, "_log_traceback", False) def _payload(value: Any) -> Any: - # The Workflow's converter runs here, in the sandbox. The plugin runs only the codec. + # The Workflow's converter runs here, in the sandbox. The backend runs only the codec. return temporalio.workflow.payload_converter().to_payload(value) @@ -285,30 +344,47 @@ class WorkflowTopic(TopicBase[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. + def publish(self, *values: T) -> asyncio.Future[int]: + """Start the write of one Record for each value, in order, and return its future. + + The write has its place in the Stream's order when the call returns, also when no + caller awaits the future. Writes to one Stream keep the order of the calls. The write + is a Local Activity, so History has its result and not the values. Publishes that wait + behind a write to the same Stream are sent as one Local Activity, and each call gets the + offset of its own last Record. A list is one value. - The write is a Local Activity, so History has its answer and not the values. Publishes - that wait behind a write to the same Stream go as one Local Activity, and each call - gets the offset of its own last Record. A list is one value. Writes to one Stream keep - the order of the calls. + A caller that awaits the future, or calls its ``result()`` or ``exception()``, gets the + write's error and owns it. An error that no caller gets raises at the end of the run, as + an error of :py:meth:`publish_nowait` does. So a ``StreamError`` fails the run. + + Returns: + The future of the write. It gives the offset of the last Record, or raises + ``StreamError`` when the backend refused the write, such as + ``StreamClosedError`` for a Stream that another caller closed. Raises: - StreamError: The backend refused the write, such as ``StreamClosedError`` for a - Stream that something else closed. StreamClosedError: This run closed the Stream. ReadOnlyContextError: The code must not change state, such as a query handler. - StreamConfigError: No Stream plugin is on the Worker, or the Stream is of a - backend that the Worker does not have. It fails the Workflow Task, not the run. + StreamConfigError: The Worker has no Stream backend, or not the Stream's backend. + It fails the Workflow Task, not the run. """ - return await self._stream._append(self._name, values) + state = _state() + future = self._stream._append(self._name, values) + + def failed(future: asyncio.Future[int]) -> None: + if _unretrieved(future): + state.failed_publishes.append(future) + + future.add_done_callback(failed) + return future def publish_nowait(self, *values: T) -> None: - """Write one Record for each value, in order, with no await. + """Write one Record for each value, in order, without an await. - It is for sync Workflow code, such as ``@workflow.init`` and sync signal and update - handlers, as ``asyncio.Queue.put_nowait`` is for code that can't await ``put``. A - backend error fails the Workflow Task. + Use it in sync Workflow code, such as ``@workflow.init`` and sync signal and update + handlers. It's like ``asyncio.Queue.put_nowait`` for code that can't await ``put``. A + backend error raises in the run as an error that the run doesn't catch, so a + ``StreamError`` fails the run. """ _write_nowait(self._stream._append(self._name, values)) @@ -317,7 +393,9 @@ class WorkflowStream(StreamHandle[E]): """A Stream in Workflow code. Workflow code writes and closes it, and can't read it. Pass it to an Activity, a child Workflow or the next run of continue-as-new as an - argument. Use it in ``async with`` to close it at the end of the block. + argument. Use it in ``async with`` to close it at the end of the block. The end of a run + closes only the Streams whose ``async with`` block deferred the close to the run end. A + Stream that the run chain never closes stays open until its ``safety_ttl``. .. warning:: This API is experimental. @@ -353,7 +431,7 @@ def topic( Args: topic: A Topic of the Stream's definition, such as ``Events.turns``, or a name. - A name of the definition takes its type. + A name in the definition gets the definition's type. type: The type of a Topic outside the definition, or of a Stream with no definition. @@ -369,7 +447,9 @@ def _check(self, action: str, *, backend_checked: bool = False) -> _RunStream: raise _errors.ReadOnlyContextError( f"Can't {action} a Stream in read-only code" ) - run_stream = _state().stream(self._model_ref()) + state = _state() + state.uses_streams = True + run_stream = state.stream(self._model_ref()) if run_stream.closed: raise _errors.StreamClosedError( f"Can't {action} the Stream {self.id}, since this run closed it", @@ -395,11 +475,18 @@ def _append(self, topic: str, values: tuple[Any, ...]) -> asyncio.Future[int]: async def close( self, *, result: Any | None = None, failure: BaseException | None = None ) -> None: - """Close the Stream with one Local Activity, after the writes before it. + """Close the Stream after the writes before it, once the Workflow Task commits. + + If the task fails first, or the Worker crashes, the Stream stays open, and the task's + retry writes to it again. On Redis the close is a regular Activity on the Workflow's + task queue, which closes the Stream as the run chain that owns it. So every Worker on + that queue must register the Redis backend. A Worker that polls Activities runs the + close, and a Worker with ``no_remote_activities=True`` doesn't. - Readers get the Records before it, then a ``Closed`` item with ``result`` or - ``failure``. A ``CancelledError`` failure closes it as canceled. Later writes in this run - raise :py:class:`StreamClosedError`. + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError`` or an ``ActivityError`` + that a cancel caused, closes the Stream as canceled. Later writes in this run raise + :py:class:`StreamClosedError`. Raises: StreamNotOwnerError: This run chain didn't create the Stream. @@ -409,26 +496,30 @@ async def close( def close_nowait( self, *, result: Any | None = None, failure: BaseException | None = None ) -> None: - """Close the Stream with no await, for sync Workflow code. + """Close the Stream without an await, for sync Workflow code. - It is to :py:meth:`close` what ``asyncio.Queue.put_nowait`` is to ``put``. A backend - error fails the Workflow Task. + It is to :py:meth:`close` what ``asyncio.Queue.put_nowait`` is to ``put``. It starts the + close, and the end of the run waits for it. A backend error raises in the run as an + error that the run doesn't catch, so a ``StreamError`` fails the run. """ _write_nowait(self._close(result, failure)) - def _close(self, result: Any, failure: BaseException | None) -> asyncio.Future[int]: + def _closing(self) -> _RunStream: info = temporalio.workflow.info() if (self.workflow_id, self.first_execution_run_id) != ( info.workflow_id, info.first_execution_run_id, ): - # Only the run chain that created a Stream ends it for every reader. The backend - # refuses it too, and this error comes before a Local Activity is spent. + # Only the run chain that created a Stream can close it. The backend refuses the + # close too, but this check fails before the Workflow schedules a command. raise _errors.StreamNotOwnerError( f"Can't close the Stream {self.id}, since this run chain didn't create it", reason="stream_not_owner", ) - self._check("close").closed = True + return self._check("close") + + def _close(self, result: Any, failure: BaseException | None) -> asyncio.Future[int]: + self._closing().closed = True request = stream_pb.CloseStreamRequest(stream=self._ref()) if result is not None: request.result.CopyFrom(_payload(result)) @@ -453,14 +544,99 @@ async def __aexit__( ) -> None: """Close the Stream, with the exception as its failure if there is one. - A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as - an eviction, leaves it open, since the block did not end in Workflow code. + A cancel of the block or of the run closes the Stream as canceled. An exception that + fails the run with no retry run closes it too. Examples are an exception in a run with + no retry policy, a non-retryable ``ApplicationError``, and an exception on the last + attempt. + + For any other exception, the close waits for the run's end. The Stream then closes with + that exception if the run returns, continues as new, or fails with no retry run. It + stays open if a retry run comes next, since that run writes to the same Stream. It also + stays open if the exception fails only the Workflow Task, since the task's retry runs + the block again. + + In a cron run, the block end leaves the Stream open, since the Stream spans the cron + chain and the next run creates it again. An eviction cancels the block, and the Worker + sends no commands for an evicted run. An explicit :py:meth:`close` always closes the + Stream. """ - if _state().stream(self._model_ref()).closed: + state = _state() + if state.stream(self._model_ref()).closed: return close, failure = exit_failure(exc) - if close: + if not close or temporalio.workflow.info().cron_schedule: + return + if _closes_at_block_end(exc): await self.close(failure=failure) + else: + # A refused close raises here, as it does for a close at the block end. + self._closing() + state.deferred_closes.append((self, failure)) + + +def _closes_at_block_end(exc: BaseException | None) -> bool: + # For other exceptions, only the run's end shows whether a retry run will write to the Stream. + if exc is None or isinstance(exc, asyncio.CancelledError) or _cancelled(exc): + return True + runtime = temporalio.workflow._Runtime.current() + return runtime.workflow_is_failure_exception(exc) and not _retried(exc) + + +def close_at_run_end(err: BaseException | None) -> None: + """Close the Streams whose block end deferred the close, unless a retry run comes next. + + ``err`` is the exception that ends the run, or None if the run returns. + """ + state = _state() + runtime = temporalio.workflow._Runtime.current() + if ( + isinstance(err, Exception) + and not _cancelled(err) + and runtime.workflow_is_failure_exception(err) + and _retried(err) + ): + return + for stream, failure in state.deferred_closes: + if not state.stream(stream._model_ref()).closed: + _write_nowait(stream._close(None, failure)) + + +def _retried(exc: BaseException) -> bool: + # Whether the server starts a retry run after the run fails with exc, by the server's rule. + info = temporalio.workflow.info() + policy = info.retry_policy + if policy is None or 0 < policy.maximum_attempts <= info.attempt: + return False + failure_type: str | None + if isinstance(exc, temporalio.exceptions.ApplicationError): + if exc.non_retryable: + return False + failure_type = exc.type + elif isinstance( + exc, + (temporalio.exceptions.CancelledError, temporalio.exceptions.TerminatedError), + ): + return False + elif isinstance(exc, temporalio.exceptions.TimeoutError): + return exc.type in ( + temporalio.exceptions.TimeoutType.START_TO_CLOSE, + temporalio.exceptions.TimeoutType.HEARTBEAT, + ) + elif isinstance(exc, temporalio.exceptions.FailureError): + return True + else: + # For any other exception, the run's failure has the class name as its type. + failure_type = type(exc).__name__ + return failure_type not in (policy.non_retryable_error_types or ()) + + +def _cancelled(exc: BaseException) -> bool: + # A cancelled run ends as canceled with no retry, even when the cancel reached the block as + # the error of an Activity, a Local Activity or a child. + return ( + temporalio.workflow.cancellation_reason() is not None + and temporalio.exceptions.is_cancelled_exception(exc) + ) def _own_handle(definition: type[E], backend: str | None) -> WorkflowStream[E]: @@ -482,47 +658,15 @@ def create_stream( definition: type[E], *, retention: timedelta | None = None, - auto_close: timedelta | None = None, - close_with_run: bool = True, backend: str | None = None, ) -> Creating[WorkflowStream[E]]: """Create a Stream that this run chain owns. - Its id is ``::``, so a client finds it from the Workflow - handle with no call. The create is one Local Activity, and its error raises here. - - Await the call for the handle, or use it in ``async with``, which closes the Stream at the - end of the block:: - - async with create_stream(streams.with_name(Events, "turn-3")) as turn: - await turn.topic(Events.tokens).publish("hi") - - Args: - definition: The Stream definition class. Its name is the Stream name. For one Stream - per turn, for example, pass ``streams.with_name(Events, f"turn-{n}")``. - retention: How long the Records stay after the close. Defaults to the backend's. - auto_close: Close the Stream with no result when no Record was appended for this - 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. 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: - StreamError: The backend refused the create, such as for another Retention. - ReadOnlyContextError: The code must not change state, such as a query handler. - 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. + See :py:func:`temporalio.workflow.create_stream`. """ async def create() -> WorkflowStream[E]: - stream, task = _create( - definition, retention, auto_close, close_with_run, backend - ) + stream, task = _create(definition, retention, backend) await task return stream @@ -533,17 +677,13 @@ def create_stream_nowait( definition: type[E], *, retention: timedelta | None = None, - auto_close: timedelta | None = None, - close_with_run: bool = True, backend: str | None = None, ) -> WorkflowStream[E]: - """Create a Stream that this run chain owns, with no await. + """Create a Stream that this run chain owns, without an await. - It is for sync Workflow code, such as ``@workflow.init``, as ``asyncio.Queue.put_nowait`` - is for code that can't await ``put``. It takes the options of :py:func:`create_stream`. A - backend error fails the Workflow Task. + See :py:func:`temporalio.workflow.create_stream_nowait`. """ - stream, task = _create(definition, retention, auto_close, close_with_run, backend) + stream, task = _create(definition, retention, backend) _write_nowait(task) return stream @@ -551,20 +691,13 @@ def create_stream_nowait( def _create( definition: type[E], retention: timedelta | None, - auto_close: timedelta | None, - close_with_run: bool, backend: str | None, ) -> tuple[WorkflowStream[E], asyncio.Future[int]]: stream = _own_handle(definition, _check_backend(backend, "create")) run_stream = stream._check("create", backend_checked=True) - if run_stream.close_with_run and not close_with_run: - raise ValueError( - f"The Stream {stream.id} already closes with the run chain, and " - "close_with_run=False can't turn that off" - ) run_id = temporalio.workflow.info().run_id - # The owner makes the Stream the run chain's, and names a retry run's withdrawal. - # The ref names the run chain, which the Store keeps for its owner check of a close. + # A retry run's withdrawal names the owner's run id. The Store keeps the ref's run chain, and + # checks each close against it. request = stream_pb.CreateStreamRequest( stream_id=stream.id, owner=stream_pb.Producer( @@ -576,9 +709,6 @@ def _create( ) if retention is not None: request.retention.FromTimedelta(retention) - if auto_close is not None: - request.auto_close.FromTimedelta(auto_close) - run_stream.close_with_run = close_with_run run_stream.backend = stream.backend return stream, _write(stream.id, CREATE_ACTIVITY, request) @@ -586,18 +716,16 @@ def _create( def stream_handle( definition: type[E], *, backend: str | None = None ) -> WorkflowStream[E]: - """A Stream that this run chain created, by its definition. It makes no call. - - A later run of the chain, such as one after continue-as-new, uses it to write the Stream - again. + """Get a Stream that this run chain created, by its definition. - Args: - definition: The Stream definition class, renamed as for the create. - backend: The backend of the Stream. Defaults to the one that the run chain created it - on, and else to the Worker's. + See :py:func:`temporalio.workflow.stream_handle`. """ stream = _own_handle(definition, backend) + if not temporalio.workflow.unsafe.is_read_only(): + _state().uses_streams = True if backend is None: known = _state().streams.get(stream.id) - stream._backend = (known.backend if known else None) or _backend_name() or None + stream._backend = (known.backend if known else None) or _check_backend( + None, "get" + ) return stream diff --git a/temporalio/contrib/streams/_writer.py b/temporalio/common/_streams/_writer.py similarity index 74% rename from temporalio/contrib/streams/_writer.py rename to temporalio/common/_streams/_writer.py index a02f419129..1ee5f2726d 100644 --- a/temporalio/contrib/streams/_writer.py +++ b/temporalio/common/_streams/_writer.py @@ -8,14 +8,14 @@ from typing import TYPE_CHECKING, Any import temporalio.api.streamservice.v1 as stream_pb -from temporalio.contrib.streams import _errors +from temporalio.common._streams import _errors if TYPE_CHECKING: - from temporalio.contrib.streams.writer_system import client as writer_client + from temporalio.common._streams.writer_system import client as writer_client class StreamWriter: - """Writes the Records of one Producer on one Stream, through the generated Writer.""" + """Writes the Records of one Producer to one Stream, through the generated Writer.""" def __init__( self, @@ -27,10 +27,10 @@ def __init__( previous_completed_task: stream_pb.WorkflowTaskEpoch | None = None, failed_tasks: Sequence[int] = (), ) -> None: - """Create a writer. It opens on its first write. + """Create a writer. It opens on the first write. - ``writers`` runs the values through the user's converter and codec. ``free`` releases - the crate's writer when Python drops this one with no close. + ``writers`` encodes the values with the user's converter and codec. ``free`` releases + the crate's writer if Python frees this one without a close. """ self._writers = writers self._resolve_id = stream_id @@ -49,20 +49,27 @@ def producer(self) -> stream_pb.Producer: return self._producer async def withdraw_earlier_attempts(self) -> None: - """Open the writer of a new attempt, which withdraws the Records of the earlier ones. + """Open the writer of a new attempt, which withdraws the Records of earlier attempts. - A Stream that does not exist has nothing to withdraw. + A Stream that doesn't exist has nothing to withdraw. A closed or expired Stream accepts + no write, so the attempt continues, and a later write raises the real error. """ try: await self.open() - except _errors.StreamNotFoundError: + except ( + _errors.StreamNotFoundError, + _errors.StreamClosedError, + _errors.StreamExpiredError, + ): pass async def open(self) -> writer_client.Writer: - """Open the crate's writer once. The open of a new attempt withdraws the earlier one's - Records, even when it writes nothing. + """Open the crate's writer once. + + The open of a new attempt withdraws the Records of earlier attempts, even if it writes + nothing. """ - from temporalio.contrib.streams.writer_system import models + from temporalio.common._streams.writer_system import models async with self._open_lock: if self._writer is None: @@ -83,18 +90,16 @@ async def open(self) -> writer_client.Writer: ) ) if self._free is not None: - # The crate keeps a writer for each open that no one closed or freed, so a - # writer that Python drops frees its open. + # The crate keeps a writer for each open that nobody closed or freed. So + # free it when Python frees this writer. self._release = weakref.finalize( self, self._free, self._writer.writer_id ) return self._writer async def write(self, records: Sequence[tuple[str, Any]]) -> int: - """Write ``(topic, value)`` pairs as Records now, in order. Returns the last Record's - offset. - """ - from temporalio.contrib.streams.writer_system import models + """Write ``(topic, value)`` pairs as Records now, in order, and return the last offset.""" + from temporalio.common._streams.writer_system import models if self._closed: raise _errors.StreamClosedError( @@ -107,7 +112,7 @@ async def write(self, records: Sequence[tuple[str, Any]]) -> int: return answer.last_offset or 0 async def close(self) -> None: - """Drop the writer. The Stream stays open.""" + """Close the writer. The Stream stays open.""" if self._closed: return self._closed = True diff --git a/temporalio/contrib/streams/_writes.py b/temporalio/common/_streams/_writes.py similarity index 62% rename from temporalio/contrib/streams/_writes.py rename to temporalio/common/_streams/_writes.py index b59c9aea5e..04e85d428b 100644 --- a/temporalio/contrib/streams/_writes.py +++ b/temporalio/common/_streams/_writes.py @@ -1,7 +1,8 @@ -"""The Local Activities that write for Workflow code. +"""The Activities that write for Workflow code. -Each Workflow ``create_stream`` and ``close`` is one of them, and the publishes that wait -behind a write to the same Stream share one. History has the answer and not the values. +Each Workflow ``create_stream`` is a Local Activity. The publishes that wait behind a write to +the same Stream share one Local Activity. History has the result and not the values. Each +Workflow ``close`` is a regular Activity, so it runs only after its Workflow Task completes. """ from __future__ import annotations @@ -13,27 +14,26 @@ import temporalio.activity import temporalio.api.streamservice.v1 as stream_pb import temporalio.common -from temporalio.contrib.streams import _activity, _errors -from temporalio.contrib.streams.system import client as service_client -from temporalio.contrib.streams.system import models -from temporalio.contrib.streams.workflow import ( +from temporalio.common._streams import _activity, _errors, _ids +from temporalio.common._streams._workflow import ( APPEND_ACTIVITY, CLOSE_ACTIVITY, CREATE_ACTIVITY, ) +from temporalio.common._streams.system import client as service_client +from temporalio.common._streams.system import models def _context() -> tuple[_activity.ActivityStreams, service_client.StreamServiceClient]: streams = _activity.current.get() if streams is None: - raise _errors.StreamConfigError("No Stream plugin runs this Local Activity") + raise _errors.StreamConfigError("No Stream backend runs this Local Activity") return streams, streams.backend.service(streams.namespace, streams.converter) def _producer(streams: _activity.ActivityStreams) -> models.Producer: - # Every attempt of one Workflow Task writes as one epoch with one sequence, so the backend - # drops the copy that a retry sends. A later task has a higher epoch, so its write - # withdraws what a failed task wrote under the same call key. + # All attempts of one Workflow Task write with one epoch and sequence, so the Store dedupes a + # retry. A later task's higher epoch withdraws a failed task's Records for the same call key. producer = streams.producer epoch = producer.epoch started = epoch.task_started_time if epoch.HasField("task_started_time") else None @@ -85,9 +85,6 @@ async def create_stream(request: stream_pb.CreateStreamRequest) -> int: retention=request.retention.ToTimedelta() if request.HasField("retention") else None, - auto_close=request.auto_close.ToTimedelta() - if request.HasField("auto_close") - else None, withdraw_run=request.withdraw_run or None, owner=models.Producer( kind=request.owner.kind, @@ -102,6 +99,10 @@ async def create_stream(request: stream_pb.CreateStreamRequest) -> int: f"The Stream {request.stream_id} was created and closed before", reason="stream_closed", ) + if response.status == stream_pb.StreamStatus.STREAM_STATUS_EXPIRED: + raise _errors.StreamExpiredError( + f"The Stream {request.stream_id} expired", reason="stream_expired" + ) return 0 @@ -112,17 +113,52 @@ async def create_stream(request: stream_pb.CreateStreamRequest) -> int: @dataclasses.dataclass class _Encoded: - # The batches of one append and the calls of each, with the bodies after the codec, and - # the answer of each batch that the Store took. + # The batches of one append, with each batch's calls and encoded bodies, and the result of + # each batch that the Store accepted. batches: list[tuple[list[int], list[stream_pb.AppendRecord]]] + size: int answers: dict[int, Any] = dataclasses.field(default_factory=dict) -# The appends that have not succeeded yet, by Producer and epoch. A retry of the Local Activity -# in this process sends the same bytes, so the Store's dedupe of a resend matches them, also -# when the codec adds random bytes. A rerun after a crash is a new epoch, so it needs none. -_PENDING: OrderedDict[bytes, _Encoded] = OrderedDict() -_PENDING_LIMIT = 1024 +_Key = tuple[str, str, str] +_Epoch = tuple[int, int, int, int] + + +class _Pending: + # The appends that haven't succeeded, one for each Producer. A retry in this process resends + # the same bytes, so the Store's dedupe matches them, even with a random codec. + + def __init__(self, limit: int) -> None: + self._limit = limit + self._size = 0 + self._entries: OrderedDict[_Key, tuple[_Epoch, _Encoded]] = OrderedDict() + + def get(self, key: _Key, epoch: _Epoch) -> _Encoded | None: + found = self._entries.get(key) + return found[1] if found is not None and found[0] == epoch else None + + def put(self, key: _Key, epoch: _Epoch, encoded: _Encoded) -> None: + found = self._entries.get(key) + if found is not None: + if found[0] > epoch: + # The Store fences this write, since the Producer has a later epoch. + return + self.pop(key, found[0]) + self._entries[key] = (epoch, encoded) + self._size += encoded.size + while self._size > self._limit and len(self._entries) > 1: + _, (_, oldest) = self._entries.popitem(last=False) + self._size -= oldest.size + + def pop(self, key: _Key, epoch: _Epoch) -> None: + found = self._entries.get(key) + if found is not None and found[0] == epoch: + del self._entries[key] + self._size -= found[1].size + + +# A run that ends while its append retries leaves its entry, so evict the oldest past this size. +_PENDING = _Pending(64 * 1024 * 1024) async def _encode( @@ -150,15 +186,15 @@ async def _encode( for record in records ] if sum(sizes) <= _MAX_BATCH_BYTES and max(sizes, default=0) <= _MAX_RECORD_BYTES: - return _Encoded([(list(range(len(counts))), records)]) - # The codec grew the merge past the Store's limit, so each call goes as its own batch, as - # it would with no merge. A call that is too large alone fails alone. + return _Encoded([(list(range(len(counts))), records)], sum(sizes)) + # The codec grew the merge past the Store's limit, so send each call as its own batch, as + # with no merge. A call that's too large by itself fails by itself. batches: list[tuple[list[int], list[stream_pb.AppendRecord]]] = [] start = 0 for call, count in enumerate(counts): batches.append(([call], records[start : start + count])) start += count - return _Encoded(batches) + return _Encoded(batches, sum(sizes)) @temporalio.activity.defn(name=APPEND_ACTIVITY) @@ -167,26 +203,25 @@ async def append_records( ) -> list[Any]: """Append the Records of a Workflow's merged publishes. - ``counts`` has the Records of each publish. Gives each publish the offset of its last - Record, or the error of a publish that the Store refused as too large. + ``counts`` has the Record count of each publish. Returns each publish's last offset, or the + error of a publish that the Store refused as too large. """ streams, service = _context() producer = _producer(streams) producer_pb = models._ProducerTransferTypeConverter().to_transfer_type(producer) - key = b"\n".join( - [ - streams.namespace.encode(), - request.stream.id.encode(), - producer_pb.SerializeToString(deterministic=True), - ] + key = (streams.namespace, request.stream.id, producer_pb.id) + task = producer_pb.epoch + epoch = ( + task.task_started_event_id, + task.task_attempt, + task.task_started_time.seconds, + task.task_started_time.nanos, ) - encoded = _PENDING.get(key) + encoded = _PENDING.get(key, epoch) if encoded is None: encoded = await _encode(streams, request, counts, producer_pb) - _PENDING[key] = encoded - while len(_PENDING) > _PENDING_LIMIT: - _PENDING.popitem(last=False) - # The bodies went through the codec once already. + _PENDING.put(key, epoch, encoded) + # The codec already ran on the bodies. service = streams.backend.service( streams.namespace, dataclasses.replace(streams.converter, payload_codec=None) ) @@ -216,6 +251,8 @@ async def append_records( answer = response.last_offset or 0 except _errors.StreamError as err: if err.reason != "stream_too_large": + if err.non_retryable: + _PENDING.pop(key, epoch) raise answer = {"reason": err.reason, "message": err.message} encoded.answers[index] = answer @@ -225,14 +262,22 @@ async def append_records( for call in calls: after -= counts[call] answers[call] = answer - after if isinstance(answer, int) else answer - _PENDING.pop(key, None) + _PENDING.pop(key, epoch) return answers @temporalio.activity.defn(name=CLOSE_ACTIVITY) async def close_stream(request: stream_pb.CloseStreamRequest) -> int: - """Close a Workflow's Stream.""" + """Close a Workflow's Stream, as the run chain that owns it.""" streams, service = _context() + # The Store refuses an Activity Producer's close of a Workflow's Stream, so the Activity + # writes the close as the run that scheduled it, and only for that run's chain. + if request.stream.workflow_id != streams.workflow_id: + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {request.stream.id} for another Workflow", + reason="stream_not_owner", + ) + run_id = streams.producer.run_id failure: Any = None if request.HasField("failure"): failure = streams.converter.failure_converter.from_failure( @@ -241,13 +286,16 @@ async def close_stream(request: stream_pb.CloseStreamRequest) -> int: await service.close_stream( models.CloseStreamRequest( stream=_ref(request.stream), - producer=_producer(streams), + producer=models.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_WORKFLOW, + id=_ids.workflow_producer_id(run_id), + run_id=run_id, + ), result=temporalio.common.RawValue(request.result) if request.HasField("result") else None, failure=failure, status=request.status or None, - **_history(streams), ) ) return 0 diff --git a/temporalio/contrib/streams/proto/__init__.py b/temporalio/common/_streams/proto/__init__.py similarity index 100% rename from temporalio/contrib/streams/proto/__init__.py rename to temporalio/common/_streams/proto/__init__.py diff --git a/temporalio/contrib/streams/proto/v1/__init__.py b/temporalio/common/_streams/proto/v1/__init__.py similarity index 100% rename from temporalio/contrib/streams/proto/v1/__init__.py rename to temporalio/common/_streams/proto/v1/__init__.py diff --git a/temporalio/contrib/streams/proto/v1/extension_pb2.py b/temporalio/common/_streams/proto/v1/extension_pb2.py similarity index 92% rename from temporalio/contrib/streams/proto/v1/extension_pb2.py rename to temporalio/common/_streams/proto/v1/extension_pb2.py index b35975f103..ea592706d7 100644 --- a/temporalio/contrib/streams/proto/v1/extension_pb2.py +++ b/temporalio/common/_streams/proto/v1/extension_pb2.py @@ -42,7 +42,7 @@ (_message.Message,), { "DESCRIPTOR": _REDISSTREAMCONFIG, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.RedisStreamConfig) }, ) @@ -53,7 +53,7 @@ (_message.Message,), { "DESCRIPTOR": _OPENWRITERREQUEST, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.OpenWriterRequest) }, ) @@ -64,7 +64,7 @@ (_message.Message,), { "DESCRIPTOR": _OPENWRITERRESPONSE, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.OpenWriterResponse) }, ) @@ -75,7 +75,7 @@ (_message.Message,), { "DESCRIPTOR": _WRITERECORDSREQUEST, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.WriteRecordsRequest) }, ) @@ -86,7 +86,7 @@ (_message.Message,), { "DESCRIPTOR": _WRITERECORDSRESPONSE, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.WriteRecordsResponse) }, ) @@ -97,7 +97,7 @@ (_message.Message,), { "DESCRIPTOR": _CLOSEWRITERREQUEST, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.CloseWriterRequest) }, ) @@ -108,7 +108,7 @@ (_message.Message,), { "DESCRIPTOR": _CLOSEWRITERRESPONSE, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.CloseWriterResponse) }, ) @@ -119,7 +119,7 @@ (_message.Message,), { "DESCRIPTOR": _FREEWRITERREQUEST, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.FreeWriterRequest) }, ) @@ -130,7 +130,7 @@ (_message.Message,), { "DESCRIPTOR": _FREEWRITERRESPONSE, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.FreeWriterResponse) }, ) @@ -141,7 +141,7 @@ (_message.Message,), { "DESCRIPTOR": _CLOSEWRITERSREQUEST, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.CloseWritersRequest) }, ) @@ -152,7 +152,7 @@ (_message.Message,), { "DESCRIPTOR": _CLOSEWRITERSRESPONSE, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.CloseWritersResponse) }, ) @@ -163,7 +163,7 @@ (_message.Message,), { "DESCRIPTOR": _STREAMERRORDETAILS, - "__module__": "temporalio.contrib.streams.proto.v1.extension_pb2", + "__module__": "temporalio.common._streams.proto.v1.extension_pb2", # @@protoc_insertion_point(class_scope:temporal.sdk.streams.v1.StreamErrorDetails) }, ) diff --git a/temporalio/contrib/streams/proto/v1/extension_pb2.pyi b/temporalio/common/_streams/proto/v1/extension_pb2.pyi similarity index 100% rename from temporalio/contrib/streams/proto/v1/extension_pb2.pyi rename to temporalio/common/_streams/proto/v1/extension_pb2.pyi diff --git a/temporalio/contrib/streams/system/__init__.py b/temporalio/common/_streams/system/__init__.py similarity index 100% rename from temporalio/contrib/streams/system/__init__.py rename to temporalio/common/_streams/system/__init__.py diff --git a/temporalio/contrib/streams/system/_resources/__init__.py b/temporalio/common/_streams/system/_resources/__init__.py similarity index 100% rename from temporalio/contrib/streams/system/_resources/__init__.py rename to temporalio/common/_streams/system/_resources/__init__.py diff --git a/temporalio/contrib/streams/system/_resources/stream.py b/temporalio/common/_streams/system/_resources/stream.py similarity index 99% rename from temporalio/contrib/streams/system/_resources/stream.py rename to temporalio/common/_streams/system/_resources/stream.py index bba5cfaeb8..1e268b96fc 100644 --- a/temporalio/contrib/streams/system/_resources/stream.py +++ b/temporalio/common/_streams/system/_resources/stream.py @@ -166,7 +166,7 @@ async def iter_read_records( async def describe( self, ) -> DescribeStreamResponse: - """Describe the Stream: its status, newest offset, Retention and auto-close.""" + """Describe the Stream: its status, newest offset and Retention.""" request = DescribeStreamRequest( stream=self.stream, ) diff --git a/temporalio/contrib/streams/system/_system_nexus_interceptor.py b/temporalio/common/_streams/system/_system_nexus_interceptor.py similarity index 100% rename from temporalio/contrib/streams/system/_system_nexus_interceptor.py rename to temporalio/common/_streams/system/_system_nexus_interceptor.py diff --git a/temporalio/contrib/streams/system/client.py b/temporalio/common/_streams/system/client.py similarity index 99% rename from temporalio/contrib/streams/system/client.py rename to temporalio/common/_streams/system/client.py index c667a86b74..3e717241c8 100644 --- a/temporalio/contrib/streams/system/client.py +++ b/temporalio/common/_streams/system/client.py @@ -483,7 +483,7 @@ async def iter_read_records( async def describe( self, ) -> DescribeStreamResponse: - """Describe the Stream: its status, newest offset, Retention and auto-close.""" + """Describe the Stream: its status, newest offset and Retention.""" request = DescribeStreamRequest( stream=self.stream, ) diff --git a/temporalio/contrib/streams/system/models.py b/temporalio/common/_streams/system/models.py similarity index 98% rename from temporalio/contrib/streams/system/models.py rename to temporalio/common/_streams/system/models.py index fb7101d8f8..9b8504ab8d 100644 --- a/temporalio/contrib/streams/system/models.py +++ b/temporalio/common/_streams/system/models.py @@ -23,12 +23,11 @@ @dataclasses.dataclass(slots=True) class CreateStreamRequest: - """Create a Stream idempotently. ALREADY_EXISTS for another Retention or auto-close.""" + """Create a Stream idempotently. ALREADY_EXISTS for another Retention.""" stream_id: str retention: datetime.timedelta | None = None withdraw_run: str | None = None - auto_close: datetime.timedelta | None = None owner: Producer | None = None close_with_run: bool | None = None ref: StreamRef | None = None @@ -60,9 +59,6 @@ def from_transfer_type( if value.HasField("retention") else None, withdraw_run=value.withdraw_run if bool(value.withdraw_run) else None, - auto_close=_support.duration_from_proto(value.auto_close) - if value.HasField("auto_close") - else None, owner=_ProducerTransferTypeConverter().from_transfer_type( value.owner, Producer ) @@ -89,8 +85,6 @@ def to_transfer_type( message.retention.CopyFrom(_support.duration_to_proto(value.retention)) if value.withdraw_run is not None: message.withdraw_run = value.withdraw_run - if value.auto_close is not None: - message.auto_close.CopyFrom(_support.duration_to_proto(value.auto_close)) if value.owner is not None: message.owner.CopyFrom( _ProducerTransferTypeConverter().to_transfer_type(value.owner) @@ -1245,7 +1239,6 @@ class DescribeStreamResponse: ) = None last_offset: int | None = None retention: datetime.timedelta | None = None - auto_close: datetime.timedelta | None = None class _DescribeStreamResponseTransferTypeConverter( @@ -1275,9 +1268,6 @@ def from_transfer_type( retention=_support.duration_from_proto(value.retention) if value.HasField("retention") else None, - auto_close=_support.duration_from_proto(value.auto_close) - if value.HasField("auto_close") - else None, ) @typing_extensions.override @@ -1296,8 +1286,6 @@ def to_transfer_type( message.last_offset = value.last_offset if value.retention is not None: message.retention.CopyFrom(_support.duration_to_proto(value.retention)) - if value.auto_close is not None: - message.auto_close.CopyFrom(_support.duration_to_proto(value.auto_close)) return message diff --git a/temporalio/contrib/streams/system/operations/__init__.py b/temporalio/common/_streams/system/operations/__init__.py similarity index 100% rename from temporalio/contrib/streams/system/operations/__init__.py rename to temporalio/common/_streams/system/operations/__init__.py diff --git a/temporalio/contrib/streams/system/operations/create_stream.py b/temporalio/common/_streams/system/operations/create_stream.py similarity index 87% rename from temporalio/contrib/streams/system/operations/create_stream.py rename to temporalio/common/_streams/system/operations/create_stream.py index f2f30d8324..10da759aad 100644 --- a/temporalio/contrib/streams/system/operations/create_stream.py +++ b/temporalio/common/_streams/system/operations/create_stream.py @@ -47,19 +47,16 @@ async def create_stream( stream_id: str, retention: datetime.timedelta | None = None, withdraw_run: str | None = None, - auto_close: datetime.timedelta | None = None, owner: Producer | None = None, close_with_run: bool | None = None, ref: StreamRef | None = None, ) -> CreateStreamResult: """ Args: - stream_id: Required from Workflow code. Empty elsewhere asks the service for an + stream_id: Required on Redis and memory. On native, empty asks the server for an id. retention: The age after which a Record is dropped. Empty is the Store default. withdraw_run: Retry run only. The previous run, whose Records this withdraws. - auto_close: Close with no result after this long with no append. Empty never - auto-closes. owner: The Workflow run that owns the Stream. The Workflow transport sets it. close_with_run: Close with the owner run chain's result when it closes. Native only. @@ -69,7 +66,6 @@ async def create_stream( stream_id=stream_id, retention=retention, withdraw_run=withdraw_run, - auto_close=auto_close, owner=owner, close_with_run=close_with_run, ref=ref, diff --git a/temporalio/contrib/streams/system/services.py b/temporalio/common/_streams/system/services.py similarity index 100% rename from temporalio/contrib/streams/system/services.py rename to temporalio/common/_streams/system/services.py diff --git a/temporalio/contrib/streams/writer_system/__init__.py b/temporalio/common/_streams/writer_system/__init__.py similarity index 100% rename from temporalio/contrib/streams/writer_system/__init__.py rename to temporalio/common/_streams/writer_system/__init__.py diff --git a/temporalio/contrib/streams/writer_system/client.py b/temporalio/common/_streams/writer_system/client.py similarity index 100% rename from temporalio/contrib/streams/writer_system/client.py rename to temporalio/common/_streams/writer_system/client.py diff --git a/temporalio/contrib/streams/writer_system/models.py b/temporalio/common/_streams/writer_system/models.py similarity index 86% rename from temporalio/contrib/streams/writer_system/models.py rename to temporalio/common/_streams/writer_system/models.py index 189ca8840c..93767c40c5 100644 --- a/temporalio/contrib/streams/writer_system/models.py +++ b/temporalio/common/_streams/writer_system/models.py @@ -10,7 +10,7 @@ import typing_extensions import temporalio.api.streamservice.v1.stream_service_pb2 -import temporalio.contrib.streams.proto.v1.extension_pb2 +import temporalio.common._streams.proto.v1.extension_pb2 import temporalio.converter import temporalio.nexus.system._support as _support @@ -30,17 +30,17 @@ class OpenWriterRequest: class _OpenWriterRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ OpenWriterRequest, - temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest, + temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest + type[temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest] | None + ) = temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest, + value: temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest, type_hint: type[OpenWriterRequest], ) -> OpenWriterRequest: return OpenWriterRequest( @@ -70,8 +70,8 @@ def from_transfer_type( def to_transfer_type( self, value: OpenWriterRequest, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest: - message = temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterRequest() + ) -> temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest: + message = temporalio.common._streams.proto.v1.extension_pb2.OpenWriterRequest() if value.stream is not None: message.stream.CopyFrom( _StreamRefTransferTypeConverter().to_transfer_type(value.stream) @@ -416,18 +416,18 @@ class WriterResult: class _WriterResultTransferTypeConverter( temporalio.converter.TransferTypeConverter[ WriterResult, - temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse, + temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse] + type[temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse + ) = temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse, + value: temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse, type_hint: type[WriterResult], ) -> WriterResult: return WriterResult( @@ -438,8 +438,8 @@ def from_transfer_type( def to_transfer_type( self, value: WriterResult, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse: - message = temporalio.contrib.streams.proto.v1.extension_pb2.OpenWriterResponse() + ) -> temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse: + message = temporalio.common._streams.proto.v1.extension_pb2.OpenWriterResponse() if value.writer_id is not None: message.writer_id = value.writer_id return message @@ -459,18 +459,18 @@ class WriteRecordsRequest: class _WriteRecordsRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ WriteRecordsRequest, - temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest, + temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest] + type[temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest + ) = temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest, + value: temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest, type_hint: type[WriteRecordsRequest], ) -> WriteRecordsRequest: return WriteRecordsRequest( @@ -487,9 +487,9 @@ def from_transfer_type( def to_transfer_type( self, value: WriteRecordsRequest, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest: + ) -> temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest: message = ( - temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsRequest() + temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsRequest() ) if value.writer_id is not None: message.writer_id = value.writer_id @@ -566,18 +566,18 @@ class WriteRecordsResponse: class _WriteRecordsResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ WriteRecordsResponse, - temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse, + temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse] + type[temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse + ) = temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse, + value: temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse, type_hint: type[WriteRecordsResponse], ) -> WriteRecordsResponse: return WriteRecordsResponse( @@ -589,9 +589,9 @@ def from_transfer_type( def to_transfer_type( self, value: WriteRecordsResponse, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse: + ) -> temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse: message = ( - temporalio.contrib.streams.proto.v1.extension_pb2.WriteRecordsResponse() + temporalio.common._streams.proto.v1.extension_pb2.WriteRecordsResponse() ) if value.first_offset is not None: message.first_offset = value.first_offset @@ -613,18 +613,18 @@ class CloseWriterRequest: class _CloseWriterRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ CloseWriterRequest, - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest, + temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest] + type[temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest + ) = temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest, + value: temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest, type_hint: type[CloseWriterRequest], ) -> CloseWriterRequest: return CloseWriterRequest( @@ -635,8 +635,8 @@ def from_transfer_type( def to_transfer_type( self, value: CloseWriterRequest, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest: - message = temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterRequest() + ) -> temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest: + message = temporalio.common._streams.proto.v1.extension_pb2.CloseWriterRequest() if value.writer_id is not None: message.writer_id = value.writer_id return message @@ -655,18 +655,18 @@ class CloseWriterResponse: class _CloseWriterResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ CloseWriterResponse, - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse, + temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse] + type[temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse + ) = temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse, + value: temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse, type_hint: type[CloseWriterResponse], ) -> CloseWriterResponse: return CloseWriterResponse( @@ -677,9 +677,9 @@ def from_transfer_type( def to_transfer_type( self, value: CloseWriterResponse, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse: + ) -> temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse: message = ( - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWriterResponse() + temporalio.common._streams.proto.v1.extension_pb2.CloseWriterResponse() ) if value.last_offset is not None: message.last_offset = value.last_offset @@ -699,17 +699,17 @@ class FreeWriterRequest: class _FreeWriterRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ FreeWriterRequest, - temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest, + temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest + type[temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest] | None + ) = temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest, + value: temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest, type_hint: type[FreeWriterRequest], ) -> FreeWriterRequest: return FreeWriterRequest( @@ -720,8 +720,8 @@ def from_transfer_type( def to_transfer_type( self, value: FreeWriterRequest, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest: - message = temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterRequest() + ) -> temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest: + message = temporalio.common._streams.proto.v1.extension_pb2.FreeWriterRequest() if value.writer_id is not None: message.writer_id = value.writer_id return message @@ -740,18 +740,18 @@ class FreeWriterResponse: class _FreeWriterResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ FreeWriterResponse, - temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse, + temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse] + type[temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse + ) = temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse, + value: temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse, type_hint: type[FreeWriterResponse], ) -> FreeWriterResponse: return FreeWriterResponse() @@ -760,8 +760,8 @@ def from_transfer_type( def to_transfer_type( self, value: FreeWriterResponse, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse: - message = temporalio.contrib.streams.proto.v1.extension_pb2.FreeWriterResponse() + ) -> temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse: + message = temporalio.common._streams.proto.v1.extension_pb2.FreeWriterResponse() return message @@ -778,18 +778,18 @@ class CloseWritersRequest: class _CloseWritersRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ CloseWritersRequest, - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest, + temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest] + type[temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest + ) = temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest, + value: temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest, type_hint: type[CloseWritersRequest], ) -> CloseWritersRequest: return CloseWritersRequest() @@ -798,9 +798,9 @@ def from_transfer_type( def to_transfer_type( self, value: CloseWritersRequest, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest: + ) -> temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest: message = ( - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersRequest() + temporalio.common._streams.proto.v1.extension_pb2.CloseWritersRequest() ) return message @@ -818,18 +818,18 @@ class CloseWritersResponse: class _CloseWritersResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ CloseWritersResponse, - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse, + temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse, ] ): transfer_type: ( - type[temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse] + type[temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse] | None - ) = temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse + ) = temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse @typing_extensions.override def from_transfer_type( self, - value: temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse, + value: temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse, type_hint: type[CloseWritersResponse], ) -> CloseWritersResponse: return CloseWritersResponse( @@ -840,9 +840,9 @@ def from_transfer_type( def to_transfer_type( self, value: CloseWritersResponse, - ) -> temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse: + ) -> temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse: message = ( - temporalio.contrib.streams.proto.v1.extension_pb2.CloseWritersResponse() + temporalio.common._streams.proto.v1.extension_pb2.CloseWritersResponse() ) if value.failed is not None: message.failed = value.failed diff --git a/temporalio/contrib/pydantic.py b/temporalio/contrib/pydantic.py index c6ea2a2da9..94e4363db0 100644 --- a/temporalio/contrib/pydantic.py +++ b/temporalio/contrib/pydantic.py @@ -75,8 +75,8 @@ def __init__( raise ValueError("max_cached_type_adapters cannot be negative") self._schema_serializer = SchemaSerializer(any_schema()) self._to_json_options = to_json_options - self._type_adapter = functools.lru_cache(maxsize=max_cached_type_adapters)( - TypeAdapter + self._type_adapter: functools._lru_cache_wrapper[TypeAdapter[Any]] = ( + functools.lru_cache(maxsize=max_cached_type_adapters)(TypeAdapter) ) def __getstate__(self) -> dict[str, Any]: diff --git a/temporalio/contrib/stream/__init__.py b/temporalio/contrib/stream/__init__.py new file mode 100644 index 0000000000..a1f132dee9 --- /dev/null +++ b/temporalio/contrib/stream/__init__.py @@ -0,0 +1,172 @@ +"""The Redis Stream backend. + +Register it on the client. :py:func:`create` creates a Stream that the client owns. +:py:mod:`temporalio.contrib.stream.workflow` has the Workflow functions that use this backend:: + + from temporalio.contrib import stream as redis + + client = await Client.connect( + "localhost:7233", stream_backends=[redis.RedisStreamBackend("redis://localhost")] + ) + +A Workflow's close is a regular Activity on the Workflow's task queue. So a Worker with a Redis +backend must poll Activities on its task queue for the close Activity. All Workers that share a +task queue must register the Redis backend. + +.. warning:: + This package is experimental. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import timedelta +from typing import TypeVar + +import temporalio.client +from temporalio.common._streams import _plugin, _stream_client +from temporalio.common._streams._errors import StreamConfigError +from temporalio.common._streams._extension import ExtensionStreamBackend +from temporalio.common._streams._handle import Creating +from temporalio.common._streams._stream_client import ClientStream +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.contrib.stream import workflow + +D = TypeVar("D") + + +class RedisStreamBackend(ExtensionStreamBackend): + """Sends Stream calls to Redis, through the temporalio-streams runtime in the bridge. + + Register it on the client with ``stream_backends=``. A Worker uses its client's backends, + so one backend serves Workflow code, Activities and the client. Each client of the backend + gets its own runtime object, and the client's Workers share it. It registers the backend + name ``redis``, and each Stream handle that it creates stores that name. A Worker serves + only one backend. A client can have more than one, and ``backend=`` picks one. + + A Workflow's close is a regular Activity on the Workflow's task queue, which runs after the + Workflow Task commits. So a Worker with this backend must poll Activities on its task queue, + and all Workers that share a task queue must register the backend. + + .. warning:: + This API is experimental. + """ + + _kind = "redis" + + def __init__( + self, + url: str | Sequence[str], + *, + cluster: bool = False, + key_prefix: str | None = None, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + response_timeout: timedelta | None = None, + name: str = "redis", + default: bool = False, + ) -> None: + """Create the backend. + + Each client of the backend connects to Redis on its first Stream call, on the client's + runtime. + + Args: + url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the + URL. + cluster: Whether the URLs are of a Redis cluster. + key_prefix: The prefix of every key. Defaults to ``temporal-streams``. + retention: The default Retention of a new Stream. Defaults to 7 days. + safety_ttl: How long an open Workflow-owned Stream is kept after its last write. + Defaults to 90 days. + response_timeout: The longest time that one Redis call can take. Defaults to 5 s. + name: The backend name that Stream handles store, so every process can route a + handle that it receives. + default: Make this the backend of a create with no backend name. With more than + one registered backend, exactly one must have it, or the client raises + ``StreamConfigError``. + """ + config = ext_pb.RedisStreamConfig( + urls=[url] if isinstance(url, str) else list(url), + cluster=cluster, + key_prefix=key_prefix or "", + ) + for field, value in ( + ("retention", retention), + ("safety_ttl", safety_ttl), + ("response_timeout", response_timeout), + ): + if value is not None: + getattr(config, field).FromTimedelta(value) + super().__init__( + "temporalio.RedisStreamBackend", False, config, None, name, default + ) + + +def create( + client: temporalio.client.Client, + definition: type[D], + *, + id: str, + retention: timedelta | None = None, + backend: str | None = None, +) -> Creating[ClientStream[D]]: + """Create a Stream that the client owns, so only a client can close it. + + The call is idempotent for the same id and Retention. Await the call to get the handle, or + use it in ``async with``, which closes the Stream at the end of the block:: + + async with redis.create(client, Events, id="chat-1") as chat: + await chat.topic(Events.turns).publish(Turn("user", "hi")) + + .. warning:: + This API is experimental. + + Args: + client: The client, with a registered Redis backend. + definition: The Stream definition class. + id: The Stream id, which the caller picks, as for a Workflow id. + retention: The age after which a Record is deleted, even while the Stream is open. + A closed Stream is kept for one Retention after its close. Defaults to the + backend's Retention. + backend: The name of the Redis backend. Defaults to the client's only Redis backend, + or the one with ``default=True``. + + Raises: + StreamConfigError: The client has no Redis backend of that name. No call is made. + StreamClosedError: A Stream of that id was created and closed before. + StreamExpiredError: A Stream of that id expired. + StreamError: The Stream exists with a different Retention, with the reason + ``stream_already_exists``. + """ + + async def create() -> ClientStream[D]: + name = _redis_backend(_plugin.client_backends(client), backend) + return await _stream_client.create( + client, definition, id=id, retention=retention, backend=name + ) + + return Creating(create) + + +def _redis_backend(backends: Sequence[_plugin.StreamBackend], name: str | None) -> str: + # The memory backend has Redis's rules, so tests can use it in place of Redis. + redis = [backend for backend in backends if backend._kind == "redis"] + if name: + if not any(backend.backend_name == name for backend in redis): + raise StreamConfigError(f"The client has no Redis Stream backend {name}") + return name + if not redis: + raise StreamConfigError( + "The client has no Redis Stream backend. Register a RedisStreamBackend with " + "stream_backends=." + ) + defaults = [backend for backend in redis if backend.default] + if len(redis) > 1 and len(defaults) != 1: + raise StreamConfigError( + "The client has more than one Redis Stream backend, so pass backend=" + ) + return (defaults if len(redis) > 1 else redis)[0].backend_name + + +__all__ = ["RedisStreamBackend", "create", "workflow"] diff --git a/temporalio/contrib/stream/testing.py b/temporalio/contrib/stream/testing.py new file mode 100644 index 0000000000..0002ec03e8 --- /dev/null +++ b/temporalio/contrib/stream/testing.py @@ -0,0 +1,60 @@ +"""A Stream backend for tests, with no Redis. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +from datetime import timedelta + +import temporalio.runtime +from temporalio.common._streams._extension import ExtensionStreamBackend +from temporalio.common._streams.proto import v1 as ext_pb + + +class MemoryStreamBackend(ExtensionStreamBackend): + """Keeps Streams in this process, with the same rules as Redis, for tests. + + Register it on the client, as with ``RedisStreamBackend``. Every Worker and client that + shares one backend shares its Streams. The Streams are lost when the process ends. It + registers the backend name ``memory``. It has Redis's rules, so the functions of + :py:mod:`temporalio.contrib.stream` accept it as a Redis backend, and tests can swap the + Store without changing Workflow code. As on Redis, a Workflow's close is a regular + Activity, so all Workers that share a task queue must register the backend. + + .. warning:: + This API is experimental. + """ + + _kind = "redis" + + def __init__( + self, + *, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + runtime: temporalio.runtime.Runtime | None = None, + name: str = "memory", + default: bool = False, + ) -> None: + """Create the backend. + + Args: + retention: The default Retention of a new Stream. Defaults to 7 days. + safety_ttl: How long an open Workflow-owned Stream is kept after its last write. + Defaults to 90 days. + runtime: The runtime of the backend's memory Store. Defaults to the default + runtime. + name: The backend name that Stream handles store. + default: Make this the backend of a create with no backend name. With more than + one registered backend, exactly one must have it. + """ + config = ext_pb.RedisStreamConfig() + if retention is not None: + config.retention.FromTimedelta(retention) + if safety_ttl is not None: + config.safety_ttl.FromTimedelta(safety_ttl) + super().__init__( + "temporalio.MemoryStreamBackend", True, config, runtime, name, default + ) diff --git a/temporalio/contrib/stream/workflow.py b/temporalio/contrib/stream/workflow.py new file mode 100644 index 0000000000..a25d3aeea9 --- /dev/null +++ b/temporalio/contrib/stream/workflow.py @@ -0,0 +1,101 @@ +"""Streams in Workflow code on the registered Redis backend. + +These functions work like :py:func:`temporalio.workflow.create_stream`, +:py:func:`temporalio.workflow.create_stream_nowait` and +:py:func:`temporalio.workflow.stream_handle`, but use the Worker's Redis backend in place of the +default one. A ``MemoryStreamBackend`` has Redis's rules, so tests can use it here without +changing Workflow code. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +from datetime import timedelta +from typing import TypeVar + +from temporalio.common._streams import _workflow +from temporalio.common._streams._errors import StreamConfigError +from temporalio.common._streams._handle import Creating +from temporalio.common._streams._workflow import WorkflowStream + +D = TypeVar("D") + + +def _backend(backend: str | None) -> str: + # This error comes from the Worker's config, so it fails the Workflow Task and not the run. + worker = _workflow.worker_backend() + if worker is None: + raise _workflow.no_worker_backend() + if worker.kind != "redis": + raise StreamConfigError( + f"The Worker's Stream backend {worker.name} is not a Redis backend" + ) + return backend or worker.name + + +def create_stream( + definition: type[D], + *, + retention: timedelta | None = None, + backend: str | None = None, +) -> Creating[WorkflowStream[D]]: + """Create a Stream that this run chain owns, on the Worker's Redis backend. + + See :py:func:`temporalio.workflow.create_stream`. On Redis, the end of the run chain doesn't + close a Stream outside an ``async with`` block, so close the Stream before the run returns. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. Its name is the Stream name. + retention: The age after which a Record is deleted, even while the Stream is open. + A closed Stream is kept for one Retention after its close. Defaults to the + backend's Retention. + backend: The name of the Redis backend. Defaults to the Worker's. + + Raises: + StreamError: The backend refused the create, such as for a different Retention. + StreamConfigError: The Worker has no Redis backend, or none of that name. It fails + the Workflow Task, not the run. + """ + + async def create() -> WorkflowStream[D]: + return await _workflow.create_stream( + definition, retention=retention, backend=_backend(backend) + ) + + return Creating(create) + + +def create_stream_nowait( + definition: type[D], + *, + retention: timedelta | None = None, + backend: str | None = None, +) -> WorkflowStream[D]: + """Create a Stream on the Worker's Redis backend, without an await. + + See :py:func:`temporalio.workflow.create_stream_nowait`. + + .. warning:: + This API is experimental. + """ + return _workflow.create_stream_nowait( + definition, retention=retention, backend=_backend(backend) + ) + + +def stream_handle( + definition: type[D], *, backend: str | None = None +) -> WorkflowStream[D]: + """Get a Stream that this run chain created on the Worker's Redis backend. It makes no call. + + See :py:func:`temporalio.workflow.stream_handle`. + + .. warning:: + This API is experimental. + """ + return _workflow.stream_handle(definition, backend=_backend(backend)) diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py deleted file mode 100644 index 7f493ef9b0..0000000000 --- a/temporalio/contrib/streams/__init__.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Streams of Records that Workflows, Activities and clients write and read. - -A Stream definition, a class with ``@streams.defn``, gives the Stream name and the type of each -Topic. Workflow code creates a Stream with -:py:func:`temporalio.contrib.streams.workflow.create_stream`. A client creates and gets Streams -with :py:class:`StreamClient`. A Stream is data: a Workflow passes it to an Activity, a child -Workflow or the next run as an argument, and the parameter's type gives the handle class of -that context, such as :py:class:`ActivityStream`. - -A Stream plugin carries the calls and registers a backend name with ``name=``, such as -:py:class:`temporalio.contrib.streams.redis.RedisStreamPlugin` (``redis``) and -:py:class:`temporalio.contrib.streams.testing.MemoryStreamPlugin` (``memory``). A handle -carries the name of its backend. A client may have two plugins, and exactly one of them has -``default=True``, the one that a create with no ``backend=`` uses. For now a Worker has one -Stream plugin, so it serves one backend. - -.. warning:: - This package is experimental. -""" - -from temporalio.contrib.streams import workflow -from temporalio.contrib.streams._activity_stream import ( - ActivityStream, - ActivityTopic, - SyncActivityStream, - SyncActivityTopic, - SyncSubscription, -) -from temporalio.contrib.streams._buffered import BufferedPublisher -from temporalio.contrib.streams._defn import ( - StreamDefinition, - TopicDefinition, - defn, - topic, - with_name, -) -from temporalio.contrib.streams._errors import ( - ReadOnlyContextError, - StreamClosedError, - StreamConfigError, - StreamError, - StreamExpiredError, - StreamFencedError, - StreamNotFoundError, - StreamNotOwnerError, - StreamOffsetError, - StreamSeqDivergentError, - StreamStorageError, - StreamUnsupportedError, -) -from temporalio.contrib.streams._handle import StreamHandle -from temporalio.contrib.streams._items import ( - END, - Closed, - CloseStatus, - Record, - Superseded, -) -from temporalio.contrib.streams._plugin import StreamPlugin -from temporalio.contrib.streams._reader import Subscription -from temporalio.contrib.streams._stream_client import ( - ClientStream, - ClientTopic, - StreamClient, -) -from temporalio.contrib.streams.workflow import WorkflowStream, WorkflowTopic - -__all__ = [ - "END", - "ActivityStream", - "ActivityTopic", - "BufferedPublisher", - "ClientStream", - "ClientTopic", - "CloseStatus", - "Closed", - "ReadOnlyContextError", - "Record", - "StreamClient", - "StreamClosedError", - "StreamConfigError", - "StreamDefinition", - "StreamError", - "StreamExpiredError", - "StreamFencedError", - "StreamHandle", - "StreamNotFoundError", - "StreamNotOwnerError", - "StreamOffsetError", - "StreamPlugin", - "StreamSeqDivergentError", - "StreamStorageError", - "StreamUnsupportedError", - "Subscription", - "Superseded", - "SyncActivityStream", - "SyncActivityTopic", - "SyncSubscription", - "TopicDefinition", - "WorkflowStream", - "WorkflowTopic", - "defn", - "topic", - "with_name", - "workflow", -] diff --git a/temporalio/contrib/streams/_plugin.py b/temporalio/contrib/streams/_plugin.py deleted file mode 100644 index 708fa139a3..0000000000 --- a/temporalio/contrib/streams/_plugin.py +++ /dev/null @@ -1,138 +0,0 @@ -"""The base of the Stream plugins.""" - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import temporalio.client -import temporalio.worker -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. With more than one - Stream plugin, exactly one has ``default=True``. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - name: str, - *, - backend_name: str, - default: bool, - stream_backend: StreamBackend | None = None, - **options: Any, - ) -> None: - """Create the plugin. ``options`` go to :py:class:`temporalio.plugin.SimplePlugin`. - - A plugin with no ``stream_backend`` gives each client its own in - :py:meth:`backend_for`. - """ - 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 - - def backend_for( - self, - client: temporalio.client.Client, # pyright: ignore[reportUnusedParameter] - ) -> StreamBackend: - """The backend that carries the calls of ``client`` and of its Workers.""" - if self._stream_backend is None: - raise NotImplementedError(f"{type(self).__name__} gives no backend") - return self._stream_backend - - def configure_client( - self, config: temporalio.client.ClientConfig - ) -> temporalio.client.ClientConfig: - """Check the default backend, then add the plugin's parts to the client.""" - check_default(config.get("plugins", []), "client") - return super().configure_client(config) - - def configure_worker( - self, config: temporalio.worker.WorkerConfig - ) -> temporalio.worker.WorkerConfig: - """Check the default backend, then add the plugin's parts to the Worker.""" - client = config.get("client") - client_plugins = client.config()["plugins"] if client is not None else [] - check_default([*client_plugins, *config.get("plugins", [])], "Worker") - return super().configure_worker(config) - - -def check_default(plugins: Iterable[object], owner: str) -> None: - """Check that one Stream plugin is the default. - - The default backend is part of Workflow code, since a replayed create uses it. So the - default must not follow the order of the plugin list. - - Raises: - StreamConfigError: More than one Stream plugin, and not exactly one with - ``default=True``. - """ - stream_plugins = list( - {id(p): p for p in plugins if isinstance(p, StreamPlugin)}.values() - ) - if len(stream_plugins) < 2: - return - defaults = [p.backend_name for p in stream_plugins if p.default] - if len(defaults) != 1: - names = ", ".join(p.backend_name for p in stream_plugins) - marked = f"{', '.join(defaults)} have" if defaults else "none has" - raise _errors.StreamConfigError( - f"The {owner} has the Stream plugins {names}, and {marked} default=True. " - "Mark exactly one with default=True." - ) - - -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 default plugin, - which is the only one or the one with ``default=True``. 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.backend_for(client), 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 deleted file mode 100644 index f5ad67d7b7..0000000000 --- a/temporalio/contrib/streams/_stream_client.py +++ /dev/null @@ -1,608 +0,0 @@ -"""Streams in client code.""" - -from __future__ import annotations - -import asyncio -import itertools -import uuid -import weakref -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence -from datetime import timedelta -from types import TracebackType -from typing import Any, TypeVar, cast, overload - -import temporalio.api.streamservice.v1 as stream_pb -import temporalio.client -from temporalio.contrib.streams import _backend, _errors, _ids -from temporalio.contrib.streams._access import ( - Access, - ReadSpec, - _AsyncTopic, - read_spec, - subscription, -) -from temporalio.contrib.streams._backend import StreamBackend -from temporalio.contrib.streams._defn import TopicDefinition, definition_of -from temporalio.contrib.streams._handle import Creating, StreamHandle, exit_failure -from temporalio.contrib.streams._items import End -from temporalio.contrib.streams._plugin import client_backend -from temporalio.contrib.streams._reader import Item, Position, Subscription -from temporalio.contrib.streams._writer import StreamWriter -from temporalio.contrib.streams.system import models - -E = TypeVar("E") -T = TypeVar("T") - -# The handles of one client and Stream write as one Producer, so the crate shares one writer -# between them. Each client gets its own key, since two clients may have two runtime objects, -# and two writers of one Producer would send clashing sequences. -_PROCESS_KEY = str(uuid.uuid4()) -_CLIENTS = itertools.count(1) -_CLIENT_KEYS: weakref.WeakKeyDictionary[temporalio.client.Client, str] = ( - weakref.WeakKeyDictionary() -) - - -def _client_producer( - client: temporalio.client.Client, named: tuple[str, int] | None -) -> stream_pb.Producer: - if named is not None: - name, epoch = named - return stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, - id=_ids.named_client_producer_id(name), - epoch=stream_pb.ProducerEpoch(attempt=epoch), - ) - key = _CLIENT_KEYS.get(client) - if key is None: - key = f"{_PROCESS_KEY}-{next(_CLIENTS)}" - _CLIENT_KEYS[client] = key - return stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_CLIENT, - id=_ids.client_producer_id(client.identity, key), - epoch=stream_pb.ProducerEpoch(attempt=1), - ) - - -class ClientTopic(_AsyncTopic[T]): - """One Topic of a :py:class:`ClientStream`. All Topics of a Stream share one order. - - .. warning:: - This API is experimental. - """ - - async def publish( - self, - *values: T, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> int: - """Write one Record for each value, in order, and return the last Record's offset. - - A list is one value. ``rpc_metadata`` and ``rpc_timeout`` are as for - :py:meth:`ClientStream.get_offset`. - """ - with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - return await super().publish(*values) - - def subscribe( - self, - *, - from_offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - create_wait: timedelta | None = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> Subscription[T]: - """Read this Topic only. Its Records have the Topic's type. See - :py:meth:`ClientStream.subscribe`. - """ - stream = self._stream - assert isinstance(stream, ClientStream) - return cast( - Subscription[T], - stream.subscribe( - self, - from_offset=from_offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - rpc_metadata=rpc_metadata, - rpc_timeout=rpc_timeout, - ), - ) - - -class ClientStream(StreamHandle[E]): - """A Stream that a client writes, reads and closes. Get one from :py:class:`StreamClient`. - - Each publish writes before it returns. Use :py:class:`BufferedPublisher` to batch. The - handles of one client and Stream write as one Producer, and :py:meth:`producer` gives a - handle that writes as a named one. Use it in ``async with`` to close it at the end of the - block. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - *, - id: str, - workflow_id: str | None = None, - first_execution_run_id: str | None = None, - backend: str | None = None, - definition: type[E] | None = None, - ) -> None: - """Create a handle. Use :py:class:`StreamClient` instead.""" - super().__init__( - id=id, - workflow_id=workflow_id, - first_execution_run_id=first_execution_run_id, - backend=backend, - definition=definition, - ) - self._client: temporalio.client.Client | None = None - self._stream_backend: StreamBackend | None = None - self._find_first_run: Callable[[], Awaitable[str]] | None = None - self._writer: StreamWriter | None = None - # 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._closed = False - # The name and the epoch of a named Producer. - self._named: tuple[str, int] | 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. A handle that looks in many backends knows it after a call.""" - return None if self._probe else self._backend - - async def _ready(self) -> None: - # The first backend that has the Stream keeps it. Two first calls may both look. - for name, backend in self._probe: - assert self._client is not None - request = models.DescribeStreamRequest( - stream=models.StreamRef(id=await self._stream_id()) - ) - try: - await backend.service(self._client.namespace).describe_stream(request) - except _errors.StreamNotFoundError: - continue - self._backend, self._stream_backend, self._probe = name, backend, [] - return - if self._probe: - raise _errors.StreamNotFoundError( - f"No backend of {[name for name, _ in self._probe]} has the Stream {self._id}", - reason="stream_not_found", - ) - - @property - def id(self) -> str: - """The Stream id. - - Raises: - RuntimeError: The Workflow handle had no first run id, and no call has found it - yet. Call :py:meth:`get_offset` first. - """ - if not self._id: - raise RuntimeError( - "The Stream id is not known yet. Call get_offset() first." - ) - return self._id - - async def _stream_id(self) -> str: - if not self._id: - # Calls that start together share one Describe call. - if self._finding is None: - assert self._find_first_run is not None - self._finding = asyncio.ensure_future(self._find_first_run()) - first_run = await asyncio.shield(self._finding) - if not self._id: - assert self._workflow_id is not None - self._first_execution_run_id = first_run - self._id = _ids.scoped_stream_id( - self._workflow_id, first_run, self._stream_name - ) - return self._id - - def _fields(self) -> dict[str, str]: - fields = super()._fields() - if not self._id: - # The Stream name tells apart the Streams of one Workflow before the id is known. - fields["name"] = self._stream_name - return fields - - def _access(self) -> Access: - client, backend = self._client, self._stream_backend - if client is None or backend is None: - raise _errors.StreamConfigError( - "This ClientStream came from a payload, so it has no client. " - "Use StreamClient.get_handle." - ) - - def writer() -> StreamWriter: - if self._writer is None: - self._writer = StreamWriter( - writers=backend.writers(client.namespace, client.data_converter), - stream_id=self._stream_id, - producer=_client_producer(client, self._named), - free=backend.free_writer, - ) - return self._writer - - return Access( - backend=backend, - converter=client.data_converter, - namespace=client.namespace, - stream_id=self._stream_id, - writer=writer, - ) - - def producer(self, name: str, epoch: int) -> ClientStream[E]: - """This Stream as a handle that writes as the Producer ``name`` at ``epoch``. - - It is for a writer that must fence its earlier self, such as a server that restarts. - Keep ``epoch`` in durable storage, and raise it at each start. After a write at a - higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The Records - of the earlier epochs stay, since they are real messages. Run one live writer for a - name. - - Args: - name: The Producer name, the same across restarts. - epoch: The writer's epoch, 1 or more. - """ - if epoch < 1: - raise ValueError("A Producer's epoch is 1 or more") - handle = ClientStream( - id=self._id, - workflow_id=self._workflow_id, - first_execution_run_id=self._first_execution_run_id, - backend=self._backend, - definition=self._definition, - ) - handle._stream_name = self._stream_name - handle._named = (name, epoch) - if self._client is not None and self._stream_backend is not None: - handle._bind( - self._client, self._stream_backend, self._find_first_run, self._probe - ) - return handle - - @overload - def topic(self, topic: TopicDefinition[T]) -> ClientTopic[T]: ... - - @overload - def topic(self, topic: str) -> ClientTopic[Any]: ... - - @overload - def topic(self, topic: str, *, type: type[T]) -> ClientTopic[T]: ... - - def topic( - self, topic: str | TopicDefinition[Any], *, type: Any = None - ) -> ClientTopic[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. - """ - name, topic_type = self._topic_spec(topic, type) - return ClientTopic(self, name, topic_type) - - def subscribe( - self, - *topics: str | TopicDefinition[Any] | ClientTopic[Any], - from_offset: int | End | None = None, - last: int | None = None, - wait_new_records: bool = True, - prefixes: Sequence[str] = (), - create_wait: timedelta | None = None, - type: Any = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> Subscription[Any]: - """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 reader that wants to stop after a quiet time sets its - own timer, and resumes later with ``from_offset=`` the read's ``next_offset``. 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. - last: Start with the newest this many Records that pass the filter, oldest - first, such as the first view of a chat. Not with ``from_offset``. A Record - that a later ``Superseded`` item withdraws doesn't count. At most 1000. - wait_new_records: Wait for new Records at the end of the Stream. - prefixes: Also read the Topics whose names start with one of these. - create_wait: How long to wait for a Stream that nobody created yet, such as the - Stream of a Workflow that just started. Defaults to no wait. - type: The type of the values of Topics outside the definition, such as Topics that - ``prefixes`` match. Topics of the definition keep their own types. Defaults to - no type, so the values decode as the converter gives them. - rpc_metadata: Headers of each call of the read. See :py:meth:`get_offset`. - rpc_timeout: The deadline of each call of the read, not of the whole read. A call - that waits for new Records answers before it. - """ - read = read_spec( - self, - topics, - prefixes=prefixes, - default_type=type, - from_offset=from_offset, - last=last, - wait_new_records=wait_new_records, - create_wait=create_wait, - ) - return subscription( - lambda position: self._items(read, position), - read, - _backend.options_of(rpc_metadata, rpc_timeout), - ) - - async def _items( - self, read: ReadSpec, position: Position - ) -> AsyncGenerator[Item, None]: - await self._ready() - items = self._access().items(read, position) - try: - async for item in items: - yield item - finally: - await items.aclose() - - async def get_offset( - self, - *, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> int: - """The next offset of the Stream, the newest offset plus one. - - Args: - rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so - they send none. - rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with - ``DEADLINE_EXCEEDED``. - """ - with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - await self._ready() - return await self._access().get_offset() - - async def close( - self, - *, - result: Any | None = None, - failure: BaseException | None = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> None: - """Write this handle's Records, then close the Stream. - - Readers get the Records before it, then a ``Closed`` item with ``result`` or - ``failure``. A ``CancelledError`` failure closes it as canceled. ``rpc_metadata`` and - ``rpc_timeout`` are as for :py:meth:`get_offset`. - - Raises: - StreamNotOwnerError: A Workflow owns the Stream, so only its run chain closes it. - """ - if self._workflow_id: - # A Workflow's Stream closes only from its run chain. The Store refuses a client's - # close too, and this error comes with no call. - raise _errors.StreamNotOwnerError( - f"Can't close the Stream {self._id or self._stream_name} from a client, since " - "a Workflow owns it", - reason="stream_not_owner", - ) - self._closed = True - with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - await self._ready() - await self._access().close(result, failure) - - async def __aenter__(self) -> ClientStream[E]: - """Use the Stream in ``async with``, which closes it at the end.""" - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - """Close the Stream, with the exception as its failure if there is one. - - A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as - ``KeyboardInterrupt``, leaves it open. A Stream that this handle closed stays as it is. - """ - if self._closed: - return - close, failure = exit_failure(exc) - if close: - await self.close(failure=failure) - - -class StreamClient: - """Creates and gets Streams over a Temporal client with a Stream plugin. - - .. warning:: - This API is experimental. - """ - - def __init__(self, client: temporalio.client.Client) -> None: - """Create a Stream client. It makes no call.""" - self._client = client - - def create( - self, - definition: type[E], - *, - id: str, - retention: timedelta | None = None, - auto_close: timedelta | None = None, - backend: str | None = None, - rpc_metadata: Mapping[str, str | bytes] = {}, - rpc_timeout: timedelta | None = None, - ) -> Creating[ClientStream[E]]: - """Create a Stream that this client owns, so only a client closes it. - - It is idempotent for one id, one Retention and one ``auto_close``. Await the call for - the handle, or use it in ``async with``, which closes the Stream at the end of the - block. - - Args: - definition: The Stream definition class. - id: The Stream id, as a Workflow takes a Workflow id. - retention: How long the Records stay after the close. Defaults to the backend's. - auto_close: Close the Stream with no result when no Record was appended for this - long, so a Stream that nobody closes does not stay open. Defaults to never. - backend: The name of the backend. Defaults to the client's default backend. - rpc_metadata: Headers of the call. See :py:meth:`ClientStream.get_offset`. - rpc_timeout: The deadline of the call. - - Raises: - 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``. - """ - - async def create() -> ClientStream[E]: - definition_of(definition) - backend_name, stream_backend = client_backend(self._client, backend) - service = stream_backend.service(self._client.namespace) - with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): - response = await service.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( - f"The Stream {id} was created and closed before", - reason="stream_closed", - ) - stream = ClientStream( - id=id, backend=backend_name or None, definition=definition - ) - return stream._bind(self._client, stream_backend) - - return Creating(create) - - def get_handle( - self, - definition: type[E], - *, - id: str | None = None, - stream: StreamHandle[Any] | None = None, - backend: str | tuple[str, ...] | None = None, - ) -> ClientStream[E]: - """A Stream by its id, 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, 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: - # An empty name is of any backend, so it needs a client with one. - backend_name, stream_backend = client_backend( - self._client, stream.backend, only=True - ) - probe: list[tuple[str, StreamBackend]] = [] - elif id is not None: - backend_name, stream_backend, probe = self._backends(backend) - else: - raise TypeError("get_handle takes id= or stream=") - handle = ClientStream( - id=stream.id if stream else id or "", - workflow_id=stream.workflow_id if stream else None, - first_execution_run_id=stream.first_execution_run_id if stream else None, - backend=backend_name or None, - definition=definition, - ) - return 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, - handle: temporalio.client.WorkflowHandle[Any, Any], - definition: type[E], - *, - backend: str | tuple[str, ...] | None = None, - ) -> ClientStream[E]: - """A Stream that a Workflow created, by the Workflow handle and the definition. - - It builds the id from the Workflow id and the first run id, so it makes no call. A - handle with no first run id makes one Describe call on the Stream's first use. - - Args: - handle: A handle of the Workflow, best the one that ``start_workflow`` returned. - definition: The Stream definition class, renamed as the Workflow's create. - 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 = definition_of(definition).name - 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) - if first_run - else "", - workflow_id=handle.id, - first_execution_run_id=first_run, - backend=backend_name or None, - definition=definition, - ) - if first_run: - return stream._bind(self._client, stream_backend, probe=probe) - - async def find_first_run() -> str: - # The first call that needs the id gives its own call options. - options = _backend.call_options.get() - description = await handle.describe( - rpc_metadata=options.metadata if options else {}, - rpc_timeout=options.timeout if options else None, - ) - first = description.raw_description.workflow_execution_info.first_run_id - if not first: - raise ValueError( - "The server gave no first run id, so the Stream id is unknown" - ) - return first - - stream._stream_name = stream_name - return stream._bind(self._client, stream_backend, find_first_run, probe) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py deleted file mode 100644 index c3cdcad1e3..0000000000 --- a/temporalio/contrib/streams/redis.py +++ /dev/null @@ -1,706 +0,0 @@ -"""The Redis Stream plugin. - -.. warning:: - This module is experimental. -""" - -from __future__ import annotations - -import asyncio -import dataclasses -import logging -import weakref -from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime, timedelta -from typing import Any, NoReturn - -import temporalio.activity -import temporalio.api.common.v1 -import temporalio.api.failure.v1 -import temporalio.api.streamservice.v1 as stream_pb -import temporalio.bridge.streams -import temporalio.client -import temporalio.common -import temporalio.converter -import temporalio.exceptions -import temporalio.runtime -import temporalio.service -import temporalio.worker -import temporalio.workflow -from temporalio.contrib.streams import _activity, _ids, _writes -from temporalio.contrib.streams import workflow as streams_workflow -from temporalio.contrib.streams._backend import StreamBackend, _ExtensionStreamBackend -from temporalio.contrib.streams._errors import StreamConfigError -from temporalio.contrib.streams._handle import StreamHandle -from temporalio.contrib.streams._items import close_status -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.worker.workflow_sandbox import SandboxedWorkflowRunner - -logger = logging.getLogger(__name__) - -# 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 run chain's open Streams to the next run in this header. -_OWNED_STREAMS_HEADER = "__temporal_owned_streams" - - -class _ExtensionStreamPlugin(StreamPlugin): - # The Redis and memory plugins: the Stream calls go to the temporalio-streams runtime. - - def __init__( - self, - name: str, - memory: bool, - config: ext_pb.RedisStreamConfig, - runtime: temporalio.runtime.Runtime | None, - backend_name: str, - default: bool, - ) -> None: - """Create the plugin over the bridge's Stream runtime, on Redis or in memory. - - ``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. - """ - config.backend = backend_name - self._config = config.SerializeToString() - self._backends: weakref.WeakSet[_ExtensionStreamBackend] = weakref.WeakSet() - # The closes of the runs that ended on this plugin's Workers, by run id, until the - # Worker evicts the run. - self._run_closes: dict[str, _RunClose] = {} - self._closing: set[asyncio.Task[None]] = set() - # Each client gets its own runtime object on Redis. The memory Store has one for the - # plugin, so every client and Worker that shares the plugin shares its Streams. - self._memory: _ExtensionStreamBackend | None = None - if memory: - self._memory = self._new_backend( - runtime or temporalio.runtime.Runtime.default(), True - ) - super().__init__( - name, - backend_name=backend_name, - default=default, - activities=_writes.ACTIVITIES, - ) - - def _new_backend( - self, runtime: temporalio.runtime.Runtime, memory: bool = False - ) -> _ExtensionStreamBackend: - if self._memory is not None and not memory: - return self._memory - backend = _ExtensionStreamBackend( - temporalio.bridge.streams.StreamRuntime( - runtime._core_runtime, memory, self._config - ) - ) - self._backends.add(backend) - return backend - - def configure_client( - self, config: temporalio.client.ClientConfig - ) -> temporalio.client.ClientConfig: - """Check the default backend, and give the client its own Stream runtime.""" - config = super().configure_client(config) - backend = self._new_backend(_runtime_of(config["service_client"])) - config["interceptors"] = [ - *config.get("interceptors", []), - _StreamInterceptor(self, backend, config["data_converter"]), - ] - return config - - def backend_for(self, client: temporalio.client.Client) -> StreamBackend: - """The Stream runtime that this plugin gave ``client``.""" - interceptor = self._interceptor_of( - client.config(active_config=True)["interceptors"] - ) - if interceptor is None or interceptor.backend is None: - raise StreamConfigError( - f"The client has no {type(self).__name__}. Put the plugin on the client." - ) - return interceptor.backend - - def _interceptor_of(self, interceptors: Sequence[Any]) -> _StreamInterceptor | None: - for interceptor in interceptors: - if ( - isinstance(interceptor, _StreamInterceptor) - and interceptor.plugin is self - ): - return interceptor - return None - - def configure_worker( - self, config: temporalio.worker.WorkerConfig - ) -> temporalio.worker.WorkerConfig: - """Add the plugin's Activities, and its interceptor when the client has none.""" - config = super().configure_worker(config) - client = config["client"] # type: ignore[reportTypedDictNotRequiredAccess] - interceptors = [ - *client.config(active_config=True)["interceptors"], - *config.get("interceptors", []), - ] - if any( - isinstance(other, _StreamInterceptor) and other.plugin is not self - for other in interceptors - ): - raise ValueError( - "The Worker has another Stream plugin. For now a Worker serves one backend." - ) - # Records get no serialization context, so Workflows, Activities and clients encode - # them alike. - mine = self._interceptor_of(interceptors) - if mine is None: - mine = _StreamInterceptor( - self, self._new_backend(_runtime_of(client.service_client)), None - ) - config["interceptors"] = [*config.get("interceptors", []), mine] - mine.converter = client.data_converter - runner = config.get("workflow_runner") - if isinstance(runner, SandboxedWorkflowRunner): - # The interceptor and the facade must share the run's Stream state, also when the - # restrictions don't pass temporalio through. - config["workflow_runner"] = dataclasses.replace( - runner, - restrictions=runner.restrictions.with_passthrough_modules( - "temporalio.contrib.streams" - ), - ) - return config - - def configure_replayer( - self, config: temporalio.worker.ReplayerConfig - ) -> temporalio.worker.ReplayerConfig: - """Add the Workflow interceptor. A replay runs no Activity, so it needs no runtime.""" - config = super().configure_replayer(config) - config["interceptors"] = [ - *config.get("interceptors", []), - _StreamInterceptor(self, None, None), - ] - return config - - def on_workflow_evicted( - self, input: temporalio.worker.WorkflowEvictedInput - ) -> None: - """Write the closes of a run that ended, once the server accepted its completion.""" - close = self._run_closes.pop(input.run_id, None) - if close is None or not input.execution_ending: - return - task = asyncio.get_running_loop().create_task(_write_closes(close)) - self._closing.add(task) - task.add_done_callback(self._closing.discard) - - async def run_worker( - self, - worker: temporalio.worker.Worker, - next: Callable[[temporalio.worker.Worker], Awaitable[None]], - ) -> None: - """Run the Worker, then wait for the closes of the runs that ended on it.""" - try: - await super().run_worker(worker, next) - finally: - if self._closing: - await asyncio.gather(*list(self._closing), return_exceptions=True) - - async def shutdown(self) -> None: - """Drop every open writer, including the client's. Call it once the process no longer - writes Streams. - """ - failed = 0 - for backend in list(self._backends): - failed += await backend.close_writers() - if failed: - logger.warning("Stream shutdown left %d failed writers", failed) - - -def _runtime_of( - service: temporalio.service.ServiceClient, -) -> temporalio.runtime.Runtime: - # The runtime of a client's connection, as the Worker finds it. - service = getattr(service, "worker_service_client", service) - runtime = getattr(service, "_runtime", None) or getattr( - getattr(service, "config", None), "runtime", None - ) - return runtime or temporalio.runtime.Runtime.default() - - -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 - code, Activities and the client. Each client of the plugin gets its own runtime object, and - the client's Workers share it. It registers the backend name ``redis``, which the Stream - handles it creates carry. For now a Worker takes one Stream plugin, so it serves one - backend. A client may have more than one, and ``backend=`` picks one. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - url: str | Sequence[str], - *, - cluster: bool = False, - key_prefix: str | None = None, - retention: timedelta | None = None, - safety_ttl: timedelta | None = None, - response_timeout: timedelta | None = None, - name: str = "redis", - default: bool = False, - ) -> None: - """Create the plugin. Each client of the plugin connects to Redis on its first Stream - call, on the client's runtime. - - Args: - url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the - URL. - cluster: Whether the URLs are of a Redis cluster. - key_prefix: The prefix of every key. Defaults to ``temporal-streams``. - retention: The default Retention of a new Stream. Defaults to 7 days. - safety_ttl: How long an open Workflow-owned Stream stays after its last write. - Defaults to 90 days. - response_timeout: How long one Redis call takes at most. Defaults to 5 s. - name: The backend name that Stream handles carry, so every process routes a - handle it receives. - default: Make this the client's default backend. A client with more than one - Stream plugin marks exactly one, else it raises ``StreamConfigError``. - """ - config = ext_pb.RedisStreamConfig( - urls=[url] if isinstance(url, str) else list(url), - cluster=cluster, - key_prefix=key_prefix or "", - ) - for field, value in ( - ("retention", retention), - ("safety_ttl", safety_ttl), - ("response_timeout", response_timeout), - ): - if value is not None: - getattr(config, field).FromTimedelta(value) - super().__init__( - "temporalio.RedisStreamPlugin", False, config, None, name, default - ) - - -class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): - # One for each client of the plugin, with that client's Stream runtime. The client's - # Workers take it from the client. - def __init__( - self, - plugin: _ExtensionStreamPlugin, - backend: _ExtensionStreamBackend | None, - converter: temporalio.converter.DataConverter | None, - ) -> None: - super().__init__() - self.plugin = plugin - self.backend = backend - self.converter = converter or temporalio.converter.DataConverter.default - - def workflow_interceptor_class( - self, input: temporalio.worker.WorkflowInterceptorClassInput - ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: - # Workflow code learns the Worker's backend name from the interceptor class, and the - # run gives its closes to the interceptor. - return type( - "_WorkflowInbound", - (_WorkflowInbound,), - { - "_backend_name": self.plugin.backend_name, - "_interceptor": self if self.backend is not None else None, - }, - ) - - def intercept_activity( - self, next: temporalio.worker.ActivityInboundInterceptor - ) -> temporalio.worker.ActivityInboundInterceptor: - return _ActivityInbound(next, self) - - -@dataclass -class _RunClose: - # The closes of a run's close-with-run Streams, which the worker writes once the server - # accepted the run's completion. - interceptor: _StreamInterceptor - namespace: str - run_id: str - streams: list[models.StreamRef] - result: temporalio.api.common.v1.Payload | None - failure: temporalio.api.failure.v1.Failure | None - status: stream_pb.CloseStatus.ValueType | None - - -class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): - _backend_name = "" - _interceptor: _StreamInterceptor | None = None - - def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: - super().init(_WorkflowOutbound(outbound, self._backend_name)) - - async def execute_workflow( - self, input: temporalio.worker.ExecuteWorkflowInput - ) -> Any: - _take_owned_streams(input.headers) - try: - result = await super().execute_workflow(input) - await streams_workflow.drain() - except temporalio.workflow.ContinueAsNewError: - # The next run continues the run chain's Streams. - await streams_workflow.drain() - raise - except BaseException as err: - if not isinstance(err, (Exception, asyncio.CancelledError)): - # A GeneratorExit or another eviction exception closes the coroutine outside - # the Workflow's loop. The run did not end, so there is nothing to close. - raise - if ( - isinstance(err, asyncio.CancelledError) - and temporalio.workflow.cancellation_reason() is not None - ): - # The run ends cancelled. - await streams_workflow.drain() - self._close_with_run( - failure=temporalio.exceptions.CancelledError("Workflow cancelled") - ) - elif _fails_run(err) and not chain_continues( - temporalio.workflow.info(), temporalio.workflow.now(), err - ): - await streams_workflow.drain() - self._close_with_run(failure=err) - raise - # A cron run that completes still goes on with the chain. - if not chain_continues( - temporalio.workflow.info(), temporalio.workflow.now(), None - ): - self._close_with_run(result=result) - return result - - def _close_with_run( - self, *, result: Any | None = None, failure: BaseException | None = None - ) -> None: - # The closes wait until the server accepts the run's completion. A last Workflow Task - # that fails then leaves the Streams open, so its retry still writes to them. - interceptor = self._interceptor - if interceptor is None or temporalio.workflow.unsafe.is_replaying(): - return - streams = [ - run_stream.ref - for run_stream in streams_workflow._state().streams.values() - if run_stream.close_with_run and not run_stream.closed - ] - if not streams: - return - converter = temporalio.workflow.payload_converter() - failure_proto = None - if failure is not None: - failure_proto = temporalio.api.failure.v1.Failure() - temporalio.converter.DataConverter.default.failure_converter.to_failure( - failure, converter, failure_proto - ) - info = temporalio.workflow.info() - interceptor.plugin._run_closes[info.run_id] = _RunClose( - interceptor=interceptor, - namespace=info.namespace, - run_id=info.run_id, - streams=streams, - result=converter.to_payload(result) if result is not None else None, - failure=failure_proto, - status=close_status(failure), - ) - - -def _fails_run(err: BaseException) -> bool: - # Any other exception fails the Workflow Task, and the run goes on. - return temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) - - -def chain_continues( - info: temporalio.workflow.Info, now: datetime, failure: BaseException | None -) -> bool: - """Whether the server starts another run of the chain after this run ends. - - Then the next run goes on with the run chain's Streams, so the runtime skips this run's - closes. A cron run always continues. A failed run continues when its retry policy has - attempts and time left for the failure. - """ - if info.cron_schedule: - return True - policy = info.retry_policy - if failure is None or policy is None: - return False - if policy.maximum_attempts != 0 and info.attempt >= policy.maximum_attempts: - return False - if _non_retryable(failure, policy.non_retryable_error_types): - return False - if info.execution_timeout is not None: - backoff = policy.initial_interval * ( - policy.backoff_coefficient ** (info.attempt - 1) - ) - if policy.maximum_interval is not None: - backoff = min(backoff, policy.maximum_interval) - # The execution timeout covers the whole chain, so this run's start is not its start. - expiration = info.execution_expiration_time or ( - info.workflow_start_time + info.execution_timeout - ) - if now + backoff >= expiration: - return False - return True - - -def _non_retryable(err: BaseException, types: Sequence[str] | None) -> bool: - if isinstance(err, temporalio.exceptions.ApplicationError): - return err.non_retryable or (err.type or "") in (types or []) - return type(err).__name__ in (types or []) - - -async def _write_closes(close: _RunClose) -> None: - backend, converter = close.interceptor.backend, close.interceptor.converter - assert backend is not None - service = backend.service(close.namespace, converter) - failure = ( - converter.failure_converter.from_failure( - close.failure, converter.payload_converter - ) - if close.failure is not None - else None - ) - for ref in close.streams: - try: - await service.close_stream( - models.CloseStreamRequest( - stream=ref, - # The run itself closes its Streams, as their owner. - producer=models.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_WORKFLOW, - id=_ids.workflow_producer_id(close.run_id), - run_id=close.run_id, - ), - result=temporalio.common.RawValue(close.result) - if close.result is not None - else None, - failure=failure, - status=close.status, - ) - ) - except Exception: - logger.warning( - "Failed closing the Stream %s with the run %s", - ref.id, - close.run_id, - exc_info=True, - ) - - -def _owned_ref(stream_id: str) -> models.StreamRef: - # A Stream that the run chain creates is the chain's own. - info = temporalio.workflow.info() - return models.StreamRef( - id=stream_id, - workflow_id=info.workflow_id, - first_execution_run_id=info.first_execution_run_id, - ) - - -def _take_owned_streams( - headers: Mapping[str, temporalio.api.common.v1.Payload], -) -> None: - payload = headers.get(_OWNED_STREAMS_HEADER) - if payload is None: - return - state = streams_workflow._state() - converter = temporalio.converter.PayloadConverter.default - 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): - _temporal_stream_transport = True - - def __init__( - self, next: temporalio.worker.WorkflowOutboundInterceptor, backend_name: str - ) -> None: - super().__init__(next) - self._temporal_stream_backend_name = backend_name - # Numbers each Activity call of the run in order, so the key is the same on replay and - # a retry keeps it. A counter, not workflow.uuid4(), leaves the user's random values - # alone. - self._activity_calls = 0 - - def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: - self._activity_calls += 1 - # Both ends use the default converter, so a custom user converter can't break the - # header. - converter = temporalio.converter.PayloadConverter.default - return {_KEY_HEADER: converter.to_payload(str(self._activity_calls))} - - def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: - # The next run needs each owned Stream's backend for stream_handle. - owned = [ - [stream_id, run_stream.backend, run_stream.close_with_run] - for stream_id, run_stream in streams_workflow._state().streams.items() - if run_stream.backend is not None and not run_stream.closed - ] - if owned: - converter = temporalio.converter.PayloadConverter.default - input.headers = { - **input.headers, - _OWNED_STREAMS_HEADER: converter.to_payload(owned), - } - super().continue_as_new(input) - - def start_activity( - self, input: temporalio.worker.StartActivityInput - ) -> temporalio.workflow.ActivityHandle[Any]: - input.headers = {**input.headers, **self._activity_headers()} - return super().start_activity(input) - - def start_local_activity( - self, input: temporalio.worker.StartLocalActivityInput - ) -> temporalio.workflow.ActivityHandle[Any]: - input.headers = {**input.headers, **self._activity_headers()} - return super().start_local_activity(input) - - -class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): - def __init__( - self, - next: temporalio.worker.ActivityInboundInterceptor, - interceptor: _StreamInterceptor, - ) -> None: - super().__init__(next) - self._interceptor = interceptor - - async def execute_activity( - self, input: temporalio.worker.ExecuteActivityInput - ) -> Any: - streams = _activity_streams(input, self._interceptor) - token = _activity.current.set(streams) - try: - if temporalio.activity.info().attempt > 1: - # The open of a new attempt withdraws the Records of the earlier ones, also - # when this attempt publishes nothing. It creates no Stream. - stream_ids = [handle.id for handle in _stream_arguments(input.args)] - for stream_id in dict.fromkeys(stream_ids): - await streams.writer(stream_id).withdraw_earlier_attempts() - result = await super().execute_activity(input) - except BaseException: - await _close_quietly(streams) - raise - else: - await streams.close() - return result - finally: - _activity.current.reset(token) - - -def _stream_arguments(args: Sequence[Any]) -> list[StreamHandle[Any]]: - # The Streams that an Activity received, which an earlier attempt may have written. They - # may sit in fields of its input object, in slots, and in lists and dicts. - handles: list[StreamHandle[Any]] = [] - seen: set[int] = set() - - def walk(value: Any, depth: int) -> None: - if depth > 4 or id(value) in seen: - return - seen.add(id(value)) - if isinstance(value, StreamHandle): - handles.append(value) - elif isinstance(value, (str, bytes, int, float, bool)) or value is None: - return - elif isinstance(value, dict): - for item in value.values(): - walk(item, depth + 1) - elif isinstance(value, (list, tuple, set, frozenset)): - for item in value: - walk(item, depth + 1) - else: - for item in _fields_of(value): - walk(item, depth + 1) - - for arg in args: - walk(arg, 0) - return handles - - -def _fields_of(value: Any) -> list[Any]: - fields = list(getattr(value, "__dict__", {}).values()) - for klass in type(value).__mro__: - slots = getattr(klass, "__slots__", ()) - for slot in (slots,) if isinstance(slots, str) else slots: - if hasattr(value, slot): - fields.append(getattr(value, slot)) - return fields - - -async def _close_quietly(streams: _activity.ActivityStreams) -> None: - try: - await asyncio.shield(streams.close()) - except Exception: - temporalio.activity.logger.warning( - "Failed flushing Stream Records of a failed Activity", exc_info=True - ) - - -def _activity_streams( - input: temporalio.worker.ExecuteActivityInput, interceptor: _StreamInterceptor -) -> _activity.ActivityStreams: - # A sync Activity has no client, so the interceptor gives the backend and the converter. - assert interceptor.backend is not None - info = temporalio.activity.info() - converter = temporalio.converter.PayloadConverter.default - key = input.headers.get(_KEY_HEADER) - key_text = converter.from_payload(key, str) if key is not None else "" - epoch = stream_pb.ProducerEpoch(attempt=info.attempt) - previous: stream_pb.WorkflowTaskEpoch | None = None - failed: tuple[int, ...] = () - if info.is_local: - task = info.workflow_task - if task is not None: - # Always set, with 0 for none completed. An unset one turns off the withdrawal of - # the run's orphaned Local Activity Records. - previous = stream_pb.WorkflowTaskEpoch( - started_event_id=task.previous_started_event_id - ) - failed = tuple(task.failed_tasks) - epoch.task_started_event_id = task.started_event_id - epoch.task_attempt = task.attempt - if task.started_time is not None: - epoch.task_started_time.FromDatetime(task.started_time) - run_id = info.workflow_run_id or "" - producer = stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_LOCAL_ACTIVITY, - id=_ids.local_activity_producer_id(run_id, key_text), - run_id=run_id, - ) - else: - epoch.attempt_scheduled_time.FromDatetime(info.current_attempt_scheduled_time) - if info.in_workflow: - run_id = info.workflow_run_id or "" - producer = stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_ACTIVITY, - id=_ids.activity_producer_id(run_id, info.activity_id, key_text), - run_id=run_id, - ) - else: - producer = stream_pb.Producer( - kind=stream_pb.ProducerKind.PRODUCER_KIND_STANDALONE_ACTIVITY, - id=_ids.standalone_activity_producer_id( - info.activity_id, info.activity_run_id or "" - ), - ) - producer.epoch.CopyFrom(epoch) - return _activity.ActivityStreams( - backend=interceptor.backend, - converter=interceptor.converter, - namespace=info.namespace, - producer=producer, - loop=asyncio.get_running_loop(), - backend_name=interceptor.plugin.backend_name, - workflow_id=info.workflow_id, - previous_completed_task=previous, - failed_tasks=failed, - ) diff --git a/temporalio/contrib/streams/testing.py b/temporalio/contrib/streams/testing.py deleted file mode 100644 index 5988c28668..0000000000 --- a/temporalio/contrib/streams/testing.py +++ /dev/null @@ -1,55 +0,0 @@ -"""A Stream plugin for tests, with no Redis. - -.. warning:: - This module is experimental. -""" - -from __future__ import annotations - -from datetime import timedelta - -import temporalio.runtime -from temporalio.contrib.streams.proto import v1 as ext_pb -from temporalio.contrib.streams.redis import _ExtensionStreamPlugin - - -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 - plugin shares its Streams, and they go when the process ends. It registers the backend - name ``memory``, so tests swap the Store with no change to Workflow code. - - .. warning:: - This API is experimental. - """ - - def __init__( - self, - *, - retention: timedelta | None = None, - safety_ttl: timedelta | None = None, - runtime: temporalio.runtime.Runtime | None = None, - name: str = "memory", - default: bool = False, - ) -> None: - """Create the plugin. - - Args: - retention: The default Retention of a new Stream. Defaults to 7 days. - safety_ttl: How long an open Workflow-owned Stream stays after its last write. - Defaults to 90 days. - runtime: The runtime of the plugin's memory Store. Defaults to the default - runtime. - name: The backend name that Stream handles carry. - default: Make this the client's default backend. A client with more than one - Stream plugin marks exactly one. - """ - config = ext_pb.RedisStreamConfig() - if retention is not None: - config.retention.FromTimedelta(retention) - if safety_ttl is not None: - config.safety_ttl.FromTimedelta(safety_ttl) - super().__init__( - "temporalio.MemoryStreamPlugin", True, config, runtime, name, default - ) diff --git a/temporalio/converter/_payload_converter.py b/temporalio/converter/_payload_converter.py index 3540d743cb..210aa33c36 100644 --- a/temporalio/converter/_payload_converter.py +++ b/temporalio/converter/_payload_converter.py @@ -110,9 +110,10 @@ def transfer_type_convertible( ) -> _TransferTypeConvertibleDecorator[ValueT, TransferTypeT]: """Decorate a class with a transfer type converter class. - A value of the class takes its transfer type at the top level of a payload, and inside a - JSON value such as a list or a dataclass field. A dataclass is the exception inside a - JSON value: ``dataclasses.asdict`` writes it as its fields, and it decodes from them. + The converter converts a value of the class to its transfer type at the top level of a + payload, and inside a JSON value such as a list item or a dataclass field. Inside a JSON + value, a dataclass is the exception. ``dataclasses.asdict`` writes it as its fields, and + the converter decodes it from them. .. warning:: This API is experimental and subject to change. @@ -132,8 +133,8 @@ def _get_transfer_type_converter( def _nested_transfer_type_converter( hint: object, ) -> TransferTypeConverter[Any, Any] | None: - # A dataclass inside another value is written as its fields, so only other classes take - # their transfer type there. + # dataclasses.asdict writes a nested dataclass as its fields, so only a class that isn't a + # dataclass is converted there. origin = typing.get_origin(hint) or hint if not inspect.isclass(origin) or dataclasses.is_dataclass(origin): return None @@ -689,9 +690,8 @@ def default(self, o: Any) -> Any: # Datetime support if isinstance(o, datetime): return o.isoformat() - # A transfer type value inside another value. A dataclass never gets here, since - # dataclasses.asdict has already turned it into a dict. - converter = _get_transfer_type_converter(type(o)) + # Nested transfer type support + converter = _nested_transfer_type_converter(type(o)) if converter is not None: value = converter.to_transfer_type(o) if isinstance(value, google.protobuf.message.Message): @@ -866,7 +866,7 @@ def value_to_type( if ret is not JSONTypeConverter.Unhandled: return ret - # A transfer type value inside another value, as AdvancedJSONEncoder writes it + # Nested transfer type values, as AdvancedJSONEncoder writes them converter = _nested_transfer_type_converter(hint) if converter is not None: transfer_type = converter.transfer_type diff --git a/temporalio/testing/_workflow.py b/temporalio/testing/_workflow.py index 6c8354162c..b4af035279 100644 --- a/temporalio/testing/_workflow.py +++ b/temporalio/testing/_workflow.py @@ -8,6 +8,7 @@ from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timedelta, timezone from typing import ( + TYPE_CHECKING, Any, cast, ) @@ -31,6 +32,9 @@ TimeSkippingConfig, ) +if TYPE_CHECKING: + from temporalio.common import StreamBackend + logger = logging.getLogger(__name__) @@ -102,6 +106,7 @@ async def start_local( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> WorkflowEnvironment: """Start a full Temporal server locally, downloading if necessary. @@ -165,6 +170,8 @@ async def start_local( become ready after spawning. Defaults to 5 seconds and does not include download time. ui_port: UI port to use if UI is enabled. + stream_backends: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started CLI dev server workflow environment. @@ -194,6 +201,7 @@ async def start_local( dev_server_download_ttl=dev_server_download_ttl, dev_server_startup_timeout=dev_server_startup_timeout, ui_port=ui_port, + stream_backends=stream_backends, ) @classmethod @@ -216,6 +224,7 @@ async def start_time_skipping( test_server_extra_args: Sequence[str] = [], test_server_download_ttl: timedelta | None = None, test_server_startup_timeout: timedelta | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> WorkflowEnvironment: """Start a V1 time skipping workflow environment. @@ -276,6 +285,8 @@ async def start_time_skipping( test_server_startup_timeout: Maximum time to wait for the test server to become ready after spawning. Defaults to 5 seconds and does not include download time. + stream_backends: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started workflow environment with time skipping. @@ -296,6 +307,7 @@ async def start_time_skipping( test_server_extra_args=test_server_extra_args, test_server_download_ttl=test_server_download_ttl, test_server_startup_timeout=test_server_startup_timeout, + stream_backends=stream_backends, ) @classmethod @@ -582,6 +594,7 @@ async def _bootstrap_dev_server( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_backends: Sequence[StreamBackend] = (), ) -> tuple[ temporalio.bridge.testing.EphemeralServer, temporalio.client.Client, @@ -649,6 +662,7 @@ async def _bootstrap_dev_server( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_backends=stream_backends, ) except: try: @@ -719,6 +733,7 @@ async def _create( test_server_extra_args: Sequence[str], test_server_download_ttl: timedelta | None, test_server_startup_timeout: timedelta | None, + stream_backends: Sequence[StreamBackend], ) -> _V1WorkflowEnvironment: runtime = runtime or temporalio.runtime.Runtime.default() download_ttl_ms = None @@ -753,6 +768,7 @@ async def _create( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_backends=stream_backends, ) except: try: diff --git a/temporalio/worker/__init__.py b/temporalio/worker/__init__.py index c322971e51..4f6efe68c3 100644 --- a/temporalio/worker/__init__.py +++ b/temporalio/worker/__init__.py @@ -25,7 +25,7 @@ WorkflowInterceptorClassInput, WorkflowOutboundInterceptor, ) -from ._plugin import Plugin, WorkflowEvictedInput +from ._plugin import Plugin from ._replayer import ( Replayer, ReplayerConfig, @@ -87,7 +87,6 @@ "WorkflowOutboundInterceptor", "NexusOperationInboundInterceptor", "Plugin", - "WorkflowEvictedInput", # Interceptor input "ContinueAsNewInput", "ExecuteActivityInput", diff --git a/temporalio/worker/_plugin.py b/temporalio/worker/_plugin.py index 6447682a67..c222f836c2 100644 --- a/temporalio/worker/_plugin.py +++ b/temporalio/worker/_plugin.py @@ -3,7 +3,6 @@ import abc from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import AbstractAsyncContextManager -from dataclasses import dataclass from typing import TYPE_CHECKING from temporalio.client import WorkflowHistory @@ -18,22 +17,6 @@ ) -@dataclass(frozen=True) -class WorkflowEvictedInput: - """Input for :py:meth:`Plugin.on_workflow_evicted`. - - .. warning:: - This API is experimental. - """ - - run_id: str - """Run ID of the Workflow.""" - - execution_ending: bool - """Whether the server accepted the run's last completion, so the run ended. False for any - other eviction, such as a full cache, a failed Workflow Task or a Worker shutdown.""" - - class Plugin(abc.ABC): """Base class for worker plugins that can intercept and modify worker behavior. @@ -81,26 +64,6 @@ async def run_worker( next: Callable to continue the worker execution. """ - def on_workflow_evicted( - self, - input: WorkflowEvictedInput, # pyright: ignore[reportUnusedParameter] - ) -> None: - """Hook called after the Worker evicts a Workflow run from its cache. - - The Worker calls it outside the sandbox, on its event loop, after Core reports the - eviction. So a plugin can act on a run that ended only once the server accepted its - completion. The hook must not block. A plugin starts its async work here and waits for - it in :py:meth:`run_worker`. The Worker logs an exception and goes on. The - :py:class:`Replayer` never calls it. The default does nothing. - - .. warning:: - This API is experimental. - - Args: - input: The run and why it was evicted. - """ - pass - @abc.abstractmethod def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig: """Hook called when creating a replayer to allow modification of configuration. diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index b3eb1a4d12..e4ae5e9d94 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass +from typing import TYPE_CHECKING from typing_extensions import TypedDict @@ -32,6 +33,9 @@ ) from .workflow_sandbox import SandboxedWorkflowRunner +if TYPE_CHECKING: + from temporalio.common import StreamBackend + logger = logging.getLogger(__name__) @@ -56,13 +60,16 @@ def __init__( runtime: temporalio.runtime.Runtime | None = None, disable_safe_workflow_eviction: bool = False, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_backends: Sequence[StreamBackend] = (), ) -> None: """Create a replayer to replay workflows from history. See :py:meth:`temporalio.worker.Worker.__init__` for a description of most of the arguments. Most of the same arguments need to be passed to the replayer that were passed to the worker when the workflow originally - ran. + ran. Pass ``stream_backends`` with the same backends as the Worker's + client, such as :py:class:`temporalio.contrib.stream.RedisStreamBackend`. + A replay makes no Stream call, so the Store doesn't need to be reachable. Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that @@ -86,6 +93,7 @@ def __init__( runtime=runtime, disable_safe_workflow_eviction=disable_safe_workflow_eviction, header_codec_behavior=header_codec_behavior, + stream_backends=list(stream_backends), ) self._initial_config = self._config.copy() self._default_workflow_logic_flags = set(_DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS) @@ -94,6 +102,10 @@ def __init__( self.plugins = plugins for plugin in plugins: self._config = plugin.configure_replayer(self._config) + if self._config.get("stream_backends"): + from temporalio.common._streams import _plugin as _stream_plugin + + self._config = _stream_plugin.configure_replayer(self._config) # Validate workflows after plugin configuration if not self._config.get("workflows"): @@ -433,6 +445,7 @@ class ReplayerConfig(TypedDict, total=False): runtime: temporalio.runtime.Runtime | None disable_safe_workflow_eviction: bool header_codec_behavior: HeaderCodecBehavior + stream_backends: Sequence[StreamBackend] @dataclass(frozen=True) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index d687faf06f..af10ffc05a 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -19,7 +19,6 @@ from typing_extensions import TypedDict -import temporalio.bridge.proto.workflow_activation import temporalio.bridge.worker import temporalio.client import temporalio.common @@ -34,7 +33,7 @@ from ._activity import SharedStateManager, _ActivityWorker from ._interceptor import Interceptor from ._nexus import _NexusWorker -from ._plugin import Plugin, WorkflowEvictedInput +from ._plugin import Plugin from ._tuning import WorkerTuner from ._workflow import ( _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, @@ -133,7 +132,6 @@ def __init__( max_task_queue_activities_per_second: float | None = None, max_eager_activity_reservations_per_workflow_task: int = 3, graceful_shutdown_timeout: timedelta = timedelta(), - plugin_shutdown_timeout: timedelta | None = timedelta(seconds=60), workflow_failure_exception_types: Sequence[type[BaseException]] = [], shared_state_manager: SharedStateManager | None = None, debug_mode: bool = False, @@ -280,11 +278,6 @@ def __init__( graceful_shutdown_timeout: Amount of time after shutdown is called that activities are given to complete before their tasks are cancelled. - plugin_shutdown_timeout: Amount of time that exiting - ``async with Worker(...)`` waits, after shutdown, for the code - that plugins run after the Worker stops in ``run_worker``. That - code is cancelled after this time. ``None`` waits with no limit. - The default is 60 seconds. workflow_failure_exception_types: The types of exceptions that, if a workflow-thrown exception extends, will cause the workflow/update to fail instead of suspending the workflow via @@ -386,7 +379,6 @@ def __init__( max_task_queue_activities_per_second=max_task_queue_activities_per_second, max_eager_activity_reservations_per_workflow_task=max_eager_activity_reservations_per_workflow_task, graceful_shutdown_timeout=graceful_shutdown_timeout, - plugin_shutdown_timeout=plugin_shutdown_timeout, workflow_failure_exception_types=workflow_failure_exception_types, shared_state_manager=shared_state_manager, debug_mode=debug_mode, @@ -418,6 +410,13 @@ def __init__( self._plugins = plugins for plugin in plugins: config = plugin.configure_worker(config) + # Configure Streams after the plugins, so a plugin that replaces the Workflow runner + # doesn't lose the Streams sandbox passthrough. + stream_backends = client.config(active_config=True).get("stream_backends") + if stream_backends: + from temporalio.common._streams import _plugin as _stream_plugin + + config = _stream_plugin.configure_worker(stream_backends, config) self._init_from_config(client, config) @@ -496,7 +495,6 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf self._async_context_inner_task: asyncio.Task | None = None self._async_context_run_task: asyncio.Task | None = None self._async_context_run_exception: BaseException | None = None - self._async_context_exiting = False self._activity_worker: _ActivityWorker | None = None self._runtime = ( @@ -577,7 +575,7 @@ def check_activity(activity: str): "disable_eager_activity_execution" ], # type: ignore[reportTypedDictNotRequiredAccess] metric_meter=self._runtime.metric_meter, - on_eviction_hook=self._eviction_hook(), + on_eviction_hook=None, disable_safe_eviction=config["disable_safe_workflow_eviction"], # type: ignore[reportTypedDictNotRequiredAccess] should_enforce_versioning_behavior=should_enforce_versioning_behavior, assert_local_activity_valid=check_activity, @@ -719,42 +717,6 @@ def check_activity(activity: str): ), ) - def _eviction_hook( - self, - ) -> ( - Callable[ - [str, temporalio.bridge.proto.workflow_activation.RemoveFromCache], None - ] - | None - ): - plugins = [ - plugin - for plugin in self._plugins - if type(plugin).on_workflow_evicted is not Plugin.on_workflow_evicted - ] - if not plugins: - return None - ending = temporalio.bridge.proto.workflow_activation.RemoveFromCache.EvictionReason.WORKFLOW_EXECUTION_ENDING - - def on_evicted( - run_id: str, - job: temporalio.bridge.proto.workflow_activation.RemoveFromCache, - ) -> None: - input = WorkflowEvictedInput( - run_id=run_id, execution_ending=job.reason == ending - ) - for plugin in plugins: - try: - plugin.on_workflow_evicted(input) - except Exception: - logger.exception( - "Plugin %s failed on the eviction of run %s", - plugin.name(), - run_id, - ) - - return on_evicted - def config(self, *, active_config: bool = False) -> WorkerConfig: """Config, as a dictionary, used to create this worker. @@ -1005,10 +967,7 @@ async def run(): await self.run() except BaseException as err: self._async_context_run_exception = err - # After __aexit__ starts, the caller is past the block, so a cancel would hit - # its later awaits. - if not self._async_context_exiting: - self._async_context_inner_task.cancel() # type: ignore[union-attr] + self._async_context_inner_task.cancel() # type: ignore[union-attr] self._async_context_run_task = asyncio.create_task(run()) return self @@ -1023,20 +982,8 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N if not self._async_context_run_task: raise RuntimeError("Never started") await self.shutdown() - # Plugins can run cleanup code after shutdown. Wait for that code instead of cancelling it. - self._async_context_exiting = True - run_task = self._async_context_run_task - timeout = self._config.get("plugin_shutdown_timeout", timedelta(seconds=60)) - try: - # wait() doesn't raise for the task's own outcome, only when this call is cancelled. - await asyncio.wait( - [run_task], timeout=timeout.total_seconds() if timeout else None - ) - except asyncio.CancelledError: - run_task.cancel() - raise - if not run_task.done(): - run_task.cancel() + # Cancel our run task + self._async_context_run_task.cancel() # Only re-raise our exception if present and exc_type is cancel if exc_type is asyncio.CancelledError and self._async_context_run_exception: raise self._async_context_run_exception @@ -1076,7 +1023,6 @@ class WorkerConfig(TypedDict, total=False): max_task_queue_activities_per_second: float | None max_eager_activity_reservations_per_workflow_task: int graceful_shutdown_timeout: timedelta - plugin_shutdown_timeout: timedelta | None workflow_failure_exception_types: Sequence[type[BaseException]] shared_state_manager: SharedStateManager | None debug_mode: bool diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index bfff26cd91..f50955c4cf 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -720,11 +720,6 @@ def _create_workflow_instance( workflow_id=init.workflow_id, workflow_type=init.workflow_type, priority=temporalio.common.Priority._from_proto(init.priority), - execution_expiration_time=init.workflow_execution_expiration_time.ToDatetime().replace( - tzinfo=timezone.utc - ) - if init.HasField("workflow_execution_expiration_time") - else None, ) last_failure = ( diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index c631fd88e8..ef1c20a37c 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -151,6 +151,11 @@ logger, unsafe, ) +from ._streams import ( + create_stream, + create_stream_nowait, + stream_handle, +) from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowConfig, @@ -258,6 +263,9 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", + "create_stream", + "create_stream_nowait", + "stream_handle", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index c86620eea9..a33759c37a 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -136,10 +136,6 @@ class Info: workflow_type: str - execution_expiration_time: datetime | None = None - """When the execution timeout of the run chain passes, across retries and continue-as-new. - None when the chain has no execution timeout.""" - def _logger_details(self) -> Mapping[str, Any]: return { # TODO(cretz): worker ID? diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py new file mode 100644 index 0000000000..856f517264 --- /dev/null +++ b/temporalio/workflow/_streams.py @@ -0,0 +1,119 @@ +"""Streams in Workflow code. + +These functions import the Streams API of :py:mod:`temporalio.common` on first use, so a +Workflow with no Stream never imports it. +""" + +from __future__ import annotations + +from datetime import timedelta +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from temporalio.common import Creating, WorkflowStream + +D = TypeVar("D") + + +def create_stream( + definition: type[D], + *, + retention: timedelta | None = None, + backend: str | None = None, +) -> Creating[WorkflowStream[D]]: + """Create a Stream that this run chain owns. + + The Stream id is ``::``, so a client can get it from the + Workflow handle without a call. The create and each publish are Local Activities, and the + create's error is raised here. Only the run chain can close the Stream. A close runs after + its Workflow Task commits, as a regular Activity on Redis, so a task that fails first leaves + the Stream open for its retry. On Redis, the end of the run chain doesn't close a Stream + outside an ``async with`` block. So a Stream that nobody closes stays open until its + ``safety_ttl``. + + Await the call to get the handle, or use it in ``async with``, which closes the Stream at + the end of the block:: + + async with workflow.create_stream(common.with_name(Events, "turn-3")) as turn: + await turn.topic(Events.tokens).publish("hi") + + An exception in the block closes the Stream with that exception as the failure. If a retry + run can come next, the close waits for the run's end. The Stream then stays open only for a + retry run, or for the retry of a failed Workflow Task. A cron run leaves the Stream open. + See :py:meth:`temporalio.common.WorkflowStream.__aexit__`. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. Its name is the Stream name. For one Stream + per turn, for example, pass ``common.with_name(Events, f"turn-{n}")``. + retention: The age after which a Record is deleted, even while the Stream is open. + A closed Stream is kept for one Retention after its close. Defaults to the + backend's Retention. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. The default is part of the Workflow's code, since a replay must + create the Stream on the same backend. A Worker serves only one backend, so the + name must be the Worker's backend name. + + Raises: + StreamError: The backend refused the create, such as for a different Retention. It's + ``StreamClosedError`` or ``StreamExpiredError`` if a Stream with that id closed or + expired. + ReadOnlyContextError: The code must not change state, such as a query handler. + StreamConfigError: The Worker has no Stream backend or none with that name, or no + backend has ``default=True`` and the call gives no backend. It fails the Workflow + Task, not the run. + """ + from temporalio.common._streams import _workflow + + return _workflow.create_stream(definition, retention=retention, backend=backend) + + +def create_stream_nowait( + definition: type[D], + *, + retention: timedelta | None = None, + backend: str | None = None, +) -> WorkflowStream[D]: + """Create a Stream that this run chain owns, without an await. + + Use it in sync Workflow code, such as ``@workflow.init``. It's like + ``asyncio.Queue.put_nowait`` for code that can't await ``put``. It takes the options of + :py:func:`create_stream`. A backend error raises in the run as an error that the run + doesn't catch, so a ``StreamError`` fails the run. + + .. warning:: + This API is experimental. + """ + from temporalio.common._streams import _workflow + + return _workflow.create_stream_nowait( + definition, retention=retention, backend=backend + ) + + +def stream_handle( + definition: type[D], *, backend: str | None = None +) -> WorkflowStream[D]: + """Get a Stream that this run chain created, by its definition. It makes no call. + + A later run of the chain, such as one after continue-as-new, uses it to write to or close + the Stream. In ``@workflow.init`` of a run that continued as new, the run hasn't read the + backends of the chain's Streams yet. So pass ``backend=`` there, or call it in the run + method. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class, with the same name as in the create. + backend: The backend of the Stream. Defaults to the backend that the run chain created + it on, or else to the registered backend with ``default=True``. + + Raises: + StreamConfigError: No backend is given, known or marked as the default. + """ + from temporalio.common._streams import _workflow + + return _workflow.stream_handle(definition, backend=backend) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index e6cd47603d..90b2e05c7e 100644 --- a/tests/contrib/streams/_support.py +++ b/tests/contrib/streams/_support.py @@ -38,4 +38,4 @@ def crate_file(relative: str) -> Path: def vectors(name: str) -> object: - return json.loads(crate_file(f"testdata/{name}").read_text()) + return json.loads(crate_file(f"testdata/{name}").read_text(encoding="utf-8")) diff --git a/tests/contrib/streams/conftest.py b/tests/contrib/streams/conftest.py index c2cde1e5e1..ba32594c13 100644 --- a/tests/contrib/streams/conftest.py +++ b/tests/contrib/streams/conftest.py @@ -9,9 +9,9 @@ import pytest from temporalio.client import Client -from temporalio.contrib.streams import StreamPlugin -from temporalio.contrib.streams.redis import RedisStreamPlugin -from temporalio.contrib.streams.testing import MemoryStreamPlugin +from temporalio.common import StreamBackend +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend # Ports that other local services own. _RESERVED = {7333, 7861, 7931, 6431, 7101, 7102, 7103} @@ -26,8 +26,8 @@ def _free_port() -> int: return port -# One scratch server serves every test module, and each test takes its own key prefix. A run -# that must use a known port gives its own server in TEMPORAL_STREAMS_REDIS_URL. +# One scratch server serves every test module, and each test uses its own key prefix. To use a +# known port, set TEMPORAL_STREAMS_REDIS_URL to your own server. @pytest.fixture(scope="session") def redis_url() -> Iterator[str]: given = os.environ.get("TEMPORAL_STREAMS_REDIS_URL") @@ -59,14 +59,16 @@ def redis_url() -> Iterator[str]: async def stream_client( client: Client, request: pytest.FixtureRequest ) -> AsyncIterator[Client]: - # Every live test runs on the memory Store and on Redis, since both must keep one set of - # rules. - plugin: StreamPlugin + # Every live test runs on the memory Store and on Redis, since both must have the same + # rules. The backend is the default, so a create with no backend name uses it. + backend: StreamBackend if request.param == "memory": - plugin = MemoryStreamPlugin() + backend = MemoryStreamBackend(default=True) else: redis_url = request.getfixturevalue("redis_url") - plugin = RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}") + backend = RedisStreamBackend( + redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True + ) config = client.config() - config["plugins"] = [plugin] + config["stream_backends"] = [backend] yield Client(**config) diff --git a/tests/contrib/streams/examples/README.md b/tests/contrib/streams/examples/README.md index 9564002ee9..d559350efa 100644 --- a/tests/contrib/streams/examples/README.md +++ b/tests/contrib/streams/examples/README.md @@ -1,6 +1,7 @@ # Examples for the Brook streams API proposal -Each file is a complete program for one section of the proposal. `test_examples.py` runs each one end to end on the memory Store and on Redis. +Each file is a complete program for one section of the proposal. `test_examples.py` runs each one +end to end on the memory Store and on Redis. | Proposal section | File | |---|---| @@ -12,4 +13,4 @@ Each file is a complete program for one section of the proposal. `test_examples. | A Workflow scopes a Stream to one turn and closes it | [`turn_scoped_stream.py`](turn_scoped_stream.py) | | Worker setup | [`worker_setup.py`](worker_setup.py) | -The proposal's Nexus examples are a later stage, so they have no file here. +The proposal's Nexus examples aren't in this directory. diff --git a/tests/contrib/streams/examples/activity_streams_tokens.py b/tests/contrib/streams/examples/activity_streams_tokens.py index 49303c502a..b68f0f3558 100644 --- a/tests/contrib/streams/examples/activity_streams_tokens.py +++ b/tests/contrib/streams/examples/activity_streams_tokens.py @@ -12,38 +12,36 @@ from dataclasses import dataclass from datetime import timedelta -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow from temporalio.worker import Worker -@streams.defn(name="events") +@common.defn(name="events") class Events: - tokens = streams.topic(str) + tokens = common.topic(str) @dataclass class CallModelInput: prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] @dataclass class SummarizeInput: text: str - stream: streams.SyncActivityStream[Events] + stream: common.SyncActivityStream[Events] @activity.defn async def call_model(input: CallModelInput) -> int: tokens = input.stream.topic(Events.tokens) words = input.prompt.split() - # Each awaited publish has written when it returns. + # Each awaited publish has written its Record when it returns. last = await tokens.publish(words[0]) # A fast model sends many small tokens, so the helper batches them. - async with streams.BufferedPublisher( + async with common.BufferedPublisher( tokens, max_delay=timedelta(milliseconds=100) ) as fast: for word in words[1:]: @@ -60,23 +58,24 @@ def summarize(input: SummarizeInput) -> int: class Model: @workflow.run async def run(self, prompt: str) -> list[int]: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) first = await workflow.execute_activity( call_model, - CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(minutes=5), ) summary = await workflow.execute_activity( summarize, - SummarizeInput(prompt, streams.SyncActivityStream.from_handle(stream)), + SummarizeInput(prompt, common.SyncActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(minutes=5), ) + await stream.close() return [first, summary] async def main( client: Client, task_queue: str -) -> tuple[list[int], list[streams.Record[str]]]: +) -> tuple[list[int], list[common.Record[str]]]: with concurrent.futures.ThreadPoolExecutor() as executor: async with Worker( client, @@ -92,8 +91,8 @@ async def main( task_queue=task_queue, ) offsets = await handle.result() - stream = streams.StreamClient(client).workflow_stream(handle, Events) + stream = handle.stream(Events) records = [ - item async for item in stream.subscribe() if isinstance(item, streams.Record) + item async for item in stream.subscribe() if isinstance(item, common.Record) ] return offsets, records diff --git a/tests/contrib/streams/examples/agent_publishes_turns.py b/tests/contrib/streams/examples/agent_publishes_turns.py index cd50961991..660ac4b9ab 100644 --- a/tests/contrib/streams/examples/agent_publishes_turns.py +++ b/tests/contrib/streams/examples/agent_publishes_turns.py @@ -1,8 +1,8 @@ """Brook API proposal, "An agent Workflow publishes its turns". The Workflow creates its Stream, publishes each turn with an awaited ``publish``, and passes the -Stream to the model Activity in its one input object. ``close_with_run`` closes the Stream with -the result of the run chain. +Stream to the model Activity in its one input object. It closes the Stream with its result +before it returns, since a run's end doesn't close the Stream. """ from __future__ import annotations @@ -11,10 +11,8 @@ from dataclasses import dataclass, field from datetime import timedelta -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow from temporalio.worker import Worker @@ -24,10 +22,10 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) @dataclass @@ -39,7 +37,7 @@ class Reply: @dataclass class CallModelInput: messages: list[str] - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] @activity.defn @@ -59,20 +57,20 @@ async def run_tool(call: str) -> str: class AgentTurns: @workflow.run async def run(self, messages: list[str]) -> str: - stream = await stream_workflow.create_stream( - Events, retention=timedelta(days=1), auto_close=timedelta(hours=1) - ) + stream = await workflow.create_stream(Events, retention=timedelta(days=1)) turns = stream.topic(Events.turns) while True: # The publish is one Local Activity. History keeps its offset, not the Turn. await turns.publish(Turn("model")) reply = await workflow.execute_activity( call_model, - CallModelInput(messages, streams.ActivityStream.from_handle(stream)), + CallModelInput(messages, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(minutes=5), ) if not reply.tool_calls: await turns.publish(Turn("done", reply.text)) + # A run's end doesn't close the Stream, so the Workflow closes it. + await stream.close(result=reply.text) return reply.text for call in reply.tool_calls: messages.append( @@ -82,9 +80,7 @@ async def run(self, messages: list[str]) -> str: ) -async def main( - client: Client, task_queue: str -) -> list[streams.Record | streams.Closed]: +async def main(client: Client, task_queue: str) -> list[common.Record | common.Closed]: async with Worker( client, task_queue=task_queue, @@ -98,9 +94,9 @@ async def main( task_queue=task_queue, ) await handle.result() - stream = streams.StreamClient(client).workflow_stream(handle, Events) + stream = handle.stream(Events) return [ item async for item in stream.subscribe() - if isinstance(item, (streams.Record, streams.Closed)) + if isinstance(item, (common.Record, common.Closed)) ] diff --git a/tests/contrib/streams/examples/client_creates_stream.py b/tests/contrib/streams/examples/client_creates_stream.py index 5894524d77..5fddc3810e 100644 --- a/tests/contrib/streams/examples/client_creates_stream.py +++ b/tests/contrib/streams/examples/client_creates_stream.py @@ -10,10 +10,9 @@ from dataclasses import dataclass from datetime import timedelta -from temporalio import workflow +from temporalio import common, workflow from temporalio.client import Client -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib import stream as redis from temporalio.worker import Worker @@ -23,16 +22,16 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) + turns = common.topic(Turn) @workflow.defn class Chat: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> None: - # The Workflow only writes. The client, or auto-close, closes the Stream. + async def run(self, stream: common.WorkflowStream[Events]) -> None: + # The Workflow only writes. The client closes the Stream. await stream.topic(Events.turns).publish( Turn("model", "hi"), Turn("last", "bye") ) @@ -42,21 +41,21 @@ async def main(client: Client, task_queue: str) -> list[Turn]: async with Worker(client, task_queue=task_queue, workflows=[Chat]): turns = [] # The block closes the Stream at its end, or with the error that ends it. - async with streams.StreamClient(client).create( + async with redis.create( + client, Events, id=f"chat-{uuid.uuid4()}", retention=timedelta(days=1), - auto_close=timedelta(hours=1), ) as chat_stream: await client.start_workflow( Chat.run, # A Stream is data, so the Workflow gets it as an argument. - stream_workflow.WorkflowStream.from_handle(chat_stream), + common.WorkflowStream.from_handle(chat_stream), id=f"chat-{uuid.uuid4()}", task_queue=task_queue, ) async for item in chat_stream.subscribe(): - if not isinstance(item, streams.Record): + if not isinstance(item, common.Record): continue turns.append(item.data) if item.data.kind == "last": diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py index 57468479f9..607b7d9565 100644 --- a/tests/contrib/streams/examples/resume_from_offset.py +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -1,8 +1,8 @@ """Brook API proposal, "A client resumes a scoped Stream from an offset". -The Workflow writes one Stream per turn, a renamed definition from ``streams.with_name``. The +The Workflow writes one Stream per turn, a renamed definition from ``common.with_name``. The Reader saves the read's ``next_offset`` after each Record, stops, and later resumes from it with -``from_offset``. A view that wants only new Records starts at ``streams.END``. +``offset``. A view that wants only new Records starts at ``common.END``. """ from __future__ import annotations @@ -10,19 +10,17 @@ import uuid from datetime import timedelta -from temporalio import workflow +from temporalio import common, workflow from temporalio.client import Client, WorkflowHandle -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow from temporalio.worker import Worker -@streams.defn(name="events") +@common.defn(name="events") class Events: - tokens = streams.topic(str) + tokens = common.topic(str) -Turn3 = streams.with_name(Events, "turn-3") +Turn3 = common.with_name(Events, "turn-3") @workflow.defn @@ -33,7 +31,7 @@ def __init__(self) -> None: @workflow.run async def run(self, first: list[str]) -> None: - turn = await stream_workflow.create_stream(Turn3) + turn = await workflow.create_stream(Turn3) tokens = turn.topic(Events.tokens) await tokens.publish(*first) await workflow.wait_condition(lambda: bool(self.more) or self.done) @@ -58,18 +56,17 @@ def save(self, key: str, offset: int) -> None: async def read_some( - client: Client, handle: WorkflowHandle[Turns, None], store: OffsetStore, count: int + handle: WorkflowHandle[Turns, None], store: OffsetStore, count: int ) -> list[str]: - turn3 = streams.StreamClient(client).workflow_stream(handle, Turn3) - tokens = turn3.topic(Events.tokens) + turn3 = handle.stream(Turn3) seen = [] - read = tokens.subscribe( - from_offset=store.load("turn-3"), create_wait=timedelta(seconds=10) + read = turn3.subscribe( + Events.tokens, offset=store.load("turn-3"), create_wait=timedelta(seconds=10) ) async for item in read: - if isinstance(item, streams.Closed): + if isinstance(item, common.Closed): break - if isinstance(item, streams.Superseded): + if isinstance(item, common.Superseded): continue seen.append(item.data) # The read's next offset is after this Record, so the next read starts there. @@ -80,15 +77,15 @@ async def read_some( return seen -async def read_new(client: Client, handle: WorkflowHandle[Turns, None]) -> list[str]: +async def read_new(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, Turn3) + turn3 = handle.stream(Turn3) return [ item.data - async for item in turn3.topic(Events.tokens).subscribe( - from_offset=streams.END, wait_new_records=False + async for item in turn3.subscribe( + Events.tokens, offset=common.END, wait_new_records=False ) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] @@ -100,9 +97,9 @@ async def main( 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) + before = await read_some(handle, store, 2) + new = await read_new(handle) await handle.signal(Turns.add, ["c", "d"]) - after = await read_some(client, handle, store, 2) + after = await read_some(handle, store, 2) await handle.result() return before, new, after diff --git a/tests/contrib/streams/examples/test_examples.py b/tests/contrib/streams/examples/test_examples.py index 4b0a6d69cf..0e41708670 100644 --- a/tests/contrib/streams/examples/test_examples.py +++ b/tests/contrib/streams/examples/test_examples.py @@ -1,9 +1,9 @@ -"""Runs each example end to end on the memory Store and on Redis, so none goes stale.""" +"""Runs each example end to end on the memory Store and on Redis, so a broken one fails.""" import uuid +from temporalio import common from temporalio.client import Client -from temporalio.contrib import streams from tests.contrib.streams.examples import ( activity_streams_tokens, agent_publishes_turns, @@ -45,12 +45,12 @@ async def test_client_creates_stream(stream_client: Client): async def test_agent_publishes_turns(stream_client: Client): items = await agent_publishes_turns.main(stream_client, task_queue()) - data = [item.data for item in items if isinstance(item, streams.Record)] + data = [item.data for item in items if isinstance(item, common.Record)] Turn = agent_publishes_turns.Turn assert data == [Turn("model"), "find", "it", Turn("model"), "search", "result"] + [ Turn("done", "the answer") ] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "the answer" + assert isinstance(items[-1], common.Closed) and items[-1].result == "the answer" async def test_activity_streams_tokens(stream_client: Client): diff --git a/tests/contrib/streams/examples/turn_scoped_stream.py b/tests/contrib/streams/examples/turn_scoped_stream.py index 826df20aa6..8c1223147e 100644 --- a/tests/contrib/streams/examples/turn_scoped_stream.py +++ b/tests/contrib/streams/examples/turn_scoped_stream.py @@ -1,6 +1,6 @@ """Brook API proposal, "A Workflow scopes a Stream to one turn and closes it". -Each turn gets its own Stream, a renamed definition from ``streams.with_name``, which +Each turn gets its own Stream, a renamed definition from ``common.with_name``, which ``async with`` closes at the end of the turn. An index Stream lists the turns, and continue-as-new passes it to the next run. """ @@ -11,33 +11,31 @@ from dataclasses import dataclass from datetime import timedelta -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow from temporalio.worker import Worker -@streams.defn(name="events") +@common.defn(name="events") class Events: - tokens = streams.topic(str) + tokens = common.topic(str) @dataclass class TurnStarted: n: int - stream: streams.ClientStream[Events] + stream: common.ClientStream[Events] -@streams.defn(name="index") +@common.defn(name="index") class TurnIndex: - started = streams.topic(TurnStarted) + started = common.topic(TurnStarted) @dataclass class CallModelInput: prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] @activity.defn @@ -51,37 +49,38 @@ class ChatTurnsInput: # The first turn number of this run. A Stream name must be unique in the run chain. next: int = 0 # From the run before continue-as-new. - index: stream_workflow.WorkflowStream[TurnIndex] | None = None + index: common.WorkflowStream[TurnIndex] | None = None @workflow.defn class ChatTurns: @workflow.run async def run(self, input: ChatTurnsInput) -> None: - index = input.index or await stream_workflow.create_stream( + index = input.index or await workflow.create_stream( TurnIndex, retention=timedelta(days=1) ) started = index.topic(TurnIndex.started) batch = input.prompts[:2] for i, prompt in enumerate(batch): - async with stream_workflow.create_stream( - streams.with_name(Events, f"turn-{input.next + i}"), + async with workflow.create_stream( + common.with_name(Events, f"turn-{input.next + i}"), retention=timedelta(days=1), - close_with_run=False, ) as turn: - # A Stream is data, so the index Record carries the turn's Stream. + # A Stream is data, so the index Record contains the turn's Stream. await started.publish( - TurnStarted(input.next + i, streams.ClientStream.from_handle(turn)) + TurnStarted(input.next + i, common.ClientStream.from_handle(turn)) ) await workflow.execute_activity( call_model_into, - CallModelInput(prompt, streams.ActivityStream.from_handle(turn)), + CallModelInput(prompt, common.ActivityStream.from_handle(turn)), start_to_close_timeout=timedelta(minutes=5), ) if len(input.prompts) > len(batch): workflow.continue_as_new( ChatTurnsInput(input.prompts[2:], input.next + len(batch), index) ) + # The last run closes the index, since a run's end doesn't close it. + await index.close() async def main(client: Client, task_queue: str) -> dict[int, list[str]]: @@ -98,17 +97,16 @@ async def main(client: Client, task_queue: str) -> dict[int, list[str]]: task_queue=task_queue, ) await handle.result() - stream_client = streams.StreamClient(client) - index = stream_client.workflow_stream(handle, TurnIndex) + index = handle.stream(TurnIndex) turns = {} async for item in index.subscribe(): - if not isinstance(item, streams.Record): + if not isinstance(item, common.Record): continue # 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) + turn = client.stream(Events, stream=item.data.stream) turns[item.data.n] = [ token.data async for token in turn.subscribe() - if isinstance(token, streams.Record) + if isinstance(token, common.Record) ] return turns diff --git a/tests/contrib/streams/examples/ui_backend_reads.py b/tests/contrib/streams/examples/ui_backend_reads.py index 006cea2651..f0d71198ab 100644 --- a/tests/contrib/streams/examples/ui_backend_reads.py +++ b/tests/contrib/streams/examples/ui_backend_reads.py @@ -1,7 +1,7 @@ """Brook API proposal, "A UI backend reads a Workflow's Stream". The agent Workflow creates the ``events`` Stream. The UI backend starts it, gets the Stream -from the Workflow handle with no call to Temporal, and renders each item as it comes. +from the Workflow handle with no call to Temporal, and renders each item as it arrives. """ from __future__ import annotations @@ -11,10 +11,8 @@ from datetime import timedelta from typing import Any -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client -from temporalio.contrib import streams -from temporalio.contrib.streams import workflow as stream_workflow from temporalio.worker import Worker @@ -24,16 +22,16 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) @dataclass class CallModelInput: prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] @activity.defn @@ -48,28 +46,28 @@ async def call_model(input: CallModelInput) -> str: class Agent: @workflow.run async def run(self, prompt: str) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) turns = stream.topic(Events.turns) await turns.publish(Turn("model")) reply = await workflow.execute_activity( call_model, - CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(minutes=5), ) await turns.publish(Turn("done", reply)) + # A run's end doesn't close the Stream, so the Workflow closes it before it returns. + await stream.close(result=reply) return reply class Screen: - """Stands in for the UI. It keeps what it shows by offset, so a withdrawal can take items - back. - """ + """A fake UI. It keeps what it shows by offset, so a withdrawal can remove items.""" def __init__(self) -> None: - self.shown: dict[int, streams.Record[Any]] = {} + self.shown: dict[int, common.Record[Any]] = {} self.result: Any = None - def render(self, record: streams.Record[Any]) -> None: + def render(self, record: common.Record[Any]) -> None: self.shown[record.offset] = record def unrender(self, offset: int) -> None: @@ -83,19 +81,19 @@ async def main(client: Client, task_queue: str) -> Screen: handle = await client.start_workflow( Agent.run, "hello there", id=f"agent-{uuid.uuid4()}", task_queue=task_queue ) - stream = streams.StreamClient(client).workflow_stream(handle, Events) + stream = handle.stream(Events) screen = Screen() # The Workflow may not have created the Stream yet, so the read waits for it. async for item in stream.subscribe( Events.turns, Events.tokens, create_wait=timedelta(seconds=10) ): - if isinstance(item, streams.Closed): + if isinstance(item, common.Closed): screen.result = item.result if item.failure is None else item.failure - elif isinstance(item, streams.Superseded): + elif isinstance(item, common.Superseded): # The offsets of the Records that this read showed and the item withdraws. for offset in item.withdrawn_offsets: screen.unrender(offset) else: - # item.data is a Turn or a str, by item.topic. + # item.data is a Turn or a str, depending on item.topic. screen.render(item) return screen diff --git a/tests/contrib/streams/examples/worker_setup.py b/tests/contrib/streams/examples/worker_setup.py index e29bda15f0..d48ec404ad 100644 --- a/tests/contrib/streams/examples/worker_setup.py +++ b/tests/contrib/streams/examples/worker_setup.py @@ -1,26 +1,25 @@ """Brook API proposal, "Worker setup". -The Stream plugin goes on the client, and the Worker takes the client's plugins. Each plugin -registers a backend name. Tests use ``MemoryStreamPlugin`` with no change to Workflow code. -Sync Workflow code, such as ``@workflow.init``, uses the ``_nowait`` calls. +Register the Stream backends on the client with ``stream_backends=``. The Worker uses the +client's backends. Each backend registers a name. A create with no backend name uses the +backend with ``default=True``. Tests use ``MemoryStreamBackend`` without changing Workflow code. +Sync Workflow code, such as ``@workflow.init``, uses the ``_nowait`` functions. """ from __future__ import annotations import uuid -from temporalio import workflow +from temporalio import common, workflow from temporalio.client import Client -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.contrib import stream as redis +from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.worker import Worker -@streams.defn(name="events") +@common.defn(name="events") class Events: - tokens = streams.topic(str) + tokens = common.topic(str) @workflow.defn @@ -28,36 +27,44 @@ class Greeter: @workflow.init def __init__(self, name: str) -> None: # The init can't await, so it creates and publishes with the _nowait calls. - self.stream = stream_workflow.create_stream_nowait(Events) + self.stream = workflow.create_stream_nowait(Events) self.stream.topic(Events.tokens).publish_nowait(f"hello {name}") @workflow.run async def run(self, name: str) -> str: await self.stream.topic(Events.tokens).publish(f"bye {name}") + await self.stream.close() return name async def connect_with_redis(target_host: str, redis_url: str) -> Client: - return await Client.connect(target_host, plugins=[RedisStreamPlugin(redis_url)]) + # With no default=True, the default is native, and Workflow code would name Redis with + # redis.workflow.create_stream. + return await Client.connect( + target_host, + stream_backends=[redis.RedisStreamBackend(redis_url, default=True)], + ) async def connect_for_tests(target_host: str) -> Client: - return await Client.connect(target_host, plugins=[MemoryStreamPlugin()]) + return await Client.connect( + target_host, stream_backends=[MemoryStreamBackend(default=True)] + ) async def connect_while_migrating(target_host: str, url_a: str, url_b: str) -> Client: - # Two backends during a move. Creates with no backend= go to the default one. + # Two backends during a migration. A create with no backend= uses the default one. return await Client.connect( target_host, - plugins=[ - RedisStreamPlugin(url_a, name="redis-a", default=True), - RedisStreamPlugin(url_b, name="redis-b"), + stream_backends=[ + redis.RedisStreamBackend(url_a, name="redis-a", default=True), + redis.RedisStreamBackend(url_b, name="redis-b"), ], ) def run_worker(client: Client, task_queue: str) -> Worker: - # The Worker takes the client's plugins, so it needs no Stream setup of its own. + # The Worker uses the client's backends, so it needs no Stream setup of its own. return Worker(client, task_queue=task_queue, workflows=[Greeter]) @@ -67,16 +74,16 @@ async def main(client: Client, task_queue: str) -> list[str]: Greeter.run, "ada", id=f"greeter-{uuid.uuid4()}", task_queue=task_queue ) await handle.result() - stream = streams.StreamClient(client).workflow_stream(handle, Events) + stream = handle.stream(Events) return [ item.data async for item in stream.subscribe() - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] async def create_on_second_backend(client: Client) -> str | None: - stream = await streams.StreamClient(client).create( - Events, id=f"chat-{uuid.uuid4()}", backend="redis-b" + stream = await redis.create( + client, Events, id=f"chat-{uuid.uuid4()}", backend="redis-b" ) return stream.backend diff --git a/tests/contrib/streams/fake_extension.py b/tests/contrib/streams/fake_extension.py index 715fa04466..152b5583bc 100644 --- a/tests/contrib/streams/fake_extension.py +++ b/tests/contrib/streams/fake_extension.py @@ -1,4 +1,4 @@ -"""An in-memory stand-in for the temporalio-streams runtime, for facade tests.""" +"""An in-memory fake of the temporalio-streams runtime, for the facade tests.""" from __future__ import annotations @@ -6,7 +6,10 @@ import temporalio.api.streamservice.v1 as stream_pb from temporalio.bridge.streams import StreamCallError -from temporalio.contrib.streams.proto import v1 as ext_pb +from temporalio.client import Client +from temporalio.common import StreamBackend +from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams.proto import v1 as ext_pb from temporalio.service import RPCStatusCode @@ -146,3 +149,27 @@ async def _read( await asyncio.wait_for(stream.changed.wait(), left) except asyncio.TimeoutError: pass + + +class _FakeBackend(StreamBackend): + # The fake has Redis's rules, so the Redis functions accept it. + _kind = "redis" + + def __init__(self, calls: _ExtensionCalls) -> None: + super().__init__("fake-streams", backend_name="fake", default=True) + self._calls = calls + + def _backend_for( + self, + client: Client, # pyright: ignore[reportUnusedParameter] + ) -> _ExtensionCalls: + return self._calls + + +def fake_client(client: Client) -> tuple[Client, FakeExtension]: + """A copy of ``client`` whose one Stream backend, ``fake``, is a fake runtime.""" + extension = FakeExtension() + backend = _FakeBackend(_ExtensionCalls(extension)) + config = client.config() + config["stream_backends"] = [backend] + return Client(**config), extension diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index b53a144a6d..9b8b395470 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -11,12 +11,11 @@ import pytest -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client, WorkflowHandle from temporalio.common import RetryPolicy -from temporalio.contrib import streams -from temporalio.contrib.streams import _activity -from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.common._streams import _activity +from temporalio.contrib import stream as redis from temporalio.exceptions import CancelledError from tests.helpers import new_worker @@ -27,35 +26,31 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) @dataclass class CallModelInput: prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] @dataclass class SyncInput: prompt: str - stream: streams.SyncActivityStream[Events] + stream: common.SyncActivityStream[Events] -async def read( - client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" -) -> list[object]: - stream = streams.StreamClient(client).workflow_stream( - handle, streams.with_name(Events, name) - ) - return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] +async def read(handle: WorkflowHandle[Any, Any], name: str = "events") -> list[object]: + stream = handle.stream(common.with_name(Events, name)) + return [item async for item in stream.subscribe(wait_new_records=False)] -def records(items: Sequence[object]) -> list[streams.Record[Any]]: - return [item for item in items if isinstance(item, streams.Record)] +def records(items: Sequence[object]) -> list[common.Record[Any]]: + return [item for item in items if isinstance(item, common.Record)] @activity.defn @@ -63,7 +58,7 @@ async def call_model(input: CallModelInput) -> list[int]: assert input.stream.definition is Events tokens = input.stream.topic(Events.tokens) offsets = [await tokens.publish(word) for word in input.prompt.split()] - async with streams.BufferedPublisher( + async with common.BufferedPublisher( input.stream.topic(Events.turns), max_delay=timedelta(seconds=10) ) as turns: turns.publish(Turn("model", "a")) @@ -80,11 +75,11 @@ def call_model_sync(input: SyncInput) -> int: @activity.defn -async def read_back(stream: streams.ActivityStream[Events]) -> list[Any]: +async def read_back(stream: common.ActivityStream[Events]) -> list[Any]: return [ item.data async for item in stream.subscribe(Events.tokens, wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] @@ -92,22 +87,23 @@ async def read_back(stream: streams.ActivityStream[Events]) -> list[Any]: class Agent: @workflow.run async def run(self, prompt: str) -> list[Any]: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) offsets = await workflow.execute_activity( call_model, - CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) last = await workflow.execute_activity( call_model_sync, - SyncInput("sync words", streams.SyncActivityStream.from_handle(stream)), + SyncInput("sync words", common.SyncActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) read = await workflow.execute_activity( read_back, - streams.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(stream), start_to_close_timeout=timedelta(seconds=30), ) + await stream.close() return [offsets, last, read] @@ -122,7 +118,7 @@ async def test_activity_publishes_and_returns_offsets(stream_client: Client): Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) offsets, last, read_tokens = await handle.result() - items = await read(stream_client, handle) + items = await read(handle) data = records(items) assert [record.data for record in data] == [ "hi", @@ -136,7 +132,7 @@ async def test_activity_publishes_and_returns_offsets(stream_client: Client): assert offsets[2] == data[3].offset + 1 assert last == data[5].offset assert read_tokens == ["hi", "there", "sync", "words"] - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[-1], common.Closed) def _executor() -> Any: @@ -160,11 +156,11 @@ async def flaky(input: CallModelInput) -> None: class Retrying: @workflow.run async def run(self, key: str) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("start") await workflow.execute_activity( flaky, - CallModelInput(key, streams.ActivityStream.from_handle(stream)), + CallModelInput(key, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) @@ -177,15 +173,74 @@ async def test_retry_withdraws_a_received_stream(stream_client: Client): Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - items = await read(stream_client, handle) - superseded = [item for item in items if isinstance(item, streams.Superseded)] + items = await read(handle) + superseded = [item for item in items if isinstance(item, common.Superseded)] # The second attempt writes nothing, and its start still withdraws the first one's Record. assert len(superseded) == 1 withdrawn = [r.data for r in records(items) if superseded[0].withdraws(r)] assert withdrawn == ["attempt-1"] -# Set by a test once its readers showed the first attempt's Record. +@dataclass +class FinishedStreams: + key: str + closed: common.ActivityStream[Events] + expired: common.ActivityStream[Events] + + +@activity.defn +async def reads_after_a_failed_attempt(input: FinishedStreams) -> list[Any]: + attempts[input.key] = attempts.get(input.key, 0) + 1 + if attempts[input.key] == 1: + raise RuntimeError("first attempt fails") + items = input.closed.subscribe(Events.tokens, wait_new_records=False) + return [item.data async for item in items if isinstance(item, common.Record)] + + +@workflow.defn +class ReadsFinishedStreams: + @workflow.run + async def run(self, key: str, expired: common.WorkflowStream[Events]) -> list[Any]: + stream = await workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("turn") + await stream.close() + return await workflow.execute_activity( + reads_after_a_failed_attempt, + FinishedStreams( + key, + common.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(expired), + ), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_retry_goes_on_with_a_closed_or_expired_stream(stream_client: Client): + # On a Stream that accepts no write, the second attempt's open withdraws nothing, and the + # attempt only reads. + expired = await redis.create( + stream_client, + Events, + id=f"expired-{uuid.uuid4()}", + retention=timedelta(milliseconds=200), + ) + await expired.topic(Events.tokens).publish("old") + await asyncio.sleep(0.5) + key = str(uuid.uuid4()) + async with new_worker( + stream_client, ReadsFinishedStreams, activities=[reads_after_a_failed_attempt] + ) as worker: + result = await stream_client.execute_workflow( + ReadsFinishedStreams.run, + args=[key, common.WorkflowStream.from_handle(expired)], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == ["turn"] and attempts[key] == 2 + + +# Set by a test after its readers showed the first attempt's Record. shown: dict[str, asyncio.Event] = {} @@ -204,13 +259,14 @@ async def fails_after_it_is_shown(input: CallModelInput) -> None: class ShownThenWithdrawn: @workflow.run async def run(self, key: str) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await workflow.execute_activity( fails_after_it_is_shown, - CallModelInput(key, streams.ActivityStream.from_handle(stream)), + CallModelInput(key, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) + await stream.close() async def test_reader_withdraws_a_shown_record_after_a_reconnect( @@ -227,49 +283,48 @@ async def test_reader_withdraws_a_shown_record_after_a_reconnect( id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) - stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) + stream = handle.stream(Events) live = stream.subscribe(Events.tokens, create_wait=timedelta(seconds=20)) reconnects = stream.subscribe(Events.tokens, create_wait=timedelta(seconds=20)) first = await anext(live) before = await anext(reconnects) - # The reader drops its connection and keeps only the Records it showed and where to - # go on from, as a UI does across a page reload. + # The reader closes its connection and keeps only the Records that it showed and its + # resume offset, as a UI does across a page reload. resume_at = reconnects.next_offset await reconnects.aclose() shown[key].set() live_items = [first, *[item async for item in live]] resumed = [ - item - async for item in stream.subscribe(Events.tokens, from_offset=resume_at) + item async for item in stream.subscribe(Events.tokens, offset=resume_at) ] - assert isinstance(first, streams.Record) and first.data == "attempt-1" - assert isinstance(before, streams.Record) and before.offset == first.offset - [live_superseded] = [i for i in live_items if isinstance(i, streams.Superseded)] + assert isinstance(first, common.Record) and first.data == "attempt-1" + assert isinstance(before, common.Record) and before.offset == first.offset + [live_superseded] = [i for i in live_items if isinstance(i, common.Superseded)] assert live_superseded.withdrawn_offsets == (first.offset,) - # A new read gave no Record of the failed attempt, so it names no offset. The kept + # A new read got no Record of the failed attempt, so the item lists no offset. The kept # Record is still withdrawn. - [superseded] = [i for i in resumed if isinstance(i, streams.Superseded)] + [superseded] = [i for i in resumed if isinstance(i, common.Superseded)] assert superseded.withdrawn_offsets == () assert superseded.withdraws(before) assert [r.data for r in records(resumed)] == ["attempt-2"] - assert isinstance(resumed[-1], streams.Closed) + assert isinstance(resumed[-1], common.Closed) @activity.defn -async def wrong_backend(stream: streams.ActivityStream[Events]) -> str: +async def wrong_backend(stream: common.ActivityStream[Events]) -> str: try: await stream.topic(Events.tokens).publish("x") - except streams.StreamConfigError: + except common.StreamConfigError: return "config" return "wrote" @activity.defn -async def close_it(stream: streams.ActivityStream[Events]) -> str: +async def close_it(stream: common.ActivityStream[Events]) -> str: await stream.topic(Events.tokens).publish("last") try: await stream.close(result="closed by activity") - except streams.StreamNotOwnerError as err: + except common.StreamNotOwnerError as err: return err.reason return "closed" @@ -278,24 +333,31 @@ async def close_it(stream: streams.ActivityStream[Events]) -> str: class Closing: @workflow.run async def run(self) -> list[str]: - stream = await stream_workflow.create_stream(Events, close_with_run=False) + stream = await workflow.create_stream(Events) backend = await workflow.execute_activity( wrong_backend, - streams.ActivityStream(id=stream.id, backend="elsewhere"), + common.ActivityStream(id=stream.id, backend="elsewhere"), start_to_close_timeout=timedelta(seconds=30), ) owner = await workflow.execute_activity( close_it, - streams.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(stream), start_to_close_timeout=timedelta(seconds=30), ) - # A Local Activity is part of the run chain, so it may close the chain's Stream. + # A Local Activity runs before its Workflow Task commits, so it can't close it either. local = await workflow.execute_local_activity( close_it, - streams.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(stream), start_to_close_timeout=timedelta(seconds=30), ) - return [backend, owner, local] + # A handle with only the id names no Workflow, so the Store refuses the close. + by_id = await workflow.execute_local_activity( + close_it, + common.ActivityStream(id=stream.id), + start_to_close_timeout=timedelta(seconds=30), + ) + await stream.close(result="closed by the run") + return [backend, owner, local, by_id] async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): @@ -305,19 +367,24 @@ async def test_activity_checks_the_backend_and_the_owner(stream_client: Client): handle = await stream_client.start_workflow( Closing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - # An Activity writes to the Workflow's Stream but can't end it for every reader. - assert await handle.result() == ["config", "stream_not_owner", "closed"] - items = await read(stream_client, handle) - assert [r.data for r in records(items)] == ["last", "last"] - assert isinstance(items[-1], streams.Closed) - assert items[-1].result == "closed by activity" + # An Activity writes to the Workflow's Stream but can't close it. + assert await handle.result() == [ + "config", + "stream_not_owner", + "stream_not_owner", + "stream_not_owner", + ] + items = await read(handle) + assert [r.data for r in records(items)] == ["last", "last", "last"] + assert isinstance(items[-1], common.Closed) + assert items[-1].result == "closed by the run" @workflow.defn class CreatesAndLeavesOpen: @workflow.run async def run(self) -> str: - stream = await stream_workflow.create_stream(Events, close_with_run=False) + stream = await workflow.create_stream(Events) return stream.id @@ -328,7 +395,7 @@ async def run(self, stream_id: str) -> str: # The handle has only the id, so the SDK can't tell that another chain owns it. return await workflow.execute_local_activity( close_it, - streams.ActivityStream(id=stream_id), + common.ActivityStream(id=stream_id), start_to_close_timeout=timedelta(seconds=30), ) @@ -355,13 +422,13 @@ async def test_store_refuses_another_chains_close(stream_client: Client): ) == "stream_not_owner" ) - stream = streams.StreamClient(stream_client).get_handle(Events, id=stream_id) + stream = stream_client.stream(Events, id=stream_id) items = [item async for item in stream.subscribe(wait_new_records=False)] - assert [type(item) for item in items] == [streams.Record] + assert [type(item) for item in items] == [common.Record] @dataclass -class _Batches(streams.ActivityTopic[int]): +class _Batches(common.ActivityTopic[int]): batches: list[tuple[int, ...]] = field(default_factory=list) fail: bool = False in_flight: int = 0 @@ -372,7 +439,7 @@ def __post_init__(self) -> None: async def publish(self, *values: int) -> int: if self.fail: - raise streams.StreamStorageError("store down") + raise common.StreamStorageError("store down") self.in_flight += 1 self.most_in_flight = max(self.most_in_flight, self.in_flight) await asyncio.sleep(0.01) @@ -383,7 +450,7 @@ async def publish(self, *values: int) -> int: async def test_buffered_publisher_batches_by_delay(): topic = _Batches() - publisher = streams.BufferedPublisher(topic, max_delay=timedelta(milliseconds=20)) + publisher = common.BufferedPublisher(topic, max_delay=timedelta(milliseconds=20)) publisher.publish(1) publisher.publish(2, 3) await asyncio.sleep(0.1) @@ -396,12 +463,12 @@ async def test_buffered_publisher_batches_by_delay(): async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight(): topic = _Batches() - publisher = streams.BufferedPublisher(topic, max_delay=timedelta(minutes=10)) + publisher = common.BufferedPublisher(topic, max_delay=timedelta(minutes=10)) publisher.publish(*range(1000)) publisher.publish(*range(1000, 2500)) await asyncio.sleep(0.2) - # A full batch goes with no wait for the delay, and the values published during its - # write go in the next one. + # A full batch is written without waiting for the delay. The values published during its + # write are in the next batch. assert [len(batch) for batch in topic.batches] == [1000, 1500] assert topic.most_in_flight == 1 assert await publisher.flush() == 2500 @@ -409,21 +476,21 @@ async def test_buffered_publisher_writes_a_full_batch_with_one_write_in_flight() async def test_buffered_publisher_raises_a_failed_batch(): topic = _Batches(fail=True) - with pytest.raises(streams.StreamStorageError): - async with streams.BufferedPublisher(topic) as publisher: + with pytest.raises(common.StreamStorageError): + async with common.BufferedPublisher(topic) as publisher: publisher.publish(1) with pytest.raises(TypeError): - streams.BufferedPublisher(topic).publish() + common.BufferedPublisher(topic).publish() @activity.defn -async def local_step(stream: streams.ActivityStream[Events]) -> list[Any]: +async def local_step(stream: common.ActivityStream[Events]) -> list[Any]: # Each Workflow write is an awaited Local Activity, so the run's create and Records are in # the Store before a later Local Activity starts. seen = [ item.data async for item in stream.subscribe(wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] await stream.topic(Events.tokens).publish("local") return seen @@ -433,15 +500,16 @@ async def local_step(stream: streams.ActivityStream[Events]) -> list[Any]: class LocalOrder: @workflow.run async def run(self) -> list[Any]: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) tokens = stream.topic(Events.tokens) await tokens.publish("before") seen = await workflow.execute_local_activity( local_step, - streams.ActivityStream.from_handle(stream), + common.ActivityStream.from_handle(stream), start_to_close_timeout=timedelta(seconds=30), ) await tokens.publish("after") + await stream.close() return seen @@ -451,32 +519,32 @@ async def test_workflow_writes_go_out_before_a_local_activity(stream_client: Cli LocalOrder.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == ["before"] - items = await read(stream_client, handle) + items = await read(handle) written = records(items) assert [record.data for record in written] == ["before", "local", "after"] producer = written[1].producer assert producer is not None and producer.id.startswith("la:") - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[-1], common.Closed) async def test_buffered_publisher_stays_failed(): topic = _Batches(fail=True) - publisher = streams.BufferedPublisher(topic) + publisher = common.BufferedPublisher(topic) publisher.publish(1, 2) - with pytest.raises(streams.StreamStorageError): + with pytest.raises(common.StreamStorageError): await publisher.flush() topic.fail = False # A later batch would leave a gap, so the publisher refuses it and keeps the error. - with pytest.raises(streams.StreamStorageError): + with pytest.raises(common.StreamStorageError): publisher.publish(3) - with pytest.raises(streams.StreamStorageError): + with pytest.raises(common.StreamStorageError): await publisher.flush() assert topic.batches == [] and publisher.unwritten == 2 @activity.defn async def never_flushes(input: CallModelInput) -> None: - publisher = streams.BufferedPublisher( + publisher = common.BufferedPublisher( input.stream.topic(Events.tokens), max_delay=timedelta(minutes=10) ) publisher.publish(*input.prompt.split()) @@ -485,8 +553,8 @@ async def never_flushes(input: CallModelInput) -> None: @dataclass class ManyStreams: key: str - targets: list[streams.ActivityStream[Events]] - maybe: streams.ActivityStream[Events] | None = None + targets: list[common.ActivityStream[Events]] + maybe: common.ActivityStream[Events] | None = None @activity.defn @@ -500,11 +568,11 @@ async def flaky_many(input: ManyStreams) -> None: @activity.defn -def sync_read(stream: streams.SyncActivityStream[Events]) -> list[Any]: +def sync_read(stream: common.SyncActivityStream[Events]) -> list[Any]: return [ item.data for item in stream.subscribe(Events.tokens, wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] @@ -512,26 +580,26 @@ def sync_read(stream: streams.SyncActivityStream[Events]) -> list[Any]: class Leftovers: @workflow.run async def run(self, key: str) -> list[Any]: - stream = await stream_workflow.create_stream(Events) - other = await stream_workflow.create_stream(streams.with_name(Events, "other")) + stream = await workflow.create_stream(Events) + other = await workflow.create_stream(common.with_name(Events, "other")) await workflow.execute_activity( never_flushes, - CallModelInput("a b", streams.ActivityStream.from_handle(stream)), + CallModelInput("a b", common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) await workflow.execute_activity( flaky_many, ManyStreams( key, - [streams.ActivityStream.from_handle(stream)], - streams.ActivityStream.from_handle(other), + [common.ActivityStream.from_handle(stream)], + common.ActivityStream.from_handle(other), ), start_to_close_timeout=timedelta(seconds=30), retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), ) return await workflow.execute_activity( sync_read, - streams.SyncActivityStream.from_handle(stream), + common.SyncActivityStream.from_handle(stream), start_to_close_timeout=timedelta(seconds=30), ) @@ -551,10 +619,10 @@ async def test_attempt_end_flushes_and_retry_finds_nested_streams( ) # The publisher was never flushed, so the end of the attempt wrote its values. assert await handle.result() == ["a", "b", "attempt-1"] - items = await read(stream_client, handle) - other = await read(stream_client, handle, "other") + items = await read(handle) + other = await read(handle, "other") for read_items in (items, other): - superseded = [i for i in read_items if isinstance(i, streams.Superseded)] + superseded = [i for i in read_items if isinstance(i, common.Superseded)] withdrawn = [r.data for r in records(read_items) if superseded[0].withdraws(r)] assert withdrawn == ["attempt-1"] @@ -573,9 +641,9 @@ async def long_poll() -> None: finally: stopped.set() - stream: streams.SyncActivityStream[Any] = streams.SyncActivityStream(id="s") + stream: common.SyncActivityStream[Any] = common.SyncActivityStream(id="s") token = _activity.current.set( - cast(Any, SimpleNamespace(loop=loop, backend_name="")) + cast(Any, SimpleNamespace(loop=loop, backend_name="", default=True)) ) errors: list[BaseException] = [] @@ -598,3 +666,106 @@ def call() -> None: assert stopped.wait(5) assert errors and isinstance(errors[0], CancelledError) loop.call_soon_threadsafe(loop.stop) + + +@activity.defn +async def plain(name: str) -> str: + return f"hello {name}" + + +@workflow.defn +class NeverTouchesStreams: + @workflow.run + async def run(self) -> str: + return await workflow.execute_activity( + plain, "ada", start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class ReceivesAStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[Events]) -> str: + # The Activity gets no Stream, but the run received one. + return await workflow.execute_activity( + plain, stream.id, start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class RetriesAfterContinueAsNew: + @workflow.run + async def run(self, key: str, stream: common.WorkflowStream[Events] | None) -> None: + if stream is None: + created = await workflow.create_stream(Events) + workflow.continue_as_new(args=[key, created]) + await workflow.execute_activity( + flaky, + CallModelInput(key, common.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def keyed(handle: WorkflowHandle[Any, Any], activity_type: str) -> list[bool]: + # Whether each schedule of the Activity has the Producer key. + history = await handle.fetch_history() + return [ + "__temporal_stream_key" in attributes.header.fields + for event in history.events + if ( + attributes := event.activity_task_scheduled_event_attributes + ).activity_type.name + == activity_type + ] + + +async def test_only_a_run_that_uses_streams_keys_its_activities(stream_client: Client): + # The key is in History, so a run with no Stream adds none. A run that created, bound or + # received a Stream adds it to each Activity, so an Activity that writes keeps its Producer id. + key = str(uuid.uuid4()) + async with new_worker( + stream_client, + NeverTouchesStreams, + ReceivesAStream, + Retrying, + RetriesAfterContinueAsNew, + activities=[plain, flaky], + ) as worker: + none = await stream_client.start_workflow( + NeverTouchesStreams.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await none.result() == "hello ada" + assert await keyed(none, "plain") == [False] + creates = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await creates.result() + assert await keyed(creates, "flaky") == [True] + receives = await stream_client.start_workflow( + ReceivesAStream.run, + common.WorkflowStream.from_handle(creates.stream(Events)), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await receives.result() + assert await keyed(receives, "plain") == [True] + # The next run gets the Stream as its argument. Its retried Activity withdraws the first + # attempt's Records, as in the run that created the Stream. + continued_key = str(uuid.uuid4()) + continued = await stream_client.start_workflow( + RetriesAfterContinueAsNew.run, + args=[continued_key, None], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await continued.result() + assert await keyed(continued, "flaky") == [True] + items = await read(continued) + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert len(superseded) == 1 + assert [r.data for r in records(items) if superseded[0].withdraws(r)] == [ + "attempt-1" + ] diff --git a/tests/contrib/streams/test_bindings.py b/tests/contrib/streams/test_bindings.py index 2780c8123e..61ae5e45d5 100644 --- a/tests/contrib/streams/test_bindings.py +++ b/tests/contrib/streams/test_bindings.py @@ -4,7 +4,7 @@ import temporalio.api.streamservice.v1 as stream_pb import temporalio.converter import temporalio.nexus.system as nexus_system -from temporalio.contrib.streams.system import client, models +from temporalio.common._streams.system import client, models class FakeTransport: diff --git a/tests/contrib/streams/test_chain.py b/tests/contrib/streams/test_chain.py deleted file mode 100644 index 06b439eddb..0000000000 --- a/tests/contrib/streams/test_chain.py +++ /dev/null @@ -1,80 +0,0 @@ -from datetime import datetime, timedelta, timezone -from types import SimpleNamespace -from typing import Any - -import pytest - -from temporalio.common import RetryPolicy -from temporalio.contrib.streams.redis import chain_continues -from temporalio.exceptions import ApplicationError - -START = datetime(2026, 10, 21, tzinfo=timezone.utc) - - -def info(**fields: Any) -> Any: - base: dict[str, Any] = { - "cron_schedule": None, - "retry_policy": None, - "attempt": 1, - "execution_timeout": None, - "workflow_start_time": START, - "execution_expiration_time": None, - } - return SimpleNamespace(**{**base, **fields}) - - -RETRY = RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3) - - -@pytest.mark.parametrize( - "case, run, failure, expected", - [ - ("success", info(), None, False), - ("cron success", info(cron_schedule="* * * * *"), None, True), - ("cron failure", info(cron_schedule="* * * * *"), RuntimeError("x"), True), - ("no retry policy", info(), RuntimeError("x"), False), - ("attempts left", info(retry_policy=RETRY), RuntimeError("x"), True), - ("last attempt", info(retry_policy=RETRY, attempt=3), RuntimeError("x"), False), - ( - "unlimited attempts", - info(retry_policy=RetryPolicy(maximum_attempts=0), attempt=50), - RuntimeError("x"), - True, - ), - ( - "non-retryable error", - info(retry_policy=RETRY), - ApplicationError("x", non_retryable=True), - False, - ), - ( - "non-retryable type", - info( - retry_policy=RetryPolicy(non_retryable_error_types=["Bad"]), - ), - ApplicationError("x", type="Bad"), - False, - ), - ( - "no time left", - info(retry_policy=RETRY, execution_timeout=timedelta(seconds=10)), - RuntimeError("x"), - False, - ), - ( - "no time left in a later run of the chain", - info( - retry_policy=RETRY, - execution_timeout=timedelta(seconds=10), - workflow_start_time=START + timedelta(seconds=9), - execution_expiration_time=START + timedelta(seconds=10), - ), - RuntimeError("x"), - False, - ), - ], -) -def test_chain_continues(case: str, run: Any, failure: Any, expected: bool): - # "no time left" runs 9.5 s into a 10 s chain, so a 1 s backoff passes the expiration. - now = START + timedelta(seconds=9.5) - assert chain_continues(run, now, failure) == expected, case diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index 3421cb56ec..c30f8ae603 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -3,22 +3,20 @@ from dataclasses import dataclass from datetime import timedelta -import google.protobuf.any_pb2 import pytest -import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb +from temporalio import common from temporalio.bridge.streams import StreamCallError from temporalio.client import Client 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.common._streams import _errors +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.common._streams.system import models +from temporalio.contrib import stream as redis from temporalio.exceptions import ApplicationError from temporalio.service import RPCError, RPCStatusCode -from tests.contrib.streams.fake_extension import FakeExtension +from tests.contrib.streams.fake_extension import fake_client @dataclass @@ -27,30 +25,17 @@ class Turn: text: str = "" -def fake_client(client: Client) -> tuple[Client, FakeExtension]: - extension = FakeExtension() - plugin = streams.StreamPlugin( - "fake-streams", - backend_name="fake", - default=False, - stream_backend=_ExtensionStreamBackend(extension), - ) - config = client.config() - config["plugins"] = [plugin] - return Client(**config), extension - - def test_runtime_errors_map_to_stream_errors(): details = ext_pb.StreamErrorDetails(reason="stream_fenced").SerializeToString() err = _errors.from_call_error( StreamCallError(int(RPCStatusCode.FAILED_PRECONDITION), "fenced", details) ) - assert isinstance(err, streams.StreamFencedError) + assert isinstance(err, common.StreamFencedError) assert err.reason == "stream_fenced" other = _errors.from_call_error( StreamCallError(int(RPCStatusCode.INTERNAL), "odd", b"") ) - assert type(other) is streams.StreamError + assert type(other) is common.StreamError def test_superseded_withdraws_earlier_attempts_of_one_producer(): @@ -62,7 +47,7 @@ def producer(attempt: int) -> models.Producer: epoch=models.ProducerEpoch(attempt=attempt), ) - item = streams.Superseded( + item = common.Superseded( offset=9, supersession=models.Supersession( producer_kind=stream_pb.ProducerKind.PRODUCER_KIND_ACTIVITY, @@ -70,21 +55,21 @@ def producer(attempt: int) -> models.Producer: epoch=models.ProducerEpoch(attempt=2), ), ) - old = streams.Record(topic="t", data=1, offset=1, producer=producer(1)) - new = streams.Record(topic="t", data=2, offset=10, producer=producer(2)) + old = common.Record(topic="t", data=1, offset=1, producer=producer(1)) + new = common.Record(topic="t", data=2, offset=10, producer=producer(2)) assert item.withdraws(old) assert not item.withdraws(new) -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) async def test_client_writes_reads_and_closes(client: Client): client, extension = fake_client(client) - stream = await streams.StreamClient(client).create(Events, id="chat-42") + stream = await redis.create(client, Events, id="chat-42") first = await stream.topic(Events.turns).publish(Turn("model")) second = await stream.topic("tokens").publish("hello") assert second == first + 1 @@ -92,50 +77,49 @@ async def test_client_writes_reads_and_closes(client: Client): await stream.close(result="bye") items = [item async for item in stream.subscribe()] assert [type(item) for item in items] == [ - streams.Record, - streams.Record, - streams.Closed, + common.Record, + common.Record, + common.Closed, ] turn, token, closed = items - assert isinstance(turn, streams.Record) and turn.data == Turn("model") - assert isinstance(token, streams.Record) and token.data == "hello" + assert isinstance(turn, common.Record) and turn.data == Turn("model") + assert isinstance(token, common.Record) and token.data == "hello" assert turn.producer is not None assert turn.producer.kind == stream_pb.ProducerKind.PRODUCER_KIND_CLIENT - assert isinstance(closed, streams.Closed) and closed.result == "bye" + assert isinstance(closed, common.Closed) and closed.result == "bye" assert extension.namespaces == {client.namespace} raw = stream.topic("raw", type=RawValue) assert raw.type is RawValue - resumed = [item async for item in stream.subscribe(from_offset=token.offset)] + resumed = [item async for item in stream.subscribe(offset=token.offset)] assert [item.offset for item in resumed] == [token.offset, closed.offset] async def test_read_waits_for_create(client: Client): client, _ = fake_client(client) - stream_client = streams.StreamClient(client) handle = client.get_workflow_handle("wf-1", first_execution_run_id="run-1") - stream = stream_client.workflow_stream(handle, Events) - with pytest.raises(streams.StreamNotFoundError): + stream = handle.stream(Events) + with pytest.raises(common.StreamNotFoundError): async for _ in stream.subscribe(): pass async def create_later() -> None: - created = await stream_client.create(Events, id=stream.id) + created = await redis.create(client, Events, id=stream.id) await created.topic(Events.tokens).publish("five") await created.close() task = asyncio.create_task(create_later()) items = [item async for item in stream.subscribe(create_wait=timedelta(seconds=5))] await task - assert isinstance(items[0], streams.Record) and items[0].data == "five" - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[0], common.Record) and items[0].data == "five" + assert isinstance(items[-1], common.Closed) def test_stream_errors_fail_the_workflow_not_the_task(): - err = streams.StreamClosedError("closed", reason="stream_closed") + err = common.StreamClosedError("closed", reason="stream_closed") assert isinstance(err, ApplicationError) assert err.non_retryable and err.type == "StreamClosedError" - assert not streams.StreamStorageError("down").non_retryable + assert not common.StreamStorageError("down").non_retryable details = ext_pb.StreamErrorDetails(reason="stream_closed").SerializeToString() mapped = _errors.from_call_error( StreamCallError(int(RPCStatusCode.FAILED_PRECONDITION), "closed", details) @@ -145,20 +129,19 @@ def test_stream_errors_fail_the_workflow_not_the_task(): async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client): client, extension = fake_client(client) - stream_client = streams.StreamClient(client) - stream = await stream_client.create(Events, id="two-writers") - other = stream_client.get_handle(Events, id="two-writers") + stream = await redis.create(client, Events, id="two-writers") + other = client.stream(Events, id="two-writers") await stream.topic(Events.tokens).publish("1") await other.topic(Events.tokens).publish("2") items = [item async for item in stream.subscribe(wait_new_records=False)] producers = { item.producer.id for item in items - if isinstance(item, streams.Record) and item.producer + if isinstance(item, common.Record) and item.producer } assert len(producers) == 1 - # A handle that goes with no close frees its writer, so a server that makes a handle per - # request keeps no writer for it. + # A handle that Python frees without a close frees its writer. So a server that makes a + # handle per request doesn't keep a writer for each. del other gc.collect() assert len(extension.freed) == 1 @@ -166,7 +149,7 @@ async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client async def test_rpc_timeout_bounds_each_call_and_not_the_read(client: Client): client, extension = fake_client(client) - stream = await streams.StreamClient(client).create(Events, id="slow") + stream = await redis.create(client, Events, id="slow") call = extension.call async def slow_describe(method: str, input: bytes) -> bytes: @@ -183,15 +166,13 @@ async def publish_later() -> None: await asyncio.sleep(1) await stream.topic(Events.tokens).publish("late") - # Each call of a read that waits answers before the deadline, so the read goes on. - reader = stream.topic(Events.tokens).subscribe( - rpc_timeout=timedelta(milliseconds=300) - ) + # Each call of a waiting read returns before the deadline, so the read continues. + reader = stream.subscribe(Events.tokens, rpc_timeout=timedelta(milliseconds=300)) publisher = asyncio.create_task(publish_later()) item = await anext(reader) await reader.aclose() await publisher - assert isinstance(item, streams.Record) and item.data == "late" + assert isinstance(item, common.Record) and item.data == "late" def test_unsupported_capability_maps_to_its_error(): @@ -201,22 +182,6 @@ def test_unsupported_capability_maps_to_its_error(): int(RPCStatusCode.UNIMPLEMENTED), "The Store lacks replica_wait", details ) ) - assert isinstance(err, streams.StreamUnsupportedError) + assert isinstance(err, common.StreamUnsupportedError) assert err.non_retryable and err.code == RPCStatusCode.UNIMPLEMENTED assert err.message.endswith("replica_wait") - - -def test_server_failure_maps_to_its_stream_error(): - detail = google.protobuf.any_pb2.Any() - detail.Pack(ext_pb.StreamErrorDetails(reason="stream_closed")) - status = temporalio.api.common.v1.GrpcStatus( - code=9, message="closed", details=[detail] - ) - err = _errors.from_rpc_error( - RPCError( - "closed", RPCStatusCode.FAILED_PRECONDITION, status.SerializeToString() - ) - ) - assert isinstance(err, streams.StreamClosedError) - plain = RPCError("down", RPCStatusCode.UNAVAILABLE, b"") - assert _errors.from_rpc_error(plain) is plain diff --git a/tests/contrib/streams/test_handles.py b/tests/contrib/streams/test_handles.py index 2d8d6267a9..6054d4bfca 100644 --- a/tests/contrib/streams/test_handles.py +++ b/tests/contrib/streams/test_handles.py @@ -1,14 +1,17 @@ import json +import subprocess +import sys +import types from dataclasses import dataclass from typing import Any, TypeVar import pydantic import pytest +from temporalio import common from temporalio.common import RawValue -from temporalio.contrib import streams +from temporalio.common._streams._defn import definition_of from temporalio.contrib.pydantic import pydantic_data_converter -from temporalio.contrib.streams._defn import definition_of from temporalio.converter import DataConverter E = TypeVar("E") @@ -19,19 +22,19 @@ class Turn: kind: str -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) - raw = streams.topic(RawValue, name="raw-values") + turns = common.topic(Turn) + tokens = common.topic(str) + raw = common.topic(RawValue, name="raw-values") -@streams.defn +@common.defn class Index: - turns = streams.topic(int) + turns = common.topic(int) -class PlainStream(streams.StreamHandle[E]): +class PlainStream(common.StreamHandle[E]): pass @@ -47,19 +50,19 @@ class PlainStream(streams.StreamHandle[E]): def test_definition_rejects_two_topics_of_one_name(): with pytest.raises(ValueError, match="two Topics named t"): - @streams.defn + @common.defn class Twice: # pyright: ignore[reportUnusedClass] - a = streams.topic(int, name="t") - b = streams.topic(str, name="t") + a = common.topic(int, name="t") + b = common.topic(str, name="t") def test_subclass_topic_replaces_the_base_topic(): class Base: - turns = streams.topic(int) + turns = common.topic(int) - @streams.defn + @common.defn class Sub(Base): - turns = streams.topic(str) + turns = common.topic(str) assert definition_of(Sub).topics["turns"].type is str @@ -73,10 +76,10 @@ def test_topic_of_the_definition_keeps_its_type(): assert OWNED._topic_spec(Events.turns, None) == ("turns", Turn) assert OWNED._topic_spec("tokens", None) == ("tokens", str) assert OWNED._topic_spec("custom", Turn) == ("custom", Turn) - with pytest.raises(streams.StreamError, match="has no Topic turns") as err: + with pytest.raises(common.StreamError, match="has no Topic turns") as err: OWNED._topic_spec(Index.turns, None) assert err.value.reason == "unknown_topic" - with pytest.raises(streams.StreamError, match="Pass type="): + with pytest.raises(common.StreamError, match="Pass type="): OWNED._topic_spec("custom", None) with pytest.raises(TypeError, match="has the type"): OWNED._topic_spec("tokens", int) @@ -106,7 +109,7 @@ def test_type_hint_picks_the_class_and_the_definition(): bare = converter.from_payload(payload, PlainStream) assert bare.definition is None and bare.backend == "redis" with pytest.raises(Exception, match="handle type of its context"): - converter.from_payload(payload, streams.StreamHandle) + converter.from_payload(payload, common.StreamHandle) @dataclass @@ -150,3 +153,32 @@ def test_any_type_argument_has_no_definition(): decoded = converter.from_payload(converter.to_payload(OWNED), PlainStream[Any]) assert type(decoded) is PlainStream and decoded.definition is None assert decoded.id == OWNED.id + + +def test_end_type_is_public(): + # subscribe takes offset: int | End | None, so the type of END is a public name. + assert isinstance(common.END, common.End) and "End" in common.__all__ + + +def test_temporalio_common_exports_the_streams_api(): + # temporalio.common imports the Stream names on first use. Each is the private package's. + import temporalio.common._streams as streams_api + + names = { + name + for name, value in vars(streams_api).items() + if not name.startswith("_") and not isinstance(value, types.ModuleType) + } + assert names == common._STREAM_NAMES and names <= set(common.__all__) + for name in common._STREAM_NAMES: + assert getattr(common, name) is getattr(streams_api, name) + + +def test_import_of_temporalio_common_loads_no_stream_code(): + # The Stream modules import temporalio.client and temporalio.workflow, which import + # temporalio.common. + code = "import sys, temporalio.common; print('temporalio.common._streams' in sys.modules)" + loaded = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ).stdout + assert loaded.strip() == "False" diff --git a/tests/contrib/streams/test_pydantic.py b/tests/contrib/streams/test_pydantic.py index 29c11e422b..f9745414c1 100644 --- a/tests/contrib/streams/test_pydantic.py +++ b/tests/contrib/streams/test_pydantic.py @@ -3,28 +3,26 @@ import pydantic -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.client import Client -from temporalio.contrib import streams from temporalio.contrib.pydantic import pydantic_data_converter -from temporalio.contrib.streams import workflow as stream_workflow -from temporalio.contrib.streams.testing import MemoryStreamPlugin +from temporalio.contrib.stream.testing import MemoryStreamBackend from tests.helpers import new_worker -@streams.defn(name="events") +@common.defn(name="events") class Events: - tokens = streams.topic(str) + tokens = common.topic(str) class CallModelInput(pydantic.BaseModel): prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] class TurnResult(pydantic.BaseModel): text: str - stream: streams.ClientStream[Events] + stream: common.ClientStream[Events] @activity.defn @@ -37,22 +35,23 @@ async def call_model(input: CallModelInput) -> int: class TurnWorkflow: @workflow.run async def run(self, prompt: str) -> TurnResult: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await workflow.execute_activity( call_model, CallModelInput( - prompt=prompt, stream=streams.ActivityStream.from_handle(stream) + prompt=prompt, stream=common.ActivityStream.from_handle(stream) ), start_to_close_timeout=timedelta(seconds=30), ) + await stream.close(result=prompt.upper()) return TurnResult( - text=prompt.upper(), stream=streams.ClientStream.from_handle(stream) + text=prompt.upper(), stream=common.ClientStream.from_handle(stream) ) async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): config = client.config() - config["plugins"] = [MemoryStreamPlugin()] + config["stream_backends"] = [MemoryStreamBackend(default=True)] config["data_converter"] = pydantic_data_converter client = Client(**config) async with new_worker(client, TurnWorkflow, activities=[call_model]) as worker: @@ -67,7 +66,7 @@ async def test_handles_in_models_round_trip_on_the_memory_store(client: Client): assert result.stream.workflow_id == handle.id assert result.stream.first_execution_run_id == handle.first_execution_run_id assert result.stream.backend == "memory" - stream = streams.StreamClient(client).get_handle(Events, id=result.stream.id) + stream = client.stream(Events, id=result.stream.id) items = [item async for item in stream.subscribe()] - assert [i.data for i in items if isinstance(i, streams.Record)] == ["hi", "there"] - assert isinstance(items[-1], streams.Closed) and items[-1].result is not None + assert [i.data for i in items if isinstance(i, common.Record)] == ["hi", "there"] + assert isinstance(items[-1], common.Closed) and items[-1].result is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index eedf145456..9b3367e964 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -10,19 +10,18 @@ import pytest -from temporalio import activity, workflow +from temporalio import activity, common, workflow from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import EventType from temporalio.client import Client, Interceptor, 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.redis import ( - RedisStreamPlugin, - _ExtensionStreamPlugin, +from temporalio.common._streams._extension import ( + ExtensionStreamBackend, _WorkflowInbound, ) -from temporalio.contrib.streams.testing import MemoryStreamPlugin +from temporalio.contrib import stream as redis +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.converter import DataConverter, PayloadCodec from temporalio.exceptions import ApplicationError from temporalio.worker import ( @@ -46,32 +45,30 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) @dataclass class CallModelInput: prompt: str - stream: streams.ActivityStream[Events] + stream: common.ActivityStream[Events] def workflow_stream( - client: Client, handle: WorkflowHandle[Any, Any], name: str | None = None -) -> streams.ClientStream[Events]: - return streams.StreamClient(client).workflow_stream( - handle, streams.with_name(Events, name) if name else Events - ) + handle: WorkflowHandle[Any, Any], name: str | None = None +) -> common.ClientStream[Events]: + return handle.stream(common.with_name(Events, name) if name else Events) -async def read_all(stream: streams.ClientStream[Any]) -> list[Any]: +async def read_all(stream: common.ClientStream[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] -def records(items: list[Any]) -> list[streams.Record[Any]]: - return [item for item in items if isinstance(item, streams.Record)] +def records(items: list[Any]) -> list[common.Record[Any]]: + return [item for item in items if isinstance(item, common.Record)] @activity.defn @@ -86,15 +83,16 @@ async def call_model(input: CallModelInput) -> str: class Agent: @workflow.run async def run(self, prompt: str) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) turns = stream.topic(Events.turns) await turns.publish(Turn("model")) reply = await workflow.execute_activity( call_model, - CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + CallModelInput(prompt, common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) await turns.publish(Turn("done", reply)) + await stream.close(result=reply) return reply @@ -103,7 +101,7 @@ async def test_workflow_and_activity_write_one_stream(stream_client: Client): handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - stream = workflow_stream(stream_client, handle) + stream = workflow_stream(handle) items = await read_all(stream) assert await handle.result() == "HI THERE" data = [(item.topic, item.data) for item in records(items)] @@ -114,12 +112,12 @@ async def test_workflow_and_activity_write_one_stream(stream_client: Client): ("turns", Turn("done", "HI THERE")), ] closed = items[-1] - assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" + assert isinstance(closed, common.Closed) and closed.result == "HI THERE" only_turns = [item async for item in stream.subscribe(Events.turns)] assert [type(item) for item in only_turns] == [ - streams.Record, - streams.Record, - streams.Closed, + common.Record, + common.Record, + common.Closed, ] @@ -130,18 +128,21 @@ async def test_workflow_and_activity_write_one_stream(stream_client: Client): class PublishThenFailTask: @workflow.run async def run(self, key: str) -> None: - turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + stream = await workflow.create_stream(Events) + turns = stream.topic(Events.tokens) await turns.publish("before") await workflow.sleep(timedelta(milliseconds=10)) await turns.publish("after") if attempts.setdefault(key, 0) == 0: attempts[key] = 1 raise RuntimeError("fail this Workflow Task once") + await stream.close() async def test_failed_workflow_task_is_withdrawn(stream_client: Client): key = str(uuid.uuid4()) - # The test's counter must survive the retry of the task, so no sandbox. + # The test's counter must keep its value across the task's retry, so the Workflow has no + # sandbox. async with new_worker( stream_client, PublishThenFailTask, @@ -154,14 +155,14 @@ async def test_failed_workflow_task_is_withdrawn(stream_client: Client): task_queue=worker.task_queue, ) await handle.result() - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) # The retried task writes "after" again, and withdraws the copy of the failed task. written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] + superseded = [item for item in items if isinstance(item, common.Superseded)] assert [record.data for record in written] == ["before", "after", "after"] withdrawn = [r for r in written if any(i.withdraws(r) for i in superseded)] assert withdrawn == [written[1]] - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[-1], common.Closed) async def test_resume_after_a_superseded_item_keeps_its_record(stream_client: Client): @@ -178,14 +179,14 @@ async def test_resume_after_a_superseded_item_keeps_its_record(stream_client: Cl task_queue=worker.task_queue, ) await handle.result() - stream = workflow_stream(stream_client, handle) + stream = workflow_stream(handle) items = [item async for item in stream.subscribe(wait_new_records=False)] # The retried task's batch has its Superseded item and its first Record in one entry. record_offsets = {r.offset for r in records(items)} [shared] = [ i for i in items - if isinstance(i, streams.Superseded) and i.offset in record_offsets + if isinstance(i, common.Superseded) and i.offset in record_offsets ] reader = stream.subscribe(wait_new_records=False) async for item in reader: @@ -196,9 +197,7 @@ async def test_resume_after_a_superseded_item_keeps_its_record(stream_client: Cl assert resume_at == shared.offset resumed = [ item - async for item in stream.subscribe( - from_offset=resume_at, wait_new_records=False - ) + async for item in stream.subscribe(offset=resume_at, wait_new_records=False) ] assert [r.data for r in records(resumed) if r.offset == shared.offset] == ["after"] @@ -221,7 +220,7 @@ async def run(self) -> bool: return same -async def test_plugin_leaves_the_workflow_random_alone(stream_client: Client): +async def test_backend_leaves_the_workflow_random_alone(stream_client: Client): async with new_worker( stream_client, RandomAroundActivity, activities=[noop] ) as worker: @@ -237,13 +236,14 @@ class WritesThenContinues: @workflow.run async def run(self, write: bool) -> str: if write: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("first run") workflow.continue_as_new(False) + await workflow.stream_handle(Events).close(result="second run") return "second run" -async def test_last_run_closes_a_stream_it_did_not_write(stream_client: Client): +async def test_later_run_closes_a_stream_it_did_not_write(stream_client: Client): async with new_worker(stream_client, WritesThenContinues) as worker: handle = await stream_client.start_workflow( WritesThenContinues.run, @@ -252,18 +252,18 @@ async def test_last_run_closes_a_stream_it_did_not_write(stream_client: Client): task_queue=worker.task_queue, ) assert await handle.result() == "second run" - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) assert [i.data for i in records(items)] == ["first run"] closed = items[-1] - assert isinstance(closed, streams.Closed) and closed.result == "second run" + assert isinstance(closed, common.Closed) and closed.result == "second run" @workflow.defn class CancelledBlock: @workflow.run async def run(self) -> None: - async with await stream_workflow.create_stream( - streams.with_name(Events, "turn"), close_with_run=False + async with await workflow.create_stream( + common.with_name(Events, "turn") ) as turn: await turn.topic(Events.tokens).publish("started") await workflow.wait_condition(lambda: False) @@ -274,33 +274,34 @@ async def test_cancelled_block_closes_with_a_failure(stream_client: Client): handle = await stream_client.start_workflow( CancelledBlock.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - turn = workflow_stream(stream_client, handle, "turn") + turn = workflow_stream(handle, "turn") reader = turn.subscribe(create_wait=timedelta(seconds=20)) - assert isinstance(await anext(reader), streams.Record) + assert isinstance(await anext(reader), common.Record) await reader.aclose() await handle.cancel() items = await read_all(turn) closed = items[-1] - assert isinstance(closed, streams.Closed) + assert isinstance(closed, common.Closed) assert closed.failure is not None - assert closed.status == streams.CloseStatus.CANCELED + assert closed.status == common.CloseStatus.CANCELED -@streams.defn(name="plain") +@common.defn(name="plain") class Plain: - t = streams.topic(str) + t = common.topic(str) @workflow.defn class PlainWriter: @workflow.run async def run(self) -> str: - stream = await stream_workflow.create_stream(Plain) + stream = await workflow.create_stream(Plain) await stream.topic(Plain.t).publish("plain") - async with await stream_workflow.create_stream( - streams.with_name(Plain, "scoped"), close_with_run=False + async with await workflow.create_stream( + common.with_name(Plain, "scoped") ) as scoped: await scoped.topic(Plain.t).publish("scoped") + await stream.close(result="done") return "done" @@ -318,15 +319,12 @@ async def test_temporal_only_passthrough_keeps_records(stream_client: Client): PlainWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - plain = streams.StreamClient(stream_client) - items = await read_all(plain.workflow_stream(handle, Plain)) - scoped = await read_all( - plain.workflow_stream(handle, streams.with_name(Plain, "scoped")) - ) + items = await read_all(handle.stream(Plain)) + scoped = await read_all(handle.stream(common.with_name(Plain, "scoped"))) assert [i.data for i in records(items)] == ["plain"] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" + assert isinstance(items[-1], common.Closed) and items[-1].result == "done" assert [i.data for i in records(scoped)] == ["scoped"] - assert isinstance(scoped[-1], streams.Closed) + assert isinstance(scoped[-1], common.Closed) @activity.defn @@ -338,12 +336,13 @@ async def local_tokens(input: CallModelInput) -> None: class LocalWriter: @workflow.run async def run(self) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await workflow.execute_local_activity( local_tokens, - CallModelInput("0 1", streams.ActivityStream.from_handle(stream)), + CallModelInput("0 1", common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) + await stream.close() async def test_local_activity_writes_with_the_task_epoch(stream_client: Client): @@ -352,7 +351,7 @@ async def test_local_activity_writes_with_the_task_epoch(stream_client: Client): LocalWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue ) await handle.result() - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) written = records(items) assert [record.data for record in written] == ["0", "1"] producer = written[0].producer @@ -366,18 +365,19 @@ async def test_local_activity_writes_with_the_task_epoch(stream_client: Client): class LocalThenFailTask: @workflow.run async def run(self, key: str) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) first = attempts.setdefault(key, 0) == 0 await workflow.execute_local_activity( local_tokens, CallModelInput( - "0 1" if first else "2 3", streams.ActivityStream.from_handle(stream) + "0 1" if first else "2 3", common.ActivityStream.from_handle(stream) ), start_to_close_timeout=timedelta(seconds=30), ) if first: attempts[key] = 1 raise RuntimeError("fail this Workflow Task once") + await stream.close() async def test_rerun_local_activity_withdraws_its_failed_task(stream_client: Client): @@ -394,12 +394,12 @@ async def test_rerun_local_activity_withdraws_its_failed_task(stream_client: Cli task_queue=worker.task_queue, ) await handle.result() - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] + superseded = [item for item in items if isinstance(item, common.Superseded)] visible = [r for r in written if not any(i.withdraws(r) for i in superseded)] assert [r.data for r in visible] == ["2", "3"] - # The rerun keeps the Producer id and moves to the retry task's epoch. + # The rerun keeps the Producer id and uses the retry task's epoch. first, second = written[0].producer, visible[0].producer assert first is not None and second is not None and first.id == second.id assert first.epoch is not None and second.epoch is not None @@ -413,21 +413,22 @@ async def test_rerun_local_activity_withdraws_its_failed_task(stream_client: Cli class LocalOnlyInFailedTask: @workflow.run async def run(self, key: str) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) if attempts.setdefault(key, 0) == 0: attempts[key] = 1 await workflow.execute_local_activity( local_tokens, - CallModelInput("0 1", streams.ActivityStream.from_handle(stream)), + CallModelInput("0 1", common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) raise RuntimeError("fail this Workflow Task once") await stream.topic(Events.tokens).publish("other") + await stream.close() async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Client): # The retry task skips the Local Activity, so only its own write can withdraw "0 1". The - # write also takes the skipped call's key. + # write also uses the skipped call's key. async with new_worker( stream_client, LocalOnlyInFailedTask, @@ -441,23 +442,24 @@ async def test_later_write_withdraws_a_skipped_local_activity(stream_client: Cli task_queue=worker.task_queue, ) await handle.result() - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] + superseded = [item for item in items if isinstance(item, common.Superseded)] visible = [r for r in written if not any(s.withdraws(r) for s in superseded)] assert [r.data for r in visible] == ["other"] - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[-1], common.Closed) @workflow.defn class FailingRun: @workflow.run async def run(self) -> str: - turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + stream = await workflow.create_stream(Events) info = workflow.info() - await turns.publish(f"run-{info.attempt}") + await stream.topic(Events.tokens).publish(f"run-{info.attempt}") if info.attempt == 1: raise ApplicationError("first run fails") + await stream.close(result="done") return "done" @@ -472,15 +474,14 @@ async def test_retry_run_withdraws_the_failed_run(stream_client: Client): ), ) assert await handle.result() == "done" - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] + superseded = [item for item in items if isinstance(item, common.Superseded)] assert [record.data for record in written] == ["run-1", "run-2"] assert len(superseded) == 1 assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] - # The failed run had retries left, so only the last run closes the Stream. closed = items[-1] - assert isinstance(closed, streams.Closed) and closed.result == "done" + assert isinstance(closed, common.Closed) and closed.result == "done" class _Rot13Codec(PayloadCodec): @@ -515,15 +516,15 @@ async def test_records_go_through_the_payload_codec(stream_client: Client): Agent.run, "a b", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - stream = workflow_stream(codec_client, handle) + stream = workflow_stream(handle) decoded = await read_all(stream) - # A Stream with no definition reads the raw Payloads. + # A read of a Topic outside the definition gets the raw Payloads. + plain = codec_client.stream(Plain, id=stream.id) raw = [ item - async for item in streams.StreamClient(codec_client) - .get_handle(Plain, id=stream.id) - .topic("tokens", type=RawValue) - .subscribe(wait_new_records=False) + async for item in plain.subscribe( + plain.topic("tokens", type=RawValue), wait_new_records=False + ) ] assert [item.data for item in records(decoded)] == [ Turn("model"), @@ -532,7 +533,7 @@ async def test_records_go_through_the_payload_codec(stream_client: Client): Turn("done", "A B"), ] first = raw[0] - assert isinstance(first, streams.Record) + assert isinstance(first, common.Record) assert first.data.payload.metadata["encoding"] == b"json/plain" @@ -541,16 +542,16 @@ async def test_client_close_goes_through_the_payload_codec(stream_client: Client 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()}") + codec_client = Client(**config) + done = await redis.create(codec_client, Events, id=f"done-{uuid.uuid4()}") await done.close(result=Turn("done")) - failed = await stream_client_.create(Events, id=f"failed-{uuid.uuid4()}") + failed = await redis.create(codec_client, 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 isinstance(closed, common.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" + assert isinstance(closed, common.Closed) and str(closed.failure) == "bad turn" @workflow.defn @@ -560,14 +561,16 @@ def __init__(self) -> None: @workflow.run async def run(self) -> str: - turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + stream = await workflow.create_stream(Events) + turns = stream.topic(Events.tokens) await turns.publish("before restart") - async with await stream_workflow.create_stream( - streams.with_name(Events, "scoped"), close_with_run=False + async with await workflow.create_stream( + common.with_name(Events, "scoped") ) as scoped: await scoped.topic(Events.tokens).publish("scoped") await workflow.wait_condition(lambda: self.go) await turns.publish("after restart") + await stream.close() return "done" @workflow.signal @@ -581,9 +584,9 @@ async def test_new_worker_replays_and_writes_each_record_once(stream_client: Cli handle = await stream_client.start_workflow( WaitsForSignal.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue ) - first = workflow_stream(stream_client, handle) + first = workflow_stream(handle) reader = first.subscribe(create_wait=timedelta(seconds=20)) - assert isinstance(await anext(reader), streams.Record) + assert isinstance(await anext(reader), common.Record) await reader.aclose() # The new Worker has no cache, so it replays the run and sends the creates and the close # of the scoped Stream again. @@ -591,14 +594,14 @@ async def test_new_worker_replays_and_writes_each_record_once(stream_client: Cli await handle.signal(WaitsForSignal.proceed) assert await handle.result() == "done" items = await read_all(first) - scoped = await read_all(workflow_stream(stream_client, handle, "scoped")) + scoped = await read_all(workflow_stream(handle, "scoped")) assert [i.data for i in records(items)] == ["before restart", "after restart"] - assert isinstance(items[-1], streams.Closed) + assert isinstance(items[-1], common.Closed) assert [i.data for i in records(scoped)] == ["scoped"] - assert [type(i) for i in scoped].count(streams.Closed) == 1 + assert [type(i) for i in scoped].count(common.Closed) == 1 -async def test_plugin_replays_a_history_made_without_it( +async def test_backend_replays_a_history_made_without_it( client: Client, stream_client: Client ): async with new_worker(client, RandomAroundActivity, activities=[noop]) as worker: @@ -611,54 +614,22 @@ async def test_plugin_replays_a_history_made_without_it( history = await handle.fetch_history() replayer = Replayer( workflows=[RandomAroundActivity], - plugins=[ - p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) - ], + stream_backends=stream_client.config()["stream_backends"], ) await replayer.replay_workflow(history) -@workflow.defn -class QuickWriter: - @workflow.run - async def run(self) -> str: - stream = await stream_workflow.create_stream(Events) - await stream.topic(Events.tokens).publish("quick") - return "done" - - -async def test_worker_shutdown_drains_closes(stream_client: Client): - 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 - ) - await handle.result() - # Calls after the Worker exits still work, and the drain waited for the run's close. - items = [ - item - async for item in workflow_stream(stream_client, handle).subscribe( - wait_new_records=False - ) - ] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" - - -async def test_plugin_shutdown_after_client_writes(stream_client: Client): - [plugin] = [ - 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()}" - ) - async with streams.BufferedPublisher(stream.topic(Events.tokens)) as tokens: +async def test_backend_shutdown_after_client_writes(stream_client: Client): + [backend] = stream_client.config()["stream_backends"] + assert isinstance(backend, ExtensionStreamBackend) + stream = await redis.create(stream_client, Events, id=f"shutdown-{uuid.uuid4()}") + async with common.BufferedPublisher(stream.topic(Events.tokens)) as tokens: tokens.publish("buffered") - await plugin.shutdown() + await backend.shutdown() written = [ item.data async for item in stream.subscribe(wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] assert written == ["buffered"] @@ -667,7 +638,7 @@ async def test_plugin_shutdown_after_client_writes(stream_client: Client): class WaitsForever: @workflow.run async def run(self) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("waiting") await workflow.wait_condition(lambda: False) @@ -686,12 +657,11 @@ async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: def test_closing_an_evicted_run_touches_no_runtime(): - # Python closes the coroutine of an evicted run with GeneratorExit, outside the - # Workflow's loop. The interceptor must let it pass with no call to the runtime, or - # Python prints "Exception ignored". + # Python closes an evicted run's coroutine with GeneratorExit, outside the Workflow's loop. + # The interceptor must let it pass without a call to the runtime. inbound = _WorkflowInbound(_NeverEnds()) run = inbound.execute_workflow( - cast(ExecuteWorkflowInput, cast(object, SimpleNamespace(headers={}))) + cast(ExecuteWorkflowInput, cast(object, SimpleNamespace(headers={}, args=[]))) ) run.send(None) run.close() @@ -712,10 +682,8 @@ async def test_eviction_prints_nothing( handle = await stream_client.start_workflow( WaitsForever.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue ) - reader = workflow_stream(stream_client, handle).subscribe( - create_wait=timedelta(seconds=20) - ) - assert isinstance(await anext(reader), streams.Record) + reader = workflow_stream(handle).subscribe(create_wait=timedelta(seconds=20)) + assert isinstance(await anext(reader), common.Record) await reader.aclose() await handle.terminate() gc.collect() @@ -728,7 +696,7 @@ class AlwaysFails: @workflow.run async def run(self, non_retryable: bool) -> None: attempt = workflow.info().attempt - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish(f"run-{attempt}") raise ApplicationError(f"attempt {attempt} fails", non_retryable=non_retryable) @@ -749,20 +717,18 @@ async def test_final_failure_under_a_retry_policy( ) with pytest.raises(WorkflowFailureError): await handle.result() - items = await read_all(workflow_stream(stream_client, handle)) + stream = workflow_stream(handle) + items = [item async for item in stream.subscribe(wait_new_records=False)] written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] - closes = [item for item in items if isinstance(item, streams.Closed)] - # Only the run that ends the chain closes the Stream, with its failure, and it is last. - assert len(closes) == 1 and items[-1] is closes[0] - assert closes[0].failure is not None + superseded = [item for item in items if isinstance(item, common.Superseded)] + # The end of the chain doesn't close the Stream. + assert not any(isinstance(item, common.Closed) for item in items) if non_retryable: - # A non-retryable failure ends the chain on attempt 1, so it closes at once. + # A non-retryable failure ends the chain on attempt 1. assert [r.data for r in written] == ["run-1"] assert superseded == [] return - # Attempts 1 and 2 left the Stream open, so attempts 2 and 3 could write. Each retry run - # withdraws the run before it. + # Each retry run withdraws the Records of the run before it. assert [r.data for r in written] == ["run-1", "run-2", "run-3"] assert len(superseded) == 2 assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] @@ -771,18 +737,16 @@ async def test_final_failure_under_a_retry_policy( async def test_reader_stops_on_its_own_timer_and_resumes(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"timer-{uuid.uuid4()}" - ) + stream = await redis.create(stream_client, Events, id=f"timer-{uuid.uuid4()}") first = await stream.topic(Events.tokens).publish("one") - seen: list[streams.Record[Any]] = [] + seen: list[common.Record[Any]] = [] async def read_until_quiet() -> None: async for item in stream.subscribe(): - assert isinstance(item, streams.Record) + assert isinstance(item, common.Record) seen.append(item) - # A read has no idle end, so a reader that wants one sets its own timer. + # A read has no idle timeout, so a reader that wants one uses its own timer. with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(read_until_quiet(), timeout=0.5) assert [item.offset for item in seen] == [first] @@ -790,34 +754,36 @@ async def read_until_quiet() -> None: resumed = [ item.data async for item in stream.subscribe( - from_offset=seen[-1].offset + 1, wait_new_records=False + offset=seen[-1].offset + 1, wait_new_records=False ) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] assert resumed == ["two"] async def test_two_backends_on_one_client(client: Client, redis_url: str): - # A client may carry two Stream plugins, such as during a move from one backend to - # another. Each handle carries its backend, so a call goes to the backend that has it. + # A client can have two Stream backends, such as during a migration to another backend. Each + # handle stores its backend name, so a call goes to the backend that has the Stream. config = client.config() - config["plugins"] = [ - RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True), - MemoryStreamPlugin(), + config["stream_backends"] = [ + RedisStreamBackend(redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True), + MemoryStreamBackend(), ] - stream_client = streams.StreamClient(Client(**config)) + stream_client = Client(**config) stream_id = f"two-{uuid.uuid4()}" - on_redis = await stream_client.create(Events, id=stream_id) - on_memory = await stream_client.create(Events, id=stream_id, backend="memory") + on_redis = await redis.create(stream_client, Events, id=stream_id) + on_memory = await redis.create( + stream_client, Events, id=stream_id, backend="memory" + ) assert (on_redis.backend, on_memory.backend) == ("redis", "memory") await on_redis.topic(Events.tokens).publish("redis") await on_memory.topic(Events.tokens).publish("memory") for backend, expected in (("redis", "redis"), ("memory", "memory")): - again = stream_client.get_handle(Events, id=stream_id, backend=backend) + again = stream_client.stream(Events, id=stream_id, backend=backend) data = [ item.data async for item in again.subscribe(wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] assert data == [expected] @@ -826,8 +792,8 @@ async def test_two_backends_on_one_client(client: Client, redis_url: str): class _FailsTheLastTaskOnce(Interceptor, WorkerInterceptor): - # On the client before the Stream plugin, so it is the outer Workflow interceptor and - # fails the task after the plugin's part of the run end. + # It's on the client before the Stream backend's interceptor, so it's the outer Workflow + # interceptor. It fails the task after the backend drained the run's writes. def workflow_interceptor_class( self, input: WorkflowInterceptorClassInput ) -> type[WorkflowInboundInterceptor]: @@ -845,52 +811,76 @@ async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: @workflow.defn -class PublishesAndReturns: +class PublishesClosesAndReturns: @workflow.run async def run(self) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("token") + await stream.close(result="done") return "done" +async def close_activity_events( + handle: WorkflowHandle[Any, Any], +) -> list[EventType.ValueType]: + # The events of the close Activity, from its schedule on. + history = await handle.fetch_history() + scheduled = { + event.event_id + for event in history.events + if event.activity_task_scheduled_event_attributes.activity_type.name + == "__temporal_streams_close" + } + return [ + event.event_type + for event in history.events + if event.event_id in scheduled + or event.activity_task_started_event_attributes.scheduled_event_id in scheduled + or event.activity_task_completed_event_attributes.scheduled_event_id + in scheduled + ] + + @pytest.mark.parametrize("cache", [0, 10]) -async def test_failed_last_task_leaves_the_close_to_its_retry( - stream_client: Client, cache: int -): - # The close with the run goes out once the server accepted the run's completion. A close - # in the failed task would make the retry's create raise StreamClosedError. +async def test_close_survives_a_failed_last_task(stream_client: Client, cache: int): + # The close is an Activity, committed with the Workflow Task that scheduled it. The last + # task fails after it, and its retry replays the close and writes nothing more. config = stream_client.config() config["interceptors"] = [_FailsTheLastTaskOnce()] client = Client(**config) async with new_worker( client, - PublishesAndReturns, + PublishesClosesAndReturns, max_cached_workflows=cache, workflow_runner=UnsandboxedWorkflowRunner(), ) as worker: handle = await client.start_workflow( - PublishesAndReturns.run, + PublishesClosesAndReturns.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) assert await handle.result() == "done" - items = await asyncio.wait_for(read_all(workflow_stream(client, handle)), 20) - written = records(items) - superseded = [item for item in items if isinstance(item, streams.Superseded)] - visible = [r for r in written if not any(s.withdraws(r) for s in superseded)] - assert [r.data for r in visible] == ["token"] - closes = [item for item in items if isinstance(item, streams.Closed)] - assert len(closes) == 1 and closes[0].result == "done" - assert closes[0].status == streams.CloseStatus.COMPLETED + stream = workflow_stream(handle) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [r.data for r in records(items)] == ["token"] + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and items[-1] is closes[0] + assert closes[0].result == "done" + assert closes[0].status == common.CloseStatus.COMPLETED + assert await close_activity_events(handle) == [ + EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED, + EventType.EVENT_TYPE_ACTIVITY_TASK_STARTED, + EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED, + ] @workflow.defn class WritesToStream: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> str: + async def run(self, stream: common.WorkflowStream[Events]) -> str: try: await stream.topic(Events.tokens).publish("from workflow") - except streams.StreamError as err: + except common.StreamError as err: return err.reason return "wrote" @@ -900,7 +890,7 @@ async def test_write_to_a_missing_stream_raises_at_the_await(stream_client: Clie assert ( await stream_client.execute_workflow( WritesToStream.run, - stream_workflow.WorkflowStream(id=f"missing-{uuid.uuid4()}"), + common.WorkflowStream(id=f"missing-{uuid.uuid4()}"), id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -909,29 +899,27 @@ async def test_write_to_a_missing_stream_raises_at_the_await(stream_client: Clie async def test_write_to_a_closed_stream_raises_at_the_await(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"closed-{uuid.uuid4()}" - ) + stream = await redis.create(stream_client, Events, id=f"closed-{uuid.uuid4()}") await stream.close() async with new_worker(stream_client, WritesToStream) as worker: assert ( await stream_client.execute_workflow( WritesToStream.run, - stream_workflow.WorkflowStream.from_handle(stream), + common.WorkflowStream.from_handle(stream), id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) == "stream_closed" ) items = [item async for item in stream.subscribe()] - assert len(items) == 1 and isinstance(items[0], streams.Closed) + assert len(items) == 1 and isinstance(items[0], common.Closed) @activity.defn async def unknown_topic(input: CallModelInput) -> str: try: input.stream.topic("missing") - except streams.StreamError as err: + except common.StreamError as err: return err.reason return "found" @@ -940,10 +928,10 @@ async def unknown_topic(input: CallModelInput) -> str: class UnknownTopic: @workflow.run async def run(self) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) return await workflow.execute_activity( unknown_topic, - CallModelInput("", streams.ActivityStream.from_handle(stream)), + CallModelInput("", common.ActivityStream.from_handle(stream)), start_to_close_timeout=timedelta(seconds=30), ) @@ -961,19 +949,20 @@ async def test_activity_refuses_a_topic_outside_the_definition(stream_client: Cl @workflow.defn class ContinuesThenRetries: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events] | None) -> str: + async def run(self, stream: common.WorkflowStream[Events] | None) -> str: if stream is None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("first") workflow.continue_as_new(stream) info = workflow.info() await stream.topic(Events.tokens).publish(f"attempt-{info.attempt}") if info.attempt == 1: raise ApplicationError("the continued run fails once") + await stream.close(result="done") return "done" -async def test_retry_of_a_continued_run_keeps_the_close_with_run(stream_client: Client): +async def test_retry_of_a_continued_run_writes_the_same_stream(stream_client: Client): async with new_worker(stream_client, ContinuesThenRetries) as worker: handle = await stream_client.start_workflow( ContinuesThenRetries.run, @@ -985,10 +974,10 @@ async def test_retry_of_a_continued_run_keeps_the_close_with_run(stream_client: ), ) assert await handle.result() == "done" - items = await read_all(workflow_stream(stream_client, handle)) + items = await read_all(workflow_stream(handle)) written = records(items) assert [r.data for r in written] == ["first", "attempt-1", "attempt-2"] - superseded = [item for item in items if isinstance(item, streams.Superseded)] + superseded = [item for item in items if isinstance(item, common.Superseded)] assert [r.data for r in written if superseded[0].withdraws(r)] == ["attempt-1"] - closes = [item for item in items if isinstance(item, streams.Closed)] + closes = [item for item in items if isinstance(item, common.Closed)] assert len(closes) == 1 and closes[0].result == "done" diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index a189031d03..97db96eaee 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -1,21 +1,27 @@ import asyncio +import dataclasses +import logging import uuid from dataclasses import dataclass from datetime import timedelta from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest -from temporalio import workflow +from temporalio import common, workflow 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.redis import RedisStreamPlugin -from temporalio.contrib.streams.testing import MemoryStreamPlugin -from temporalio.converter import DataConverter -from tests.contrib.streams.fake_extension import FakeExtension +from temporalio.common._streams import _stream_client +from temporalio.contrib import stream as redis +from temporalio.contrib.stream import RedisStreamBackend +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.converter import DataConverter, DefaultFailureConverter +from temporalio.exceptions import ApplicationError +from temporalio.plugin import SimplePlugin +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Replayer +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner +from tests.contrib.streams.fake_extension import fake_client from tests.helpers import new_worker @@ -25,51 +31,39 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) - - -def fake_client(client: Client) -> tuple[Client, FakeExtension]: - extension = FakeExtension() - plugin = streams.StreamPlugin( - "fake-streams", - backend_name="fake", - default=False, - stream_backend=_ExtensionStreamBackend(extension), - ) - config = client.config() - config["plugins"] = [plugin] - return Client(**config), extension + turns = common.topic(Turn) + tokens = common.topic(str) 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.StreamConfigError, match="It has fake"): - await stream_client.create(Events, id="chat-1", backend="redis") - with pytest.raises(streams.StreamConfigError): - stream_client.get_handle(Events, id="chat-1", backend="redis") + with pytest.raises(common.StreamConfigError, match="no Redis Stream backend redis"): + await redis.create(client, Events, id="chat-1", backend="redis") + with pytest.raises(common.StreamConfigError, match="It has fake"): + client.stream(Events, id="chat-1", backend="redis") assert not extension.streams and not extension.namespaces - stream = await stream_client.create(Events, id="chat-1") + stream = await redis.create(client, Events, id="chat-1") assert stream.backend == "fake" -async def test_client_without_plugin_has_no_transport(client: Client): - with pytest.raises(streams.StreamConfigError): - streams.StreamClient(client).get_handle(Events, id="chat-1") +async def test_client_without_backend_has_no_transport(client: Client): + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + client.stream(Events, id="chat-1") + with pytest.raises(common.StreamConfigError, match="no Redis Stream backend"): + await redis.create(client, Events, id="chat-1") async def test_decoded_client_stream_has_no_client(client: Client): client, _ = fake_client(client) - stream = await streams.StreamClient(client).create(Events, id="chat-2") + stream = await redis.create(client, Events, id="chat-2") converter = DataConverter.default.payload_converter decoded = converter.from_payload( - converter.to_payload(stream), streams.ClientStream[Events] + converter.to_payload(stream), common.ClientStream[Events] ) assert decoded == stream - with pytest.raises(streams.StreamConfigError, match="get_handle"): + with pytest.raises(common.StreamConfigError, match="Client.stream"): await decoded.get_offset() @@ -93,28 +87,86 @@ class Description: return Description() setattr(handle, "describe", describe) - stream_client = streams.StreamClient(client) - await stream_client.create(Events, id="agent-1:first:n") - stream = stream_client.workflow_stream(handle, streams.with_name(Events, "n")) + await redis.create(client, Events, id="agent-1:first:n") + stream = handle.stream(common.with_name(Events, "n")) with pytest.raises(RuntimeError, match="not known yet"): stream.id assert await stream.get_offset() == 0 assert stream.id == "agent-1:first:n" assert stream.first_execution_run_id == "first" assert await stream.get_offset() == 0 + # The handle keeps the first run id, so its other Streams make no call. + other = handle.stream(common.with_name(Events, "n")) + assert await other.get_offset() == 0 assert calls == [1] +async def test_failed_first_run_lookup_is_tried_again(client: Client): + client, _ = fake_client(client) + handle = client.get_workflow_handle("agent-4") + describes = [] + + async def describe(**_: Any) -> Any: + describes.append(1) + if len(describes) == 1: + raise RuntimeError("the Workflow has not started yet") + info = SimpleNamespace(first_run_id="run-4") + return SimpleNamespace( + raw_description=SimpleNamespace(workflow_execution_info=info) + ) + + setattr(handle, "describe", describe) + await redis.create(client, Events, id="agent-4:run-4:events") + stream = handle.stream(Events) + with pytest.raises(RuntimeError, match="not started"): + await stream.get_offset() + assert await stream.get_offset() == 0 + assert stream.id == "agent-4:run-4:events" + # The handle keeps the id it found, so its other Streams make no call. + assert await handle.stream(Events).get_offset() == 0 + assert describes == [1, 1] + + +async def test_cancelled_first_run_lookup_is_tried_again(client: Client): + client, _ = fake_client(client) + handle = client.get_workflow_handle("agent-5") + started = asyncio.Event() + + async def hangs(**_: Any) -> Any: + started.set() + await asyncio.Event().wait() + + setattr(handle, "describe", hangs) + stream = handle.stream(Events) + offset = asyncio.ensure_future(stream.get_offset()) + await started.wait() + # The Describe call is cancelled, such as when the loop that ran it stops. + lookup = _stream_client._FIRST_RUNS[handle] + assert isinstance(lookup, asyncio.Future) + lookup.cancel() + with pytest.raises(asyncio.CancelledError): + await offset + + async def describe(**_: Any) -> Any: + info = SimpleNamespace(first_run_id="run-5") + return SimpleNamespace( + raw_description=SimpleNamespace(workflow_execution_info=info) + ) + + setattr(handle, "describe", describe) + await redis.create(client, Events, id="agent-5:run-5:events") + assert await stream.get_offset() == 0 + + @workflow.defn class WriteToStream: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> None: + async def run(self, stream: common.WorkflowStream[Events]) -> None: await stream.topic(Events.turns).publish(Turn("workflow")) async def test_client_stream_takes_workflow_writes(stream_client: Client): - client = streams.StreamClient(stream_client) - stream = await client.create(Events, id=f"chat-{uuid.uuid4()}") + stream = await redis.create(stream_client, Events, id=f"chat-{uuid.uuid4()}") assert await stream.get_offset() == 0 first = await stream.topic(Events.turns).publish(Turn("client")) last = await stream.topic(Events.tokens).publish("a", "b") @@ -122,35 +174,35 @@ async def test_client_stream_takes_workflow_writes(stream_client: Client): async with new_worker(stream_client, WriteToStream) as worker: await stream_client.execute_workflow( WriteToStream.run, - stream_workflow.WorkflowStream.from_handle(stream), + common.WorkflowStream.from_handle(stream), id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) await stream.close(result="bye") items = [item async for item in stream.subscribe()] - data = [item.data for item in items if isinstance(item, streams.Record)] + data = [item.data for item in items if isinstance(item, common.Record)] assert data == [Turn("client"), "a", "b", Turn("workflow")] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "bye" + assert isinstance(items[-1], common.Closed) and items[-1].result == "bye" # Another handle of the same Stream reads it, from an offset and by Topic. - again = client.get_handle(Events, id=stream.id) + again = stream_client.stream(Events, id=stream.id) resumed = [ item.data - async for item in again.subscribe(Events.tokens, from_offset=first + 1) - if isinstance(item, streams.Record) + async for item in again.subscribe(Events.tokens, offset=first + 1) + if isinstance(item, common.Record) ] assert resumed == ["a", "b"] by_prefix = [ item.data - async for item in again.subscribe(prefixes=["tok"]) - if isinstance(item, streams.Record) + async for item in again.subscribe(topic_prefixes=["tok"]) + if isinstance(item, common.Record) ] assert by_prefix == ["a", "b"] - only_turns = [item async for item in again.topic(Events.turns).subscribe()] + only_turns = [item async for item in again.subscribe(again.topic(Events.turns))] assert [type(item) for item in only_turns] == [ - streams.Record, - streams.Record, - streams.Closed, + common.Record, + common.Record, + common.Closed, ] @@ -159,8 +211,9 @@ class Agent: @workflow.run async def run(self, prompt: str) -> str: await workflow.sleep(timedelta(milliseconds=200)) - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish(*prompt.split()) + await stream.close(result=prompt.upper()) return prompt.upper() @@ -169,43 +222,44 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) + stream = handle.stream(Events) # The Workflow creates its Stream later, so the read waits for the create. items = [ item async for item in stream.subscribe(create_wait=timedelta(seconds=20)) ] - assert [item.data for item in items if isinstance(item, streams.Record)] == [ + assert [item.data for item in items if isinstance(item, common.Record)] == [ "hi", "there", ] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "HI THERE" + assert isinstance(items[-1], common.Closed) and items[-1].result == "HI THERE" assert stream.workflow_id == handle.id - with pytest.raises(streams.StreamNotOwnerError): + with pytest.raises(common.StreamNotOwnerError): await stream.close() async def test_async_with_closes_a_client_stream(stream_client: Client): - client = streams.StreamClient(stream_client) - async with client.create(Events, id=f"block-{uuid.uuid4()}") as done: + async with redis.create(stream_client, Events, id=f"block-{uuid.uuid4()}") as done: await done.topic(Events.tokens).publish("a") - failed = client.get_handle(Events, id=f"block-{uuid.uuid4()}") + failed = stream_client.stream(Events, id=f"block-{uuid.uuid4()}") with pytest.raises(ValueError): - async with client.create(Events, id=failed.id): + async with redis.create(stream_client, Events, id=failed.id): raise ValueError("turn failed") - async with await client.create(Events, id=f"block-{uuid.uuid4()}") as early: + async with await redis.create( + stream_client, Events, id=f"block-{uuid.uuid4()}" + ) as early: await early.close(result="early") - cancelled = client.get_handle(Events, id=f"block-{uuid.uuid4()}") + cancelled = stream_client.stream(Events, id=f"block-{uuid.uuid4()}") with pytest.raises(asyncio.CancelledError): - async with client.create(Events, id=cancelled.id): + async with redis.create(stream_client, Events, id=cancelled.id): raise asyncio.CancelledError() for stream, result, failure, status in ( - (done, None, None, streams.CloseStatus.COMPLETED), - (failed, None, "turn failed", streams.CloseStatus.FAILED), - (early, "early", None, streams.CloseStatus.COMPLETED), - (cancelled, None, "Cancelled", streams.CloseStatus.CANCELED), + (done, None, None, common.CloseStatus.COMPLETED), + (failed, None, "turn failed", common.CloseStatus.FAILED), + (early, "early", None, common.CloseStatus.COMPLETED), + (cancelled, None, "Cancelled", common.CloseStatus.CANCELED), ): items = [item async for item in stream.subscribe()] - closes = [item for item in items if isinstance(item, streams.Closed)] + closes = [item for item in items if isinstance(item, common.Closed)] assert len(closes) == 1 and closes[0].result == result assert closes[0].status == status assert (failure is None) == (closes[0].failure is None) @@ -213,85 +267,114 @@ async def test_async_with_closes_a_client_stream(stream_client: Client): async def test_next_offset_resumes_after_the_last_item(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"resume-{uuid.uuid4()}" - ) + stream = await redis.create(stream_client, Events, id=f"resume-{uuid.uuid4()}") await stream.topic(Events.tokens).publish("a") turn = await stream.topic(Events.turns).publish(Turn("model")) - tokens = stream.topic(Events.tokens).subscribe(wait_new_records=False) + tokens = stream.subscribe(Events.tokens, wait_new_records=False) assert tokens.next_offset == 0 - assert [item.data async for item in tokens if isinstance(item, streams.Record)] == [ + assert [item.data async for item in tokens if isinstance(item, common.Record)] == [ "a" ] # The page ends past the Record that the Topic filter left out. assert tokens.next_offset == turn + 1 - new = stream.subscribe(from_offset=streams.END, wait_new_records=False) + new = stream.subscribe(offset=common.END, wait_new_records=False) assert new.next_offset is None assert [item async for item in new] == [] assert new.next_offset == turn + 1 await stream.topic(Events.tokens).publish("b") - resumed = stream.subscribe(from_offset=tokens.next_offset, wait_new_records=False) - assert [ - item.data async for item in resumed if isinstance(item, streams.Record) - ] == ["b"] + resumed = stream.subscribe(offset=tokens.next_offset, wait_new_records=False) + assert [item.data async for item in resumed if isinstance(item, common.Record)] == [ + "b" + ] -async def test_last_reads_the_newest_records_then_follows(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"last-{uuid.uuid4()}" +class _EncodesCommonAttributes(DefaultFailureConverter): + def __init__(self) -> None: + super().__init__(encode_common_attributes=True) + + +async def test_closed_failure_decodes_with_the_clients_converter(stream_client: Client): + # The failure converter puts the message in the encoded attributes, which only the payload + # converter reads. + config = stream_client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, failure_converter_class=_EncodesCommonAttributes ) + client = Client(**config) + stream = await redis.create(client, Events, id=f"failed-{uuid.uuid4()}") + await stream.close( + failure=ApplicationError("tool failed", "detail", {"k": 1}, type="Tool") + ) + [closed] = [item async for item in stream.subscribe()] + assert isinstance(closed, common.Closed) + assert isinstance(closed.failure, ApplicationError) + assert closed.failure.message == "tool failed" and closed.failure.type == "Tool" + assert closed.failure.details == ("detail", {"k": 1}) + + +async def test_no_wait_read_goes_past_records_of_other_topics(stream_client: Client): + # A Store scans at most 10,000 entries per call, so the first pages of this read are empty. + stream = await redis.create(stream_client, Events, id=f"scan-{uuid.uuid4()}") + for _ in range(11): + await stream.topic(Events.tokens).publish(*["x"] * 1000) + turn = await stream.topic(Events.turns).publish(Turn("model")) + turns = stream.subscribe(Events.turns, wait_new_records=False) + assert [item.data async for item in turns if isinstance(item, common.Record)] == [ + Turn("model") + ] + assert turns.next_offset == turn + 1 + + +async def test_last_reads_the_newest_records_then_follows(stream_client: Client): + stream = await redis.create(stream_client, Events, id=f"last-{uuid.uuid4()}") tokens = stream.topic(Events.tokens) await tokens.publish("a", "b", "c") await stream.topic(Events.turns).publish(Turn("model")) - # A chat view shows the last few Records first, then follows. - tail = tokens.subscribe(last=2) + # A chat view shows the last few Records first, then waits for new ones. + tail = stream.subscribe(tokens, last=2) assert tail.next_offset is None shown = [await anext(tail), await anext(tail)] await tokens.publish("d") shown.append(await anext(tail)) await tail.aclose() - assert [item.data for item in shown if isinstance(item, streams.Record)] == [ + assert [item.data for item in shown if isinstance(item, common.Record)] == [ "b", "c", "d", ] newest = stream.subscribe(last=1, wait_new_records=False) - assert [item.data async for item in newest if isinstance(item, streams.Record)] == [ + assert [item.data async for item in newest if isinstance(item, common.Record)] == [ "d" ] - with pytest.raises(ValueError, match="no from_offset"): - stream.subscribe(last=2, from_offset=0) + with pytest.raises(ValueError, match="no offset"): + stream.subscribe(last=2, offset=0) async def test_read_type_decodes_topics_outside_the_definition(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"typed-{uuid.uuid4()}" - ) + stream = await redis.create(stream_client, Events, id=f"typed-{uuid.uuid4()}") await stream.topic("extra", type=Turn).publish(Turn("extra")) await stream.topic(Events.tokens).publish("a") async def read(**options: Any) -> list[object]: items = stream.subscribe(wait_new_records=False, **options) - return [item.data async for item in items if isinstance(item, streams.Record)] + return [item.data async for item in items if isinstance(item, common.Record)] assert await read() == [{"kind": "extra", "text": ""}, "a"] assert await read(type=Turn) == [Turn("extra"), "a"] async def test_buffered_publisher_keeps_the_order_across_topics(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"buffered-{uuid.uuid4()}" - ) - async with streams.BufferedPublisher(stream) as publisher: + stream = await redis.create(stream_client, Events, id=f"buffered-{uuid.uuid4()}") + async with common.BufferedPublisher(stream) as publisher: publisher.publish(Events.turns, Turn("user")) publisher.publish(Events.tokens, "a", "b") publisher.publish("turns", Turn("model")) - with pytest.raises(streams.StreamError, match="no Topic"): + with pytest.raises(common.StreamError, match="no Topic"): publisher.publish("turn", Turn("typo")) written = [ (item.topic, item.data) async for item in stream.subscribe(wait_new_records=False) - if isinstance(item, streams.Record) + if isinstance(item, common.Record) ] assert written == [ ("turns", Turn("user")), @@ -301,45 +384,111 @@ async def test_buffered_publisher_keeps_the_order_across_topics(stream_client: C ] -async def test_named_producer_fences_its_earlier_epoch(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"named-{uuid.uuid4()}" +@common.defn(name="events") +class LooseEvents: + turns = common.topic(str) + tokens = common.topic(str) + + +async def test_record_that_fails_to_decode_raises_and_resumes_past_it( + stream_client: Client, +): + stream_id = f"poison-{uuid.uuid4()}" + writer = await redis.create(stream_client, LooseEvents, id=stream_id) + bad = await writer.topic(LooseEvents.turns).publish("not a turn") + await writer.topic(LooseEvents.tokens).publish("after") + reader = stream_client.stream(Events, id=stream_id) + read = reader.subscribe(wait_new_records=False) + with pytest.raises(common.StreamError, match=f"offset {bad} of the Topic turns"): + await anext(read) + assert read.next_offset == bad + 1 + resumed = reader.subscribe(offset=read.next_offset, wait_new_records=False) + assert [item.data async for item in resumed if isinstance(item, common.Record)] == [ + "after" + ] + + +async def test_async_with_refuses_a_workflow_stream_before_the_block( + stream_client: Client, +): + # The close at the block end would raise and hide the block's own error. + handle = stream_client.get_workflow_handle( + "agent-1", first_execution_run_id="run-1" ) - # A server that restarts writes as one name with a higher epoch, so the zombie it + ran = False + with pytest.raises(common.StreamNotOwnerError): + async with handle.stream(Events): + ran = True + assert not ran + + +async def test_record_over_the_store_limit_raises_too_large(stream_client: Client): + # The client's writer splits a batch at the Store's limit, so only one Record can be over it. + stream = await redis.create(stream_client, Events, id=f"large-{uuid.uuid4()}") + with pytest.raises(common.StreamTooLargeError) as refused: + await stream.topic(Events.tokens).publish("x" * 2_000_000) + assert refused.value.reason == "stream_too_large" and refused.value.non_retryable + assert await stream.get_offset() == 0 + + +async def test_read_keeps_no_offsets_of_client_records(stream_client: Client): + # Nothing withdraws a client's Records, so a long UI read of them keeps no state. + stream = await redis.create(stream_client, Events, id=f"chat-{uuid.uuid4()}") + await stream.topic(Events.tokens).publish("a", "b", "c") + read = stream.subscribe(wait_new_records=False) + assert [item.data async for item in read if isinstance(item, common.Record)] == [ + "a", + "b", + "c", + ] + assert read._shown == {} + + +async def test_empty_stream_id_raises_with_no_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(ValueError, match="id is empty"): + client.stream(Events, id="") + with pytest.raises(ValueError, match="id is empty"): + client.stream(Events, stream=common.WorkflowStream(id="")) + assert not extension.namespaces + + +async def test_create_of_an_expired_stream_raises(stream_client: Client): + # An open client Stream expires one Retention after its last write. + stream_id = f"expired-{uuid.uuid4()}" + retention = timedelta(milliseconds=200) + stream = await redis.create( + stream_client, Events, id=stream_id, retention=retention + ) + await stream.topic(Events.tokens).publish("old") + await asyncio.sleep(0.5) + with pytest.raises(common.StreamExpiredError): + await redis.create(stream_client, Events, id=stream_id, retention=retention) + + +async def test_named_producer_fences_its_earlier_epoch(stream_client: Client): + stream = await redis.create(stream_client, Events, id=f"named-{uuid.uuid4()}") + # A server that restarts writes as the same name with a higher epoch, so the zombie that it # replaces can't write after it. before = stream.producer("ingest", epoch=1) await before.topic(Events.tokens).publish("before restart") after = stream.producer("ingest", epoch=2) await after.topic(Events.tokens).publish("after restart") - with pytest.raises(streams.StreamFencedError): + with pytest.raises(common.StreamFencedError): await before.topic(Events.tokens).publish("zombie") # The earlier epoch's Records are real messages, so nothing withdraws them. items = [item async for item in stream.subscribe(wait_new_records=False)] - assert [type(item) for item in items] == [streams.Record, streams.Record] + assert [type(item) for item in items] == [common.Record, common.Record] with pytest.raises(ValueError): stream.producer("ingest", epoch=0) -async def test_auto_close_of_a_client_stream(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"quiet-{uuid.uuid4()}", auto_close=timedelta(seconds=1) - ) - await stream.topic(Events.tokens).publish("only") - await asyncio.sleep(1.5) - items = [item async for item in stream.subscribe(wait_new_records=False)] - assert [item.data for item in items if isinstance(item, streams.Record)] == ["only"] - assert isinstance(items[-1], streams.Closed) and items[-1].result is None - with pytest.raises(streams.StreamClosedError): - await stream.topic(Events.tokens).publish("late") - - async def test_unresolved_workflow_stream_compares_and_prints(client: Client): client, _ = fake_client(client) - stream_client = streams.StreamClient(client) handle = client.get_workflow_handle("agent-2") - first = stream_client.workflow_stream(handle, Events) - same = stream_client.workflow_stream(handle, Events) - other = stream_client.workflow_stream(handle, streams.with_name(Events, "other")) + first = handle.stream(Events) + same = handle.stream(Events) + other = handle.stream(common.with_name(Events, "other")) key = hash(first) assert first == same and first != other describes = [] @@ -353,50 +502,146 @@ async def describe(**_: Any) -> Any: ) setattr(handle, "describe", describe) - await stream_client.create(Events, id="agent-2:run-2:events") - # Calls that start together share one Describe call. + await redis.create(client, Events, id="agent-2:run-2:events") + # Concurrent calls share one Describe call. assert await asyncio.gather(first.get_offset(), first.get_offset()) == [0, 0] assert describes == [1] assert first.id == "agent-2:run-2:events" assert hash(first) == key -async def test_default_plugin_is_the_default_backend(client: Client): +async def test_registration_rules(client: Client): config = client.config() - config["plugins"] = [ - MemoryStreamPlugin(), - MemoryStreamPlugin(name="b", default=True), + config["stream_backends"] = [ + MemoryStreamBackend(), + MemoryStreamBackend(name="b", default=True), ] - stream = await streams.StreamClient(Client(**config)).create(Events, id="d-1") - assert stream.backend == "b" - # The default is part of Workflow code, so it must not follow the order of the list. - for plugins in ( - [MemoryStreamPlugin(), MemoryStreamPlugin(name="b")], - [MemoryStreamPlugin(default=True), MemoryStreamPlugin(name="b", default=True)], + two = Client(**config) + assert (await redis.create(two, Events, id="d-1")).backend == "b" + assert two.stream(Events, id="d-1").backend == "b" + # The default is part of Workflow code, so it must not depend on the order of the list. + for backends in ( + [MemoryStreamBackend(), MemoryStreamBackend(name="b")], + [ + MemoryStreamBackend(default=True), + MemoryStreamBackend(name="b", default=True), + ], ): - config["plugins"] = plugins - with pytest.raises(streams.StreamConfigError, match="exactly one"): + config["stream_backends"] = backends + with pytest.raises(common.StreamConfigError, match="exactly one"): Client(**config) - config["plugins"] = [MemoryStreamPlugin()] - with pytest.raises(streams.StreamConfigError, match="Worker has"): - new_worker(Client(**config), Agent, plugins=[RedisStreamPlugin("redis://-")]) + config["stream_backends"] = [MemoryStreamBackend(), MemoryStreamBackend()] + with pytest.raises(common.StreamConfigError, match="two Stream backends"): + Client(**config) + # A Worker serves only one backend, and so does a Replayer. + config["stream_backends"] = [ + MemoryStreamBackend(), + MemoryStreamBackend(name="b", default=True), + ] + with pytest.raises(common.StreamConfigError, match="one backend"): + new_worker(Client(**config), Agent) + with pytest.raises(common.StreamConfigError, match="one backend"): + Replayer(workflows=[Agent], stream_backends=config["stream_backends"]) + # A generator of backends works too, and a copy of the client keeps them. + config["stream_backends"] = cast( + Any, (b for b in [MemoryStreamBackend(default=True)]) + ) + from_generator = Client(**config) + assert from_generator.stream(Events, id="d-3").backend == "memory" + assert ( + Client(**from_generator.config()).stream(Events, id="d-3").backend == "memory" + ) + + +async def test_worker_that_polls_no_activities_warns_of_the_close( + client: Client, caplog: pytest.LogCaptureFixture +): + # The close of a Workflow's Stream is a regular Activity, which another Worker can run. + config = client.config() + config["stream_backends"] = [MemoryStreamBackend(default=True)] + with caplog.at_level(logging.WARNING, logger="temporalio.common._streams"): + async with new_worker(Client(**config), Agent, no_remote_activities=True): + pass + assert "polls no Activities" in caplog.text + + +async def test_backend_is_not_a_plugin(client: Client): + # The connect refuses it before it connects, so the unreachable target doesn't matter. + with pytest.raises(common.StreamConfigError, match="not a plugin"): + await Client.connect("localhost:1", plugins=[MemoryStreamBackend()]) + config = client.config() + config["plugins"] = [MemoryStreamBackend()] + with pytest.raises(common.StreamConfigError, match="not a plugin"): + Client(**config) + with pytest.raises(common.StreamConfigError, match="not a plugin"): + new_worker(client, Agent, plugins=[RedisStreamBackend("redis://-")]) + with pytest.raises(common.StreamConfigError, match="not a plugin"): + Replayer(workflows=[Agent], plugins=[MemoryStreamBackend()]) + + +async def test_worker_plugin_keeps_the_streams_passthrough(client: Client): + # The backends configure the Worker after its plugins, so a plugin that replaces the + # Workflow runner doesn't lose the Streams sandbox passthrough. + config = client.config() + config["stream_backends"] = [MemoryStreamBackend(default=True)] + plugin = SimplePlugin("replaces-runner", workflow_runner=SandboxedWorkflowRunner()) + worker = new_worker(Client(**config), Agent, plugins=[plugin]) + runner = worker.config(active_config=True).get("workflow_runner") + assert isinstance(runner, SandboxedWorkflowRunner) + assert "temporalio.common._streams" in runner.restrictions.passthrough_modules -async def test_handle_looks_in_each_backend_during_a_move(client: Client): +async def test_default_is_native_with_no_default_backend(client: Client): + config = client.config() + config["stream_backends"] = [MemoryStreamBackend()] + one = Client(**config) + with pytest.raises(common.StreamConfigError, match="default is native"): + one.stream(Events, id="d-2") + with pytest.raises(common.StreamConfigError, match="default is native"): + one.get_workflow_handle("agent-3").stream(Events) + assert one.stream(Events, id="d-2", backend="memory").backend == "memory" + # The Redis create picks the client's only Redis backend. The memory backend counts as one. + assert (await redis.create(one, Events, id="d-2")).backend == "memory" + + +async def test_stream_takes_the_handle_backend_or_the_default(client: Client): config = client.config() - config["plugins"] = [ - MemoryStreamPlugin(name="a", default=True), - MemoryStreamPlugin(name="b"), + config["stream_backends"] = [ + MemoryStreamBackend(name="a", default=True), + MemoryStreamBackend(name="b"), ] - stream_client = streams.StreamClient(Client(**config)) - created = await stream_client.create(Events, id="moving", backend="b") + two = Client(**config) + created = await redis.create(two, Events, id="moving", backend="b") await created.topic(Events.tokens).publish("on b") - moving = stream_client.get_handle(Events, id="moving", backend=("a", "b")) - assert await moving.get_offset() > 0 - assert moving.backend == "b" - nowhere = stream_client.get_handle(Events, id="nowhere", backend=("a", "b")) - with pytest.raises(streams.StreamNotFoundError): - await nowhere.get_offset() - # A handle with no backend name is ambiguous when the client has two backends. - with pytest.raises(streams.StreamConfigError): - stream_client.get_handle(Events, stream=streams.ClientStream(id="moving")) + assert await two.stream(Events, stream=created).get_offset() > 0 + # A handle with no backend name uses the default backend. + assert two.stream(Events, stream=common.ClientStream(id="moving")).backend == "a" + with pytest.raises(ValueError): + two.stream(Events, stream=created, backend="a") + with pytest.raises(TypeError): + two.stream(Events) + + +@workflow.defn +class CreatesAndCloses: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("from the environment") + await stream.close() + + +async def test_workflow_environment_takes_stream_backends(): + async with await WorkflowEnvironment.start_local( + stream_backends=[MemoryStreamBackend(default=True)] + ) as env: + async with new_worker(env.client, CreatesAndCloses) as worker: + handle = await env.client.start_workflow( + CreatesAndCloses.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = [item async for item in handle.stream(Events).subscribe()] + data = [item.data for item in items if isinstance(item, common.Record)] + assert data == ["from the environment"] diff --git a/tests/contrib/streams/test_vectors.py b/tests/contrib/streams/test_vectors.py index a8ae5dfa45..72bf764fe4 100644 --- a/tests/contrib/streams/test_vectors.py +++ b/tests/contrib/streams/test_vectors.py @@ -4,8 +4,8 @@ import pytest import temporalio.api.streamservice.v1 as stream_pb -from temporalio.contrib.streams import _ids, _items -from temporalio.contrib.streams.system import models +from temporalio.common._streams import _ids, _items +from temporalio.common._streams.system import models from tests.contrib.streams._support import crate_file, vectors _KINDS = { @@ -27,8 +27,8 @@ ], ) def test_vendored_protos_match_the_crate(relative: str): - crate = crate_file(f"protos/{relative}").read_text() - assert (_SCRIPTS / relative).read_text() == crate + crate = crate_file(f"protos/{relative}").read_text(encoding="utf-8") + assert (_SCRIPTS / relative).read_text(encoding="utf-8") == crate def test_stream_ids_match_the_crate(): diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index cbb78e93c3..f32d1fc244 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -9,7 +9,7 @@ import pytest -from temporalio import workflow +from temporalio import activity, common, workflow from temporalio.api.common.v1 import Payload from temporalio.api.enums.v1 import EventType from temporalio.client import ( @@ -17,13 +17,15 @@ WorkflowFailureError, WorkflowHandle, ) -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._defn import definition_of +from temporalio.common import RetryPolicy +from temporalio.common._streams import _workflow, _writes +from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams._defn import definition_of +from temporalio.contrib import stream as redis +from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.converter import DataConverter, PayloadCodec from temporalio.exceptions import ApplicationError -from temporalio.worker import Replayer +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner from tests.helpers import assert_task_fail_eventually, new_worker @@ -33,46 +35,44 @@ class Turn: text: str = "" -@streams.defn(name="events") +@common.defn(name="events") class Events: - turns = streams.topic(Turn) - tokens = streams.topic(str) + turns = common.topic(Turn) + tokens = common.topic(str) -async def read( - client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" -) -> list[object]: - stream = streams.StreamClient(client).workflow_stream( - handle, streams.with_name(Events, name) - ) - return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] +async def read(handle: WorkflowHandle[Any, Any], name: str = "events") -> list[object]: + stream = handle.stream(common.with_name(Events, name)) + return [item async for item in stream.subscribe(wait_new_records=False)] def data(items: Sequence[object]) -> list[object]: - return [item.data for item in items if isinstance(item, streams.Record)] + return [item.data for item in items if isinstance(item, common.Record)] @workflow.defn class Agent: @workflow.run async def run(self, prompt: str) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) turns = stream.topic(Events.turns) await turns.publish(Turn("user", prompt), Turn("model")) await stream.topic("tokens").publish(*prompt.split()) + # A run's end doesn't close this Stream, so the run closes it before it returns. + await stream.close(result=prompt.upper()) return prompt.upper() -async def test_create_publish_and_close_with_run(stream_client: Client): +async def test_create_publish_and_close(stream_client: Client): async with new_worker(stream_client, Agent) as worker: handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == "HI THERE" - items = await read(stream_client, handle) + items = await read(handle) assert data(items) == [Turn("user", "hi there"), Turn("model"), "hi", "there"] closed = items[-1] - assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" + assert isinstance(closed, common.Closed) and closed.result == "HI THERE" @workflow.defn @@ -81,8 +81,8 @@ class Turns: async def run(self, fail_at: int) -> None: for n in range(2): try: - async with stream_workflow.create_stream( - streams.with_name(Events, f"turn-{n}"), close_with_run=False + async with workflow.create_stream( + common.with_name(Events, f"turn-{n}") ) as turn: await turn.topic(Events.tokens).publish(f"token-{n}") if n == fail_at: @@ -98,26 +98,27 @@ async def test_async_with_closes_the_stream(stream_client: Client): Turns.run, 1, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - first = await read(stream_client, handle, "turn-0") - second = await read(stream_client, handle, "turn-1") + first = await read(handle, "turn-0") + second = await read(handle, "turn-1") assert data(first) == ["token-0"] - assert isinstance(first[-1], streams.Closed) and first[-1].failure is None + assert isinstance(first[-1], common.Closed) and first[-1].failure is None assert data(second) == ["token-1"] - assert isinstance(second[-1], streams.Closed) + assert isinstance(second[-1], common.Closed) assert str(second[-1].failure) == "turn failed" @workflow.defn class Continuing: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events] | None) -> str: + async def run(self, stream: common.WorkflowStream[Events] | None) -> str: if stream is None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("run 1") workflow.continue_as_new(stream) assert stream.definition is Events - assert stream == stream_workflow.stream_handle(Events) + assert stream == workflow.stream_handle(Events) await stream.topic(Events.tokens).publish("run 2") + await stream.close(result="done") return "done" @@ -125,14 +126,15 @@ async def run(self, stream: stream_workflow.WorkflowStream[Events] | None) -> st class RenamedAcrossRuns: @workflow.run async def run(self, first: bool) -> list[str]: - renamed = streams.with_name(Events, "turn-1") + renamed = common.with_name(Events, "turn-1") if first: - turn = await stream_workflow.create_stream(renamed) + turn = await workflow.create_stream(renamed) await turn.topic(Events.tokens).publish("run 1") workflow.continue_as_new(False) - turn = stream_workflow.stream_handle(renamed) + turn = workflow.stream_handle(renamed) await turn.topic(Events.tokens).publish("run 2") - return [stream_workflow.stream_handle(Events).id, turn.id] + await turn.close() + return [workflow.stream_handle(Events).id, turn.id] async def test_renamed_definition_is_another_stream(stream_client: Client): @@ -144,13 +146,13 @@ async def test_renamed_definition_is_another_stream(stream_client: Client): task_queue=worker.task_queue, ) plain_id, turn_id = await handle.result() - turn = await read(stream_client, handle, "turn-1") + turn = await read(handle, "turn-1") assert plain_id.endswith(":events") and turn_id.endswith(":turn-1") assert data(turn) == ["run 1", "run 2"] def test_with_name_leaves_the_definition_unchanged(): - renamed = streams.with_name(Events, "turn-1") + renamed = common.with_name(Events, "turn-1") assert definition_of(Events).name == "events" assert definition_of(renamed).name == "turn-1" assert definition_of(renamed).topics == definition_of(Events).topics @@ -162,34 +164,34 @@ async def test_continue_as_new_passes_the_stream(stream_client: Client): Continuing.run, None, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) assert await handle.result() == "done" - items = await read(stream_client, handle) + items = await read(handle) assert data(items) == ["run 1", "run 2"] - closed = [item for item in items if isinstance(item, streams.Closed)] + closed = [item for item in items if isinstance(item, common.Closed)] assert len(closed) == 1 and closed[0].result == "done" @workflow.defn class Mistakes: def __init__(self) -> None: - self._stream: stream_workflow.WorkflowStream[Events] | None = None + self._stream: common.WorkflowStream[Events] | None = None @workflow.run async def run(self) -> list[str]: errors = [] - stream = await stream_workflow.create_stream(Events, close_with_run=False) + stream = await workflow.create_stream(Events) self._stream = stream try: stream.topic("unknown") - except streams.StreamError as err: + except common.StreamError as err: errors.append(err.reason) await stream.close(result="early") try: await stream.topic(Events.tokens).publish("late") - except streams.StreamClosedError as err: + except common.StreamClosedError as err: errors.append(err.reason) try: - await stream_workflow.create_stream(Events) - except streams.StreamClosedError as err: + await workflow.create_stream(Events) + except common.StreamClosedError as err: errors.append(err.reason) return errors @@ -198,7 +200,7 @@ def publish_in_query(self) -> str: assert self._stream is not None try: self._stream.topic(Events.tokens)._stream._append("tokens", ("x",)) - except streams.ReadOnlyContextError as err: + except common.ReadOnlyContextError as err: return type(err).__name__ return "wrote" @@ -216,37 +218,35 @@ async def test_replay_reproducible_errors_raise_in_workflow_code( "stream_closed", ] assert await handle.query(Mistakes.publish_in_query) == "ReadOnlyContextError" - items = await read(stream_client, handle) - assert isinstance(items[-1], streams.Closed) and items[-1].result == "early" + items = await read(handle) + assert isinstance(items[-1], common.Closed) and items[-1].result == "early" @workflow.defn class ClosesAReceivedStream: @workflow.run - async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> str: + async def run(self, stream: common.WorkflowStream[Events]) -> str: await stream.topic(Events.tokens).publish("from workflow") try: await stream.close(result="ended by a writer") - except streams.StreamNotOwnerError as err: + except common.StreamNotOwnerError as err: return err.reason return "closed" async def test_workflow_closes_only_its_own_streams(stream_client: Client): - stream = await streams.StreamClient(stream_client).create( - Events, id=f"owned-{uuid.uuid4()}" - ) + stream = await redis.create(stream_client, Events, id=f"owned-{uuid.uuid4()}") async with new_worker(stream_client, ClosesAReceivedStream) as worker: assert ( await stream_client.execute_workflow( ClosesAReceivedStream.run, - stream_workflow.WorkflowStream.from_handle(stream), + common.WorkflowStream.from_handle(stream), id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) == "stream_not_owner" ) - # The writer can't end the Stream for every reader, so it stays open. + # The Workflow that received the Stream can't close it, so it stays open. items = [item async for item in stream.subscribe(wait_new_records=False)] assert data(items) == ["from workflow"] @@ -255,72 +255,279 @@ async def test_workflow_closes_only_its_own_streams(stream_client: Client): class Failing: @workflow.run async def run(self) -> None: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("before") raise ApplicationError("run failed", non_retryable=True) -async def test_close_with_run_takes_the_failure(stream_client: Client): +async def test_run_end_closes_nothing(stream_client: Client): async with new_worker(stream_client, Failing) as worker: handle = await stream_client.start_workflow( Failing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) with pytest.raises(WorkflowFailureError): await handle.result() - items = await read(stream_client, handle) + items = await read(handle) + # The Stream stays open, so a later reader still waits for a close. assert data(items) == ["before"] - assert isinstance(items[-1], streams.Closed) - assert str(items[-1].failure) == "run failed" - assert items[-1].status == streams.CloseStatus.FAILED + assert not any(isinstance(item, common.Closed) for item in items) @workflow.defn -class LeavesOpen: +class FailsInItsBlockOnce: @workflow.run - async def run(self) -> None: - stream = await stream_workflow.create_stream( - Events, auto_close=timedelta(seconds=1), close_with_run=False + async def run(self) -> str: + attempt = workflow.info().attempt + async with workflow.create_stream(Events) as stream: + await stream.topic(Events.tokens).publish(f"run-{attempt}") + if attempt == 1: + raise ApplicationError("the first run fails") + return f"done on attempt {attempt}" + + +async def test_block_end_leaves_the_stream_open_for_a_retry_run(stream_client: Client): + async with new_worker(stream_client, FailsInItsBlockOnce) as worker: + handle = await stream_client.start_workflow( + FailsInItsBlockOnce.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=100), maximum_attempts=3 + ), + ) + assert await handle.result() == "done on attempt 2" + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert [r.data for r in written] == ["run-1", "run-2"] + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] + # Only the block end of the run that succeeds closes the Stream. + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and items[-1] is closes[0] + assert closes[0].failure is None + + +@workflow.defn +class EndsItsBlock: + @workflow.run + async def run(self, mode: str) -> str: + try: + async with workflow.create_stream(Events) as stream: + await stream.topic(Events.tokens).publish(mode) + if mode == "non_retryable": + raise ApplicationError("block failed", non_retryable=True) + if mode == "listed_type": + raise ApplicationError("block failed", type="Listed") + if mode == "caught_task_failure": + raise ValueError("block failed") + raise ApplicationError("block failed") + except (ApplicationError, ValueError) as err: + if mode.startswith("caught"): + return "caught" + if mode == "then_non_retryable": + raise ApplicationError("run failed", non_retryable=True) from err + raise + + +async def test_block_end_closes_when_no_retry_run_follows(stream_client: Client): + # Each run ends with no retry run, so each block end closes its Stream, at once or at the + # run's end, with the block's exception. + retries = RetryPolicy( + initial_interval=timedelta(milliseconds=100), maximum_attempts=3 + ) + cases = { + "non_retryable": retries, + "last_attempt": RetryPolicy(maximum_attempts=1), + "listed_type": dataclasses.replace( + retries, non_retryable_error_types=["Listed"] + ), + "caught": retries, + "caught_task_failure": None, + "then_non_retryable": retries, + } + async with new_worker(stream_client, EndsItsBlock) as worker: + handles = { + mode: await stream_client.start_workflow( + EndsItsBlock.run, + mode, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=policy, + ) + for mode, policy in cases.items() + } + for mode, handle in handles.items(): + if mode.startswith("caught"): + assert await handle.result() == "caught" + else: + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + assert data(items) == [mode] + closed = items[-1] + assert isinstance(closed, common.Closed), mode + assert closed.status == common.CloseStatus.FAILED + assert isinstance(closed.failure, ApplicationError) + assert closed.failure.message == "block failed" + # A close at the run's end is a command after the run's code, even on replay. + replayer = Replayer( + workflows=[EndsItsBlock], + stream_backends=stream_client.config()["stream_backends"], + ) + for mode in ("caught", "then_non_retryable"): + await replayer.replay_workflow(await handles[mode].fetch_history()) + + +# The Activity of a test's run sets this when it starts, so the test cancels the run then. +waiting: dict[str, asyncio.Event] = {} + + +@activity.defn +async def waits_for_cancel() -> None: + waiting[activity.info().workflow_id or ""].set() + while True: + activity.heartbeat() + await asyncio.sleep(0.1) + + +@workflow.defn +class CancelledInItsBlock: + @workflow.run + async def run(self, local: bool) -> None: + async with workflow.create_stream(Events) as stream: + await stream.topic(Events.tokens).publish("started") + if local: + await workflow.execute_local_activity( + waits_for_cancel, start_to_close_timeout=timedelta(minutes=1) + ) + else: + await workflow.execute_activity( + waits_for_cancel, + start_to_close_timeout=timedelta(minutes=1), + heartbeat_timeout=timedelta(seconds=1), + ) + + +async def test_cancel_during_an_activity_closes_as_canceled(stream_client: Client): + # The cancel reaches the block as an ActivityError, and a cancelled run gets no retry run. + async with new_worker( + stream_client, CancelledInItsBlock, activities=[waits_for_cancel] + ) as worker: + for local, policy in ((False, RetryPolicy(maximum_attempts=3)), (True, None)): + workflow_id = f"wf-{uuid.uuid4()}" + waiting[workflow_id] = asyncio.Event() + handle = await stream_client.start_workflow( + CancelledInItsBlock.run, + local, + id=workflow_id, + task_queue=worker.task_queue, + retry_policy=policy, + ) + await asyncio.wait_for(waiting[workflow_id].wait(), 20) + await handle.cancel() + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + assert data(items) == ["started"] + closed = items[-1] + assert isinstance(closed, common.Closed), f"local={local}" + assert closed.status == common.CloseStatus.CANCELED + + +@workflow.defn +class CronBlock: + @workflow.run + async def run(self) -> str: + async with workflow.create_stream(Events) as stream: + await stream.topic(Events.tokens).publish("turn") + return "done" + + +async def test_block_end_leaves_a_cron_stream_open(stream_client: Client): + # The Stream spans the cron chain, so a close would fail the next run's create. + async with new_worker(stream_client, CronBlock) as worker: + handle = await stream_client.start_workflow( + CronBlock.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + cron_schedule="@every 1s", ) - await stream.topic(Events.tokens).publish("only") + reader = handle.stream(Events).subscribe(create_wait=timedelta(seconds=20)) + first = await asyncio.wait_for(anext(reader), 20) + second = await asyncio.wait_for(anext(reader), 20) + await reader.aclose() + await handle.terminate() + assert isinstance(first, common.Record) and isinstance(second, common.Record) + items = await read(handle) + assert not any(isinstance(item, common.Closed) for item in items) + + +failed_close_tasks: set[str] = set() -async def test_auto_close_closes_a_quiet_stream(stream_client: Client): - async with new_worker(stream_client, LeavesOpen) as worker: +@workflow.defn +class SchedulesACloseThenFailsItsTask: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("token") + stream.close_nowait(result="closed") + workflow_id = workflow.info().workflow_id + if workflow_id not in failed_close_tasks: + failed_close_tasks.add(workflow_id) + raise RuntimeError("fail the Workflow Task that schedules the close") + + +async def test_failed_task_writes_no_close(stream_client: Client): + # The test's set must keep its value across the task's retry, so the Workflow has no + # sandbox. + async with new_worker( + stream_client, + SchedulesACloseThenFailsItsTask, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: handle = await stream_client.start_workflow( - LeavesOpen.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + SchedulesACloseThenFailsItsTask.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, ) await handle.result() - await asyncio.sleep(1.5) - stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) - items = [item async for item in stream.subscribe(wait_new_records=False)] - assert data(items) == ["only"] - assert isinstance(items[-1], streams.Closed) - assert items[-1].result is None and items[-1].failure is None + items = await read(handle) + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and closes[0].result == "closed" + history = await handle.fetch_history() + # The failed task's close was never sent, so History has one close Activity. + scheduled = [ + event + for event in history.events + if event.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED + ] + assert [ + event.activity_task_scheduled_event_attributes.activity_type.name + for event in scheduled + ] == ["__temporal_streams_close"] @workflow.defn class CreatesOn: @workflow.run - async def run(self, backend: str | None, again_without_close: bool) -> None: - await stream_workflow.create_stream(Events, backend=backend) - if again_without_close: - await stream_workflow.create_stream(Events, close_with_run=False) + async def run(self, backend: str | None) -> None: + await workflow.create_stream(Events, backend=backend) async def test_worker_config_errors_fail_the_task( client: Client, stream_client: Client ): # A Worker with no plugin, or with another backend, fails the Workflow Task, so a Worker - # with the right config can take the run on. - for task_client, args, message in ( - (client, [None, False], "No Stream plugin"), - (stream_client, ["elsewhere", False], "backend elsewhere"), - (stream_client, [None, True], "close_with_run=False"), + # with the right config can continue the run. + for task_client, backend, message in ( + (client, None, "No Stream backend"), + (stream_client, "elsewhere", "backend elsewhere"), ): async with new_worker(task_client, CreatesOn) as worker: handle = await task_client.start_workflow( CreatesOn.run, - args=args, + backend, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -332,14 +539,13 @@ async def test_worker_config_errors_fail_the_task( class CronTurn: @workflow.run async def run(self) -> str: - stream = await stream_workflow.create_stream(Events) + stream = await workflow.create_stream(Events) await stream.topic(Events.tokens).publish("turn") return "done" async def test_cron_runs_share_an_open_stream(stream_client: Client): - # A cron chain goes on after each run, so no run closes the chain's Stream, and each run - # writes to it again. + # No run end closes the chain's Stream, so each cron run writes to it again. async with new_worker(stream_client, CronTurn) as worker: handle = await stream_client.start_workflow( CronTurn.run, @@ -347,26 +553,24 @@ async def test_cron_runs_share_an_open_stream(stream_client: Client): task_queue=worker.task_queue, cron_schedule="@every 1s", ) - stream = streams.StreamClient(stream_client).workflow_stream(handle, Events) + stream = handle.stream(Events) reader = stream.subscribe(create_wait=timedelta(seconds=20)) first = await reader.__anext__() second = await reader.__anext__() await reader.aclose() await handle.terminate() - assert isinstance(first, streams.Record) and first.data == "turn" - assert isinstance(second, streams.Record) and second.data == "turn" + assert isinstance(first, common.Record) and first.data == "turn" + assert isinstance(second, common.Record) and second.data == "turn" items = [item async for item in stream.subscribe(wait_new_records=False)] - assert not any(isinstance(item, streams.Closed) for item in items) + assert not any(isinstance(item, common.Closed) for item in items) @workflow.defn class SyncHandlers: @workflow.init def __init__(self, prompt: str) -> None: - # The create goes first, so the signal's writes find the Stream. - self._stream = stream_workflow.create_stream_nowait( - Events, close_with_run=False - ) + # The create runs first, so the Stream exists for the signal's writes. + self._stream = workflow.create_stream_nowait(Events) self._prompt = prompt self._done = False @@ -384,7 +588,7 @@ def finish(self) -> str: self._stream.close_nowait(result="finished") try: self._stream.topic(Events.tokens).publish_nowait("late") - except streams.StreamClosedError as err: + except common.StreamClosedError as err: return err.reason finally: self._done = True @@ -399,21 +603,22 @@ async def test_sync_handlers_publish_and_close_with_no_await(stream_client: Clie await handle.signal(SyncHandlers.token, "a") assert await handle.execute_update(SyncHandlers.finish) == "stream_closed" await handle.result() - items = await read(stream_client, handle) + items = await read(handle) assert data(items) == ["a", "p"] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "finished" + assert isinstance(items[-1], common.Closed) and items[-1].result == "finished" @workflow.defn class CreatesInInit: @workflow.init def __init__(self, prompt: str) -> None: - self._stream = stream_workflow.create_stream_nowait(Events) + self._stream = workflow.create_stream_nowait(Events) self._stream.topic(Events.tokens).publish_nowait(prompt) @workflow.run async def run(self, prompt: str) -> str: await self._stream.topic(Events.tokens).publish(prompt.upper()) + await self._stream.close(result="done") return "done" @@ -426,16 +631,113 @@ async def test_init_creates_and_publishes_with_no_await(stream_client: Client): task_queue=worker.task_queue, ) assert await handle.result() == "done" - items = await read(stream_client, handle) + items = await read(handle) assert data(items) == ["hi", "HI"] - assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" + assert isinstance(items[-1], common.Closed) and items[-1].result == "done" + + +@workflow.defn +class PublishesWithNoAwait: + @workflow.run + async def run(self) -> list[int]: + tokens = (await workflow.create_stream(Events)).topic(Events.tokens) + first = tokens.publish("a") + second = tokens.publish("b") + last = await tokens.publish("c") + # Nobody awaits this write, and the run's end still waits for it. + tokens.publish("d") + return [first.result(), second.result(), last] + + +async def test_publish_starts_the_write_and_returns_its_future(stream_client: Client): + async with new_worker(stream_client, PublishesWithNoAwait) as worker: + handle = await stream_client.start_workflow( + PublishesWithNoAwait.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + offsets = await handle.result() + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] + assert [r.data for r in written] == ["a", "b", "c", "d"] + assert offsets == [r.offset for r in written[:3]] + + +@workflow.defn +class LeavesAFailedPublish: + @workflow.run + async def run(self, stream_id: str, retrieve: bool) -> str: + # Nobody created the Stream, so the backend refuses the write. + missing = common.WorkflowStream(id=stream_id).topic(Events.tokens) + future = missing.publish("lost") + errors: list[BaseException | None] = [] + if retrieve: + future.add_done_callback(lambda done: errors.append(done.exception())) + await workflow.wait_condition(lambda: bool(errors)) + return "done" + + +async def test_publish_error_that_no_caller_gets_raises_at_the_run_end( + stream_client: Client, +): + # The error raises as if the run didn't catch it, as with publish_nowait. + async with new_worker(stream_client, LeavesAFailedPublish) as worker: + lost = await stream_client.start_workflow( + LeavesAFailedPublish.run, + args=[f"missing-{uuid.uuid4()}", False], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + with pytest.raises(WorkflowFailureError) as failed: + await lost.result() + assert isinstance(failed.value.cause, ApplicationError) + assert failed.value.cause.type == "StreamNotFoundError" + # A done callback that reads the error owns it, so the run completes. + owned = await stream_client.start_workflow( + LeavesAFailedPublish.run, + args=[f"missing-{uuid.uuid4()}", True], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await owned.result() == "done" + + +def test_asyncio_flags_a_future_error_that_no_caller_got(): + # The run's end reads this asyncio flag, so pin it for both Future classes. + loop = asyncio.new_event_loop() + try: + for future_class in (asyncio.Future, getattr(asyncio.futures, "_PyFuture")): + unread: asyncio.Future[int] = future_class(loop=loop) + unread.set_exception(ValueError("refused")) + assert _workflow._unretrieved(unread) + unread.exception() + assert not _workflow._unretrieved(unread) + by_result: asyncio.Future[int] = future_class(loop=loop) + by_result.set_exception(ValueError("refused")) + with pytest.raises(ValueError): + by_result.result() + assert not _workflow._unretrieved(by_result) + awaited: asyncio.Future[int] = future_class(loop=loop) + loop.call_soon(awaited.set_exception, ValueError("refused")) + + async def wait(future: asyncio.Future[int]) -> None: + with pytest.raises(ValueError): + await future + + loop.run_until_complete(wait(awaited)) + assert not _workflow._unretrieved(awaited) + cancelled: asyncio.Future[int] = future_class(loop=loop) + cancelled.cancel() + assert not _workflow._unretrieved(cancelled) + finally: + loop.close() @workflow.defn class Gathers: @workflow.run async def run(self) -> list[int]: - tokens = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + tokens = (await workflow.create_stream(Events)).topic(Events.tokens) tokens.publish_nowait("0") # The first write starts, so the next ones wait behind it. await asyncio.sleep(0) @@ -450,13 +752,13 @@ async def test_concurrent_publishes_keep_their_order(stream_client: Client): Gathers.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) offsets = await handle.result() - items = await read(stream_client, handle) + items = await read(handle) assert data(items) == ["0", "1", "2", "3", "4"] - # The four writes that waited went as one Local Activity, and each call got the offset of - # its own Record. + # The four writes that waited were sent as one Local Activity, and each call got the offset + # of its own Record. assert ( offsets - == [item.offset for item in items if isinstance(item, streams.Record)][1:] + == [item.offset for item in items if isinstance(item, common.Record)][1:] ) history = await handle.fetch_history() markers = [ @@ -464,17 +766,13 @@ async def test_concurrent_publishes_keep_their_order(stream_client: Client): for e in history.events if e.event_type == EventType.EVENT_TYPE_MARKER_RECORDED ] - # The create, the first write and the four that waited. The worker writes the close with - # the run after the run, so it has no marker. + # The create, the first write and the four that waited. assert len(markers) == 3 - # Replay answers from the markers and writes nothing. - plugins = [ - p - for p in stream_client.config()["plugins"] - if isinstance(p, streams.StreamPlugin) - ] - await Replayer(workflows=[Gathers], plugins=plugins).replay_workflow(history) - assert data(await read(stream_client, handle)) == ["0", "1", "2", "3", "4"] + # Replay uses the markers and writes nothing. + await Replayer( + workflows=[Gathers], stream_backends=stream_client.config()["stream_backends"] + ).replay_workflow(history) + assert data(await read(handle)) == ["0", "1", "2", "3", "4"] class _RandomCodec(PayloadCodec): @@ -495,7 +793,7 @@ async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: class _GrowingCodec(PayloadCodec): - # Triples each payload, so a merge under the Workflow's cap passes the Store's batch limit + # Triples each payload, so a merge under the Workflow's cap exceeds the Store's batch limit # after the codec. async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: return [ @@ -539,7 +837,7 @@ def with_codec(client: Client, codec: PayloadCodec) -> Client: async def test_retried_write_goes_in_once( stream_client: Client, monkeypatch: pytest.MonkeyPatch ): - original = _ExtensionStreamBackend.call_async + original = _ExtensionCalls.call_async sent: list[list[bytes]] = [] async def loses_the_first_answer( @@ -552,25 +850,70 @@ async def loses_the_first_answer( raise RuntimeError("the answer is lost") return answer - monkeypatch.setattr(_ExtensionStreamBackend, "call_async", loses_the_first_answer) + monkeypatch.setattr(_ExtensionCalls, "call_async", loses_the_first_answer) client = with_codec(stream_client, _RandomCodec()) async with new_worker(client, Gathers) as worker: handle = await client.start_workflow( Gathers.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) await handle.result() - items = await read(client, handle) + items = await read(handle) assert data(items) == ["0", "1", "2", "3", "4"] - # The retry sends the bytes of the first attempt, so the Store's dedupe takes it as the - # same batch, though the codec adds random bytes. + # The retry sends the bytes of the first attempt, so the Store's dedupe matches it as the + # same batch, even though the codec adds random bytes. assert sent[0] == sent[1] +@workflow.defn +class PublishesToAMissingStream: + @workflow.run + async def run(self) -> None: + # No create, so the Store refuses the append, and no retry sends it again. + stream = workflow.stream_handle(Events) + await stream.topic(Events.tokens).publish("x" * 100_000) + + +async def test_refused_write_keeps_no_pending_bytes(stream_client: Client): + async with new_worker(stream_client, PublishesToAMissingStream) as worker: + handle = await stream_client.start_workflow( + PublishesToAMissingStream.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + with pytest.raises(WorkflowFailureError): + await handle.result() + stream_id = handle.stream(Events).id + assert not [key for key in _writes._PENDING._entries if key[1] == stream_id] + + +def _entry(size: int) -> Any: + return _writes._Encoded([], size) + + +def test_pending_appends_keep_the_newest_epoch_and_a_bounded_size(): + pending = _writes._Pending(100) + one, two = ("ns", "s", "la:1"), ("ns", "s", "la:2") + first, retried = _entry(10), _entry(10) + pending.put(one, (5, 1, 0, 0), first) + # A rerun under a later epoch replaces the earlier one. + pending.put(one, (9, 1, 0, 0), retried) + assert pending.get(one, (5, 1, 0, 0)) is None + assert pending.get(one, (9, 1, 0, 0)) is retried + # A write of an earlier epoch is fenced, so its put and its pop don't change the later entry. + pending.put(one, (5, 1, 0, 0), first) + pending.pop(one, (5, 1, 0, 0)) + assert pending.get(one, (9, 1, 0, 0)) is retried + # Past the size limit, the oldest entry is evicted. + pending.put(two, (9, 1, 0, 0), _entry(95)) + assert pending.get(one, (9, 1, 0, 0)) is None + assert pending.get(two, (9, 1, 0, 0)) is not None + + @workflow.defn class PublishesLargeValues: @workflow.run async def run(self, size: int) -> list[int]: - tokens = (await stream_workflow.create_stream(Events)).topic(Events.tokens) + tokens = (await workflow.create_stream(Events)).topic(Events.tokens) return list( await asyncio.gather( tokens.publish(*["a" * size] * 3), tokens.publish(*["b" * size] * 3) @@ -579,8 +922,8 @@ async def run(self, size: int) -> list[int]: async def test_merge_that_the_codec_grows_past_the_batch_limit(stream_client: Client): - # The two publishes merge under the Workflow's cap, and the codec makes the merge larger - # than one batch may be, so each publish goes as its own batch. + # The two publishes merge under the Workflow's cap. The codec makes the merge larger than the + # batch limit, so each publish is sent as its own batch. client = with_codec(stream_client, _GrowingCodec()) async with new_worker(client, PublishesLargeValues) as worker: handle = await client.start_workflow( @@ -590,7 +933,124 @@ async def test_merge_that_the_codec_grows_past_the_batch_limit(stream_client: Cl task_queue=worker.task_queue, ) offsets = await handle.result() - items = await read(client, handle) - written = [item for item in items if isinstance(item, streams.Record)] + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] assert [len(r.data) for r in written] == [300_000] * 6 assert offsets == [written[2].offset, written[5].offset] + + +@workflow.defn +class PublishesTooMuch: + @workflow.run + async def run(self) -> str: + tokens = (await workflow.create_stream(Events)).topic(Events.tokens) + try: + await tokens.publish(*["x" * 1_400_000] * 4) + except common.StreamTooLargeError as err: + return err.reason + return "written" + + +async def test_batch_over_the_store_limit_raises_too_large(stream_client: Client): + # Each Record is under the Store's Record limit, and the publish's batch is over 4 MiB. + async with new_worker(stream_client, PublishesTooMuch) as worker: + handle = await stream_client.start_workflow( + PublishesTooMuch.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == "stream_too_large" + assert data(await read(handle)) == [] + + +@workflow.defn +class CreatesOnTheDefault: + @workflow.run + async def run(self) -> None: + await workflow.create_stream(Events) + + +@workflow.defn +class CreatesOnRedis: + @workflow.run + async def run(self) -> str: + stream = await redis.workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("on redis") + await redis.workflow.stream_handle(Events).close(result="closed") + return stream.backend or "" + + +@workflow.defn +class CreatesAfterAnIdleTime: + @workflow.run + async def run(self) -> str: + await workflow.create_stream(Events) + await workflow.sleep(timedelta(seconds=1.5)) + try: + await workflow.create_stream(Events) + except common.StreamExpiredError as err: + return err.reason + return "created" + + +async def test_create_of_an_expired_workflow_stream_raises(client: Client): + # An open Workflow Stream expires one safety TTL after its last write. + config = client.config() + config["stream_backends"] = [ + MemoryStreamBackend(default=True, safety_ttl=timedelta(milliseconds=500)) + ] + expiring = Client(**config) + async with new_worker(expiring, CreatesAfterAnIdleTime) as worker: + assert ( + await expiring.execute_workflow( + CreatesAfterAnIdleTime.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "stream_expired" + ) + + +async def test_default_is_native_with_no_default_backend(client: Client): + # With no backend marked default=True, the default is native, which this SDK lacks. The + # Redis functions still pick the registered Redis backend, and the memory backend counts as one. + config = client.config() + config["stream_backends"] = [MemoryStreamBackend()] + memory_client = Client(**config) + async with new_worker(memory_client, CreatesOnTheDefault, CreatesOnRedis) as worker: + handle = await memory_client.start_workflow( + CreatesOnTheDefault.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually(handle, message_contains="default is native") + await handle.terminate() + on_redis = await memory_client.start_workflow( + CreatesOnRedis.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await on_redis.result() == "memory" + with pytest.raises(common.StreamConfigError, match="default=True"): + on_redis.stream(Events) + stream = on_redis.stream(Events, backend="memory") + items = [item async for item in stream.subscribe()] + assert data(items) == ["on redis"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "closed" + + +class _RecordsCloses(MemoryStreamBackend): + # A backend that closes another way overrides the close of a Workflow's Stream. + closes: list[str] = [] + + async def _close_in_workflow(self, request: Any) -> int: + self.closes.append(request.stream.id) + return await super()._close_in_workflow(request) + + +async def test_backend_carries_the_workflow_close(client: Client): + config = client.config() + config["stream_backends"] = [_RecordsCloses(default=True)] + recording = Client(**config) + async with new_worker(recording, Agent) as worker: + handle = await recording.start_workflow( + Agent.run, "hi", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + assert _RecordsCloses.closes == [handle.stream(Events).id] diff --git a/tests/test_converter.py b/tests/test_converter.py index d21d9193ab..97618e2904 100644 --- a/tests/test_converter.py +++ b/tests/test_converter.py @@ -512,6 +512,18 @@ def test_transfer_type_values_nested_in_json(): assert decoded.proto == TemporalTransferTypeValueWithoutHint("c") +def test_transfer_type_dataclass_in_a_top_level_list_keeps_its_fields(): + # The decode of a dataclass inside a JSON value reads its fields, so the encode writes them. + converter = DataConverter.default.payload_converter + value = [TemporalTransferTypeValueWithoutHint("c")] + payload = converter.to_payload(value) + assert json.loads(payload.data) == [{"value": "c"}] + decoded = converter.from_payload( + payload, list[TemporalTransferTypeValueWithoutHint] + ) + assert decoded == value + + def test_transfer_type_convertible_rejects_existing_converter(): with pytest.raises(TypeError, match="already has a transfer type converter"): transfer_type_convertible(TemporalTransferTypeValueConverter)( diff --git a/tests/worker/test_worker_exit.py b/tests/worker/test_worker_exit.py deleted file mode 100644 index 737639f475..0000000000 --- a/tests/worker/test_worker_exit.py +++ /dev/null @@ -1,80 +0,0 @@ -import asyncio -import uuid -from collections.abc import Awaitable, Callable -from datetime import timedelta - -from temporalio import workflow -from temporalio.client import Client -from temporalio.plugin import SimplePlugin -from temporalio.worker import Worker - - -@workflow.defn -class NothingWorkflow: - @workflow.run - async def run(self) -> None: - pass - - -class CleanupPlugin(SimplePlugin): - def __init__(self, client: Client) -> None: - super().__init__("test-cleanup") - self.client = client - self.cleanup_result: int | None = None - - async def run_worker( - self, worker: Worker, next: Callable[[Worker], Awaitable[None]] - ) -> None: - try: - await super().run_worker(worker, next) - finally: - # A call after the Worker stops, as a plugin's cleanup makes. - self.cleanup_result = (await self.client.count_workflows()).count - - -async def test_plugin_cleanup_awaits_after_worker_exit(client: Client): - plugin = CleanupPlugin(client) - async with Worker( - client, - task_queue=f"tq-{uuid.uuid4()}", - workflows=[NothingWorkflow], - plugins=[plugin], - ): - pass - assert plugin.cleanup_result is not None - # The caller can still await after the block. - await client.count_workflows() - - -class StuckCleanupPlugin(SimplePlugin): - def __init__(self) -> None: - super().__init__("test-stuck-cleanup") - self.cancelled = asyncio.Event() - - async def run_worker( - self, worker: Worker, next: Callable[[Worker], Awaitable[None]] - ) -> None: - try: - await super().run_worker(worker, next) - finally: - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - self.cancelled.set() - raise - - -async def test_plugin_cleanup_is_cancelled_after_plugin_shutdown_timeout( - client: Client, -): - plugin = StuckCleanupPlugin() - async with Worker( - client, - task_queue=f"tq-{uuid.uuid4()}", - workflows=[NothingWorkflow], - plugins=[plugin], - plugin_shutdown_timeout=timedelta(milliseconds=100), - ): - pass - # The cancel reaches the cleanup code on a later loop iteration, so wait for it. - await asyncio.wait_for(plugin.cancelled.wait(), timeout=10)