From 50537f75466f274a97ed4c4d15dcc100214e3457 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 14:41:16 -0700 Subject: [PATCH 1/2] Generated the Stream service and writer bindings with nexgen. No nexgen release has the long-poll, client and client-only features, so the generator installs nexgen at a fork revision. `poe gen-protos` runs it, so `check-protos` checks the bindings for drift. --- .github/workflows/ci.yml | 2 + .gitignore | 3 + pyproject.toml | 11 +- scripts/_wit/stream-service.wit | 435 +++++ scripts/_wit/stream-writer.wit | 118 ++ scripts/gen_stream_service.py | 188 +++ temporalio/common/_streams/system/__init__.py | 85 + .../_streams/system/_resources/__init__.py | 9 + .../_streams/system/_resources/stream.py | 341 ++++ .../system/_system_nexus_interceptor.py | 265 +++ temporalio/common/_streams/system/client.py | 439 +++++ temporalio/common/_streams/system/models.py | 1439 +++++++++++++++++ .../_streams/system/operations/__init__.py | 3 + .../system/operations/create_stream.py | 64 + temporalio/common/_streams/system/services.py | 51 + .../common/_streams/writer_system/__init__.py | 49 + .../common/_streams/writer_system/client.py | 256 +++ .../common/_streams/writer_system/models.py | 893 ++++++++++ .../common/_streams/writer_system/services.py | 51 + temporalio/nexus/system/_support.py | 45 + tests/contrib/streams/__init__.py | 0 tests/contrib/streams/test_bindings.py | 91 ++ 22 files changed, 4837 insertions(+), 1 deletion(-) create mode 100644 scripts/_wit/stream-service.wit create mode 100644 scripts/_wit/stream-writer.wit create mode 100644 scripts/gen_stream_service.py create mode 100644 temporalio/common/_streams/system/__init__.py create mode 100644 temporalio/common/_streams/system/_resources/__init__.py create mode 100644 temporalio/common/_streams/system/_resources/stream.py create mode 100644 temporalio/common/_streams/system/_system_nexus_interceptor.py create mode 100644 temporalio/common/_streams/system/client.py create mode 100644 temporalio/common/_streams/system/models.py create mode 100644 temporalio/common/_streams/system/operations/__init__.py create mode 100644 temporalio/common/_streams/system/operations/create_stream.py create mode 100644 temporalio/common/_streams/system/services.py create mode 100644 temporalio/common/_streams/writer_system/__init__.py create mode 100644 temporalio/common/_streams/writer_system/client.py create mode 100644 temporalio/common/_streams/writer_system/models.py create mode 100644 temporalio/common/_streams/writer_system/services.py create mode 100644 tests/contrib/streams/__init__.py create mode 100644 tests/contrib/streams/test_bindings.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 32f405d5d9..f498f08f09 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -230,6 +230,8 @@ jobs: - run: uv add --python 3.10 "protobuf<4" - run: uv sync --all-extras - run: cargo install --locked nexgen --version 0.2.8 --features advanced --force + # The Stream service bindings need a nexgen revision that no release has yet. + - run: uv run scripts/gen_stream_service.py --install-nexgen - run: poe build-develop - run: poe gen-protos - name: Check generation unchanged diff --git a/.gitignore b/.gitignore index 8cd439e052..f16269ea96 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,6 @@ temporalio/bridge/temporal_sdk_bridge* tags /.claude tmpclaude-* + +# nexgen builds that scripts/gen_stream_service.py installs +/.nexgen-streams/ diff --git a/pyproject.toml b/pyproject.toml index 1300699a27..23f366dbd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -109,9 +109,11 @@ format = [ ] gen-docs = "uv run scripts/gen_docs.py" gen-nexus-system-api = "uv run scripts/gen_nexus_system_api.py" +gen-stream-service = "uv run scripts/gen_stream_service.py" gen-protos = [ { cmd = "uv run scripts/gen_protos.py" }, { ref = "gen-nexus-system-api" }, + { ref = "gen-stream-service" }, { cmd = "uv run scripts/gen_payload_visitor.py" }, { cmd = "uv run scripts/gen_bridge_client.py" }, { ref = "format" }, @@ -119,6 +121,7 @@ gen-protos = [ gen-protos-docker = [ { cmd = "uv run scripts/gen_protos_docker.py" }, { ref = "gen-nexus-system-api" }, + { ref = "gen-stream-service" }, { cmd = "uv run scripts/gen_payload_visitor.py" }, { cmd = "uv run scripts/gen_bridge_client.py" }, { ref = "format" }, @@ -184,6 +187,8 @@ exclude = [ 'temporalio/bridge/proto', 'temporalio/nexus/system/workflow_service', 'temporalio/common/_streams/proto', + 'temporalio/common/_streams/system', + 'temporalio/common/_streams/writer_system', ] [[tool.mypy.overrides]] @@ -193,13 +198,15 @@ ignore_errors = true [[tool.mypy.overrides]] module = [ "temporalio.common._streams.proto.*", + "temporalio.common._streams.system.*", + "temporalio.common._streams.writer_system.*", ] ignore_errors = true [tool.pydocstyle] convention = "google" # https://github.com/PyCQA/pydocstyle/issues/363#issuecomment-625563088 -match_dir = "^(?!(docs|scripts|tests|api|proto|system|_support|notifications|\\.)).*" +match_dir = "^(?!(docs|scripts|tests|api|proto|system|writer_system|_support|notifications|\\.)).*" add_ignore = [ # We like to wrap at a certain number of chars, even long summary sentences. # https://github.com/PyCQA/pydocstyle/issues/184 @@ -232,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" diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit new file mode 100644 index 0000000000..480e8ddfed --- /dev/null +++ b/scripts/_wit/stream-service.wit @@ -0,0 +1,435 @@ +// Draft. The Stream service as a System Nexus service, for the Brook API proposal. Each record maps +// to a message of the protos under ../protos/temporal/api (stream/v1, streamservice/v1 and the +// Stream enums in enums/v1). This file is the source. Each SDK copies it and generates its Workflow +// API and its client with nexgen (`--targets workflow,client --system-nexus`). +package temporal:nexus@1.0.0; + +world system { + export stream-service; +} + +/// @nexus.endpoint "__temporal_system" +/// @nexus.service-name "temporal.api.streamservice.v1.StreamService" +/// @nexus.experimental +interface stream-service { + use nexus:temporal-types/model@1.0.0.{ + duration, + failure, + payload, + placeholder, + }; + + /// @nexus.proto "google.protobuf.Timestamp" + record timestamp { + seconds: option, + nanos: option, + } + + // A Stream that a Workflow creates has the id ::. Each part is + // percent-encoded. So a caller that knows the first run id can build the id with no call. + /// @nexus.doc "Identifies one Stream. The id is opaque to callers." + /// @nexus.proto "temporal.api.stream.v1.StreamRef" + record stream-ref { + /// @nexus.doc "The Stream id. 1 to 4096 bytes after encoding." + id: string, + /// @nexus.doc "The owning Workflow run chain or Standalone Activity. Empty on a client ref." + owner: option, + /// @nexus.doc "Registered backend name, empty for the default. The Store never reads it." + backend: option, + } + + /// @nexus.proto "temporal.api.enums.v1.ExecutionType" + enum execution-type { + execution-type-unspecified, + execution-type-workflow, + execution-type-activity, + execution-type-nexus-operation, + } + + /// @nexus.doc "A Workflow run chain or a Standalone Activity run." + /// @nexus.proto "temporal.api.common.v1.Execution" + record execution { + %type: option, + /// @nexus.doc "The Workflow id, or the Standalone Activity's id." + business-id: option, + /// @nexus.doc "The run chain's first run id, or the Standalone Activity's run id." + run-id: option, + } + + /// @nexus.proto "temporal.api.enums.v1.StreamProducerKind" + enum stream-producer-kind { + stream-producer-kind-unspecified, + stream-producer-kind-workflow, + stream-producer-kind-local-activity, + stream-producer-kind-activity, + stream-producer-kind-client, + stream-producer-kind-standalone-activity, + } + + // Epochs compare field by field, in the order of the fields. A field that a kind does not use is + // empty and compares as zero. A timestamp compares seconds, then nanos. A Workflow epoch has + // task-started-event-id, task-attempt and task-started-time. A Local Activity epoch has the same + // three of the dispatching Workflow Task, then attempt. An Activity epoch has + // attempt-scheduled-time, then attempt. A client epoch has attempt. + /// @nexus.doc "The ordered attempt of a Producer. A higher epoch supersedes a lower one." + /// @nexus.proto "temporal.api.stream.v1.ProducerEpoch" + record producer-epoch { + /// @nexus.doc "The started event id of the Workflow Task." + task-started-event-id: option, + /// @nexus.doc "The attempt of the Workflow Task." + task-attempt: option, + /// @nexus.doc "When the Workflow Task attempt started. It tells a reused started id apart." + task-started-time: option, + /// @nexus.doc "When the server scheduled this Activity attempt. A reset makes it later." + attempt-scheduled-time: option, + /// @nexus.doc "The Activity, Local Activity or client attempt." + attempt: option, + } + + /// @nexus.doc "The Producer of a batch." + /// @nexus.proto "temporal.api.stream.v1.Producer" + record producer { + /// @nexus.doc "Stable across attempts. The proposal gives the form per kind." + id: string, + kind: option, + epoch: option, + /// @nexus.doc "The run of the Producer, so nobody parses the id. Empty for a client." + run-id: option, + } + + /// @nexus.doc "A Workflow Task attempt of the run." + /// @nexus.proto "temporal.api.stream.v1.WorkflowTaskEpoch" + record workflow-task-epoch { + started-event-id: option, + attempt: option, + started-time: option, + } + + /// @nexus.doc "One Record to append." + /// @nexus.proto "temporal.api.stream.v1.AppendRecord" + record append-record { + /// @nexus.doc "The value. It goes through the payload codec." + body: payload, + } + + // An Activity or Local Activity item withdraws the Records of its Producer with a lower epoch. A + // Workflow item, with this task t in epoch and the previous completed task c, withdraws the + // Records of the run and of its Local Activities whose task-started-event-id is in failed-tasks + // or strictly between c and t. It also withdraws those at the started event id of c or of t with + // a lower epoch than c or t. No task in that range can have completed. + /// @nexus.doc "Which earlier Records a Superseded item withdraws." + /// @nexus.proto "temporal.api.stream.v1.Supersession" + record supersession { + producer-id: string, + /// @nexus.doc "Workflow only. Started event ids of failed or timed out Workflow Tasks." + failed-tasks: list, + producer-kind: option, + epoch: option, + /// @nexus.doc "Workflow only. The previous completed Workflow Task of the run." + previous-completed-task: option, + /// @nexus.doc "Workflow retry only. The previous run, all of whose Records are withdrawn." + withdraw-run: option, + /// @nexus.doc "The withdrawing Producer's run, so no Reader parses ids. Empty for a client." + run-id: option, + } + + /// @nexus.doc "How a Stream closed. Only the owner's end on the server terminates or times out." + /// @nexus.proto "temporal.api.enums.v1.StreamCloseStatus" + enum stream-close-status { + stream-close-status-unspecified, + stream-close-status-completed, + stream-close-status-failed, + stream-close-status-canceled, + stream-close-status-terminated, + stream-close-status-timed-out, + } + + /// @nexus.doc "One Record that a Producer appended." + /// @nexus.proto "temporal.api.stream.v1.StreamData" + record stream-data { + /// @nexus.doc "The value. It goes through the payload codec." + body: payload, + producer: option, + } + + /// @nexus.doc "The close of a Stream. A read gives it last." + /// @nexus.proto "temporal.api.stream.v1.StreamClose" + record stream-close { + /// @nexus.doc "The Producer that closed the Stream." + producer: option, + %result: option, + failure: option, + /// @nexus.doc "Set together with the result or the failure." + status: option, + } + + // The read builds a Superseded item from a flagged batch. + /// @nexus.doc "A read item: a Record, a Superseded item or the close." + /// @nexus.proto "temporal.api.stream.v1.StreamRecord.item" + variant stream-item { + data(stream-data), + superseded(supersession), + close(stream-close), + } + + // An offset is the Redis entry id packed as ms << 20 | seq. It is ordered and unique in one + // Stream, and sparse. Offsets have gaps, so they don't count Records. A read names its Stream, so + // the read checks the offset against that Stream. + /// @nexus.doc "One item that a read gives." + /// @nexus.proto "temporal.api.stream.v1.StreamRecord" + record stream-record { + item: stream-item, + /// @nexus.doc "The place of this item in the Stream. Ordered and unique, with gaps." + offset: option, + } + + /// @nexus.doc "Create a Stream, or get the one that exists, which keeps its Retention." + /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamRequest" + record create-stream-request { + /// @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 "The Workflow run that owns the Stream. The Workflow transport sets it." + owner: option, + /// @nexus.doc "The Stream's ref, with the owner fields." + %ref: option, + // Native only. Brook's Stores don't read it. + /// @nexus.omit + close-with-run: placeholder, + /// @nexus.omit + namespace: placeholder, + // Native only. Brook's Stores don't read it. + /// @nexus.omit + task-token: placeholder, + } + + /// @nexus.proto "temporal.api.enums.v1.StreamStatus" + enum stream-status { + stream-status-unspecified, + stream-status-open, + stream-status-closed, + stream-status-expired, + } + + // The meta key of each Stream has a format version v. Lua refuses a write to a Stream with a + // higher v with FAILED_PRECONDITION and the detail stream_format_too_new. + /// @nexus.doc + /// "Append. FAILED_PRECONDITION if closed, fenced or format too new. NOT_FOUND if absent." + /// @nexus.proto "temporal.api.streamservice.v1.AppendRecordsRequest" + record append-records-request { + %stream: stream-ref, + records: list, + /// @nexus.doc "Local Activity only. The run's failed or timed out Workflow Tasks." + failed-tasks: list, + /// @nexus.doc "Required outside a Workflow. In a Workflow the transport sets it." + producer: option, + /// @nexus.doc "The batch number in this epoch, 1 or more. It is part of the dedupe key." + sequence: option, + /// @nexus.doc "Local Activity only. The run's last completed Workflow Task." + previous-completed-task: option, + /// @nexus.doc "Retry run only. The previous run, whose Records this withdraws." + withdraw-run: option, + /// @nexus.omit + namespace: placeholder, + // Native only. Brook's Stores don't read them. + /// @nexus.omit + defer-withdrawal: placeholder, + /// @nexus.omit + task-token: placeholder, + } + + /// @nexus.proto "temporal.api.streamservice.v1.AppendRecordsResponse" + record append-records-response { + first-offset: option, + last-offset: option, + } + + /// @nexus.doc "Close a Stream. Readers drain what is left, then end. Later appends are refused." + /// @nexus.proto "temporal.api.streamservice.v1.CloseStreamRequest" + record close-stream-request { + %stream: stream-ref, + /// @nexus.doc "Workflow only. The run's failed or timed out Workflow Tasks." + failed-tasks: list, + producer: option, + %result: option, + failure: option, + /// @nexus.doc "Empty is completed, or failed with a failure. A caller may set canceled." + status: option, + /// @nexus.doc "Workflow only. The previous completed Workflow Task of the run." + previous-completed-task: option, + /// @nexus.omit + namespace: placeholder, + // Native only. Brook's Stores don't read them. + /// @nexus.omit + withdraw-run: placeholder, + /// @nexus.omit + task-token: placeholder, + } + + /// @nexus.proto "temporal.api.streamservice.v1.CloseStreamResponse" + record close-stream-response { + } + + /// @nexus.proto "temporal.api.enums.v1.StreamReadStart" + enum stream-read-start { + stream-read-start-unspecified, + stream-read-start-beginning, + stream-read-start-end, + } + + /// @nexus.doc "Read Records from an offset." + /// @nexus.proto "temporal.api.streamservice.v1.ReadRecordsRequest" + record read-records-request { + %stream: stream-ref, + /// @nexus.doc "Start offset. Empty uses start. INVALID_ARGUMENT past the newest offset + 1." + from-offset: option, + /// @nexus.doc "Where to start when from-offset is empty. Unspecified is the beginning." + start: option, + /// @nexus.doc "At most this many Records. 1 to 1000." + max-records: option, + /// @nexus.doc "How long to wait for a new Record before an empty answer. At most 60 s." + wait: option, + /// @nexus.doc "How long to wait for an uncreated Stream. Empty fails at once. At most 60 s." + create-wait: option, + /// @nexus.doc "Read the newest this many Records, oldest first. No from-offset. No wait." + last: option, + /// @nexus.doc "The maximum Record bytes in a page. A larger Record is still returned alone." + max-bytes: option, + /// @nexus.omit + namespace: placeholder, + // Native only. Brook's Stores refuse to-offset and wake-only. + /// @nexus.omit + to-offset: placeholder, + /// @nexus.omit + wake-only: placeholder, + /// @nexus.omit + floor-from: placeholder, + } + + /// @nexus.proto "temporal.api.streamservice.v1.ReadRecordsResponse" + record read-records-response { + records: list, + /// @nexus.doc "The from-offset of the next read, which is the newest offset read + 1." + next-offset: option, + /// @nexus.doc "True after the close Record. No later read gives a Record." + done: option, + // Native only. Brook's Stores never set it. + /// @nexus.omit + continue-as-new-suggested: placeholder, + } + + /// @nexus.proto "temporal.api.streamservice.v1.DescribeStreamRequest" + record describe-stream-request { + %stream: stream-ref, + /// @nexus.omit + namespace: placeholder, + } + + /// @nexus.proto "temporal.api.streamservice.v1.DescribeStreamResponse" + record describe-stream-response { + status: option, + /// @nexus.doc "The offset of the newest Record. Empty when the Stream has no Record." + last-offset: option, + retention: option, + } + + // A handle never changes, since offsets are values that its methods return. Only a create's + // response sets created and status, so a handle built from a ref has neither. + /// @nexus.doc "A handle on one Stream. Build it from a stream-ref, or get it from a create." + resource %stream { + constructor(%stream: stream-ref, created: option, status: option); + + /// @nexus.doc "Publish Records. The answer has the first and the last offset." + /// @nexus.operation "append-records" + publish: func( + producer: option, + sequence: option, + records: list, + previous-completed-task: option, + failed-tasks: list, + withdraw-run: option, + ) -> append-records-response; + + /// @nexus.doc "Close the Stream. Readers drain what is left, then end." + /// @nexus.operation "close-stream" + close: func( + producer: option, + %result: option, + failure: option, + status: option, + previous-completed-task: option, + failed-tasks: list, + ); + + /// @nexus.doc "Read one page of Records. The poll iterator reads them all." + /// @nexus.operation "read-records" + read: func( + from-offset: option, + start: option, + max-records: option, + wait: option, + create-wait: option, + ) -> read-records-response; + + /// @nexus.doc "Describe the Stream's status, newest offset and Retention." + /// @nexus.operation "describe-stream" + describe: func() -> describe-stream-response; + } + + // created is false when the Stream existed, whatever its Retention. status is closed when an + // earlier create and close used the id. + /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" + type create-stream-result = own<%stream>; + + /// @nexus.operation name="CreateStream" + create-stream: func(request: create-stream-request) -> create-stream-result; + + /// @nexus.doc "Append one batch. Workflow code makes the call from a Local Activity." + /// @nexus.operation name="AppendRecords" + append-records: func(request: append-records-request) -> append-records-response; + + /// @nexus.operation name="CloseStream" + close-stream: func(request: close-stream-request) -> close-stream-response; + + // A WIT stream can't cross a Nexus call, so a read is a request and a response. The SDK's support + // helper builds the next request from a page for the poll iterator. + /// @nexus.doc + /// "A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local Activity." + /// @nexus.operation name="ReadRecords" + /// @nexus.poll python="next_read_request" + read-records: func(request: read-records-request) -> read-records-response; + + /// @nexus.operation name="DescribeStream" + describe-stream: func(request: describe-stream-request) -> describe-stream-response; + + // No operation returns these. A failed call carries them, so each SDK can raise one error type + // per reason. + /// @nexus.doc "Why a Stream call failed. Each reason has one gRPC status code." + /// @nexus.proto "temporal.api.enums.v1.StreamErrorReason" + enum stream-error-reason { + stream-error-reason-unspecified, + stream-error-reason-not-found, + stream-error-reason-expired, + stream-error-reason-closed, + stream-error-reason-fenced, + stream-error-reason-sequence, + stream-error-reason-seq-divergent, + stream-error-reason-not-owner, + stream-error-reason-format-too-new, + stream-error-reason-offset, + stream-error-reason-too-large, + stream-error-reason-invalid, + stream-error-reason-storage, + stream-error-reason-unknown-writer, + stream-error-reason-unsupported, + } + + /// @nexus.doc "The details of a failed Stream call." + /// @nexus.proto "temporal.api.stream.v1.StreamErrorDetails" + record stream-error-details { + reason: option, + } +} diff --git a/scripts/_wit/stream-writer.wit b/scripts/_wit/stream-writer.wit new file mode 100644 index 0000000000..5f03e746ea --- /dev/null +++ b/scripts/_wit/stream-writer.wit @@ -0,0 +1,118 @@ +// Internal. This isn't a Temporal API, and no server serves it. The crate's runtime object serves +// these calls in process, with the messages of ../protos/temporal/sdk/streams/v1/runtime.proto. +// Lang calls them through the bridge, from Activities and clients only, so the interface has no +// endpoint. This file is the source. Each SDK copies it next to stream-service.wit and generates +// its client with nexgen's client target (`--targets client`, which needs `--system-nexus`). +package temporal:streams@1.0.0; + +world writer { + export stream-writer; +} + +/// @nexus.service-name "temporal.sdk.streams.v1.StreamWriter" +/// @nexus.experimental +interface stream-writer { + use nexus:temporal-types/model@1.0.0.{duration, placeholder}; + use temporal:nexus/stream-service@1.0.0.{ + append-record, producer, stream-ref, workflow-task-epoch, + }; + + /// @nexus.doc "Open a writer for one Producer of one Stream." + /// @nexus.proto "temporal.sdk.streams.v1.OpenWriterRequest" + record open-writer-request { + %stream: stream-ref, + /// @nexus.doc "Local Activity only. The run's failed or timed out Workflow Tasks." + failed-tasks: list, + producer: option, + /// @nexus.doc "A write sends at most this many Records in one batch. Zero is 1000." + max-batch-size: option, + /// @nexus.doc "How long a write retries a Store failure. Empty is 10 s." + max-retry-duration: option, + /// @nexus.doc "Local Activity only. The run's last completed Workflow Task." + previous-completed-task: option, + /// @nexus.omit + namespace: placeholder, + } + + /// @nexus.proto "temporal.sdk.streams.v1.WriteRecordsRequest" + record write-records-request { + records: list, + writer-id: option, + } + + /// @nexus.proto "temporal.sdk.streams.v1.WriteRecordsResponse" + record write-records-response { + /// @nexus.doc "The offset of the first Record. Zero when the request has no Record." + first-offset: option, + /// @nexus.doc "The offset of the last Record. Zero when the request has no Record." + last-offset: option, + } + + /// @nexus.proto "temporal.sdk.streams.v1.CloseWriterRequest" + record close-writer-request { + writer-id: option, + } + + /// @nexus.proto "temporal.sdk.streams.v1.CloseWriterResponse" + record close-writer-response { + /// @nexus.doc "The offset of the newest Record the writer wrote. Zero if it wrote none." + last-offset: option, + } + + /// @nexus.proto "temporal.sdk.streams.v1.FreeWriterRequest" + record free-writer-request { + writer-id: option, + } + + /// @nexus.proto "temporal.sdk.streams.v1.FreeWriterResponse" + record free-writer-response { + } + + /// @nexus.proto "temporal.sdk.streams.v1.CloseWritersRequest" + record close-writers-request { + } + + /// @nexus.proto "temporal.sdk.streams.v1.CloseWritersResponse" + record close-writers-response { + /// @nexus.doc "The writers that had failed." + failed: option, + } + + // A writer never changes. Offsets are values that its methods return. The handles of one + // Producer epoch on one Stream share one writer, and each close or free releases one handle. + /// @nexus.doc "An open writer of one Producer. The crate keeps its epoch and sequence." + resource writer { + constructor(writer-id: u64); + + /// @nexus.doc "Write Records at once. The answer has the first and the last offset." + /// @nexus.operation "write-records" + write-records: func(records: list) -> write-records-response; + + /// @nexus.doc "Release the handle. The last one waits for a write, then drops the writer." + /// @nexus.operation "close-writer" + close: func() -> close-writer-response; + + /// @nexus.doc "Release the handle with no wait, for a handle that lang drops." + /// @nexus.operation "free-writer" + free: func() -> free-writer-response; + } + + /// @nexus.proto "temporal.sdk.streams.v1.OpenWriterResponse" + type writer-result = own; + + /// @nexus.operation name="OpenWriter" + open-writer: func(request: open-writer-request) -> writer-result; + + /// @nexus.operation name="WriteRecords" + write-records: func(request: write-records-request) -> write-records-response; + + /// @nexus.operation name="CloseWriter" + close-writer: func(request: close-writer-request) -> close-writer-response; + + /// @nexus.operation name="FreeWriter" + free-writer: func(request: free-writer-request) -> free-writer-response; + + /// @nexus.doc "Drop every open writer, for example at shutdown." + /// @nexus.operation name="CloseWriters" + close-writers: func(request: close-writers-request) -> close-writers-response; +} diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py new file mode 100644 index 0000000000..8abc5a8618 --- /dev/null +++ b/scripts/gen_stream_service.py @@ -0,0 +1,188 @@ +"""Generates the Stream service bindings of the Streams API in temporalio.common. + +The WIT and the runtime proto are copies of the temporalio-streams crate's, in scripts/_wit and +scripts/_proto. The Stream service protos are copies of the API repo's. Run scripts/gen_protos.py +first for the proto modules. No nexgen release has the client target and the poll directive, +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 +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import gen_protos + +base_dir = Path(__file__).parent.parent +wit_path = base_dir / "scripts" / "_wit" / "stream-service.wit" +writer_wit_path = base_dir / "scripts" / "_wit" / "stream-writer.wit" +model_wit_path = ( + gen_protos.api_proto_dir / "nexus" / "deps" / "nexus-temporal-types" / "model.wit" +) +# It imports the Stream messages and enums, so --include_imports puts them in the descriptor. +proto_path = ( + gen_protos.additional_proto_dir + / "temporal" + / "api" + / "streamservice" + / "v1" + / "request_response.proto" +) +output_dir = base_dir / "temporalio" / "common" / "_streams" / "system" +# Only Activities and 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 fork chain of nexgen with the client target and the poll directive. +NEX_GEN_REVISION = "d70f5f1754c2e694b859f4a38b24c73fab27d8da" +nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION + + +def nex_gen_command() -> list[str]: + if bin_path := os.environ.get("NEX_GEN_STREAMS_BIN"): + return [bin_path] + binary = nex_gen_root / "bin" / "nexgen" + if not binary.exists(): + subprocess.check_call( + [ + "cargo", + "install", + "--locked", + "nexgen", + "--git", + NEX_GEN_REPOSITORY, + "--rev", + NEX_GEN_REVISION, + "--features", + "advanced", + "--root", + str(nex_gen_root), + ], + # The git CLI honours the user's URL rewrites and credentials, which cargo's + # built-in client does not. + env={**os.environ, "CARGO_NET_GIT_FETCH_WITH_CLI": "true"}, + ) + return [str(binary)] + + +def api_protos() -> list[Path]: + return sorted( + path + for path in (gen_protos.api_proto_dir / "temporal" / "api").glob("**/*.proto") + if "cloud" not in path.parts and "operatorservice" not in path.parts + ) + + +def generate_stream_service() -> None: + for path in (wit_path, writer_wit_path, model_wit_path, proto_path): + if not path.exists(): + 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 generated runtime messages are in temporalio.common._streams.proto, so the + # descriptor must name their file by that module path. + runtime_root = Path(temp_dir) / "runtime" + runtime_path = ( + runtime_root + / "temporalio" + / "common" + / "_streams" + / "proto" + / "v1" + / "runtime.proto" + ) + runtime_path.parent.mkdir(parents=True) + shutil.copy( + gen_protos.additional_proto_dir + / "temporal" + / "sdk" + / "streams" + / "v1" + / "runtime.proto", + runtime_path, + ) + subprocess.check_call( + [ + sys.executable, + "-mgrpc_tools.protoc", + f"--proto_path={gen_protos.api_proto_dir}", + f"--proto_path={gen_protos.additional_proto_dir}", + f"--proto_path={runtime_root}", + "--include_imports", + f"--descriptor_set_out={descriptor_path}", + str(proto_path), + str(runtime_path), + # The model WIT names API types, so their descriptors must be in the set. + *map(str, api_protos()), + ] + ) + shutil.rmtree(output_dir, ignore_errors=True) + subprocess.check_call( + [ + *nex_gen_command(), + "python", + str(wit_path), + str(model_wit_path), + "--system-nexus", + "--native-api", + "--targets", + "workflow,client", + "--support-package", + "temporalio.nexus.system._support", + "--descriptors", + str(descriptor_path), + "--output", + str(output_dir), + ] + ) + shutil.rmtree(writer_output_dir, ignore_errors=True) + subprocess.check_call( + [ + *nex_gen_command(), + "python", + str(writer_wit_path), + str(wit_path), + str(model_wit_path), + "--system-nexus", + "--native-api", + "--targets", + "client", + "--support-package", + "temporalio.nexus.system._support", + "--descriptors", + str(descriptor_path), + "--output", + str(writer_output_dir), + ] + ) + for directory in (output_dir, writer_output_dir): + subprocess.check_call( + [ + sys.executable, + "-m", + "ruff", + "check", + "--select", + "I", + "--fix", + str(directory), + ] + ) + subprocess.check_call([sys.executable, "-m", "ruff", "format", str(directory)]) + + +if __name__ == "__main__": + if sys.argv[1:] == ["--install-nexgen"]: + # CI runs the install as its own step, so an install failure shows apart from drift. + nex_gen_command() + sys.exit() + print( + f"Generating the Stream service bindings with nexgen {NEX_GEN_REPOSITORY}" + f"@{NEX_GEN_REVISION}...", + file=sys.stderr, + ) + generate_stream_service() + print("Done", file=sys.stderr) diff --git a/temporalio/common/_streams/system/__init__.py b/temporalio/common/_streams/system/__init__.py new file mode 100644 index 0000000000..8d916ce5b3 --- /dev/null +++ b/temporalio/common/_streams/system/__init__.py @@ -0,0 +1,85 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import typing + +import nexusrpc + +import temporalio.converter + +from . import services as _services +from ._resources.stream import ( + append_records, + close_stream, + describe_stream, + read_records, +) +from .operations.create_stream import create_stream + +__all__ = [ + "append_records", + "close_stream", + "create_stream", + "describe_stream", + "read_records", +] + + +_InputT = typing.TypeVar("_InputT") +_OutputT = typing.TypeVar("_OutputT") + + +_SerializationContextFactory = collections.abc.Callable[ + [_InputT], temporalio.converter.SerializationContext +] + + +class _NexusOperationInfo(typing.Generic[_InputT, _OutputT]): + def __init__( + self, + *, + operation: nexusrpc.Operation[_InputT, _OutputT], + serialization_context: _SerializationContextFactory[_InputT] | None = None, + ) -> None: + self.operation: nexusrpc.Operation[_InputT, _OutputT] = operation + self.serialization_context: _SerializationContextFactory[_InputT] | None = ( + serialization_context + ) + + +__nexus_operation_registry__: dict[ + tuple[str, str], _NexusOperationInfo[typing.Any, typing.Any] +] = { + ( + "temporal.api.streamservice.v1.StreamService", + "CreateStream", + ): _NexusOperationInfo( + operation=_services.StreamService.create_stream, + ), + ( + "temporal.api.streamservice.v1.StreamService", + "AppendRecords", + ): _NexusOperationInfo( + operation=_services.StreamService.append_records, + ), + ( + "temporal.api.streamservice.v1.StreamService", + "CloseStream", + ): _NexusOperationInfo( + operation=_services.StreamService.close_stream, + ), + ( + "temporal.api.streamservice.v1.StreamService", + "ReadRecords", + ): _NexusOperationInfo( + operation=_services.StreamService.read_records, + ), + ( + "temporal.api.streamservice.v1.StreamService", + "DescribeStream", + ): _NexusOperationInfo( + operation=_services.StreamService.describe_stream, + ), +} diff --git a/temporalio/common/_streams/system/_resources/__init__.py b/temporalio/common/_streams/system/_resources/__init__.py new file mode 100644 index 0000000000..efbe663e21 --- /dev/null +++ b/temporalio/common/_streams/system/_resources/__init__.py @@ -0,0 +1,9 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +from .stream import Stream + +__all__ = [ + "Stream", +] diff --git a/temporalio/common/_streams/system/_resources/stream.py b/temporalio/common/_streams/system/_resources/stream.py new file mode 100644 index 0000000000..663d423aec --- /dev/null +++ b/temporalio/common/_streams/system/_resources/stream.py @@ -0,0 +1,341 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import dataclasses +import datetime +import typing + +import temporalio.nexus.system._support as _support +import temporalio.workflow + +from ..models import ( + AppendRecord, + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + DescribeStreamRequest, + DescribeStreamResponse, + Producer, + ReadRecordsRequest, + ReadRecordsResponse, + StreamCloseStatus, + StreamReadStart, + StreamRef, + StreamStatus, + WorkflowTaskEpoch, +) + + +@dataclasses.dataclass +class Stream: + """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" + + stream: StreamRef + created: bool | None + status: StreamStatus | None + + async def publish( + self, + producer: Producer | None = None, + sequence: int | None = None, + records: list[AppendRecord] | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + withdraw_run: str | None = None, + ) -> AppendRecordsResponse: + """Publish Records. The answer has the first and the last offset.""" + request = AppendRecordsRequest( + stream=self.stream, + producer=producer, + sequence=sequence, + records=records or [], + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + withdraw_run=withdraw_run, + ) + handle = await _append_records(request) + return await handle + + async def close( + self, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + status: StreamCloseStatus | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + ) -> None: + """Close the Stream. Readers drain what is left, then end.""" + request = CloseStreamRequest( + stream=self.stream, + producer=producer, + result=result, + failure=failure, + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + status=status, + ) + handle = await _close_stream(request) + await handle + + async def read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> ReadRecordsResponse: + """Read one page of Records. The poll iterator reads them all.""" + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + ) + handle = await _read_records(request) + return await handle + + async def iter_read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: + """Yield each response of ``read``. The next request comes from the poll helper, + and the iterator ends when the helper returns ``None``. + """ + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + ) + while True: + handle = await _read_records(request) + response = await handle + yield response + next_request = _support.next_read_request(request, response) + if next_request is None: + return + request = next_request + + async def describe( + self, + ) -> DescribeStreamResponse: + """Describe the Stream's status, newest offset and Retention.""" + request = DescribeStreamRequest( + stream=self.stream, + ) + handle = await _describe_stream(request) + return await handle + + +async def _append_records( + request: AppendRecordsRequest, +) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="AppendRecords", + input=request, + output_type=AppendRecordsResponse, + ) + + +async def append_records( + *, + stream: StreamRef, + records: list[AppendRecord] | None = None, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + sequence: int | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + withdraw_run: str | None = None, +) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: + """Append one batch. Workflow code makes the call from a Local Activity. + + Args: + failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. + producer: Required outside a Workflow. In a Workflow the transport sets it. + sequence: The batch number in this epoch, 1 or more. It is part of the dedupe + key. + previous_completed_task: Local Activity only. The run's last completed Workflow + Task. + withdraw_run: Retry run only. The previous run, whose Records this withdraws. + """ + request = AppendRecordsRequest( + stream=stream, + records=records or [], + failed_tasks=failed_tasks or [], + producer=producer, + sequence=sequence, + previous_completed_task=previous_completed_task, + withdraw_run=withdraw_run, + ) + return await _append_records(request) + + +async def _close_stream( + request: CloseStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="CloseStream", + input=request, + output_type=CloseStreamResponse, + ) + + +async def close_stream( + *, + stream: StreamRef, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + status: StreamCloseStatus | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, +) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: + """ + Args: + failed_tasks: Workflow only. The run's failed or timed out Workflow Tasks. + status: Empty is completed, or failed with a failure. A caller may set canceled. + previous_completed_task: Workflow only. The previous completed Workflow Task of + the run. + """ + request = CloseStreamRequest( + stream=stream, + failed_tasks=failed_tasks or [], + producer=producer, + result=result, + failure=failure, + status=status, + previous_completed_task=previous_completed_task, + ) + return await _close_stream(request) + + +async def _read_records( + request: ReadRecordsRequest, +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="ReadRecords", + input=request, + output_type=ReadRecordsResponse, + ) + + +async def read_records( + *, + stream: StreamRef, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: + """A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local + Activity. + + Args: + from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the beginning. + max_records: At most this many Records. 1 to 1000. + wait: How long to wait for a new Record before an empty answer. At most 60 s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. At + most 60 s. + last: Read the newest this many Records, oldest first. No from-offset. No wait. + max_bytes: The maximum Record bytes in a page. A larger Record is still returned + alone. + """ + request = ReadRecordsRequest( + stream=stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + return await _read_records(request) + + +async def iter_read_records( + *, + stream: StreamRef, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, +) -> collections.abc.AsyncIterator[ReadRecordsResponse]: + """Yield each response of ``read_records``. The next request comes from the poll + helper, and the iterator ends when the helper returns ``None``. + """ + request = ReadRecordsRequest( + stream=stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + while True: + handle = await _read_records(request) + response = await handle + yield response + next_request = _support.next_read_request(request, response) + if next_request is None: + return + request = next_request + + +async def _describe_stream( + request: DescribeStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="DescribeStream", + input=request, + output_type=DescribeStreamResponse, + ) + + +async def describe_stream( + *, + stream: StreamRef, +) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: + request = DescribeStreamRequest( + stream=stream, + ) + return await _describe_stream(request) diff --git a/temporalio/common/_streams/system/_system_nexus_interceptor.py b/temporalio/common/_streams/system/_system_nexus_interceptor.py new file mode 100644 index 0000000000..647dd384fe --- /dev/null +++ b/temporalio/common/_streams/system/_system_nexus_interceptor.py @@ -0,0 +1,265 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import abc +import typing + +from . import models + +if typing.TYPE_CHECKING: + import temporalio.workflow + from temporalio.worker._interceptor import StartNexusOperationInput + + +__all__ = [ + "_start_system_nexus_operation", + "_SystemNexusWorkflowOutboundInterceptorBase", + "_SystemNexusWorkflowOutboundInterceptorTerminal", +] + + +_InputT = typing.TypeVar("_InputT") +_OutputT = typing.TypeVar("_OutputT") + + +async def _start_system_nexus_operation( + interceptor: _SystemNexusWorkflowOutboundInterceptorBase, + input: StartNexusOperationInput[_InputT, _OutputT], +) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + if ( + input.service == "temporal.api.streamservice.v1.StreamService" + and input.operation_name == "CreateStream" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.CreateStreamRequest, models.CreateStreamResult]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_create_stream(typed_input.input), + ) + if ( + input.service == "temporal.api.streamservice.v1.StreamService" + and input.operation_name == "AppendRecords" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.AppendRecordsRequest, models.AppendRecordsResponse]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_append_records(typed_input.input), + ) + if ( + input.service == "temporal.api.streamservice.v1.StreamService" + and input.operation_name == "CloseStream" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.CloseStreamRequest, models.CloseStreamResponse]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_close_stream(typed_input.input), + ) + if ( + input.service == "temporal.api.streamservice.v1.StreamService" + and input.operation_name == "ReadRecords" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.ReadRecordsRequest, models.ReadRecordsResponse]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_read_records(typed_input.input), + ) + if ( + input.service == "temporal.api.streamservice.v1.StreamService" + and input.operation_name == "DescribeStream" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.DescribeStreamRequest, models.DescribeStreamResponse]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_describe_stream(typed_input.input), + ) + return await interceptor.start_system_nexus_operation(input) + + +class _SystemNexusWorkflowOutboundInterceptorBase(abc.ABC): + @abc.abstractmethod + def _next_system_nexus_interceptor( + self, + ) -> _SystemNexusWorkflowOutboundInterceptorBase: ... + + @abc.abstractmethod + async def start_system_nexus_operation( + self, + input: StartNexusOperationInput[_InputT, _OutputT], + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + """Intercept a System Nexus operation through the generic hook.""" + ... + + async def start_create_stream( + self, request: models.CreateStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + """Intercept the CreateStream operation.""" + return await self._next_system_nexus_interceptor().start_create_stream(request) + + async def start_append_records( + self, request: models.AppendRecordsRequest + ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: + """Intercept the AppendRecords operation.""" + return await self._next_system_nexus_interceptor().start_append_records(request) + + async def start_close_stream( + self, request: models.CloseStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CloseStreamResponse]: + """Intercept the CloseStream operation.""" + return await self._next_system_nexus_interceptor().start_close_stream(request) + + async def start_read_records( + self, request: models.ReadRecordsRequest + ) -> temporalio.workflow.NexusOperationHandle[models.ReadRecordsResponse]: + """Intercept the ReadRecords operation.""" + return await self._next_system_nexus_interceptor().start_read_records(request) + + async def start_describe_stream( + self, request: models.DescribeStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.DescribeStreamResponse]: + """Intercept the DescribeStream operation.""" + return await self._next_system_nexus_interceptor().start_describe_stream( + request + ) + + +class _SystemNexusWorkflowOutboundInterceptorTerminal(abc.ABC): + @abc.abstractmethod + async def _intercept_system_nexus_operation( + self, + input: StartNexusOperationInput[_InputT, _OutputT], + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: ... + + async def start_create_stream( + self, request: models.CreateStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._intercept_system_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.streamservice.v1.StreamService", + operation="CreateStream", + input=request, + output_type=models.CreateStreamResult, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) + + async def start_append_records( + self, request: models.AppendRecordsRequest + ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._intercept_system_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.streamservice.v1.StreamService", + operation="AppendRecords", + input=request, + output_type=models.AppendRecordsResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) + + async def start_close_stream( + self, request: models.CloseStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CloseStreamResponse]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._intercept_system_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.streamservice.v1.StreamService", + operation="CloseStream", + input=request, + output_type=models.CloseStreamResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) + + async def start_read_records( + self, request: models.ReadRecordsRequest + ) -> temporalio.workflow.NexusOperationHandle[models.ReadRecordsResponse]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._intercept_system_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.streamservice.v1.StreamService", + operation="ReadRecords", + input=request, + output_type=models.ReadRecordsResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) + + async def start_describe_stream( + self, request: models.DescribeStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.DescribeStreamResponse]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._intercept_system_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.streamservice.v1.StreamService", + operation="DescribeStream", + input=request, + output_type=models.DescribeStreamResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) diff --git a/temporalio/common/_streams/system/client.py b/temporalio/common/_streams/system/client.py new file mode 100644 index 0000000000..8a229d37f9 --- /dev/null +++ b/temporalio/common/_streams/system/client.py @@ -0,0 +1,439 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import abc +import collections.abc +import dataclasses +import datetime +import typing + +import temporalio.nexus.system._support as _support + +from .models import ( + AppendRecord, + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + CreateStreamRequest, + CreateStreamResult, + DescribeStreamRequest, + DescribeStreamResponse, + Producer, + ReadRecordsRequest, + ReadRecordsResponse, + StreamCloseStatus, + StreamReadStart, + StreamRef, + StreamStatus, + WorkflowTaskEpoch, +) + +_OutputT = typing.TypeVar("_OutputT") + + +class _NexusClient(typing.Protocol): + """Starts the Nexus operations of one service from code outside a Workflow.""" + + @typing.overload + async def start_operation( + self, + operation: str, + input: typing.Any, + *, + output_type: type[_OutputT], + ) -> collections.abc.Awaitable[_OutputT]: ... + + @typing.overload + async def start_operation( + self, + operation: str, + input: typing.Any, + ) -> collections.abc.Awaitable[None]: ... + + +class StreamServiceClientExtension(abc.ABC): + """Starts the operations of ``StreamService`` from code outside a Workflow. + + A client inherits this class and implements ``_create_nexus_client``. + + .. warning:: + This API is experimental and subject to change. + """ + + @abc.abstractmethod + def _create_nexus_client( + self, + *, + service: str, + endpoint: str | None, + ) -> _NexusClient: + """Return the Nexus client that starts this service's operations. + + ``endpoint`` is ``None`` for a service with no endpoint, which the client serves. + """ + ... + + async def _create_stream( + self, + request: CreateStreamRequest, + ) -> Stream: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="CreateStream", + input=request, + output_type=CreateStreamResult, + ) + result = await handle + if result.stream is None: + raise RuntimeError("CreateStream answered without `stream`") + return Stream( + _client=self, + stream=result.stream, + created=result.created or None, + status=result.status or None, + ) + + async def create_stream( + self, + *, + stream_id: str, + retention: datetime.timedelta | None = None, + withdraw_run: str | None = None, + owner: Producer | None = None, + ref: StreamRef | None = None, + ) -> Stream: + """ + Args: + 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. + owner: The Workflow run that owns the Stream. The Workflow transport sets it. + ref: The Stream's ref, with the owner fields. + """ + request = CreateStreamRequest( + stream_id=stream_id, + retention=retention, + withdraw_run=withdraw_run, + owner=owner, + ref=ref, + ) + return await self._create_stream(request) + + async def _append_records( + self, + request: AppendRecordsRequest, + ) -> collections.abc.Awaitable[AppendRecordsResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="AppendRecords", + input=request, + output_type=AppendRecordsResponse, + ) + + async def append_records( + self, + *, + stream: StreamRef, + records: list[AppendRecord] | None = None, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + sequence: int | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + withdraw_run: str | None = None, + ) -> collections.abc.Awaitable[AppendRecordsResponse,]: + """Append one batch. Workflow code makes the call from a Local Activity. + + Args: + failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. + producer: Required outside a Workflow. In a Workflow the transport sets it. + sequence: The batch number in this epoch, 1 or more. It is part of the dedupe + key. + previous_completed_task: Local Activity only. The run's last completed Workflow + Task. + withdraw_run: Retry run only. The previous run, whose Records this withdraws. + """ + request = AppendRecordsRequest( + stream=stream, + records=records or [], + failed_tasks=failed_tasks or [], + producer=producer, + sequence=sequence, + previous_completed_task=previous_completed_task, + withdraw_run=withdraw_run, + ) + return await self._append_records(request) + + async def _close_stream( + self, + request: CloseStreamRequest, + ) -> collections.abc.Awaitable[CloseStreamResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="CloseStream", + input=request, + output_type=CloseStreamResponse, + ) + + async def close_stream( + self, + *, + stream: StreamRef, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + status: StreamCloseStatus | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + ) -> collections.abc.Awaitable[CloseStreamResponse,]: + """ + Args: + failed_tasks: Workflow only. The run's failed or timed out Workflow Tasks. + status: Empty is completed, or failed with a failure. A caller may set canceled. + previous_completed_task: Workflow only. The previous completed Workflow Task of + the run. + """ + request = CloseStreamRequest( + stream=stream, + failed_tasks=failed_tasks or [], + producer=producer, + result=result, + failure=failure, + status=status, + previous_completed_task=previous_completed_task, + ) + return await self._close_stream(request) + + async def _read_records( + self, + request: ReadRecordsRequest, + ) -> collections.abc.Awaitable[ReadRecordsResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="ReadRecords", + input=request, + output_type=ReadRecordsResponse, + ) + + async def read_records( + self, + *, + stream: StreamRef, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, + ) -> collections.abc.Awaitable[ReadRecordsResponse,]: + """A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local + Activity. + + Args: + from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the beginning. + max_records: At most this many Records. 1 to 1000. + wait: How long to wait for a new Record before an empty answer. At most 60 s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. At + most 60 s. + last: Read the newest this many Records, oldest first. No from-offset. No wait. + max_bytes: The maximum Record bytes in a page. A larger Record is still returned + alone. + """ + request = ReadRecordsRequest( + stream=stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + return await self._read_records(request) + + async def iter_read_records( + self, + *, + stream: StreamRef, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, + ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: + """Yield each response of ``read_records``. The next request comes from the poll + helper, and the iterator ends when the helper returns ``None``. + """ + request = ReadRecordsRequest( + stream=stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + while True: + handle = await self._read_records(request) + response = await handle + yield response + next_request = _support.next_read_request(request, response) + if next_request is None: + return + request = next_request + + async def _describe_stream( + self, + request: DescribeStreamRequest, + ) -> collections.abc.Awaitable[DescribeStreamResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + return await nexus_client.start_operation( + operation="DescribeStream", + input=request, + output_type=DescribeStreamResponse, + ) + + async def describe_stream( + self, + *, + stream: StreamRef, + ) -> collections.abc.Awaitable[DescribeStreamResponse,]: + request = DescribeStreamRequest( + stream=stream, + ) + return await self._describe_stream(request) + + +@dataclasses.dataclass +class Stream: + """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" + + _client: StreamServiceClientExtension = dataclasses.field(repr=False, compare=False) + stream: StreamRef + created: bool | None + status: StreamStatus | None + + async def publish( + self, + producer: Producer | None = None, + sequence: int | None = None, + records: list[AppendRecord] | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + withdraw_run: str | None = None, + ) -> AppendRecordsResponse: + """Publish Records. The answer has the first and the last offset.""" + request = AppendRecordsRequest( + stream=self.stream, + producer=producer, + sequence=sequence, + records=records or [], + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + withdraw_run=withdraw_run, + ) + handle = await self._client._append_records(request) + return await handle + + async def close( + self, + producer: Producer | None = None, + result: typing.Any | None = None, + failure: BaseException | None = None, + status: StreamCloseStatus | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + failed_tasks: list[int] | None = None, + ) -> None: + """Close the Stream. Readers drain what is left, then end.""" + request = CloseStreamRequest( + stream=self.stream, + producer=producer, + result=result, + failure=failure, + previous_completed_task=previous_completed_task, + failed_tasks=failed_tasks or [], + status=status, + ) + handle = await self._client._close_stream(request) + await handle + + async def read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> ReadRecordsResponse: + """Read one page of Records. The poll iterator reads them all.""" + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + ) + handle = await self._client._read_records(request) + return await handle + + async def iter_read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: + """Yield each response of ``read``. The next request comes from the poll helper, + and the iterator ends when the helper returns ``None``. + """ + request = ReadRecordsRequest( + stream=self.stream, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + ) + while True: + handle = await self._client._read_records(request) + response = await handle + yield response + next_request = _support.next_read_request(request, response) + if next_request is None: + return + request = next_request + + async def describe( + self, + ) -> DescribeStreamResponse: + """Describe the Stream's status, newest offset and Retention.""" + request = DescribeStreamRequest( + stream=self.stream, + ) + handle = await self._client._describe_stream(request) + return await handle diff --git a/temporalio/common/_streams/system/models.py b/temporalio/common/_streams/system/models.py new file mode 100644 index 0000000000..b937cdfd28 --- /dev/null +++ b/temporalio/common/_streams/system/models.py @@ -0,0 +1,1439 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import dataclasses +import datetime +import enum +import typing + +import google.protobuf.timestamp_pb2 +import typing_extensions + +import temporalio.api.common.v1.message_pb2 +import temporalio.api.enums.v1.common_pb2 +import temporalio.api.enums.v1.stream_pb2 +import temporalio.api.stream.v1.message_pb2 +import temporalio.api.streamservice.v1.request_response_pb2 +import temporalio.converter +import temporalio.nexus.system._support as _support + + +class StreamProducerKind(enum.IntEnum): + StreamProducerKindUnspecified = 0 + StreamProducerKindWorkflow = 1 + StreamProducerKindLocalActivity = 2 + StreamProducerKindActivity = 3 + StreamProducerKindClient = 4 + StreamProducerKindStandaloneActivity = 5 + + +class ExecutionType(enum.IntEnum): + ExecutionTypeUnspecified = 0 + ExecutionTypeWorkflow = 1 + ExecutionTypeActivity = 2 + ExecutionTypeNexusOperation = 3 + + +class StreamStatus(enum.IntEnum): + StreamStatusUnspecified = 0 + StreamStatusOpen = 1 + StreamStatusClosed = 2 + StreamStatusExpired = 3 + + +class StreamCloseStatus(enum.IntEnum): + StreamCloseStatusUnspecified = 0 + StreamCloseStatusCompleted = 1 + StreamCloseStatusFailed = 2 + StreamCloseStatusCanceled = 3 + StreamCloseStatusTerminated = 4 + StreamCloseStatusTimedOut = 5 + + +class StreamReadStart(enum.IntEnum): + StreamReadStartUnspecified = 0 + StreamReadStartBeginning = 1 + StreamReadStartEnd = 2 + + +@dataclasses.dataclass(slots=True) +class CreateStreamRequest: + """Create a Stream, or get the one that exists, which keeps its Retention.""" + + stream_id: str + retention: datetime.timedelta | None = None + withdraw_run: str | None = None + owner: Producer | None = None + ref: StreamRef | None = None + + +class _CreateStreamRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CreateStreamRequest, + temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest, + type_hint: type[CreateStreamRequest], + ) -> CreateStreamRequest: + stream_id = value.stream_id + return CreateStreamRequest( + stream_id=stream_id, + retention=_support.duration_from_proto(value.retention) + if value.HasField("retention") + else None, + withdraw_run=value.withdraw_run if bool(value.withdraw_run) else None, + owner=_ProducerTransferTypeConverter().from_transfer_type( + value.owner, Producer + ) + if value.HasField("owner") + else None, + ref=_StreamRefTransferTypeConverter().from_transfer_type( + value.ref, StreamRef + ) + if value.HasField("ref") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CreateStreamRequest, + ) -> temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest() + ) + message.stream_id = value.stream_id + if value.retention is not None: + message.retention.CopyFrom(_support.duration_to_proto(value.retention)) + if value.withdraw_run is not None: + message.withdraw_run = value.withdraw_run + if value.owner is not None: + message.owner.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.owner) + ) + if value.ref is not None: + message.ref.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.ref) + ) + return message + + +temporalio.converter.transfer_type_convertible( + _CreateStreamRequestTransferTypeConverter +)(CreateStreamRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Producer: + """The Producer of a batch.""" + + id: str + kind: StreamProducerKind | None = None + epoch: ProducerEpoch | None = None + run_id: str | None = None + + +class _ProducerTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Producer, temporalio.api.stream.v1.message_pb2.Producer + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.Producer] | None = ( + temporalio.api.stream.v1.message_pb2.Producer + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.Producer, + type_hint: type[Producer], + ) -> Producer: + id = value.id + return Producer( + id=id, + kind=StreamProducerKind(value.kind) if value.kind != 0 else None, + epoch=_ProducerEpochTransferTypeConverter().from_transfer_type( + value.epoch, ProducerEpoch + ) + if value.HasField("epoch") + else None, + run_id=value.run_id if bool(value.run_id) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Producer, + ) -> temporalio.api.stream.v1.message_pb2.Producer: + message = temporalio.api.stream.v1.message_pb2.Producer() + message.id = value.id + if value.kind is not None: + message.kind = ( + temporalio.api.enums.v1.stream_pb2.StreamProducerKind.ValueType( + value.kind + ) + ) + if value.epoch is not None: + message.epoch.CopyFrom( + _ProducerEpochTransferTypeConverter().to_transfer_type(value.epoch) + ) + if value.run_id is not None: + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_ProducerTransferTypeConverter)(Producer) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class ProducerEpoch: + """The ordered attempt of a Producer. A higher epoch supersedes a lower one.""" + + task_started_event_id: int | None = None + task_attempt: int | None = None + task_started_time: Timestamp | None = None + attempt_scheduled_time: Timestamp | None = None + attempt: int | None = None + + +class _ProducerEpochTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + ProducerEpoch, temporalio.api.stream.v1.message_pb2.ProducerEpoch + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.ProducerEpoch] | None = ( + temporalio.api.stream.v1.message_pb2.ProducerEpoch + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.ProducerEpoch, + type_hint: type[ProducerEpoch], + ) -> ProducerEpoch: + return ProducerEpoch( + task_started_event_id=value.task_started_event_id + if value.task_started_event_id != 0 + else None, + task_attempt=value.task_attempt if value.task_attempt != 0 else None, + task_started_time=_TimestampTransferTypeConverter().from_transfer_type( + value.task_started_time, Timestamp + ) + if value.HasField("task_started_time") + else None, + attempt_scheduled_time=_TimestampTransferTypeConverter().from_transfer_type( + value.attempt_scheduled_time, Timestamp + ) + if value.HasField("attempt_scheduled_time") + else None, + attempt=value.attempt if value.attempt != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: ProducerEpoch, + ) -> temporalio.api.stream.v1.message_pb2.ProducerEpoch: + message = temporalio.api.stream.v1.message_pb2.ProducerEpoch() + if value.task_started_event_id is not None: + message.task_started_event_id = value.task_started_event_id + if value.task_attempt is not None: + message.task_attempt = value.task_attempt + if value.task_started_time is not None: + message.task_started_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type( + value.task_started_time + ) + ) + if value.attempt_scheduled_time is not None: + message.attempt_scheduled_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type( + value.attempt_scheduled_time + ) + ) + if value.attempt is not None: + message.attempt = value.attempt + return message + + +temporalio.converter.transfer_type_convertible(_ProducerEpochTransferTypeConverter)( + ProducerEpoch +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Timestamp: + seconds: int | None = None + nanos: int | None = None + + +class _TimestampTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Timestamp, google.protobuf.timestamp_pb2.Timestamp + ] +): + transfer_type: type[google.protobuf.timestamp_pb2.Timestamp] | None = ( + google.protobuf.timestamp_pb2.Timestamp + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: google.protobuf.timestamp_pb2.Timestamp, + type_hint: type[Timestamp], + ) -> Timestamp: + return Timestamp( + seconds=value.seconds if value.seconds != 0 else None, + nanos=value.nanos if value.nanos != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Timestamp, + ) -> google.protobuf.timestamp_pb2.Timestamp: + message = google.protobuf.timestamp_pb2.Timestamp() + if value.seconds is not None: + message.seconds = value.seconds + if value.nanos is not None: + message.nanos = value.nanos + return message + + +temporalio.converter.transfer_type_convertible(_TimestampTransferTypeConverter)( + Timestamp +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamRef: + """Identifies one Stream. The id is opaque to callers.""" + + id: str + owner: Execution | None = None + backend: str | None = None + + +class _StreamRefTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamRef, temporalio.api.stream.v1.message_pb2.StreamRef + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRef] | None = ( + temporalio.api.stream.v1.message_pb2.StreamRef + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamRef, + type_hint: type[StreamRef], + ) -> StreamRef: + id = value.id + return StreamRef( + id=id, + owner=_ExecutionTransferTypeConverter().from_transfer_type( + value.owner, Execution + ) + if value.HasField("owner") + else None, + backend=value.backend if bool(value.backend) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamRef, + ) -> temporalio.api.stream.v1.message_pb2.StreamRef: + message = temporalio.api.stream.v1.message_pb2.StreamRef() + message.id = value.id + if value.owner is not None: + message.owner.CopyFrom( + _ExecutionTransferTypeConverter().to_transfer_type(value.owner) + ) + if value.backend is not None: + message.backend = value.backend + return message + + +temporalio.converter.transfer_type_convertible(_StreamRefTransferTypeConverter)( + StreamRef +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Execution: + """A Workflow run chain or a Standalone Activity run.""" + + type: ExecutionType | None = None + business_id: str | None = None + run_id: str | None = None + + +class _ExecutionTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Execution, temporalio.api.common.v1.message_pb2.Execution + ] +): + transfer_type: type[temporalio.api.common.v1.message_pb2.Execution] | None = ( + temporalio.api.common.v1.message_pb2.Execution + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.common.v1.message_pb2.Execution, + type_hint: type[Execution], + ) -> Execution: + return Execution( + type=ExecutionType(value.type) if value.type != 0 else None, + business_id=value.business_id if bool(value.business_id) else None, + run_id=value.run_id if bool(value.run_id) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Execution, + ) -> temporalio.api.common.v1.message_pb2.Execution: + message = temporalio.api.common.v1.message_pb2.Execution() + if value.type is not None: + message.type = temporalio.api.enums.v1.common_pb2.ExecutionType.ValueType( + value.type + ) + if value.business_id is not None: + message.business_id = value.business_id + if value.run_id is not None: + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_ExecutionTransferTypeConverter)( + Execution +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CreateStreamResult: + stream: StreamRef | None = None + created: bool | None = None + status: StreamStatus | None = None + + +class _CreateStreamResultTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CreateStreamResult, + temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse, + type_hint: type[CreateStreamResult], + ) -> CreateStreamResult: + return CreateStreamResult( + stream=_StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + if value.HasField("stream") + else None, + created=value.created if bool(value.created) else None, + status=StreamStatus(value.status) if value.status != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CreateStreamResult, + ) -> temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse() + ) + if value.stream is not None: + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + if value.created is not None: + message.created = value.created + if value.status is not None: + message.status = temporalio.api.enums.v1.stream_pb2.StreamStatus.ValueType( + value.status + ) + return message + + +temporalio.converter.transfer_type_convertible( + _CreateStreamResultTransferTypeConverter +)(CreateStreamResult) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class AppendRecordsRequest: + """Append. FAILED_PRECONDITION if closed, fenced or format too new. NOT_FOUND if + absent. + """ + + stream: StreamRef + records: list[AppendRecord] | None = dataclasses.field(default_factory=list) + failed_tasks: list[int] | None = dataclasses.field(default_factory=list) + producer: Producer | None = None + sequence: int | None = None + previous_completed_task: WorkflowTaskEpoch | None = None + withdraw_run: str | None = None + + +class _AppendRecordsRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + AppendRecordsRequest, + temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest, + type_hint: type[AppendRecordsRequest], + ) -> AppendRecordsRequest: + if not value.HasField("stream"): + raise ValueError("missing required field AppendRecordsRequest.stream") + stream = _StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + return AppendRecordsRequest( + stream=stream, + records=[ + _AppendRecordTransferTypeConverter().from_transfer_type( + item, AppendRecord + ) + for item in value.records + ], + failed_tasks=list(value.failed_tasks), + producer=_ProducerTransferTypeConverter().from_transfer_type( + value.producer, Producer + ) + if value.HasField("producer") + else None, + sequence=value.sequence if value.sequence != 0 else None, + previous_completed_task=_WorkflowTaskEpochTransferTypeConverter().from_transfer_type( + value.previous_completed_task, WorkflowTaskEpoch + ) + if value.HasField("previous_completed_task") + else None, + withdraw_run=value.withdraw_run if bool(value.withdraw_run) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: AppendRecordsRequest, + ) -> temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest() + ) + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + if value.records: + for element in value.records: + item = message.records.add() + item.CopyFrom( + _AppendRecordTransferTypeConverter().to_transfer_type(element) + ) + if value.failed_tasks: + message.failed_tasks.extend(value.failed_tasks) + if value.producer is not None: + message.producer.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.producer) + ) + if value.sequence is not None: + message.sequence = value.sequence + if value.previous_completed_task is not None: + message.previous_completed_task.CopyFrom( + _WorkflowTaskEpochTransferTypeConverter().to_transfer_type( + value.previous_completed_task + ) + ) + if value.withdraw_run is not None: + message.withdraw_run = value.withdraw_run + return message + + +temporalio.converter.transfer_type_convertible( + _AppendRecordsRequestTransferTypeConverter +)(AppendRecordsRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class AppendRecord: + """One Record to append.""" + + body: typing.Any + + +class _AppendRecordTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + AppendRecord, temporalio.api.stream.v1.message_pb2.AppendRecord + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.AppendRecord] | None = ( + temporalio.api.stream.v1.message_pb2.AppendRecord + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.AppendRecord, + type_hint: type[AppendRecord], + ) -> AppendRecord: + if not value.HasField("body"): + raise ValueError("missing required field AppendRecord.body") + body = _support.payload_from_proto(value.body) + return AppendRecord( + body=body, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: AppendRecord, + ) -> temporalio.api.stream.v1.message_pb2.AppendRecord: + message = temporalio.api.stream.v1.message_pb2.AppendRecord() + message.body.CopyFrom(_support.payload_to_proto(value.body)) + return message + + +temporalio.converter.transfer_type_convertible(_AppendRecordTransferTypeConverter)( + AppendRecord +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WorkflowTaskEpoch: + """A Workflow Task attempt of the run.""" + + started_event_id: int | None = None + attempt: int | None = None + started_time: Timestamp | None = None + + +class _WorkflowTaskEpochTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WorkflowTaskEpoch, temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch + ] +): + transfer_type: ( + type[temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch] | None + ) = temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch, + type_hint: type[WorkflowTaskEpoch], + ) -> WorkflowTaskEpoch: + return WorkflowTaskEpoch( + started_event_id=value.started_event_id + if value.started_event_id != 0 + else None, + attempt=value.attempt if value.attempt != 0 else None, + started_time=_TimestampTransferTypeConverter().from_transfer_type( + value.started_time, Timestamp + ) + if value.HasField("started_time") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WorkflowTaskEpoch, + ) -> temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch: + message = temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch() + if value.started_event_id is not None: + message.started_event_id = value.started_event_id + if value.attempt is not None: + message.attempt = value.attempt + if value.started_time is not None: + message.started_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type(value.started_time) + ) + return message + + +temporalio.converter.transfer_type_convertible(_WorkflowTaskEpochTransferTypeConverter)( + WorkflowTaskEpoch +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class AppendRecordsResponse: + first_offset: int | None = None + last_offset: int | None = None + + +class _AppendRecordsResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + AppendRecordsResponse, + temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse, + type_hint: type[AppendRecordsResponse], + ) -> AppendRecordsResponse: + return AppendRecordsResponse( + first_offset=value.first_offset if value.first_offset != 0 else None, + last_offset=value.last_offset if value.last_offset != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: AppendRecordsResponse, + ) -> temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsResponse() + ) + if value.first_offset is not None: + message.first_offset = value.first_offset + if value.last_offset is not None: + message.last_offset = value.last_offset + return message + + +temporalio.converter.transfer_type_convertible( + _AppendRecordsResponseTransferTypeConverter +)(AppendRecordsResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseStreamRequest: + """Close a Stream. Readers drain what is left, then end. Later appends are refused.""" + + stream: StreamRef + failed_tasks: list[int] | None = dataclasses.field(default_factory=list) + producer: Producer | None = None + result: typing.Any | None = None + failure: BaseException | None = None + status: StreamCloseStatus | None = None + previous_completed_task: WorkflowTaskEpoch | None = None + + +class _CloseStreamRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseStreamRequest, + temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest, + type_hint: type[CloseStreamRequest], + ) -> CloseStreamRequest: + if not value.HasField("stream"): + raise ValueError("missing required field CloseStreamRequest.stream") + stream = _StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + return CloseStreamRequest( + stream=stream, + failed_tasks=list(value.failed_tasks), + producer=_ProducerTransferTypeConverter().from_transfer_type( + value.producer, Producer + ) + if value.HasField("producer") + else None, + result=_support.payload_from_proto(value.result) + if value.HasField("result") + else None, + failure=_support.failure_from_proto(value.failure) + if value.HasField("failure") + else None, + status=StreamCloseStatus(value.status) if value.status != 0 else None, + previous_completed_task=_WorkflowTaskEpochTransferTypeConverter().from_transfer_type( + value.previous_completed_task, WorkflowTaskEpoch + ) + if value.HasField("previous_completed_task") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseStreamRequest, + ) -> temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest() + ) + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + if value.failed_tasks: + message.failed_tasks.extend(value.failed_tasks) + if value.producer is not None: + message.producer.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.producer) + ) + if value.result is not None: + message.result.CopyFrom(_support.payload_to_proto(value.result)) + if value.failure is not None: + message.failure.CopyFrom(_support.failure_to_proto(value.failure)) + if value.status is not None: + message.status = ( + temporalio.api.enums.v1.stream_pb2.StreamCloseStatus.ValueType( + value.status + ) + ) + if value.previous_completed_task is not None: + message.previous_completed_task.CopyFrom( + _WorkflowTaskEpochTransferTypeConverter().to_transfer_type( + value.previous_completed_task + ) + ) + return message + + +temporalio.converter.transfer_type_convertible( + _CloseStreamRequestTransferTypeConverter +)(CloseStreamRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseStreamResponse: + pass + + +class _CloseStreamResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseStreamResponse, + temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse, + type_hint: type[CloseStreamResponse], + ) -> CloseStreamResponse: + return CloseStreamResponse() + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseStreamResponse, + ) -> temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.CloseStreamResponse() + ) + return message + + +temporalio.converter.transfer_type_convertible( + _CloseStreamResponseTransferTypeConverter +)(CloseStreamResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class ReadRecordsRequest: + """Read Records from an offset.""" + + stream: StreamRef + from_offset: int | None = None + start: StreamReadStart | None = None + max_records: int | None = None + wait: datetime.timedelta | None = None + create_wait: datetime.timedelta | None = None + last: int | None = None + max_bytes: int | None = None + + +class _ReadRecordsRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + ReadRecordsRequest, + temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest, + type_hint: type[ReadRecordsRequest], + ) -> ReadRecordsRequest: + if not value.HasField("stream"): + raise ValueError("missing required field ReadRecordsRequest.stream") + stream = _StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + return ReadRecordsRequest( + stream=stream, + from_offset=value.from_offset if value.from_offset != 0 else None, + start=StreamReadStart(value.start) if value.start != 0 else None, + max_records=value.max_records if value.max_records != 0 else None, + wait=_support.duration_from_proto(value.wait) + if value.HasField("wait") + else None, + create_wait=_support.duration_from_proto(value.create_wait) + if value.HasField("create_wait") + else None, + last=value.last if value.last != 0 else None, + max_bytes=value.max_bytes if value.max_bytes != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: ReadRecordsRequest, + ) -> temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest() + ) + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + if value.from_offset is not None: + message.from_offset = value.from_offset + if value.start is not None: + message.start = ( + temporalio.api.enums.v1.stream_pb2.StreamReadStart.ValueType( + value.start + ) + ) + if value.max_records is not None: + message.max_records = value.max_records + if value.wait is not None: + message.wait.CopyFrom(_support.duration_to_proto(value.wait)) + if value.create_wait is not None: + message.create_wait.CopyFrom(_support.duration_to_proto(value.create_wait)) + if value.last is not None: + message.last = value.last + if value.max_bytes is not None: + message.max_bytes = value.max_bytes + return message + + +temporalio.converter.transfer_type_convertible( + _ReadRecordsRequestTransferTypeConverter +)(ReadRecordsRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class ReadRecordsResponse: + records: list[StreamRecord] | None = dataclasses.field(default_factory=list) + next_offset: int | None = None + done: bool | None = None + + +class _ReadRecordsResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + ReadRecordsResponse, + temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse, + type_hint: type[ReadRecordsResponse], + ) -> ReadRecordsResponse: + return ReadRecordsResponse( + records=[ + _StreamRecordTransferTypeConverter().from_transfer_type( + item, StreamRecord + ) + for item in value.records + ], + next_offset=value.next_offset if value.next_offset != 0 else None, + done=value.done if bool(value.done) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: ReadRecordsResponse, + ) -> temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse() + ) + if value.records: + for element in value.records: + item = message.records.add() + item.CopyFrom( + _StreamRecordTransferTypeConverter().to_transfer_type(element) + ) + if value.next_offset is not None: + message.next_offset = value.next_offset + if value.done is not None: + message.done = value.done + return message + + +temporalio.converter.transfer_type_convertible( + _ReadRecordsResponseTransferTypeConverter +)(ReadRecordsResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamRecord: + """One item that a read gives.""" + + item: StreamItem + offset: int | None = None + + +class _StreamRecordTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamRecord, temporalio.api.stream.v1.message_pb2.StreamRecord + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRecord] | None = ( + temporalio.api.stream.v1.message_pb2.StreamRecord + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamRecord, + type_hint: type[StreamRecord], + ) -> StreamRecord: + _oneof_item: StreamItemData | StreamItemSuperseded | StreamItemClose + match value.WhichOneof("item"): + case "data": + _oneof_item = StreamItemData( + _StreamDataTransferTypeConverter().from_transfer_type( + value.data, StreamData + ) + ) + case "superseded": + _oneof_item = StreamItemSuperseded( + _SupersessionTransferTypeConverter().from_transfer_type( + value.superseded, Supersession + ) + ) + case "close": + _oneof_item = StreamItemClose( + _StreamCloseTransferTypeConverter().from_transfer_type( + value.close, StreamClose + ) + ) + case None: + raise ValueError("missing required field StreamRecord.item") + return StreamRecord( + item=_oneof_item, + offset=value.offset if value.offset != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamRecord, + ) -> temporalio.api.stream.v1.message_pb2.StreamRecord: + message = temporalio.api.stream.v1.message_pb2.StreamRecord() + match value.item: + case StreamItemData(): + message.data.CopyFrom( + _StreamDataTransferTypeConverter().to_transfer_type( + value.item.value + ) + ) + case StreamItemSuperseded(): + message.superseded.CopyFrom( + _SupersessionTransferTypeConverter().to_transfer_type( + value.item.value + ) + ) + case StreamItemClose(): + message.close.CopyFrom( + _StreamCloseTransferTypeConverter().to_transfer_type( + value.item.value + ) + ) + if value.offset is not None: + message.offset = value.offset + return message + + +temporalio.converter.transfer_type_convertible(_StreamRecordTransferTypeConverter)( + StreamRecord +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamData: + """One Record that a Producer appended.""" + + body: typing.Any + producer: Producer | None = None + + +class _StreamDataTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamData, temporalio.api.stream.v1.message_pb2.StreamData + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamData] | None = ( + temporalio.api.stream.v1.message_pb2.StreamData + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamData, + type_hint: type[StreamData], + ) -> StreamData: + if not value.HasField("body"): + raise ValueError("missing required field StreamData.body") + body = _support.payload_from_proto(value.body) + return StreamData( + body=body, + producer=_ProducerTransferTypeConverter().from_transfer_type( + value.producer, Producer + ) + if value.HasField("producer") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamData, + ) -> temporalio.api.stream.v1.message_pb2.StreamData: + message = temporalio.api.stream.v1.message_pb2.StreamData() + message.body.CopyFrom(_support.payload_to_proto(value.body)) + if value.producer is not None: + message.producer.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.producer) + ) + return message + + +temporalio.converter.transfer_type_convertible(_StreamDataTransferTypeConverter)( + StreamData +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Supersession: + """Which earlier Records a Superseded item withdraws.""" + + producer_id: str + failed_tasks: list[int] | None = dataclasses.field(default_factory=list) + producer_kind: StreamProducerKind | None = None + epoch: ProducerEpoch | None = None + previous_completed_task: WorkflowTaskEpoch | None = None + withdraw_run: str | None = None + run_id: str | None = None + + +class _SupersessionTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Supersession, temporalio.api.stream.v1.message_pb2.Supersession + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.Supersession] | None = ( + temporalio.api.stream.v1.message_pb2.Supersession + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.Supersession, + type_hint: type[Supersession], + ) -> Supersession: + producer_id = value.producer_id + return Supersession( + producer_id=producer_id, + failed_tasks=list(value.failed_tasks), + producer_kind=StreamProducerKind(value.producer_kind) + if value.producer_kind != 0 + else None, + epoch=_ProducerEpochTransferTypeConverter().from_transfer_type( + value.epoch, ProducerEpoch + ) + if value.HasField("epoch") + else None, + previous_completed_task=_WorkflowTaskEpochTransferTypeConverter().from_transfer_type( + value.previous_completed_task, WorkflowTaskEpoch + ) + if value.HasField("previous_completed_task") + else None, + withdraw_run=value.withdraw_run if bool(value.withdraw_run) else None, + run_id=value.run_id if bool(value.run_id) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Supersession, + ) -> temporalio.api.stream.v1.message_pb2.Supersession: + message = temporalio.api.stream.v1.message_pb2.Supersession() + message.producer_id = value.producer_id + if value.failed_tasks: + message.failed_tasks.extend(value.failed_tasks) + if value.producer_kind is not None: + message.producer_kind = ( + temporalio.api.enums.v1.stream_pb2.StreamProducerKind.ValueType( + value.producer_kind + ) + ) + if value.epoch is not None: + message.epoch.CopyFrom( + _ProducerEpochTransferTypeConverter().to_transfer_type(value.epoch) + ) + if value.previous_completed_task is not None: + message.previous_completed_task.CopyFrom( + _WorkflowTaskEpochTransferTypeConverter().to_transfer_type( + value.previous_completed_task + ) + ) + if value.withdraw_run is not None: + message.withdraw_run = value.withdraw_run + if value.run_id is not None: + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_SupersessionTransferTypeConverter)( + Supersession +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamClose: + """The close of a Stream. A read gives it last.""" + + producer: Producer | None = None + result: typing.Any | None = None + failure: BaseException | None = None + status: StreamCloseStatus | None = None + + +class _StreamCloseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamClose, temporalio.api.stream.v1.message_pb2.StreamClose + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamClose] | None = ( + temporalio.api.stream.v1.message_pb2.StreamClose + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamClose, + type_hint: type[StreamClose], + ) -> StreamClose: + return StreamClose( + producer=_ProducerTransferTypeConverter().from_transfer_type( + value.producer, Producer + ) + if value.HasField("producer") + else None, + result=_support.payload_from_proto(value.result) + if value.HasField("result") + else None, + failure=_support.failure_from_proto(value.failure) + if value.HasField("failure") + else None, + status=StreamCloseStatus(value.status) if value.status != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamClose, + ) -> temporalio.api.stream.v1.message_pb2.StreamClose: + message = temporalio.api.stream.v1.message_pb2.StreamClose() + if value.producer is not None: + message.producer.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.producer) + ) + if value.result is not None: + message.result.CopyFrom(_support.payload_to_proto(value.result)) + if value.failure is not None: + message.failure.CopyFrom(_support.failure_to_proto(value.failure)) + if value.status is not None: + message.status = ( + temporalio.api.enums.v1.stream_pb2.StreamCloseStatus.ValueType( + value.status + ) + ) + return message + + +temporalio.converter.transfer_type_convertible(_StreamCloseTransferTypeConverter)( + StreamClose +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class DescribeStreamRequest: + stream: StreamRef + + +class _DescribeStreamRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + DescribeStreamRequest, + temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest, + ] +): + transfer_type: ( + type[temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest, + type_hint: type[DescribeStreamRequest], + ) -> DescribeStreamRequest: + if not value.HasField("stream"): + raise ValueError("missing required field DescribeStreamRequest.stream") + stream = _StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + return DescribeStreamRequest( + stream=stream, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: DescribeStreamRequest, + ) -> temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest: + message = ( + temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest() + ) + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + return message + + +temporalio.converter.transfer_type_convertible( + _DescribeStreamRequestTransferTypeConverter +)(DescribeStreamRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class DescribeStreamResponse: + status: StreamStatus | None = None + last_offset: int | None = None + retention: datetime.timedelta | None = None + + +class _DescribeStreamResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + DescribeStreamResponse, + temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse, + ] +): + transfer_type: ( + type[ + temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse + ] + | None + ) = temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse, + type_hint: type[DescribeStreamResponse], + ) -> DescribeStreamResponse: + return DescribeStreamResponse( + status=StreamStatus(value.status) if value.status != 0 else None, + last_offset=value.last_offset if value.last_offset != 0 else None, + retention=_support.duration_from_proto(value.retention) + if value.HasField("retention") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: DescribeStreamResponse, + ) -> temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse: + message = temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamResponse() + if value.status is not None: + message.status = temporalio.api.enums.v1.stream_pb2.StreamStatus.ValueType( + value.status + ) + if value.last_offset is not None: + message.last_offset = value.last_offset + if value.retention is not None: + message.retention.CopyFrom(_support.duration_to_proto(value.retention)) + return message + + +temporalio.converter.transfer_type_convertible( + _DescribeStreamResponseTransferTypeConverter +)(DescribeStreamResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamItemData: + value: StreamData + + +@dataclasses.dataclass(slots=True) +class StreamItemSuperseded: + value: Supersession + + +@dataclasses.dataclass(slots=True) +class StreamItemClose: + value: StreamClose + + +StreamItem = StreamItemData | StreamItemSuperseded | StreamItemClose +"""A read item: a Record, a Superseded item or the close.""" diff --git a/temporalio/common/_streams/system/operations/__init__.py b/temporalio/common/_streams/system/operations/__init__.py new file mode 100644 index 0000000000..0dd909df4d --- /dev/null +++ b/temporalio/common/_streams/system/operations/__init__.py @@ -0,0 +1,3 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations diff --git a/temporalio/common/_streams/system/operations/create_stream.py b/temporalio/common/_streams/system/operations/create_stream.py new file mode 100644 index 0000000000..94fc4a4926 --- /dev/null +++ b/temporalio/common/_streams/system/operations/create_stream.py @@ -0,0 +1,64 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import datetime + +import temporalio.workflow + +from .._resources import Stream +from ..models import ( + CreateStreamRequest, + CreateStreamResult, + Producer, + StreamRef, +) + + +async def _create_stream( + request: CreateStreamRequest, +) -> Stream: + nexus_client = temporalio.workflow.create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="CreateStream", + input=request, + output_type=CreateStreamResult, + ) + result = await handle + if result.stream is None: + raise RuntimeError("CreateStream answered without `stream`") + return Stream( + stream=result.stream, + created=result.created or None, + status=result.status or None, + ) + + +async def create_stream( + *, + stream_id: str, + retention: datetime.timedelta | None = None, + withdraw_run: str | None = None, + owner: Producer | None = None, + ref: StreamRef | None = None, +) -> Stream: + """ + Args: + 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. + owner: The Workflow run that owns the Stream. The Workflow transport sets it. + ref: The Stream's ref, with the owner fields. + """ + request = CreateStreamRequest( + stream_id=stream_id, + retention=retention, + withdraw_run=withdraw_run, + owner=owner, + ref=ref, + ) + return await _create_stream(request) diff --git a/temporalio/common/_streams/system/services.py b/temporalio/common/_streams/system/services.py new file mode 100644 index 0000000000..bbb5f522cc --- /dev/null +++ b/temporalio/common/_streams/system/services.py @@ -0,0 +1,51 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +from nexusrpc import Operation, service + +from .models import ( + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + CreateStreamRequest, + CreateStreamResult, + DescribeStreamRequest, + DescribeStreamResponse, + ReadRecordsRequest, + ReadRecordsResponse, +) + + +@service(name="temporal.api.streamservice.v1.StreamService") +class StreamService: + """ + .. warning:: + This API is experimental and subject to change. + """ + + create_stream: Operation[ + CreateStreamRequest, + CreateStreamResult, + ] = Operation(name="CreateStream") + + append_records: Operation[ + AppendRecordsRequest, + AppendRecordsResponse, + ] = Operation(name="AppendRecords") + + close_stream: Operation[ + CloseStreamRequest, + CloseStreamResponse, + ] = Operation(name="CloseStream") + + read_records: Operation[ + ReadRecordsRequest, + ReadRecordsResponse, + ] = Operation(name="ReadRecords") + + describe_stream: Operation[ + DescribeStreamRequest, + DescribeStreamResponse, + ] = Operation(name="DescribeStream") diff --git a/temporalio/common/_streams/writer_system/__init__.py b/temporalio/common/_streams/writer_system/__init__.py new file mode 100644 index 0000000000..d1849f5471 --- /dev/null +++ b/temporalio/common/_streams/writer_system/__init__.py @@ -0,0 +1,49 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +from .models import ( + AppendRecord, + CloseWriterRequest, + CloseWriterResponse, + CloseWritersRequest, + CloseWritersResponse, + Execution, + ExecutionType, + FreeWriterRequest, + FreeWriterResponse, + OpenWriterRequest, + Producer, + ProducerEpoch, + StreamProducerKind, + StreamRef, + Timestamp, + WorkflowTaskEpoch, + WriteRecordsRequest, + WriteRecordsResponse, + WriterResult, +) +from .services import StreamWriter + +__all__ = [ + "AppendRecord", + "CloseWriterRequest", + "CloseWriterResponse", + "CloseWritersRequest", + "CloseWritersResponse", + "Execution", + "ExecutionType", + "FreeWriterRequest", + "FreeWriterResponse", + "OpenWriterRequest", + "Producer", + "ProducerEpoch", + "StreamProducerKind", + "StreamRef", + "Timestamp", + "WorkflowTaskEpoch", + "WriteRecordsRequest", + "WriteRecordsResponse", + "WriterResult", + "StreamWriter", +] diff --git a/temporalio/common/_streams/writer_system/client.py b/temporalio/common/_streams/writer_system/client.py new file mode 100644 index 0000000000..7046efe0c2 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client.py @@ -0,0 +1,256 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import abc +import collections.abc +import dataclasses +import datetime +import typing + +from .models import ( + AppendRecord, + CloseWriterRequest, + CloseWriterResponse, + CloseWritersRequest, + CloseWritersResponse, + FreeWriterRequest, + FreeWriterResponse, + OpenWriterRequest, + Producer, + StreamRef, + WorkflowTaskEpoch, + WriteRecordsRequest, + WriteRecordsResponse, + WriterResult, +) + +_OutputT = typing.TypeVar("_OutputT") + + +class _NexusClient(typing.Protocol): + """Starts the Nexus operations of one service from code outside a Workflow.""" + + @typing.overload + async def start_operation( + self, + operation: str, + input: typing.Any, + *, + output_type: type[_OutputT], + ) -> collections.abc.Awaitable[_OutputT]: ... + + @typing.overload + async def start_operation( + self, + operation: str, + input: typing.Any, + ) -> collections.abc.Awaitable[None]: ... + + +class StreamWriterClientExtension(abc.ABC): + """Starts the operations of ``StreamWriter`` from code outside a Workflow. + + A client inherits this class and implements ``_create_nexus_client``. + + .. warning:: + This API is experimental and subject to change. + """ + + @abc.abstractmethod + def _create_nexus_client( + self, + *, + service: str, + endpoint: str | None, + ) -> _NexusClient: + """Return the Nexus client that starts this service's operations. + + ``endpoint`` is ``None`` for a service with no endpoint, which the client serves. + """ + ... + + async def _open_writer( + self, + request: OpenWriterRequest, + ) -> Writer: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint=None, + ) + handle = await nexus_client.start_operation( + operation="OpenWriter", + input=request, + output_type=WriterResult, + ) + result = await handle + if result.writer_id is None: + raise RuntimeError("OpenWriter answered without `writer_id`") + return Writer( + _client=self, + writer_id=result.writer_id, + ) + + async def open_writer( + self, + *, + stream: StreamRef, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + max_batch_size: int | None = None, + max_retry_duration: datetime.timedelta | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + ) -> Writer: + """ + Args: + failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. + max_batch_size: A write sends at most this many Records in one batch. Zero is + 1000. + max_retry_duration: How long a write retries a Store failure. Empty is 10 s. + previous_completed_task: Local Activity only. The run's last completed Workflow + Task. + """ + request = OpenWriterRequest( + stream=stream, + failed_tasks=failed_tasks or [], + producer=producer, + max_batch_size=max_batch_size, + max_retry_duration=max_retry_duration, + previous_completed_task=previous_completed_task, + ) + return await self._open_writer(request) + + async def _write_records( + self, + request: WriteRecordsRequest, + ) -> collections.abc.Awaitable[WriteRecordsResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint=None, + ) + return await nexus_client.start_operation( + operation="WriteRecords", + input=request, + output_type=WriteRecordsResponse, + ) + + async def write_records( + self, + *, + records: list[AppendRecord] | None = None, + writer_id: int | None = None, + ) -> collections.abc.Awaitable[WriteRecordsResponse,]: + request = WriteRecordsRequest( + records=records or [], + writer_id=writer_id, + ) + return await self._write_records(request) + + async def _close_writer( + self, + request: CloseWriterRequest, + ) -> collections.abc.Awaitable[CloseWriterResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint=None, + ) + return await nexus_client.start_operation( + operation="CloseWriter", + input=request, + output_type=CloseWriterResponse, + ) + + async def close_writer( + self, + *, + writer_id: int | None = None, + ) -> collections.abc.Awaitable[CloseWriterResponse,]: + request = CloseWriterRequest( + writer_id=writer_id, + ) + return await self._close_writer(request) + + async def _free_writer( + self, + request: FreeWriterRequest, + ) -> collections.abc.Awaitable[FreeWriterResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint=None, + ) + return await nexus_client.start_operation( + operation="FreeWriter", + input=request, + output_type=FreeWriterResponse, + ) + + async def free_writer( + self, + *, + writer_id: int | None = None, + ) -> collections.abc.Awaitable[FreeWriterResponse,]: + request = FreeWriterRequest( + writer_id=writer_id, + ) + return await self._free_writer(request) + + async def _close_writers( + self, + request: CloseWritersRequest, + ) -> collections.abc.Awaitable[CloseWritersResponse,]: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint=None, + ) + return await nexus_client.start_operation( + operation="CloseWriters", + input=request, + output_type=CloseWritersResponse, + ) + + async def close_writers( + self, + ) -> collections.abc.Awaitable[CloseWritersResponse,]: + """Drop every open writer, for example at shutdown.""" + request = CloseWritersRequest() + return await self._close_writers(request) + + +@dataclasses.dataclass +class Writer: + """An open writer of one Producer. The crate keeps its epoch and sequence.""" + + _client: StreamWriterClientExtension = dataclasses.field(repr=False, compare=False) + writer_id: int + + async def write_records( + self, + records: list[AppendRecord] | None = None, + ) -> WriteRecordsResponse: + """Write Records at once. The answer has the first and the last offset.""" + request = WriteRecordsRequest( + writer_id=self.writer_id, + records=records or [], + ) + handle = await self._client._write_records(request) + return await handle + + async def close( + self, + ) -> CloseWriterResponse: + """Release the handle. The last one waits for a write, then drops the writer.""" + request = CloseWriterRequest( + writer_id=self.writer_id, + ) + handle = await self._client._close_writer(request) + return await handle + + async def free( + self, + ) -> FreeWriterResponse: + """Release the handle with no wait, for a handle that lang drops.""" + request = FreeWriterRequest( + writer_id=self.writer_id, + ) + handle = await self._client._free_writer(request) + return await handle diff --git a/temporalio/common/_streams/writer_system/models.py b/temporalio/common/_streams/writer_system/models.py new file mode 100644 index 0000000000..833405207f --- /dev/null +++ b/temporalio/common/_streams/writer_system/models.py @@ -0,0 +1,893 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +import dataclasses +import datetime +import enum +import typing + +import google.protobuf.timestamp_pb2 +import typing_extensions + +import temporalio.api.common.v1.message_pb2 +import temporalio.api.enums.v1.common_pb2 +import temporalio.api.enums.v1.stream_pb2 +import temporalio.api.stream.v1.message_pb2 +import temporalio.common._streams.proto.v1.runtime_pb2 +import temporalio.converter +import temporalio.nexus.system._support as _support + + +class ExecutionType(enum.IntEnum): + ExecutionTypeUnspecified = 0 + ExecutionTypeWorkflow = 1 + ExecutionTypeActivity = 2 + ExecutionTypeNexusOperation = 3 + + +class StreamProducerKind(enum.IntEnum): + StreamProducerKindUnspecified = 0 + StreamProducerKindWorkflow = 1 + StreamProducerKindLocalActivity = 2 + StreamProducerKindActivity = 3 + StreamProducerKindClient = 4 + StreamProducerKindStandaloneActivity = 5 + + +@dataclasses.dataclass(slots=True) +class OpenWriterRequest: + """Open a writer for one Producer of one Stream.""" + + stream: StreamRef + failed_tasks: list[int] | None = dataclasses.field(default_factory=list) + producer: Producer | None = None + max_batch_size: int | None = None + max_retry_duration: datetime.timedelta | None = None + previous_completed_task: WorkflowTaskEpoch | None = None + + +class _OpenWriterRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + OpenWriterRequest, + temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest, + type_hint: type[OpenWriterRequest], + ) -> OpenWriterRequest: + if not value.HasField("stream"): + raise ValueError("missing required field OpenWriterRequest.stream") + stream = _StreamRefTransferTypeConverter().from_transfer_type( + value.stream, StreamRef + ) + return OpenWriterRequest( + stream=stream, + failed_tasks=list(value.failed_tasks), + producer=_ProducerTransferTypeConverter().from_transfer_type( + value.producer, Producer + ) + if value.HasField("producer") + else None, + max_batch_size=value.max_batch_size if value.max_batch_size != 0 else None, + max_retry_duration=_support.duration_from_proto(value.max_retry_duration) + if value.HasField("max_retry_duration") + else None, + previous_completed_task=_WorkflowTaskEpochTransferTypeConverter().from_transfer_type( + value.previous_completed_task, WorkflowTaskEpoch + ) + if value.HasField("previous_completed_task") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: OpenWriterRequest, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest: + message = temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest() + message.stream.CopyFrom( + _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + ) + if value.failed_tasks: + message.failed_tasks.extend(value.failed_tasks) + if value.producer is not None: + message.producer.CopyFrom( + _ProducerTransferTypeConverter().to_transfer_type(value.producer) + ) + if value.max_batch_size is not None: + message.max_batch_size = value.max_batch_size + if value.max_retry_duration is not None: + message.max_retry_duration.CopyFrom( + _support.duration_to_proto(value.max_retry_duration) + ) + if value.previous_completed_task is not None: + message.previous_completed_task.CopyFrom( + _WorkflowTaskEpochTransferTypeConverter().to_transfer_type( + value.previous_completed_task + ) + ) + return message + + +temporalio.converter.transfer_type_convertible(_OpenWriterRequestTransferTypeConverter)( + OpenWriterRequest +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamRef: + """Identifies one Stream. The id is opaque to callers.""" + + id: str + owner: Execution | None = None + backend: str | None = None + + +class _StreamRefTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamRef, temporalio.api.stream.v1.message_pb2.StreamRef + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRef] | None = ( + temporalio.api.stream.v1.message_pb2.StreamRef + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamRef, + type_hint: type[StreamRef], + ) -> StreamRef: + id = value.id + return StreamRef( + id=id, + owner=_ExecutionTransferTypeConverter().from_transfer_type( + value.owner, Execution + ) + if value.HasField("owner") + else None, + backend=value.backend if bool(value.backend) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamRef, + ) -> temporalio.api.stream.v1.message_pb2.StreamRef: + message = temporalio.api.stream.v1.message_pb2.StreamRef() + message.id = value.id + if value.owner is not None: + message.owner.CopyFrom( + _ExecutionTransferTypeConverter().to_transfer_type(value.owner) + ) + if value.backend is not None: + message.backend = value.backend + return message + + +temporalio.converter.transfer_type_convertible(_StreamRefTransferTypeConverter)( + StreamRef +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Execution: + """A Workflow run chain or a Standalone Activity run.""" + + type: ExecutionType | None = None + business_id: str | None = None + run_id: str | None = None + + +class _ExecutionTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Execution, temporalio.api.common.v1.message_pb2.Execution + ] +): + transfer_type: type[temporalio.api.common.v1.message_pb2.Execution] | None = ( + temporalio.api.common.v1.message_pb2.Execution + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.common.v1.message_pb2.Execution, + type_hint: type[Execution], + ) -> Execution: + return Execution( + type=ExecutionType(value.type) if value.type != 0 else None, + business_id=value.business_id if bool(value.business_id) else None, + run_id=value.run_id if bool(value.run_id) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Execution, + ) -> temporalio.api.common.v1.message_pb2.Execution: + message = temporalio.api.common.v1.message_pb2.Execution() + if value.type is not None: + message.type = temporalio.api.enums.v1.common_pb2.ExecutionType.ValueType( + value.type + ) + if value.business_id is not None: + message.business_id = value.business_id + if value.run_id is not None: + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_ExecutionTransferTypeConverter)( + Execution +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Producer: + """The Producer of a batch.""" + + id: str + kind: StreamProducerKind | None = None + epoch: ProducerEpoch | None = None + run_id: str | None = None + + +class _ProducerTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Producer, temporalio.api.stream.v1.message_pb2.Producer + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.Producer] | None = ( + temporalio.api.stream.v1.message_pb2.Producer + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.Producer, + type_hint: type[Producer], + ) -> Producer: + id = value.id + return Producer( + id=id, + kind=StreamProducerKind(value.kind) if value.kind != 0 else None, + epoch=_ProducerEpochTransferTypeConverter().from_transfer_type( + value.epoch, ProducerEpoch + ) + if value.HasField("epoch") + else None, + run_id=value.run_id if bool(value.run_id) else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Producer, + ) -> temporalio.api.stream.v1.message_pb2.Producer: + message = temporalio.api.stream.v1.message_pb2.Producer() + message.id = value.id + if value.kind is not None: + message.kind = ( + temporalio.api.enums.v1.stream_pb2.StreamProducerKind.ValueType( + value.kind + ) + ) + if value.epoch is not None: + message.epoch.CopyFrom( + _ProducerEpochTransferTypeConverter().to_transfer_type(value.epoch) + ) + if value.run_id is not None: + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_ProducerTransferTypeConverter)(Producer) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class ProducerEpoch: + """The ordered attempt of a Producer. A higher epoch supersedes a lower one.""" + + task_started_event_id: int | None = None + task_attempt: int | None = None + task_started_time: Timestamp | None = None + attempt_scheduled_time: Timestamp | None = None + attempt: int | None = None + + +class _ProducerEpochTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + ProducerEpoch, temporalio.api.stream.v1.message_pb2.ProducerEpoch + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.ProducerEpoch] | None = ( + temporalio.api.stream.v1.message_pb2.ProducerEpoch + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.ProducerEpoch, + type_hint: type[ProducerEpoch], + ) -> ProducerEpoch: + return ProducerEpoch( + task_started_event_id=value.task_started_event_id + if value.task_started_event_id != 0 + else None, + task_attempt=value.task_attempt if value.task_attempt != 0 else None, + task_started_time=_TimestampTransferTypeConverter().from_transfer_type( + value.task_started_time, Timestamp + ) + if value.HasField("task_started_time") + else None, + attempt_scheduled_time=_TimestampTransferTypeConverter().from_transfer_type( + value.attempt_scheduled_time, Timestamp + ) + if value.HasField("attempt_scheduled_time") + else None, + attempt=value.attempt if value.attempt != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: ProducerEpoch, + ) -> temporalio.api.stream.v1.message_pb2.ProducerEpoch: + message = temporalio.api.stream.v1.message_pb2.ProducerEpoch() + if value.task_started_event_id is not None: + message.task_started_event_id = value.task_started_event_id + if value.task_attempt is not None: + message.task_attempt = value.task_attempt + if value.task_started_time is not None: + message.task_started_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type( + value.task_started_time + ) + ) + if value.attempt_scheduled_time is not None: + message.attempt_scheduled_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type( + value.attempt_scheduled_time + ) + ) + if value.attempt is not None: + message.attempt = value.attempt + return message + + +temporalio.converter.transfer_type_convertible(_ProducerEpochTransferTypeConverter)( + ProducerEpoch +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class Timestamp: + seconds: int | None = None + nanos: int | None = None + + +class _TimestampTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Timestamp, google.protobuf.timestamp_pb2.Timestamp + ] +): + transfer_type: type[google.protobuf.timestamp_pb2.Timestamp] | None = ( + google.protobuf.timestamp_pb2.Timestamp + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: google.protobuf.timestamp_pb2.Timestamp, + type_hint: type[Timestamp], + ) -> Timestamp: + return Timestamp( + seconds=value.seconds if value.seconds != 0 else None, + nanos=value.nanos if value.nanos != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: Timestamp, + ) -> google.protobuf.timestamp_pb2.Timestamp: + message = google.protobuf.timestamp_pb2.Timestamp() + if value.seconds is not None: + message.seconds = value.seconds + if value.nanos is not None: + message.nanos = value.nanos + return message + + +temporalio.converter.transfer_type_convertible(_TimestampTransferTypeConverter)( + Timestamp +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WorkflowTaskEpoch: + """A Workflow Task attempt of the run.""" + + started_event_id: int | None = None + attempt: int | None = None + started_time: Timestamp | None = None + + +class _WorkflowTaskEpochTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WorkflowTaskEpoch, temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch + ] +): + transfer_type: ( + type[temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch] | None + ) = temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch, + type_hint: type[WorkflowTaskEpoch], + ) -> WorkflowTaskEpoch: + return WorkflowTaskEpoch( + started_event_id=value.started_event_id + if value.started_event_id != 0 + else None, + attempt=value.attempt if value.attempt != 0 else None, + started_time=_TimestampTransferTypeConverter().from_transfer_type( + value.started_time, Timestamp + ) + if value.HasField("started_time") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WorkflowTaskEpoch, + ) -> temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch: + message = temporalio.api.stream.v1.message_pb2.WorkflowTaskEpoch() + if value.started_event_id is not None: + message.started_event_id = value.started_event_id + if value.attempt is not None: + message.attempt = value.attempt + if value.started_time is not None: + message.started_time.CopyFrom( + _TimestampTransferTypeConverter().to_transfer_type(value.started_time) + ) + return message + + +temporalio.converter.transfer_type_convertible(_WorkflowTaskEpochTransferTypeConverter)( + WorkflowTaskEpoch +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WriterResult: + writer_id: int | None = None + + +class _WriterResultTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WriterResult, temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse, + type_hint: type[WriterResult], + ) -> WriterResult: + return WriterResult( + writer_id=value.writer_id if value.writer_id != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WriterResult, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse: + message = temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterResponse() + if value.writer_id is not None: + message.writer_id = value.writer_id + return message + + +temporalio.converter.transfer_type_convertible(_WriterResultTransferTypeConverter)( + WriterResult +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WriteRecordsRequest: + records: list[AppendRecord] | None = dataclasses.field(default_factory=list) + writer_id: int | None = None + + +class _WriteRecordsRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WriteRecordsRequest, + temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest, + type_hint: type[WriteRecordsRequest], + ) -> WriteRecordsRequest: + return WriteRecordsRequest( + records=[ + _AppendRecordTransferTypeConverter().from_transfer_type( + item, AppendRecord + ) + for item in value.records + ], + writer_id=value.writer_id if value.writer_id != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WriteRecordsRequest, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest: + message = temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsRequest() + if value.records: + for element in value.records: + item = message.records.add() + item.CopyFrom( + _AppendRecordTransferTypeConverter().to_transfer_type(element) + ) + if value.writer_id is not None: + message.writer_id = value.writer_id + return message + + +temporalio.converter.transfer_type_convertible( + _WriteRecordsRequestTransferTypeConverter +)(WriteRecordsRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class AppendRecord: + """One Record to append.""" + + body: typing.Any + + +class _AppendRecordTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + AppendRecord, temporalio.api.stream.v1.message_pb2.AppendRecord + ] +): + transfer_type: type[temporalio.api.stream.v1.message_pb2.AppendRecord] | None = ( + temporalio.api.stream.v1.message_pb2.AppendRecord + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.AppendRecord, + type_hint: type[AppendRecord], + ) -> AppendRecord: + if not value.HasField("body"): + raise ValueError("missing required field AppendRecord.body") + body = _support.payload_from_proto(value.body) + return AppendRecord( + body=body, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: AppendRecord, + ) -> temporalio.api.stream.v1.message_pb2.AppendRecord: + message = temporalio.api.stream.v1.message_pb2.AppendRecord() + message.body.CopyFrom(_support.payload_to_proto(value.body)) + return message + + +temporalio.converter.transfer_type_convertible(_AppendRecordTransferTypeConverter)( + AppendRecord +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WriteRecordsResponse: + first_offset: int | None = None + last_offset: int | None = None + + +class _WriteRecordsResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WriteRecordsResponse, + temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse] + | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse, + type_hint: type[WriteRecordsResponse], + ) -> WriteRecordsResponse: + return WriteRecordsResponse( + first_offset=value.first_offset if value.first_offset != 0 else None, + last_offset=value.last_offset if value.last_offset != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WriteRecordsResponse, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse: + message = temporalio.common._streams.proto.v1.runtime_pb2.WriteRecordsResponse() + if value.first_offset is not None: + message.first_offset = value.first_offset + if value.last_offset is not None: + message.last_offset = value.last_offset + return message + + +temporalio.converter.transfer_type_convertible( + _WriteRecordsResponseTransferTypeConverter +)(WriteRecordsResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseWriterRequest: + writer_id: int | None = None + + +class _CloseWriterRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseWriterRequest, + temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest, + type_hint: type[CloseWriterRequest], + ) -> CloseWriterRequest: + return CloseWriterRequest( + writer_id=value.writer_id if value.writer_id != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseWriterRequest, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest: + message = temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterRequest() + if value.writer_id is not None: + message.writer_id = value.writer_id + return message + + +temporalio.converter.transfer_type_convertible( + _CloseWriterRequestTransferTypeConverter +)(CloseWriterRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseWriterResponse: + last_offset: int | None = None + + +class _CloseWriterResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseWriterResponse, + temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse, + type_hint: type[CloseWriterResponse], + ) -> CloseWriterResponse: + return CloseWriterResponse( + last_offset=value.last_offset if value.last_offset != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseWriterResponse, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse: + message = temporalio.common._streams.proto.v1.runtime_pb2.CloseWriterResponse() + if value.last_offset is not None: + message.last_offset = value.last_offset + return message + + +temporalio.converter.transfer_type_convertible( + _CloseWriterResponseTransferTypeConverter +)(CloseWriterResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class FreeWriterRequest: + writer_id: int | None = None + + +class _FreeWriterRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + FreeWriterRequest, + temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest, + type_hint: type[FreeWriterRequest], + ) -> FreeWriterRequest: + return FreeWriterRequest( + writer_id=value.writer_id if value.writer_id != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: FreeWriterRequest, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest: + message = temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterRequest() + if value.writer_id is not None: + message.writer_id = value.writer_id + return message + + +temporalio.converter.transfer_type_convertible(_FreeWriterRequestTransferTypeConverter)( + FreeWriterRequest +) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class FreeWriterResponse: + pass + + +class _FreeWriterResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + FreeWriterResponse, + temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse, + type_hint: type[FreeWriterResponse], + ) -> FreeWriterResponse: + return FreeWriterResponse() + + @typing_extensions.override + def to_transfer_type( + self, + value: FreeWriterResponse, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse: + message = temporalio.common._streams.proto.v1.runtime_pb2.FreeWriterResponse() + return message + + +temporalio.converter.transfer_type_convertible( + _FreeWriterResponseTransferTypeConverter +)(FreeWriterResponse) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseWritersRequest: + pass + + +class _CloseWritersRequestTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseWritersRequest, + temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest] | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest, + type_hint: type[CloseWritersRequest], + ) -> CloseWritersRequest: + return CloseWritersRequest() + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseWritersRequest, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest: + message = temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersRequest() + return message + + +temporalio.converter.transfer_type_convertible( + _CloseWritersRequestTransferTypeConverter +)(CloseWritersRequest) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class CloseWritersResponse: + failed: int | None = None + + +class _CloseWritersResponseTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + CloseWritersResponse, + temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse, + ] +): + transfer_type: ( + type[temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse] + | None + ) = temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse, + type_hint: type[CloseWritersResponse], + ) -> CloseWritersResponse: + return CloseWritersResponse( + failed=value.failed if value.failed != 0 else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: CloseWritersResponse, + ) -> temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse: + message = temporalio.common._streams.proto.v1.runtime_pb2.CloseWritersResponse() + if value.failed is not None: + message.failed = value.failed + return message + + +temporalio.converter.transfer_type_convertible( + _CloseWritersResponseTransferTypeConverter +)(CloseWritersResponse) # pyright: ignore[reportUnusedCallResult] diff --git a/temporalio/common/_streams/writer_system/services.py b/temporalio/common/_streams/writer_system/services.py new file mode 100644 index 0000000000..9f388ca0e3 --- /dev/null +++ b/temporalio/common/_streams/writer_system/services.py @@ -0,0 +1,51 @@ +# Generated by nexgen v0.2.10. DO NOT EDIT! + +from __future__ import annotations + +from nexusrpc import Operation, service + +from .models import ( + CloseWriterRequest, + CloseWriterResponse, + CloseWritersRequest, + CloseWritersResponse, + FreeWriterRequest, + FreeWriterResponse, + OpenWriterRequest, + WriteRecordsRequest, + WriteRecordsResponse, + WriterResult, +) + + +@service(name="temporal.sdk.streams.v1.StreamWriter") +class StreamWriter: + """ + .. warning:: + This API is experimental and subject to change. + """ + + open_writer: Operation[ + OpenWriterRequest, + WriterResult, + ] = Operation(name="OpenWriter") + + write_records: Operation[ + WriteRecordsRequest, + WriteRecordsResponse, + ] = Operation(name="WriteRecords") + + close_writer: Operation[ + CloseWriterRequest, + CloseWriterResponse, + ] = Operation(name="CloseWriter") + + free_writer: Operation[ + FreeWriterRequest, + FreeWriterResponse, + ] = Operation(name="FreeWriter") + + close_writers: Operation[ + CloseWritersRequest, + CloseWritersResponse, + ] = Operation(name="CloseWriters") diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 15a8e80102..0b3d0be51f 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -1,4 +1,5 @@ import collections.abc +import dataclasses import typing from datetime import timedelta @@ -298,3 +299,47 @@ def versioning_override_from_proto( if proto.auto_upgrade: return temporalio.common.AutoUpgradeVersioningOverride() raise ValueError("unknown versioning override proto shape") + + +class _ReadRequest(typing.Protocol): + __dataclass_fields__: typing.ClassVar[dict[str, typing.Any]] + + @property + def from_offset(self) -> int | None: ... + + @property + def wait(self) -> timedelta | None: ... + + +class _ReadResponse(typing.Protocol): + @property + def records(self) -> collections.abc.Sized | None: ... + + @property + def next_offset(self) -> int | None: ... + + @property + def done(self) -> bool | None: ... + + +_ReadRequestT = typing.TypeVar("_ReadRequestT", bound=_ReadRequest) + + +def next_read_request( + request: _ReadRequestT, response: _ReadResponse +) -> _ReadRequestT | None: + """The request after one page of a Stream read, or None to stop. + + A read ends at a close. A read with no wait also ends at an empty page that didn't move, + since the Store had nothing past the newest Record. + """ + if response.done or response.next_offset is None: + return None + moved = response.next_offset > (request.from_offset or 0) + if not response.records and not moved and request.wait is None: + return None + # The next page starts where this one ended, so the first page's start and last no longer + # apply. + return dataclasses.replace( + request, from_offset=response.next_offset, start=None, last=None + ) diff --git a/tests/contrib/streams/__init__.py b/tests/contrib/streams/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/streams/test_bindings.py b/tests/contrib/streams/test_bindings.py new file mode 100644 index 0000000000..923800d361 --- /dev/null +++ b/tests/contrib/streams/test_bindings.py @@ -0,0 +1,91 @@ +from collections.abc import Awaitable +from datetime import timedelta +from typing import Any + +from typing_extensions import override + +from temporalio.common._streams.system import client, models + + +def _data(offset: int) -> models.StreamRecord: + return models.StreamRecord( + item=models.StreamItemData(models.StreamData(body=None)), offset=offset + ) + + +class FakeNexusClient: + def __init__(self, *answers: models.ReadRecordsResponse) -> None: + self.requests: list[models.ReadRecordsRequest] = [] + self.answers = list(answers) + + async def start_operation( + self, operation: str, input: Any, *, output_type: type[Any] | None = None + ) -> Awaitable[Any]: + assert operation == "ReadRecords" + assert output_type is models.ReadRecordsResponse + self.requests.append(input) + answer = self.answers.pop(0) + + async def done() -> Any: + return answer + + return done() + + +class FakeService(client.StreamServiceClientExtension): + def __init__(self, nexus_client: FakeNexusClient) -> None: + self.nexus_client = nexus_client + + @override + def _create_nexus_client( + self, *, service: str, endpoint: str | None + ) -> FakeNexusClient: + assert service == "temporal.api.streamservice.v1.StreamService" + assert endpoint == "__temporal_system" + return self.nexus_client + + +async def test_a_poll_carries_the_next_offset_until_the_close(): + nexus_client = FakeNexusClient( + models.ReadRecordsResponse(next_offset=8), + models.ReadRecordsResponse( + records=[_data(9)], + next_offset=10, + ), + models.ReadRecordsResponse(next_offset=10, done=True), + ) + stream = client.Stream( + _client=FakeService(nexus_client), + stream=models.StreamRef(id="wf:run"), + created=None, + status=None, + ) + pages = [ + page + async for page in stream.iter_read(from_offset=1, wait=timedelta(seconds=5)) + ] + assert [page.next_offset for page in pages] == [8, 10, 10] + assert [request.from_offset for request in nexus_client.requests] == [1, 8, 10] + assert all(request.wait for request in nexus_client.requests) + + +async def test_a_read_with_no_wait_ends_at_an_empty_page_that_didnt_move(): + nexus_client = FakeNexusClient( + models.ReadRecordsResponse( + records=[_data(4)], + next_offset=5, + ), + models.ReadRecordsResponse(next_offset=5), + ) + pages = [ + page + async for page in FakeService(nexus_client).iter_read_records( + stream=models.StreamRef(id="wf:run"), + start=models.StreamReadStart.StreamReadStartBeginning, + ) + ] + assert len(pages) == 2 + first, second = nexus_client.requests + assert first.start == models.StreamReadStart.StreamReadStartBeginning + # The second page starts where the first ended, so the start no longer applies. + assert (second.from_offset, second.start) == (5, None) From 681d912a99f2942cb22db09100acdff1f4f22f4b Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 31 Oct 2026 01:37:25 -0700 Subject: [PATCH 2/2] Generated the Stream bindings from the API's WIT, and wrote the writer. The writer has no WIT, so its client keeps the shape of nexgen's output by hand and raises the service's errors. The bindings need `StreamDefinition` and the Stream helpers of the support module, so this layer adds them. The Workflow sources fail until Workflow code can write. The Workflow Service bindings move to nexgen 0.2.11, whose dispatch gives an unknown System Nexus call to the generic hook. --- scripts/_wit/stream-service.wit | 202 ++++--- scripts/_wit/stream-writer.wit | 118 ---- scripts/gen_nexus_system_api.py | 2 +- scripts/gen_stream_service.py | 90 +--- temporalio/common/__init__.py | 1 + temporalio/common/_streams/__init__.py | 6 + temporalio/common/_streams/_defn.py | 15 + temporalio/common/_streams/system/__init__.py | 18 +- temporalio/common/_streams/system/_poll.py | 65 +++ .../common/_streams/system/_resource_call.py | 93 ++++ .../_streams/system/_resources/__init__.py | 2 +- .../_streams/system/_resources/stream.py | 485 +++++++++++------ .../system/_system_nexus_interceptor.py | 23 +- temporalio/common/_streams/system/client.py | 439 --------------- .../common/_streams/system/client/__init__.py | 11 + .../_streams/system/client/_extensions.py | 497 +++++++++++++++++ .../_streams/system/client/_nexus_client.py | 50 ++ .../system/client/_resources/__init__.py | 9 + .../system/client/_resources/stream.py | 335 ++++++++++++ temporalio/common/_streams/system/errors.py | 146 +++++ temporalio/common/_streams/system/models.py | 508 +++++++++++------- .../_streams/system/operations/__init__.py | 2 +- .../system/operations/create_stream.py | 103 ++-- temporalio/common/_streams/system/services.py | 8 +- .../common/_streams/writer_system/__init__.py | 7 +- .../writer_system/_resources/__init__.py | 10 + .../writer_system/_resources/writer.py | 13 + .../common/_streams/writer_system/client.py | 256 --------- .../_streams/writer_system/client/__init__.py | 12 + .../writer_system/client/_extensions.py | 338 ++++++++++++ .../writer_system/client/_nexus_client.py | 51 ++ .../client/_resources/__init__.py | 10 + .../writer_system/client/_resources/writer.py | 139 +++++ .../common/_streams/writer_system/models.py | 55 +- .../common/_streams/writer_system/services.py | 3 +- temporalio/nexus/notifications/__init__.py | 2 +- temporalio/nexus/notifications/models.py | 2 +- temporalio/nexus/system/_support.py | 206 +++++++ .../nexus/system/workflow_service/__init__.py | 6 +- .../_system_nexus_interceptor.py | 14 +- .../nexus/system/workflow_service/models.py | 11 +- .../workflow_service/operations/__init__.py | 2 +- .../operations/signal_with_start_workflow.py | 2 +- .../nexus/system/workflow_service/services.py | 2 +- tests/contrib/streams/test_bindings.py | 108 ++-- 45 files changed, 3042 insertions(+), 1435 deletions(-) delete mode 100644 scripts/_wit/stream-writer.wit create mode 100644 temporalio/common/_streams/_defn.py create mode 100644 temporalio/common/_streams/system/_poll.py create mode 100644 temporalio/common/_streams/system/_resource_call.py delete mode 100644 temporalio/common/_streams/system/client.py create mode 100644 temporalio/common/_streams/system/client/__init__.py create mode 100644 temporalio/common/_streams/system/client/_extensions.py create mode 100644 temporalio/common/_streams/system/client/_nexus_client.py create mode 100644 temporalio/common/_streams/system/client/_resources/__init__.py create mode 100644 temporalio/common/_streams/system/client/_resources/stream.py create mode 100644 temporalio/common/_streams/system/errors.py create mode 100644 temporalio/common/_streams/writer_system/_resources/__init__.py create mode 100644 temporalio/common/_streams/writer_system/_resources/writer.py delete mode 100644 temporalio/common/_streams/writer_system/client.py create mode 100644 temporalio/common/_streams/writer_system/client/__init__.py create mode 100644 temporalio/common/_streams/writer_system/client/_extensions.py create mode 100644 temporalio/common/_streams/writer_system/client/_nexus_client.py create mode 100644 temporalio/common/_streams/writer_system/client/_resources/__init__.py create mode 100644 temporalio/common/_streams/writer_system/client/_resources/writer.py diff --git a/scripts/_wit/stream-service.wit b/scripts/_wit/stream-service.wit index 480e8ddfed..c2973c4055 100644 --- a/scripts/_wit/stream-service.wit +++ b/scripts/_wit/stream-service.wit @@ -1,5 +1,5 @@ -// Draft. The Stream service as a System Nexus service, for the Brook API proposal. Each record maps -// to a message of the protos under ../protos/temporal/api (stream/v1, streamservice/v1 and the +// Draft. The Stream service as a System Nexus service, for the Streams API proposal. Each record +// maps to a message of the protos under ../temporal/api (stream/v1, streamservice/v1 and the // Stream enums in enums/v1). This file is the source. Each SDK copies it and generates its Workflow // API and its client with nexgen (`--targets workflow,client --system-nexus`). package temporal:nexus@1.0.0; @@ -19,6 +19,10 @@ interface stream-service { placeholder, }; + /// @nexus.doc "The value type of a Stream's Records. A Stream definition sets it." + /// @nexus.type-parameter + type body-t = placeholder; + /// @nexus.proto "google.protobuf.Timestamp" record timestamp { seconds: option, @@ -28,11 +32,11 @@ interface stream-service { // A Stream that a Workflow creates has the id ::. Each part is // percent-encoded. So a caller that knows the first run id can build the id with no call. /// @nexus.doc "Identifies one Stream. The id is opaque to callers." - /// @nexus.proto "temporal.api.stream.v1.StreamRef" - record stream-ref { - /// @nexus.doc "The Stream id. 1 to 4096 bytes after encoding." + /// @nexus.proto "temporal.api.stream.v1.StreamHandle" + record stream-handle { + /// @nexus.doc "The Stream id, 1 to 4096 bytes." id: string, - /// @nexus.doc "The owning Workflow run chain or Standalone Activity. Empty on a client ref." + /// @nexus.doc "The owning Workflow run chain or Standalone Activity. Empty for a client." owner: option, /// @nexus.doc "Registered backend name, empty for the default. The Store never reads it." backend: option, @@ -40,6 +44,7 @@ interface stream-service { /// @nexus.proto "temporal.api.enums.v1.ExecutionType" enum execution-type { + /// @nexus.fallback execution-type-unspecified, execution-type-workflow, execution-type-activity, @@ -58,6 +63,7 @@ interface stream-service { /// @nexus.proto "temporal.api.enums.v1.StreamProducerKind" enum stream-producer-kind { + /// @nexus.fallback stream-producer-kind-unspecified, stream-producer-kind-workflow, stream-producer-kind-local-activity, @@ -107,9 +113,10 @@ interface stream-service { /// @nexus.doc "One Record to append." /// @nexus.proto "temporal.api.stream.v1.AppendRecord" + /// @nexus.unwrap record append-record { - /// @nexus.doc "The value. It goes through the payload codec." - body: payload, + /// @nexus.doc "The value." + body: body-t, } // An Activity or Local Activity item withdraws the Records of its Producer with a lower epoch. A @@ -118,8 +125,8 @@ interface stream-service { // or strictly between c and t. It also withdraws those at the started event id of c or of t with // a lower epoch than c or t. No task in that range can have completed. /// @nexus.doc "Which earlier Records a Superseded item withdraws." - /// @nexus.proto "temporal.api.stream.v1.Supersession" - record supersession { + /// @nexus.proto "temporal.api.stream.v1.StreamSuperseded" + record stream-superseded { producer-id: string, /// @nexus.doc "Workflow only. Started event ids of failed or timed out Workflow Tasks." failed-tasks: list, @@ -136,6 +143,7 @@ interface stream-service { /// @nexus.doc "How a Stream closed. Only the owner's end on the server terminates or times out." /// @nexus.proto "temporal.api.enums.v1.StreamCloseStatus" enum stream-close-status { + /// @nexus.fallback stream-close-status-unspecified, stream-close-status-completed, stream-close-status-failed, @@ -147,8 +155,8 @@ interface stream-service { /// @nexus.doc "One Record that a Producer appended." /// @nexus.proto "temporal.api.stream.v1.StreamData" record stream-data { - /// @nexus.doc "The value. It goes through the payload codec." - body: payload, + /// @nexus.doc "The value." + body: body-t, producer: option, } @@ -168,46 +176,64 @@ interface stream-service { /// @nexus.proto "temporal.api.stream.v1.StreamRecord.item" variant stream-item { data(stream-data), - superseded(supersession), + superseded(stream-superseded), close(stream-close), } - // An offset is the Redis entry id packed as ms << 20 | seq. It is ordered and unique in one - // Stream, and sparse. Offsets have gaps, so they don't count Records. A read names its Stream, so - // the read checks the offset against that Stream. + // An offset is ordered and unique in one Stream, and sparse. Offsets have gaps, so they don't + // count Records. A read names its Stream, so the read checks the offset against that Stream. /// @nexus.doc "One item that a read gives." /// @nexus.proto "temporal.api.stream.v1.StreamRecord" record stream-record { - item: stream-item, + // A newer server can send an item kind that this SDK doesn't know. Then the item is empty. + item: option, /// @nexus.doc "The place of this item in the Stream. Ordered and unique, with gaps." offset: option, } - /// @nexus.doc "Create a Stream, or get the one that exists, which keeps its Retention." + /// @nexus.doc + /// "Create a Stream, or get the open one. A closed, expired or another owner's Stream fails." /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamRequest" record create-stream-request { - /// @nexus.doc "Required on Redis and memory. On native, empty asks the server for an id." + /// @nexus.doc "The Stream id. Required, and never empty." + /// @nexus.name python="id" typescript="id" dotnet="Id" + /// @nexus.source targets="workflow" python="workflow_stream_id(definition)" + /// typescript="workflowStreamId(definition)" + /// dotnet="StreamWorkflowContext.StreamId(definition)" stream-id: string, /// @nexus.doc "The age after which a Record is dropped. Empty is the Store default." retention: option, + // A wrong run would silently withdraw real Records, so it is never a parameter. The helper + // gives the previous run in a retry run's Workflow, and nothing elsewhere. /// @nexus.doc "Retry run only. The previous run, whose Records this withdraws." + /// @nexus.source targets="workflow" python="withdraw_run()" typescript="withdrawRun()" + /// dotnet="StreamWorkflowContext.WithdrawRun()" + /// @nexus.source targets="client" python="withdraw_run()" typescript="withdrawRun()" + /// dotnet="StreamWorkflowContext.WithdrawRun()" withdraw-run: option, /// @nexus.doc "The Workflow run that owns the Stream. The Workflow transport sets it." + /// @nexus.source targets="workflow" python="workflow_owner()" typescript="workflowOwner()" + /// dotnet="StreamWorkflowContext.Owner()" + /// @nexus.api-omit targets="client" owner: option, - /// @nexus.doc "The Stream's ref, with the owner fields." - %ref: option, - // Native only. Brook's Stores don't read it. + /// @nexus.doc "The handle, whose id is stream-id. With an owner, it needs the handle's owner." + /// @nexus.flatten-in-api fields="backend" fill="id=stream-id" + /// @nexus.source field="owner" targets="workflow" python="workflow_handle_owner()" + /// typescript="workflowHandleOwner()" dotnet="StreamWorkflowContext.HandleOwner()" + stream-handle: option, + // Native only. The SDK Stores don't read it. /// @nexus.omit close-with-run: placeholder, /// @nexus.omit namespace: placeholder, - // Native only. Brook's Stores don't read it. + // Native only. The SDK Stores don't read it. /// @nexus.omit task-token: placeholder, } /// @nexus.proto "temporal.api.enums.v1.StreamStatus" enum stream-status { + /// @nexus.fallback stream-status-unspecified, stream-status-open, stream-status-closed, @@ -220,21 +246,42 @@ interface stream-service { /// "Append. FAILED_PRECONDITION if closed, fenced or format too new. NOT_FOUND if absent." /// @nexus.proto "temporal.api.streamservice.v1.AppendRecordsRequest" record append-records-request { - %stream: stream-ref, + handle: stream-handle, records: list, /// @nexus.doc "Local Activity only. The run's failed or timed out Workflow Tasks." + /// @nexus.api-omit targets="workflow" + /// @nexus.source targets="client" python="local_activity_failed_tasks()" + /// typescript="localActivityFailedTasks()" + /// dotnet="StreamActivityContext.LocalActivityFailedTasks()" failed-tasks: list, /// @nexus.doc "Required outside a Workflow. In a Workflow the transport sets it." + /// @nexus.source targets="workflow" python="workflow_producer()" + /// typescript="workflowProducer()" dotnet="StreamWorkflowContext.Producer()" + /// @nexus.source targets="client" python="client_producer()" typescript="clientProducer()" + /// dotnet="StreamActivityContext.Producer()" producer: option, - /// @nexus.doc "The batch number in this epoch, 1 or more. It is part of the dedupe key." + // A call with no writer state, such as publish, sends no sequence. The Store then numbers the + // batch, so nothing dedupes it. + /// @nexus.doc + /// "The batch number in this epoch. Empty takes the next one, and a retry writes it again." sequence: option, /// @nexus.doc "Local Activity only. The run's last completed Workflow Task." + /// @nexus.api-omit targets="workflow" + /// @nexus.source targets="client" python="local_activity_previous_completed_task()" + /// typescript="localActivityPreviousCompletedTask()" + /// dotnet="StreamActivityContext.LocalActivityPreviousCompletedTask()" previous-completed-task: option, + // A wrong run would silently withdraw real Records, so it is never a parameter. The helper + // gives the previous run in a retry run's Workflow, and nothing elsewhere. /// @nexus.doc "Retry run only. The previous run, whose Records this withdraws." + /// @nexus.source targets="workflow" python="withdraw_run()" typescript="withdrawRun()" + /// dotnet="StreamWorkflowContext.WithdrawRun()" + /// @nexus.source targets="client" python="withdraw_run()" typescript="withdrawRun()" + /// dotnet="StreamWorkflowContext.WithdrawRun()" withdraw-run: option, /// @nexus.omit namespace: placeholder, - // Native only. Brook's Stores don't read them. + // Native only. The SDK Stores don't read them. /// @nexus.omit defer-withdrawal: placeholder, /// @nexus.omit @@ -250,19 +297,26 @@ interface stream-service { /// @nexus.doc "Close a Stream. Readers drain what is left, then end. Later appends are refused." /// @nexus.proto "temporal.api.streamservice.v1.CloseStreamRequest" record close-stream-request { - %stream: stream-ref, - /// @nexus.doc "Workflow only. The run's failed or timed out Workflow Tasks." - failed-tasks: list, + handle: stream-handle, + // Native only. The SDK Stores don't read it. + /// @nexus.omit + failed-tasks: placeholder, + /// @nexus.source targets="workflow" python="workflow_producer()" + /// typescript="workflowProducer()" dotnet="StreamWorkflowContext.Producer()" + /// @nexus.source targets="client" python="client_producer()" typescript="clientProducer()" + /// dotnet="StreamActivityContext.Producer()" producer: option, %result: option, + // Each SDK gives a canceled Activity, child Workflow or Nexus Operation as a failure of that + // call whose cause is the cancel. So such a failure also cancels. + /// @nexus.doc "Empty completes the Stream. A cancel failure cancels it. Any other fails it." failure: option, - /// @nexus.doc "Empty is completed, or failed with a failure. A caller may set canceled." - status: option, - /// @nexus.doc "Workflow only. The previous completed Workflow Task of the run." - previous-completed-task: option, + // Native only. The SDK Stores don't read it. + /// @nexus.omit + previous-completed-task: placeholder, /// @nexus.omit namespace: placeholder, - // Native only. Brook's Stores don't read them. + // Native only. The SDK Stores don't read them. /// @nexus.omit withdraw-run: placeholder, /// @nexus.omit @@ -283,24 +337,27 @@ interface stream-service { /// @nexus.doc "Read Records from an offset." /// @nexus.proto "temporal.api.streamservice.v1.ReadRecordsRequest" record read-records-request { - %stream: stream-ref, - /// @nexus.doc "Start offset. Empty uses start. INVALID_ARGUMENT past the newest offset + 1." + handle: stream-handle, + /// @nexus.doc "Start offset. Empty uses start. OUT_OF_RANGE past the newest offset + 1." from-offset: option, /// @nexus.doc "Where to start when from-offset is empty. Unspecified is the beginning." start: option, - /// @nexus.doc "At most this many Records. 1 to 1000." + /// @nexus.doc + /// "About this many items per page, 1 to 1000. A Superseded item stays with its Record." max-records: option, /// @nexus.doc "How long to wait for a new Record before an empty answer. At most 60 s." wait: option, /// @nexus.doc "How long to wait for an uncreated Stream. Empty fails at once. At most 60 s." create-wait: option, + // Replay needs a fixed range, so the Workflow API omits `last`. /// @nexus.doc "Read the newest this many Records, oldest first. No from-offset. No wait." + /// @nexus.api-omit targets="workflow" last: option, /// @nexus.doc "The maximum Record bytes in a page. A larger Record is still returned alone." max-bytes: option, /// @nexus.omit namespace: placeholder, - // Native only. Brook's Stores refuse to-offset and wake-only. + // Native only. The SDK Stores refuse to-offset and wake-only. /// @nexus.omit to-offset: placeholder, /// @nexus.omit @@ -316,14 +373,14 @@ interface stream-service { next-offset: option, /// @nexus.doc "True after the close Record. No later read gives a Record." done: option, - // Native only. Brook's Stores never set it. + // Native only. The SDK Stores never set it. /// @nexus.omit continue-as-new-suggested: placeholder, } /// @nexus.proto "temporal.api.streamservice.v1.DescribeStreamRequest" record describe-stream-request { - %stream: stream-ref, + handle: stream-handle, /// @nexus.omit namespace: placeholder, } @@ -331,38 +388,32 @@ interface stream-service { /// @nexus.proto "temporal.api.streamservice.v1.DescribeStreamResponse" record describe-stream-response { status: option, - /// @nexus.doc "The offset of the newest Record. Empty when the Stream has no Record." + /// @nexus.doc "The newest offset, of a Record, a Superseded item or the close. Empty for none." last-offset: option, retention: option, } - // A handle never changes, since offsets are values that its methods return. Only a create's - // response sets created and status, so a handle built from a ref has neither. - /// @nexus.doc "A handle on one Stream. Build it from a stream-ref, or get it from a create." + // A stream never changes, since offsets are values that its methods return. Only a create's + // response sets created, so a stream built from a handle has none. + /// @nexus.doc "One Stream. Build it from a stream-handle, or get it from a create." + /// @nexus.type-parameter "body-t" + /// @nexus.definition python="temporalio.common.StreamDefinition" + /// typescript="StreamDefinition" typescript-import="@temporalio/common" + /// dotnet="Temporalio.Common.StreamDefinition" + /// @nexus.flatten-in-api param="handle" fields="id backend" keep=true resource %stream { - constructor(%stream: stream-ref, created: option, status: option); + constructor(handle: stream-handle, created: option); - /// @nexus.doc "Publish Records. The answer has the first and the last offset." + /// @nexus.doc "Publish values, and return the offset of the last one." /// @nexus.operation "append-records" - publish: func( - producer: option, - sequence: option, - records: list, - previous-completed-task: option, - failed-tasks: list, - withdraw-run: option, - ) -> append-records-response; + /// @nexus.varargs name="values" + /// @nexus.result-field "last-offset" + publish: func(records: list) -> append-records-response; /// @nexus.doc "Close the Stream. Readers drain what is left, then end." /// @nexus.operation "close-stream" - close: func( - producer: option, - %result: option, - failure: option, - status: option, - previous-completed-task: option, - failed-tasks: list, - ); + /// @nexus.close failure="failure" + close: func(%result: option, failure: option); /// @nexus.doc "Read one page of Records. The poll iterator reads them all." /// @nexus.operation "read-records" @@ -372,6 +423,7 @@ interface stream-service { max-records: option, wait: option, create-wait: option, + max-bytes: option, ) -> read-records-response; /// @nexus.doc "Describe the Stream's status, newest offset and Retention." @@ -379,8 +431,7 @@ interface stream-service { describe: func() -> describe-stream-response; } - // created is false when the Stream existed, whatever its Retention. status is closed when an - // earlier create and close used the id. + // created is false when the Stream existed, whatever its Retention. /// @nexus.proto "temporal.api.streamservice.v1.CreateStreamResponse" type create-stream-result = own<%stream>; @@ -399,31 +450,48 @@ interface stream-service { /// @nexus.doc /// "A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local Activity." /// @nexus.operation name="ReadRecords" - /// @nexus.poll python="next_read_request" + /// @nexus.poll python="next_read_request" typescript="nextReadRequest" + /// dotnet="StreamPoll.NextReadRequest" items="records" python-resume="resume_read_request" + /// targets="client" read-records: func(request: read-records-request) -> read-records-response; /// @nexus.operation name="DescribeStream" describe-stream: func(request: describe-stream-request) -> describe-stream-response; - // No operation returns these. A failed call carries them, so each SDK can raise one error type - // per reason. + // No operation returns these. A failed call has one in its details. Each SDK raises the class + // of its reason, or the base class for a reason with no class. /// @nexus.doc "Why a Stream call failed. Each reason has one gRPC status code." /// @nexus.proto "temporal.api.enums.v1.StreamErrorReason" + /// @nexus.error details="stream-error-details" field="reason" name="stream-error" + /// python-extends="temporalio.exceptions.ApplicationError" + /// typescript-extends="ApplicationFailure" typescript-extends-import="@temporalio/common" + /// dotnet-extends="Temporalio.Exceptions.ApplicationFailureException" enum stream-error-reason { + /// @nexus.fallback stream-error-reason-unspecified, + /// @nexus.error stream-error-reason-not-found, + /// @nexus.error stream-error-reason-expired, + /// @nexus.error stream-error-reason-closed, + /// @nexus.error stream-error-reason-fenced, stream-error-reason-sequence, + /// @nexus.error stream-error-reason-seq-divergent, + /// @nexus.error stream-error-reason-not-owner, stream-error-reason-format-too-new, + /// @nexus.error stream-error-reason-offset, + /// @nexus.error stream-error-reason-too-large, stream-error-reason-invalid, + /// @nexus.error retryable=true stream-error-reason-storage, stream-error-reason-unknown-writer, + /// @nexus.error stream-error-reason-unsupported, } diff --git a/scripts/_wit/stream-writer.wit b/scripts/_wit/stream-writer.wit deleted file mode 100644 index 5f03e746ea..0000000000 --- a/scripts/_wit/stream-writer.wit +++ /dev/null @@ -1,118 +0,0 @@ -// Internal. This isn't a Temporal API, and no server serves it. The crate's runtime object serves -// these calls in process, with the messages of ../protos/temporal/sdk/streams/v1/runtime.proto. -// Lang calls them through the bridge, from Activities and clients only, so the interface has no -// endpoint. This file is the source. Each SDK copies it next to stream-service.wit and generates -// its client with nexgen's client target (`--targets client`, which needs `--system-nexus`). -package temporal:streams@1.0.0; - -world writer { - export stream-writer; -} - -/// @nexus.service-name "temporal.sdk.streams.v1.StreamWriter" -/// @nexus.experimental -interface stream-writer { - use nexus:temporal-types/model@1.0.0.{duration, placeholder}; - use temporal:nexus/stream-service@1.0.0.{ - append-record, producer, stream-ref, workflow-task-epoch, - }; - - /// @nexus.doc "Open a writer for one Producer of one Stream." - /// @nexus.proto "temporal.sdk.streams.v1.OpenWriterRequest" - record open-writer-request { - %stream: stream-ref, - /// @nexus.doc "Local Activity only. The run's failed or timed out Workflow Tasks." - failed-tasks: list, - producer: option, - /// @nexus.doc "A write sends at most this many Records in one batch. Zero is 1000." - max-batch-size: option, - /// @nexus.doc "How long a write retries a Store failure. Empty is 10 s." - max-retry-duration: option, - /// @nexus.doc "Local Activity only. The run's last completed Workflow Task." - previous-completed-task: option, - /// @nexus.omit - namespace: placeholder, - } - - /// @nexus.proto "temporal.sdk.streams.v1.WriteRecordsRequest" - record write-records-request { - records: list, - writer-id: option, - } - - /// @nexus.proto "temporal.sdk.streams.v1.WriteRecordsResponse" - record write-records-response { - /// @nexus.doc "The offset of the first Record. Zero when the request has no Record." - first-offset: option, - /// @nexus.doc "The offset of the last Record. Zero when the request has no Record." - last-offset: option, - } - - /// @nexus.proto "temporal.sdk.streams.v1.CloseWriterRequest" - record close-writer-request { - writer-id: option, - } - - /// @nexus.proto "temporal.sdk.streams.v1.CloseWriterResponse" - record close-writer-response { - /// @nexus.doc "The offset of the newest Record the writer wrote. Zero if it wrote none." - last-offset: option, - } - - /// @nexus.proto "temporal.sdk.streams.v1.FreeWriterRequest" - record free-writer-request { - writer-id: option, - } - - /// @nexus.proto "temporal.sdk.streams.v1.FreeWriterResponse" - record free-writer-response { - } - - /// @nexus.proto "temporal.sdk.streams.v1.CloseWritersRequest" - record close-writers-request { - } - - /// @nexus.proto "temporal.sdk.streams.v1.CloseWritersResponse" - record close-writers-response { - /// @nexus.doc "The writers that had failed." - failed: option, - } - - // A writer never changes. Offsets are values that its methods return. The handles of one - // Producer epoch on one Stream share one writer, and each close or free releases one handle. - /// @nexus.doc "An open writer of one Producer. The crate keeps its epoch and sequence." - resource writer { - constructor(writer-id: u64); - - /// @nexus.doc "Write Records at once. The answer has the first and the last offset." - /// @nexus.operation "write-records" - write-records: func(records: list) -> write-records-response; - - /// @nexus.doc "Release the handle. The last one waits for a write, then drops the writer." - /// @nexus.operation "close-writer" - close: func() -> close-writer-response; - - /// @nexus.doc "Release the handle with no wait, for a handle that lang drops." - /// @nexus.operation "free-writer" - free: func() -> free-writer-response; - } - - /// @nexus.proto "temporal.sdk.streams.v1.OpenWriterResponse" - type writer-result = own; - - /// @nexus.operation name="OpenWriter" - open-writer: func(request: open-writer-request) -> writer-result; - - /// @nexus.operation name="WriteRecords" - write-records: func(request: write-records-request) -> write-records-response; - - /// @nexus.operation name="CloseWriter" - close-writer: func(request: close-writer-request) -> close-writer-response; - - /// @nexus.operation name="FreeWriter" - free-writer: func(request: free-writer-request) -> free-writer-response; - - /// @nexus.doc "Drop every open writer, for example at shutdown." - /// @nexus.operation name="CloseWriters" - close-writers: func(request: close-writers-request) -> close-writers-response; -} diff --git a/scripts/gen_nexus_system_api.py b/scripts/gen_nexus_system_api.py index be72cf37a9..c37df01812 100644 --- a/scripts/gen_nexus_system_api.py +++ b/scripts/gen_nexus_system_api.py @@ -25,7 +25,7 @@ workflow_output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service" notification_output_dir = base_dir / "temporalio" / "nexus" / "notifications" workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py" -NEX_GEN_VERSION = "0.2.8" +NEX_GEN_VERSION = "0.2.11" proto_files = [ gen_protos.api_proto_dir / "temporal" diff --git a/scripts/gen_stream_service.py b/scripts/gen_stream_service.py index 8abc5a8618..95fce8adab 100644 --- a/scripts/gen_stream_service.py +++ b/scripts/gen_stream_service.py @@ -1,10 +1,9 @@ """Generates the Stream service bindings of the Streams API in temporalio.common. -The WIT and the runtime proto are copies of the temporalio-streams crate's, in scripts/_wit and -scripts/_proto. The Stream service protos are copies of the API repo's. Run scripts/gen_protos.py -first for the proto modules. No nexgen release has the client target and the poll directive, -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. +The WIT in scripts/_wit and the Stream service protos in scripts/_proto are copies of the API +repo's. Run scripts/gen_protos.py first for the proto modules. No nexgen release has the +directives of the WIT, 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 @@ -18,7 +17,6 @@ base_dir = Path(__file__).parent.parent wit_path = base_dir / "scripts" / "_wit" / "stream-service.wit" -writer_wit_path = base_dir / "scripts" / "_wit" / "stream-writer.wit" model_wit_path = ( gen_protos.api_proto_dir / "nexus" / "deps" / "nexus-temporal-types" / "model.wit" ) @@ -32,12 +30,10 @@ / "request_response.proto" ) output_dir = base_dir / "temporalio" / "common" / "_streams" / "system" -# Only Activities and 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 fork chain of nexgen with the client target and the poll directive. -NEX_GEN_REVISION = "d70f5f1754c2e694b859f4a38b24c73fab27d8da" +# The fork chain of nexgen that reads the directives of the WIT. +NEX_GEN_REVISION = "2670394236dde95fe6fd1c90543cbfc97848f5e9" nex_gen_root = base_dir / ".nexgen-streams" / NEX_GEN_REVISION @@ -77,44 +73,20 @@ def api_protos() -> list[Path]: def generate_stream_service() -> None: - for path in (wit_path, writer_wit_path, model_wit_path, proto_path): + for path in (wit_path, model_wit_path, proto_path): if not path.exists(): 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 generated runtime messages are in temporalio.common._streams.proto, so the - # descriptor must name their file by that module path. - runtime_root = Path(temp_dir) / "runtime" - runtime_path = ( - runtime_root - / "temporalio" - / "common" - / "_streams" - / "proto" - / "v1" - / "runtime.proto" - ) - runtime_path.parent.mkdir(parents=True) - shutil.copy( - gen_protos.additional_proto_dir - / "temporal" - / "sdk" - / "streams" - / "v1" - / "runtime.proto", - runtime_path, - ) subprocess.check_call( [ sys.executable, "-mgrpc_tools.protoc", f"--proto_path={gen_protos.api_proto_dir}", f"--proto_path={gen_protos.additional_proto_dir}", - f"--proto_path={runtime_root}", "--include_imports", f"--descriptor_set_out={descriptor_path}", str(proto_path), - str(runtime_path), # The model WIT names API types, so their descriptors must be in the set. *map(str, api_protos()), ] @@ -127,7 +99,6 @@ def generate_stream_service() -> None: str(wit_path), str(model_wit_path), "--system-nexus", - "--native-api", "--targets", "workflow,client", "--support-package", @@ -138,40 +109,19 @@ def generate_stream_service() -> None: str(output_dir), ] ) - shutil.rmtree(writer_output_dir, ignore_errors=True) - subprocess.check_call( - [ - *nex_gen_command(), - "python", - str(writer_wit_path), - str(wit_path), - str(model_wit_path), - "--system-nexus", - "--native-api", - "--targets", - "client", - "--support-package", - "temporalio.nexus.system._support", - "--descriptors", - str(descriptor_path), - "--output", - str(writer_output_dir), - ] - ) - for directory in (output_dir, writer_output_dir): - subprocess.check_call( - [ - sys.executable, - "-m", - "ruff", - "check", - "--select", - "I", - "--fix", - str(directory), - ] - ) - subprocess.check_call([sys.executable, "-m", "ruff", "format", str(directory)]) + subprocess.check_call( + [ + sys.executable, + "-m", + "ruff", + "check", + "--select", + "I", + "--fix", + str(output_dir), + ] + ) + subprocess.check_call([sys.executable, "-m", "ruff", "format", str(output_dir)]) if __name__ == "__main__": diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 578d714c6a..e6260ee000 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -31,6 +31,7 @@ import temporalio.api.enums.v1 import temporalio.api.workflow.v1 import temporalio.types +from temporalio.common._streams._defn import StreamDefinition as StreamDefinition @dataclass diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 3888aaa785..39ddb83f52 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1 +1,7 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" + +from temporalio.common._streams._defn import StreamDefinition + +__all__ = [ + "StreamDefinition", +] diff --git a/temporalio/common/_streams/_defn.py b/temporalio/common/_streams/_defn.py new file mode 100644 index 0000000000..1141e9dd42 --- /dev/null +++ b/temporalio/common/_streams/_defn.py @@ -0,0 +1,15 @@ +"""Stream definitions, which give the Stream name and its one payload type.""" + +from __future__ import annotations + +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class StreamDefinition(Generic[T]): + """Base class of a Stream definition. Its type argument is the Stream's payload type. + + .. warning:: + This API is experimental. + """ diff --git a/temporalio/common/_streams/system/__init__.py b/temporalio/common/_streams/system/__init__.py index 8d916ce5b3..f869c2dbde 100644 --- a/temporalio/common/_streams/system/__init__.py +++ b/temporalio/common/_streams/system/__init__.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -15,8 +15,16 @@ close_stream, describe_stream, read_records, + start_append_records, + start_close_stream, + start_describe_stream, + start_read_records, + stream, +) +from .operations.create_stream import ( + create_stream, + start_create_stream, ) -from .operations.create_stream import create_stream __all__ = [ "append_records", @@ -24,6 +32,12 @@ "create_stream", "describe_stream", "read_records", + "start_append_records", + "start_close_stream", + "start_create_stream", + "start_describe_stream", + "start_read_records", + "stream", ] diff --git a/temporalio/common/_streams/system/_poll.py b/temporalio/common/_streams/system/_poll.py new file mode 100644 index 0000000000..73d321309b --- /dev/null +++ b/temporalio/common/_streams/system/_poll.py @@ -0,0 +1,65 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import typing + +_PollRequestT = typing.TypeVar("_PollRequestT") +_PollItemT = typing.TypeVar("_PollItemT") + + +@typing.final +class PollIterator(typing.Generic[_PollRequestT, _PollItemT]): + """Yields each item that a poll reads, page by page. + + ``next_request`` is the request to resume with. With a ``resume`` helper, it resumes after + the last item that the iterator yielded. Without one, it reads that item's page again, + until the iterator yields the page's last item. It is ``None`` after the last page. + """ + + def __init__( + self, + request: _PollRequestT, + read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ], + resume: collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] + | None, + ) -> None: + self.next_request: _PollRequestT | None = request + self._read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ] = read + self._resume: ( + collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] | None + ) = resume + self._page: list[_PollItemT] = [] + self._page_request: _PollRequestT = request + self._page_next: _PollRequestT | None = None + + def __aiter__(self) -> PollIterator[_PollRequestT, _PollItemT]: + return self + + async def __anext__(self) -> _PollItemT: + while not self._page: + if self.next_request is None: + raise StopAsyncIteration + self._page_request = self.next_request + self._page_next, items = await self._read(self._page_request) + self._page = list(items) + if not self._page: + self.next_request = self._page_next + item = self._page.pop(0) + # After a page's last item, the next read starts where the page helper says. + if not self._page: + self.next_request = self._page_next + elif self._resume is not None: + self.next_request = self._resume(self._page_request, item) + return item diff --git a/temporalio/common/_streams/system/_resource_call.py b/temporalio/common/_streams/system/_resource_call.py new file mode 100644 index 0000000000..ee12545a4c --- /dev/null +++ b/temporalio/common/_streams/system/_resource_call.py @@ -0,0 +1,93 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import functools +import types +import typing + +import typing_extensions + + +class _Closable(typing.Protocol): + async def __aenter__(self) -> object: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: types.TracebackType | None, + ) -> None: ... + + +_ResourceT = typing.TypeVar("_ResourceT", bound=_Closable) +_P = typing_extensions.ParamSpec("_P") + + +@typing.final +class ResourceCall( + collections.abc.Coroutine[typing.Any, typing.Any, _ResourceT], + typing.Generic[_ResourceT], +): + """A call that returns a resource. Await it for the resource, or enter it with + ``async with``, which closes the resource at the block's end. + """ + + def __init__( + self, call: collections.abc.Coroutine[typing.Any, typing.Any, _ResourceT] + ) -> None: + self._call: collections.abc.Coroutine[typing.Any, typing.Any, _ResourceT] = call + self._resource: _ResourceT | None = None + self.__name__: str = getattr(call, "__name__", type(self).__name__) + self.__qualname__: str = getattr(call, "__qualname__", self.__name__) + + @typing_extensions.override + def __await__(self) -> collections.abc.Generator[typing.Any, None, _ResourceT]: + return self._call.__await__() + + @typing_extensions.override + def send(self, value: typing.Any) -> typing.Any: + return self._call.send(value) + + @typing_extensions.override + def throw( + self, + typ: typing.Any, + val: typing.Any = None, + tb: types.TracebackType | None = None, + ) -> typing.Any: + return self._call.throw(typ, val, tb) + + @typing_extensions.override + def close(self) -> None: + self._call.close() + + async def __aenter__(self) -> _ResourceT: + resource = await self._call + _ = await resource.__aenter__() + self._resource = resource + return resource + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: types.TracebackType | None, + ) -> None: + if self._resource is not None: + await self._resource.__aexit__(exc_type, exc, tb) + + +def resource_call( + function: collections.abc.Callable[ + _P, collections.abc.Coroutine[typing.Any, typing.Any, _ResourceT] + ], +) -> collections.abc.Callable[_P, ResourceCall[_ResourceT]]: + """Makes a call that returns a closable resource return a ``ResourceCall``.""" + + @functools.wraps(function) + def call(*args: _P.args, **kwargs: _P.kwargs) -> ResourceCall[_ResourceT]: + return ResourceCall(function(*args, **kwargs)) + + return call diff --git a/temporalio/common/_streams/system/_resources/__init__.py b/temporalio/common/_streams/system/_resources/__init__.py index efbe663e21..e2897447a6 100644 --- a/temporalio/common/_streams/system/_resources/__init__.py +++ b/temporalio/common/_streams/system/_resources/__init__.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/common/_streams/system/_resources/stream.py b/temporalio/common/_streams/system/_resources/stream.py index 663d423aec..0a4d50f2b3 100644 --- a/temporalio/common/_streams/system/_resources/stream.py +++ b/temporalio/common/_streams/system/_resources/stream.py @@ -1,15 +1,20 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations -import collections.abc -import dataclasses +import asyncio import datetime +import types import typing +import typing_extensions + +import temporalio.api.stream.v1.message_pb2 +import temporalio.converter import temporalio.nexus.system._support as _support import temporalio.workflow +from ..errors import StreamError, stream_error from ..models import ( AppendRecord, AppendRecordsRequest, @@ -18,68 +23,114 @@ CloseStreamResponse, DescribeStreamRequest, DescribeStreamResponse, - Producer, ReadRecordsRequest, ReadRecordsResponse, - StreamCloseStatus, + StreamErrorDetails, + StreamHandle, StreamReadStart, - StreamRef, - StreamStatus, - WorkflowTaskEpoch, + _StreamHandleTransferTypeConverter, ) +if typing.TYPE_CHECKING: + import temporalio.common + + +BodyT = typing.TypeVar("BodyT") -@dataclasses.dataclass -class Stream: - """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" - stream: StreamRef +class Stream(typing.Generic[BodyT]): + """One Stream. Build it from a stream-handle, or get it from a create.""" + + handle: StreamHandle created: bool | None - status: StreamStatus | None + + def __init__( + self, + handle: StreamHandle, + created: bool | None = None, + *, + body_type: type[BodyT] | None = None, + ) -> None: + self.handle = handle + self.created = created + # Not a dataclass field, so the converter doesn't encode it. A + # decoded handle decodes its generic values as `typing.Any`. + self.body_type: type[BodyT] = body_type or typing.cast( + "type[BodyT]", typing.Any + ) + + @typing_extensions.override + def __eq__(self, other: object) -> bool: + if not isinstance(other, Stream): + return NotImplemented + return (self.handle, self.created) == (other.handle, other.created) + + __hash__: typing.ClassVar[None] = None # pyright: ignore[reportIncompatibleMethodOverride] + + @typing_extensions.override + def __repr__(self) -> str: + return f"Stream(handle={self.handle!r}, created={self.created!r})" + + async def __aenter__(self) -> typing_extensions.Self: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: types.TracebackType | None, + ) -> None: + close, failure = _support.exit_close(self, exc) + if close: + await self.close(failure=failure) async def publish( self, - producer: Producer | None = None, - sequence: int | None = None, - records: list[AppendRecord] | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, - withdraw_run: str | None = None, - ) -> AppendRecordsResponse: - """Publish Records. The answer has the first and the last offset.""" + *values: BodyT, + ) -> int: + """Publish values, and return the offset of the last one.""" request = AppendRecordsRequest( - stream=self.stream, - producer=producer, - sequence=sequence, - records=records or [], - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - withdraw_run=withdraw_run, + handle=self.handle, + records=[AppendRecord(body=value) for value in values], + producer=_support.workflow_producer(), + withdraw_run=_support.withdraw_run(), ) - handle = await _append_records(request) - return await handle + try: + handle = await _append_records(request) + response = await handle + return response.last_offset or 0 + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def close( self, - producer: Producer | None = None, result: typing.Any | None = None, failure: BaseException | None = None, - status: StreamCloseStatus | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, ) -> None: """Close the Stream. Readers drain what is left, then end.""" request = CloseStreamRequest( - stream=self.stream, - producer=producer, + handle=self.handle, result=result, failure=failure, - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - status=status, + producer=_support.workflow_producer(), ) - handle = await _close_stream(request) - await handle + try: + handle = await _close_stream(request) + await handle + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def read( self, @@ -88,105 +139,170 @@ async def read( max_records: int | None = None, wait: datetime.timedelta | None = None, create_wait: datetime.timedelta | None = None, - ) -> ReadRecordsResponse: + max_bytes: int | None = None, + ) -> ReadRecordsResponse[BodyT]: """Read one page of Records. The poll iterator reads them all.""" request = ReadRecordsRequest( - stream=self.stream, + handle=self.handle, from_offset=from_offset, start=start, max_records=max_records, wait=wait, create_wait=create_wait, + max_bytes=max_bytes, ) - handle = await _read_records(request) - return await handle - - async def iter_read( - self, - from_offset: int | None = None, - start: StreamReadStart | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: - """Yield each response of ``read``. The next request comes from the poll helper, - and the iterator ends when the helper returns ``None``. - """ - request = ReadRecordsRequest( - stream=self.stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - ) - while True: - handle = await _read_records(request) - response = await handle - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request + try: + handle = await _read_records(request, self.body_type) + return await handle + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def describe( self, ) -> DescribeStreamResponse: """Describe the Stream's status, newest offset and Retention.""" request = DescribeStreamRequest( - stream=self.stream, + handle=self.handle, + ) + try: + handle = await _describe_stream(request) + return await handle + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + +def stream( + definition: type[temporalio.common.StreamDefinition[BodyT]], + *, + id: str | None = None, + handle: StreamHandle | None = None, + backend: str | None = None, + created: bool | None = None, +) -> Stream[BodyT]: + """A handle on an existing ``Stream``, with no call. Pass ``id`` or ``handle``.""" + (body_type,) = _support.definition_type_arguments(definition) + if handle is None: + if id is None: + raise TypeError("pass id or handle") + handle = StreamHandle(id=id, backend=backend) + elif id is not None or backend is not None: + raise TypeError("pass handle or its fields, not both") + return Stream( + handle=handle, + created=created, + body_type=body_type, + ) + + +class _StreamTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Stream[BodyT], temporalio.api.stream.v1.message_pb2.StreamHandle + ] +): + """Passes a ``Stream`` as its ``handle``. A decoded handle calls the Workflow + operations. + """ + + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamHandle] | None = ( + temporalio.api.stream.v1.message_pb2.StreamHandle + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamHandle, + type_hint: type[Stream[BodyT]], + ) -> Stream[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) + return Stream( + handle=_StreamHandleTransferTypeConverter().from_transfer_type( + value, StreamHandle + ), + body_type=typing.cast(type[BodyT], body_type), ) - handle = await _describe_stream(request) - return await handle + + @typing_extensions.override + def to_transfer_type( + self, value: Stream[BodyT] + ) -> temporalio.api.stream.v1.message_pb2.StreamHandle: + return _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) + + +temporalio.converter.transfer_type_convertible(_StreamTransferTypeConverter)(Stream) # pyright: ignore[reportUnusedCallResult] async def _append_records( - request: AppendRecordsRequest, + request: AppendRecordsRequest[BodyT], ) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: nexus_client = temporalio.workflow.create_nexus_client( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - return await nexus_client.start_operation( - operation="AppendRecords", - input=request, - output_type=AppendRecordsResponse, - ) + try: + return await nexus_client.start_operation( + operation="AppendRecords", + input=request, + output_type=AppendRecordsResponse, + ) + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def append_records( *, - stream: StreamRef, - records: list[AppendRecord] | None = None, - failed_tasks: list[int] | None = None, - producer: Producer | None = None, + handle: StreamHandle, + records: list[AppendRecord[BodyT]] | None = None, sequence: int | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - withdraw_run: str | None = None, ) -> temporalio.workflow.NexusOperationHandle[AppendRecordsResponse,]: """Append one batch. Workflow code makes the call from a Local Activity. Args: - failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. - producer: Required outside a Workflow. In a Workflow the transport sets it. - sequence: The batch number in this epoch, 1 or more. It is part of the dedupe - key. - previous_completed_task: Local Activity only. The run's last completed Workflow - Task. - withdraw_run: Retry run only. The previous run, whose Records this withdraws. + sequence: The batch number in this epoch. Empty takes the next one, and a retry + writes it again. """ request = AppendRecordsRequest( - stream=stream, + handle=handle, records=records or [], - failed_tasks=failed_tasks or [], - producer=producer, + producer=_support.workflow_producer(), sequence=sequence, - previous_completed_task=previous_completed_task, - withdraw_run=withdraw_run, + withdraw_run=_support.withdraw_run(), ) return await _append_records(request) +def start_append_records( + *, + handle: StreamHandle, + records: list[AppendRecord[BodyT]] | None = None, + sequence: int | None = None, +) -> asyncio.Task[temporalio.workflow.NexusOperationHandle[AppendRecordsResponse]]: + """Start ``append_records`` at the call, as ``workflow.start_activity`` places an + Activity, so sync Workflow code keeps the operation's place in the command order. + The task gives what ``append_records`` returns. + """ + return _support.start_now( + append_records(handle=handle, records=records, sequence=sequence) + ) + + async def _close_stream( request: CloseStreamRequest, ) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: @@ -194,127 +310,146 @@ async def _close_stream( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - return await nexus_client.start_operation( - operation="CloseStream", - input=request, - output_type=CloseStreamResponse, - ) + try: + return await nexus_client.start_operation( + operation="CloseStream", + input=request, + output_type=CloseStreamResponse, + ) + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def close_stream( *, - stream: StreamRef, - failed_tasks: list[int] | None = None, - producer: Producer | None = None, + handle: StreamHandle, result: typing.Any | None = None, failure: BaseException | None = None, - status: StreamCloseStatus | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, ) -> temporalio.workflow.NexusOperationHandle[CloseStreamResponse,]: """ Args: - failed_tasks: Workflow only. The run's failed or timed out Workflow Tasks. - status: Empty is completed, or failed with a failure. A caller may set canceled. - previous_completed_task: Workflow only. The previous completed Workflow Task of - the run. + failure: Empty completes the Stream. A cancel failure cancels it. Any other + fails it. """ request = CloseStreamRequest( - stream=stream, - failed_tasks=failed_tasks or [], - producer=producer, + handle=handle, + producer=_support.workflow_producer(), result=result, failure=failure, - status=status, - previous_completed_task=previous_completed_task, ) return await _close_stream(request) +def start_close_stream( + *, + handle: StreamHandle, + result: typing.Any | None = None, + failure: BaseException | None = None, +) -> asyncio.Task[temporalio.workflow.NexusOperationHandle[CloseStreamResponse]]: + """Start ``close_stream`` at the call, as ``workflow.start_activity`` places an + Activity, so sync Workflow code keeps the operation's place in the command order. + The task gives what ``close_stream`` returns. + """ + return _support.start_now( + close_stream(handle=handle, result=result, failure=failure) + ) + + async def _read_records( request: ReadRecordsRequest, -) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: + body_type: type[BodyT], +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse[BodyT],]: nexus_client = temporalio.workflow.create_nexus_client( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - return await nexus_client.start_operation( - operation="ReadRecords", - input=request, - output_type=ReadRecordsResponse, - ) + try: + return await nexus_client.start_operation( + operation="ReadRecords", + input=request, + output_type=ReadRecordsResponse[body_type], + ) + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def read_records( + body_type: type[BodyT], *, - stream: StreamRef, + handle: StreamHandle, from_offset: int | None = None, start: StreamReadStart | None = None, max_records: int | None = None, wait: datetime.timedelta | None = None, create_wait: datetime.timedelta | None = None, - last: int | None = None, max_bytes: int | None = None, -) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse,]: +) -> temporalio.workflow.NexusOperationHandle[ReadRecordsResponse[BodyT],]: """A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local Activity. Args: - from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest - offset + 1. + from_offset: Start offset. Empty uses start. OUT_OF_RANGE past the newest offset + + 1. start: Where to start when from-offset is empty. Unspecified is the beginning. - max_records: At most this many Records. 1 to 1000. + max_records: About this many items per page, 1 to 1000. A Superseded item stays + with its Record. wait: How long to wait for a new Record before an empty answer. At most 60 s. create_wait: How long to wait for an uncreated Stream. Empty fails at once. At most 60 s. - last: Read the newest this many Records, oldest first. No from-offset. No wait. max_bytes: The maximum Record bytes in a page. A larger Record is still returned alone. """ request = ReadRecordsRequest( - stream=stream, + handle=handle, from_offset=from_offset, start=start, max_records=max_records, wait=wait, create_wait=create_wait, - last=last, max_bytes=max_bytes, ) - return await _read_records(request) + return await _read_records(request, body_type) -async def iter_read_records( +def start_read_records( + body_type: type[BodyT], *, - stream: StreamRef, + handle: StreamHandle, from_offset: int | None = None, start: StreamReadStart | None = None, max_records: int | None = None, wait: datetime.timedelta | None = None, create_wait: datetime.timedelta | None = None, - last: int | None = None, max_bytes: int | None = None, -) -> collections.abc.AsyncIterator[ReadRecordsResponse]: - """Yield each response of ``read_records``. The next request comes from the poll - helper, and the iterator ends when the helper returns ``None``. +) -> asyncio.Task[temporalio.workflow.NexusOperationHandle[ReadRecordsResponse[BodyT]]]: + """Start ``read_records`` at the call, as ``workflow.start_activity`` places an + Activity, so sync Workflow code keeps the operation's place in the command order. + The task gives what ``read_records`` returns. """ - request = ReadRecordsRequest( - stream=stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - last=last, - max_bytes=max_bytes, + return _support.start_now( + read_records( + body_type, + handle=handle, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + max_bytes=max_bytes, + ) ) - while True: - handle = await _read_records(request) - response = await handle - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request async def _describe_stream( @@ -324,18 +459,38 @@ async def _describe_stream( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - return await nexus_client.start_operation( - operation="DescribeStream", - input=request, - output_type=DescribeStreamResponse, - ) + try: + return await nexus_client.start_operation( + operation="DescribeStream", + input=request, + output_type=DescribeStreamResponse, + ) + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err async def describe_stream( *, - stream: StreamRef, + handle: StreamHandle, ) -> temporalio.workflow.NexusOperationHandle[DescribeStreamResponse,]: request = DescribeStreamRequest( - stream=stream, + handle=handle, ) return await _describe_stream(request) + + +def start_describe_stream( + *, + handle: StreamHandle, +) -> asyncio.Task[temporalio.workflow.NexusOperationHandle[DescribeStreamResponse]]: + """Start ``describe_stream`` at the call, as ``workflow.start_activity`` places an + Activity, so sync Workflow code keeps the operation's place in the command order. + The task gives what ``describe_stream`` returns. + """ + return _support.start_now(describe_stream(handle=handle)) diff --git a/temporalio/common/_streams/system/_system_nexus_interceptor.py b/temporalio/common/_streams/system/_system_nexus_interceptor.py index 647dd384fe..33450f91a8 100644 --- a/temporalio/common/_streams/system/_system_nexus_interceptor.py +++ b/temporalio/common/_streams/system/_system_nexus_interceptor.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -45,7 +45,7 @@ async def _start_system_nexus_operation( and input.operation_name == "AppendRecords" ): typed_input = typing.cast( - "StartNexusOperationInput[models.AppendRecordsRequest, models.AppendRecordsResponse]", + "StartNexusOperationInput[models.AppendRecordsRequest[typing.Any], models.AppendRecordsResponse]", input, ) # The dispatch check above establishes that this operation's response type is _OutputT. @@ -71,7 +71,7 @@ async def _start_system_nexus_operation( and input.operation_name == "ReadRecords" ): typed_input = typing.cast( - "StartNexusOperationInput[models.ReadRecordsRequest, models.ReadRecordsResponse]", + "StartNexusOperationInput[models.ReadRecordsRequest, models.ReadRecordsResponse[typing.Any]]", input, ) # The dispatch check above establishes that this operation's response type is _OutputT. @@ -116,7 +116,7 @@ async def start_create_stream( return await self._next_system_nexus_interceptor().start_create_stream(request) async def start_append_records( - self, request: models.AppendRecordsRequest + self, request: models.AppendRecordsRequest[typing.Any] ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: """Intercept the AppendRecords operation.""" return await self._next_system_nexus_interceptor().start_append_records(request) @@ -129,7 +129,9 @@ async def start_close_stream( async def start_read_records( self, request: models.ReadRecordsRequest - ) -> temporalio.workflow.NexusOperationHandle[models.ReadRecordsResponse]: + ) -> temporalio.workflow.NexusOperationHandle[ + models.ReadRecordsResponse[typing.Any] + ]: """Intercept the ReadRecords operation.""" return await self._next_system_nexus_interceptor().start_read_records(request) @@ -173,7 +175,7 @@ async def start_create_stream( ) async def start_append_records( - self, request: models.AppendRecordsRequest + self, request: models.AppendRecordsRequest[typing.Any] ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput @@ -220,7 +222,9 @@ async def start_close_stream( async def start_read_records( self, request: models.ReadRecordsRequest - ) -> temporalio.workflow.NexusOperationHandle[models.ReadRecordsResponse]: + ) -> temporalio.workflow.NexusOperationHandle[ + models.ReadRecordsResponse[typing.Any] + ]: from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput from temporalio.workflow import NexusOperationCancellationType @@ -231,7 +235,10 @@ async def start_read_records( service="temporal.api.streamservice.v1.StreamService", operation="ReadRecords", input=request, - output_type=models.ReadRecordsResponse, + output_type=typing.cast( + "type[models.ReadRecordsResponse[typing.Any]]", + models.ReadRecordsResponse, + ), schedule_to_close_timeout=None, schedule_to_start_timeout=None, start_to_close_timeout=None, diff --git a/temporalio/common/_streams/system/client.py b/temporalio/common/_streams/system/client.py deleted file mode 100644 index 8a229d37f9..0000000000 --- a/temporalio/common/_streams/system/client.py +++ /dev/null @@ -1,439 +0,0 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! - -from __future__ import annotations - -import abc -import collections.abc -import dataclasses -import datetime -import typing - -import temporalio.nexus.system._support as _support - -from .models import ( - AppendRecord, - AppendRecordsRequest, - AppendRecordsResponse, - CloseStreamRequest, - CloseStreamResponse, - CreateStreamRequest, - CreateStreamResult, - DescribeStreamRequest, - DescribeStreamResponse, - Producer, - ReadRecordsRequest, - ReadRecordsResponse, - StreamCloseStatus, - StreamReadStart, - StreamRef, - StreamStatus, - WorkflowTaskEpoch, -) - -_OutputT = typing.TypeVar("_OutputT") - - -class _NexusClient(typing.Protocol): - """Starts the Nexus operations of one service from code outside a Workflow.""" - - @typing.overload - async def start_operation( - self, - operation: str, - input: typing.Any, - *, - output_type: type[_OutputT], - ) -> collections.abc.Awaitable[_OutputT]: ... - - @typing.overload - async def start_operation( - self, - operation: str, - input: typing.Any, - ) -> collections.abc.Awaitable[None]: ... - - -class StreamServiceClientExtension(abc.ABC): - """Starts the operations of ``StreamService`` from code outside a Workflow. - - A client inherits this class and implements ``_create_nexus_client``. - - .. warning:: - This API is experimental and subject to change. - """ - - @abc.abstractmethod - def _create_nexus_client( - self, - *, - service: str, - endpoint: str | None, - ) -> _NexusClient: - """Return the Nexus client that starts this service's operations. - - ``endpoint`` is ``None`` for a service with no endpoint, which the client serves. - """ - ... - - async def _create_stream( - self, - request: CreateStreamRequest, - ) -> Stream: - nexus_client = self._create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - handle = await nexus_client.start_operation( - operation="CreateStream", - input=request, - output_type=CreateStreamResult, - ) - result = await handle - if result.stream is None: - raise RuntimeError("CreateStream answered without `stream`") - return Stream( - _client=self, - stream=result.stream, - created=result.created or None, - status=result.status or None, - ) - - async def create_stream( - self, - *, - stream_id: str, - retention: datetime.timedelta | None = None, - withdraw_run: str | None = None, - owner: Producer | None = None, - ref: StreamRef | None = None, - ) -> Stream: - """ - Args: - 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. - owner: The Workflow run that owns the Stream. The Workflow transport sets it. - ref: The Stream's ref, with the owner fields. - """ - request = CreateStreamRequest( - stream_id=stream_id, - retention=retention, - withdraw_run=withdraw_run, - owner=owner, - ref=ref, - ) - return await self._create_stream(request) - - async def _append_records( - self, - request: AppendRecordsRequest, - ) -> collections.abc.Awaitable[AppendRecordsResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="AppendRecords", - input=request, - output_type=AppendRecordsResponse, - ) - - async def append_records( - self, - *, - stream: StreamRef, - records: list[AppendRecord] | None = None, - failed_tasks: list[int] | None = None, - producer: Producer | None = None, - sequence: int | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - withdraw_run: str | None = None, - ) -> collections.abc.Awaitable[AppendRecordsResponse,]: - """Append one batch. Workflow code makes the call from a Local Activity. - - Args: - failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. - producer: Required outside a Workflow. In a Workflow the transport sets it. - sequence: The batch number in this epoch, 1 or more. It is part of the dedupe - key. - previous_completed_task: Local Activity only. The run's last completed Workflow - Task. - withdraw_run: Retry run only. The previous run, whose Records this withdraws. - """ - request = AppendRecordsRequest( - stream=stream, - records=records or [], - failed_tasks=failed_tasks or [], - producer=producer, - sequence=sequence, - previous_completed_task=previous_completed_task, - withdraw_run=withdraw_run, - ) - return await self._append_records(request) - - async def _close_stream( - self, - request: CloseStreamRequest, - ) -> collections.abc.Awaitable[CloseStreamResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="CloseStream", - input=request, - output_type=CloseStreamResponse, - ) - - async def close_stream( - self, - *, - stream: StreamRef, - failed_tasks: list[int] | None = None, - producer: Producer | None = None, - result: typing.Any | None = None, - failure: BaseException | None = None, - status: StreamCloseStatus | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - ) -> collections.abc.Awaitable[CloseStreamResponse,]: - """ - Args: - failed_tasks: Workflow only. The run's failed or timed out Workflow Tasks. - status: Empty is completed, or failed with a failure. A caller may set canceled. - previous_completed_task: Workflow only. The previous completed Workflow Task of - the run. - """ - request = CloseStreamRequest( - stream=stream, - failed_tasks=failed_tasks or [], - producer=producer, - result=result, - failure=failure, - status=status, - previous_completed_task=previous_completed_task, - ) - return await self._close_stream(request) - - async def _read_records( - self, - request: ReadRecordsRequest, - ) -> collections.abc.Awaitable[ReadRecordsResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="ReadRecords", - input=request, - output_type=ReadRecordsResponse, - ) - - async def read_records( - self, - *, - stream: StreamRef, - from_offset: int | None = None, - start: StreamReadStart | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, - last: int | None = None, - max_bytes: int | None = None, - ) -> collections.abc.Awaitable[ReadRecordsResponse,]: - """A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local - Activity. - - Args: - from_offset: Start offset. Empty uses start. INVALID_ARGUMENT past the newest - offset + 1. - start: Where to start when from-offset is empty. Unspecified is the beginning. - max_records: At most this many Records. 1 to 1000. - wait: How long to wait for a new Record before an empty answer. At most 60 s. - create_wait: How long to wait for an uncreated Stream. Empty fails at once. At - most 60 s. - last: Read the newest this many Records, oldest first. No from-offset. No wait. - max_bytes: The maximum Record bytes in a page. A larger Record is still returned - alone. - """ - request = ReadRecordsRequest( - stream=stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - last=last, - max_bytes=max_bytes, - ) - return await self._read_records(request) - - async def iter_read_records( - self, - *, - stream: StreamRef, - from_offset: int | None = None, - start: StreamReadStart | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, - last: int | None = None, - max_bytes: int | None = None, - ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: - """Yield each response of ``read_records``. The next request comes from the poll - helper, and the iterator ends when the helper returns ``None``. - """ - request = ReadRecordsRequest( - stream=stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - last=last, - max_bytes=max_bytes, - ) - while True: - handle = await self._read_records(request) - response = await handle - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request - - async def _describe_stream( - self, - request: DescribeStreamRequest, - ) -> collections.abc.Awaitable[DescribeStreamResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.api.streamservice.v1.StreamService", - endpoint="__temporal_system", - ) - return await nexus_client.start_operation( - operation="DescribeStream", - input=request, - output_type=DescribeStreamResponse, - ) - - async def describe_stream( - self, - *, - stream: StreamRef, - ) -> collections.abc.Awaitable[DescribeStreamResponse,]: - request = DescribeStreamRequest( - stream=stream, - ) - return await self._describe_stream(request) - - -@dataclasses.dataclass -class Stream: - """A handle on one Stream. Build it from a stream-ref, or get it from a create.""" - - _client: StreamServiceClientExtension = dataclasses.field(repr=False, compare=False) - stream: StreamRef - created: bool | None - status: StreamStatus | None - - async def publish( - self, - producer: Producer | None = None, - sequence: int | None = None, - records: list[AppendRecord] | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, - withdraw_run: str | None = None, - ) -> AppendRecordsResponse: - """Publish Records. The answer has the first and the last offset.""" - request = AppendRecordsRequest( - stream=self.stream, - producer=producer, - sequence=sequence, - records=records or [], - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - withdraw_run=withdraw_run, - ) - handle = await self._client._append_records(request) - return await handle - - async def close( - self, - producer: Producer | None = None, - result: typing.Any | None = None, - failure: BaseException | None = None, - status: StreamCloseStatus | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - failed_tasks: list[int] | None = None, - ) -> None: - """Close the Stream. Readers drain what is left, then end.""" - request = CloseStreamRequest( - stream=self.stream, - producer=producer, - result=result, - failure=failure, - previous_completed_task=previous_completed_task, - failed_tasks=failed_tasks or [], - status=status, - ) - handle = await self._client._close_stream(request) - await handle - - async def read( - self, - from_offset: int | None = None, - start: StreamReadStart | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, - ) -> ReadRecordsResponse: - """Read one page of Records. The poll iterator reads them all.""" - request = ReadRecordsRequest( - stream=self.stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - ) - handle = await self._client._read_records(request) - return await handle - - async def iter_read( - self, - from_offset: int | None = None, - start: StreamReadStart | None = None, - max_records: int | None = None, - wait: datetime.timedelta | None = None, - create_wait: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ReadRecordsResponse]: - """Yield each response of ``read``. The next request comes from the poll helper, - and the iterator ends when the helper returns ``None``. - """ - request = ReadRecordsRequest( - stream=self.stream, - from_offset=from_offset, - start=start, - max_records=max_records, - wait=wait, - create_wait=create_wait, - ) - while True: - handle = await self._client._read_records(request) - response = await handle - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request - - async def describe( - self, - ) -> DescribeStreamResponse: - """Describe the Stream's status, newest offset and Retention.""" - request = DescribeStreamRequest( - stream=self.stream, - ) - handle = await self._client._describe_stream(request) - return await handle diff --git a/temporalio/common/_streams/system/client/__init__.py b/temporalio/common/_streams/system/client/__init__.py new file mode 100644 index 0000000000..40e08eb489 --- /dev/null +++ b/temporalio/common/_streams/system/client/__init__.py @@ -0,0 +1,11 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +from ._extensions import StreamServiceClientExtension +from ._resources import Stream + +__all__ = [ + "StreamServiceClientExtension", + "Stream", +] diff --git a/temporalio/common/_streams/system/client/_extensions.py b/temporalio/common/_streams/system/client/_extensions.py new file mode 100644 index 0000000000..c84a198bee --- /dev/null +++ b/temporalio/common/_streams/system/client/_extensions.py @@ -0,0 +1,497 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import abc +import collections.abc +import datetime +import typing + +import temporalio.nexus.system._support as _support + +from .._poll import PollIterator +from .._resource_call import resource_call +from ..errors import StreamError, stream_error +from ..models import ( + AppendRecord, + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + CreateStreamRequest, + CreateStreamResult, + DescribeStreamRequest, + DescribeStreamResponse, + ReadRecordsRequest, + ReadRecordsResponse, + StreamErrorDetails, + StreamHandle, + StreamReadStart, + StreamRecord, +) +from ._nexus_client import ( + _NexusClient, + _operation_id, +) +from ._resources import Stream + +if typing.TYPE_CHECKING: + import temporalio.common + + +BodyT = typing.TypeVar("BodyT") + + +class StreamServiceClientExtension(abc.ABC): + """Starts the operations of ``StreamService`` from code outside a Workflow. + + A client class inherits this class and implements ``_create_nexus_client``. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init_subclass__(cls, **kwargs: typing.Any) -> None: + super().__init_subclass__(**kwargs) + # The MRO picks one base for a name that two bases have, with no error. + for base in cls.__mro__: + if StreamServiceClientExtension in base.__mro__: + continue + for name in ( + "create_stream", + "append_records", + "close_stream", + "read_records", + "iter_read_records", + "describe_stream", + "get_stream", + ): + if name in base.__dict__: + raise TypeError( + f"{cls.__qualname__} gets {name} from {base.__qualname__} and from StreamServiceClientExtension" + ) + + @abc.abstractmethod + def _create_nexus_client( + self, + *, + service: str, + endpoint: str, + ) -> _NexusClient: + """Return the Nexus client that starts the operations of ``service``.""" + ... + + def get_stream( + self, + definition: type[temporalio.common.StreamDefinition[BodyT]], + *, + id: str | None = None, + handle: StreamHandle | None = None, + backend: str | None = None, + created: bool | None = None, + ) -> Stream[BodyT]: + """A handle on an existing ``Stream``, with no call. Pass ``id`` or ``handle``.""" + (body_type,) = _support.definition_type_arguments(definition) + if handle is None: + if id is None: + raise TypeError("pass id or handle") + handle = StreamHandle(id=id, backend=backend) + elif id is not None or backend is not None: + raise TypeError("pass handle or its fields, not both") + return Stream( + self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ), + handle=handle, + created=created, + body_type=body_type, + ) + + async def _stream_service_create_stream( + self, + request: CreateStreamRequest, + body_type: type[BodyT], + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> Stream[BodyT]: + try: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="CreateStream", + arg=request, + id=_operation_id(), + result_type=CreateStreamResult, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + result = await handle.result() + if result.handle is None: + raise RuntimeError("CreateStream answered without `handle`") + resource = Stream( + nexus_client, + body_type=body_type, + handle=result.handle, + created=result.created or None, + ) + return resource + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + @resource_call + async def create_stream( + self, + definition: type[temporalio.common.StreamDefinition[BodyT]], + *, + id: str, + retention: datetime.timedelta | None = None, + backend: str | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> Stream[BodyT]: + """ + Args: + id: The Stream id. Required, and never empty. + retention: The age after which a Record is dropped. Empty is the Store + default. + backend: Registered backend name, empty for the default. The Store never + reads it. + """ + (body_type,) = _support.definition_type_arguments(definition) + stream_handle = StreamHandle( + backend=backend, + id=id, + ) + request = CreateStreamRequest( + id=id, + retention=retention, + withdraw_run=_support.withdraw_run(), + stream_handle=stream_handle, + ) + return await self._stream_service_create_stream( + request, body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_service_append_records( + self, + request: AppendRecordsRequest[BodyT], + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> AppendRecordsResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="AppendRecords", + arg=request, + id=_operation_id(), + result_type=AppendRecordsResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def append_records( + self, + *, + handle: StreamHandle, + records: list[AppendRecord[BodyT]] | None = None, + sequence: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> AppendRecordsResponse: + """Append one batch. Workflow code makes the call from a Local Activity. + + Args: + sequence: The batch number in this epoch. Empty takes the next one, and a + retry writes it again. + """ + request = AppendRecordsRequest( + handle=handle, + records=records or [], + failed_tasks=_support.local_activity_failed_tasks(), + producer=_support.client_producer(), + sequence=sequence, + previous_completed_task=_support.local_activity_previous_completed_task(), + withdraw_run=_support.withdraw_run(), + ) + return await self._stream_service_append_records( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_service_close_stream( + self, + request: CloseStreamRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseStreamResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="CloseStream", + arg=request, + id=_operation_id(), + result_type=CloseStreamResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def close_stream( + self, + *, + handle: StreamHandle, + result: typing.Any | None = None, + failure: BaseException | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseStreamResponse: + """ + Args: + failure: Empty completes the Stream. A cancel failure cancels it. Any other + fails it. + """ + request = CloseStreamRequest( + handle=handle, + producer=_support.client_producer(), + result=result, + failure=failure, + ) + return await self._stream_service_close_stream( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_service_read_records( + self, + request: ReadRecordsRequest, + body_type: type[BodyT], + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> ReadRecordsResponse[BodyT]: + try: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="ReadRecords", + arg=request, + id=_operation_id(), + result_type=ReadRecordsResponse[body_type], + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def read_records( + self, + body_type: type[BodyT], + *, + handle: StreamHandle, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> ReadRecordsResponse[BodyT]: + """A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local + Activity. + + Args: + from_offset: Start offset. Empty uses start. OUT_OF_RANGE past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the + beginning. + max_records: About this many items per page, 1 to 1000. A Superseded item + stays with its Record. + wait: How long to wait for a new Record before an empty answer. At most 60 + s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. + At most 60 s. + last: Read the newest this many Records, oldest first. No from-offset. No + wait. + max_bytes: The maximum Record bytes in a page. A larger Record is still + returned alone. + """ + request = ReadRecordsRequest( + handle=handle, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + return await self._stream_service_read_records( + request, body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + def iter_read_records( + self, + body_type: type[BodyT], + *, + handle: StreamHandle, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + last: int | None = None, + max_bytes: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> PollIterator[ReadRecordsRequest, StreamRecord[BodyT]]: + """Yield each item that ``read_records`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. + + A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local + Activity. + + Args: + from_offset: Start offset. Empty uses start. OUT_OF_RANGE past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the + beginning. + max_records: About this many items per page, 1 to 1000. A Superseded item + stays with its Record. + wait: How long to wait for a new Record before an empty answer. At most 60 + s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. + At most 60 s. + last: Read the newest this many Records, oldest first. No from-offset. No + wait. + max_bytes: The maximum Record bytes in a page. A larger Record is still + returned alone. + """ + request = ReadRecordsRequest( + handle=handle, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + last=last, + max_bytes=max_bytes, + ) + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + + async def read( + request: ReadRecordsRequest, + ) -> tuple[ + ReadRecordsRequest | None, collections.abc.Sequence[StreamRecord[BodyT]] + ]: + try: + handle = await nexus_client.start_operation( + operation="ReadRecords", + arg=request, + id=_operation_id(), + result_type=ReadRecordsResponse[body_type], + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + response = await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + return _support.next_read_request(request, response), response.records or [] + + return PollIterator(request, read, _support.resume_read_request) + + async def _stream_service_describe_stream( + self, + request: DescribeStreamRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> DescribeStreamResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.api.streamservice.v1.StreamService", + endpoint="__temporal_system", + ) + handle = await nexus_client.start_operation( + operation="DescribeStream", + arg=request, + id=_operation_id(), + result_type=DescribeStreamResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def describe_stream( + self, + *, + handle: StreamHandle, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> DescribeStreamResponse: + request = DescribeStreamRequest( + handle=handle, + ) + return await self._stream_service_describe_stream( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) diff --git a/temporalio/common/_streams/system/client/_nexus_client.py b/temporalio/common/_streams/system/client/_nexus_client.py new file mode 100644 index 0000000000..8a36e5f693 --- /dev/null +++ b/temporalio/common/_streams/system/client/_nexus_client.py @@ -0,0 +1,50 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import datetime +import typing +import uuid + +_OutputT = typing.TypeVar("_OutputT") +_OutputT_co = typing.TypeVar("_OutputT_co", covariant=True) + + +class _NexusOperationHandle(typing.Protocol[_OutputT_co]): + """A started Nexus operation.""" + + async def result(self) -> _OutputT_co: ... + + +# The client modules import the names below, so they aren't unused. +class _NexusClient(typing.Protocol): # pyright: ignore[reportUnusedClass] + """Starts the operations of one Nexus service, like ``temporalio.client.NexusClient``.""" + + @typing.overload + async def start_operation( + self, + operation: str, + arg: typing.Any, + *, + id: str, + result_type: type[_OutputT], + rpc_metadata: collections.abc.Mapping[str, str | bytes] = ..., + rpc_timeout: datetime.timedelta | None = ..., + ) -> _NexusOperationHandle[_OutputT]: ... + + @typing.overload + async def start_operation( + self, + operation: str, + arg: typing.Any, + *, + id: str, + rpc_metadata: collections.abc.Mapping[str, str | bytes] = ..., + rpc_timeout: datetime.timedelta | None = ..., + ) -> _NexusOperationHandle[typing.Any]: ... + + +def _operation_id() -> str: # pyright: ignore[reportUnusedFunction] + """Return a new random operation ID.""" + return str(uuid.uuid4()) diff --git a/temporalio/common/_streams/system/client/_resources/__init__.py b/temporalio/common/_streams/system/client/_resources/__init__.py new file mode 100644 index 0000000000..e2897447a6 --- /dev/null +++ b/temporalio/common/_streams/system/client/_resources/__init__.py @@ -0,0 +1,9 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +from .stream import Stream + +__all__ = [ + "Stream", +] diff --git a/temporalio/common/_streams/system/client/_resources/stream.py b/temporalio/common/_streams/system/client/_resources/stream.py new file mode 100644 index 0000000000..bc19eda109 --- /dev/null +++ b/temporalio/common/_streams/system/client/_resources/stream.py @@ -0,0 +1,335 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import datetime +import types +import typing + +import typing_extensions + +import temporalio.api.stream.v1.message_pb2 +import temporalio.converter +import temporalio.nexus.system._support as _support + +from ..._poll import PollIterator +from ...errors import StreamError, stream_error +from ...models import ( + AppendRecord, + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + DescribeStreamRequest, + DescribeStreamResponse, + ReadRecordsRequest, + ReadRecordsResponse, + StreamErrorDetails, + StreamHandle, + StreamReadStart, + StreamRecord, + _StreamHandleTransferTypeConverter, +) +from .._nexus_client import ( + _NexusClient, + _operation_id, +) + +BodyT = typing.TypeVar("BodyT") + + +class Stream(typing.Generic[BodyT]): + """One Stream. Build it from a stream-handle, or get it from a create.""" + + handle: StreamHandle + created: bool | None + + def __init__( + self, + nexus_client: _NexusClient, + /, + *, + handle: StreamHandle, + created: bool | None = None, + body_type: type[BodyT] | None = None, + ) -> None: + self._nexus_client: _NexusClient = nexus_client + self.handle = handle + self.created = created + # Not a dataclass field, so the converter doesn't encode it. A + # decoded handle decodes its generic values as `typing.Any`. + self.body_type: type[BodyT] = body_type or typing.cast( + "type[BodyT]", typing.Any + ) + + @typing_extensions.override + def __eq__(self, other: object) -> bool: + if not isinstance(other, Stream): + return NotImplemented + return (self.handle, self.created) == (other.handle, other.created) + + __hash__: typing.ClassVar[None] = None # pyright: ignore[reportIncompatibleMethodOverride] + + @typing_extensions.override + def __repr__(self) -> str: + return f"Stream(handle={self.handle!r}, created={self.created!r})" + + async def __aenter__(self) -> typing_extensions.Self: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: types.TracebackType | None, + ) -> None: + close, failure = _support.exit_close(self, exc) + if close: + await self.close(failure=failure) + + async def publish( + self, + *values: BodyT, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> int: + """Publish values, and return the offset of the last one.""" + request = AppendRecordsRequest( + handle=self.handle, + records=[AppendRecord(body=value) for value in values], + failed_tasks=_support.local_activity_failed_tasks(), + producer=_support.client_producer(), + previous_completed_task=_support.local_activity_previous_completed_task(), + withdraw_run=_support.withdraw_run(), + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="AppendRecords", + arg=request, + id=_operation_id(), + result_type=AppendRecordsResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + response = await handle.result() + return response.last_offset or 0 + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def close( + self, + result: typing.Any | None = None, + failure: BaseException | None = None, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> None: + """Close the Stream. Readers drain what is left, then end.""" + request = CloseStreamRequest( + handle=self.handle, + result=result, + failure=failure, + producer=_support.client_producer(), + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="CloseStream", + arg=request, + id=_operation_id(), + result_type=CloseStreamResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + _ = await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + max_bytes: int | None = None, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> ReadRecordsResponse[BodyT]: + """Read one page of Records. The poll iterator reads them all.""" + request = ReadRecordsRequest( + handle=self.handle, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + max_bytes=max_bytes, + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="ReadRecords", + arg=request, + id=_operation_id(), + result_type=ReadRecordsResponse[self.body_type], + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + def iter_read( + self, + from_offset: int | None = None, + start: StreamReadStart | None = None, + max_records: int | None = None, + wait: datetime.timedelta | None = None, + create_wait: datetime.timedelta | None = None, + max_bytes: int | None = None, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> PollIterator[ReadRecordsRequest, StreamRecord[BodyT]]: + """Yield each item that ``read`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. + + Read one page of Records. The poll iterator reads them all. + + Args: + from_offset: Start offset. Empty uses start. OUT_OF_RANGE past the newest + offset + 1. + start: Where to start when from-offset is empty. Unspecified is the + beginning. + max_records: About this many items per page, 1 to 1000. A Superseded item + stays with its Record. + wait: How long to wait for a new Record before an empty answer. At most 60 + s. + create_wait: How long to wait for an uncreated Stream. Empty fails at once. + At most 60 s. + max_bytes: The maximum Record bytes in a page. A larger Record is still + returned alone. + """ + request = ReadRecordsRequest( + handle=self.handle, + from_offset=from_offset, + start=start, + max_records=max_records, + wait=wait, + create_wait=create_wait, + max_bytes=max_bytes, + ) + nexus_client = self._nexus_client + + async def read( + request: ReadRecordsRequest, + ) -> tuple[ + ReadRecordsRequest | None, collections.abc.Sequence[StreamRecord[BodyT]] + ]: + try: + handle = await nexus_client.start_operation( + operation="ReadRecords", + arg=request, + id=_operation_id(), + result_type=ReadRecordsResponse[self.body_type], + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + response = await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + return _support.next_read_request(request, response), response.records or [] + + return PollIterator(request, read, _support.resume_read_request) + + async def describe( + self, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> DescribeStreamResponse: + """Describe the Stream's status, newest offset and Retention.""" + request = DescribeStreamRequest( + handle=self.handle, + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="DescribeStream", + arg=request, + id=_operation_id(), + result_type=DescribeStreamResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + +class _StreamTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + Stream[BodyT], temporalio.api.stream.v1.message_pb2.StreamHandle + ] +): + """Passes a ``Stream`` as its ``handle``. A client can't decode a handle, since the + handle needs a client. Decode the ``handle`` and pass it to ``get_stream``. + """ + + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamHandle] | None = ( + temporalio.api.stream.v1.message_pb2.StreamHandle + ) + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamHandle, + type_hint: type[Stream[BodyT]], + ) -> Stream[BodyT]: + raise TypeError("a client Stream needs a client. Pass its handle to get_stream") + + @typing_extensions.override + def to_transfer_type( + self, value: Stream[BodyT] + ) -> temporalio.api.stream.v1.message_pb2.StreamHandle: + return _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) + + +temporalio.converter.transfer_type_convertible(_StreamTransferTypeConverter)(Stream) # pyright: ignore[reportUnusedCallResult] + + +__all__ = ["Stream"] diff --git a/temporalio/common/_streams/system/errors.py b/temporalio/common/_streams/system/errors.py new file mode 100644 index 0000000000..3f5e479bf9 --- /dev/null +++ b/temporalio/common/_streams/system/errors.py @@ -0,0 +1,146 @@ +# Generated by nexgen v0.2.11. DO NOT EDIT! + +from __future__ import annotations + +import typing + +import temporalio.exceptions + +from .models import StreamErrorDetails, StreamErrorReason + + +class StreamError(temporalio.exceptions.ApplicationError): + """A failed call, with its ``StreamErrorReason``. A reason with no class of its own + raises this class. + """ + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonUnspecified + ) + _failure_type: typing.ClassVar[str] = "StreamError" + _retryable: typing.ClassVar[bool] = False + + def __init__( + self, message: str, *, reason: StreamErrorReason | None = None + ) -> None: + reason = self._reason if reason is None else reason + super().__init__( + message, + StreamErrorDetails(reason=reason), + type=self._failure_type, + non_retryable=not self._retryable, + ) + self.reason: StreamErrorReason = reason + + +class StreamNotFoundError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonNotFound``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonNotFound + ) + _failure_type: typing.ClassVar[str] = "StreamNotFoundError" + + +class StreamExpiredError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonExpired``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonExpired + ) + _failure_type: typing.ClassVar[str] = "StreamExpiredError" + + +class StreamClosedError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonClosed``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonClosed + ) + _failure_type: typing.ClassVar[str] = "StreamClosedError" + + +class StreamFencedError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonFenced``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonFenced + ) + _failure_type: typing.ClassVar[str] = "StreamFencedError" + + +class StreamSeqDivergentError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonSeqDivergent``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonSeqDivergent + ) + _failure_type: typing.ClassVar[str] = "StreamSeqDivergentError" + + +class StreamNotOwnerError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonNotOwner``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonNotOwner + ) + _failure_type: typing.ClassVar[str] = "StreamNotOwnerError" + + +class StreamOffsetError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonOffset``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonOffset + ) + _failure_type: typing.ClassVar[str] = "StreamOffsetError" + + +class StreamTooLargeError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonTooLarge``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonTooLarge + ) + _failure_type: typing.ClassVar[str] = "StreamTooLargeError" + + +class StreamStorageError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonStorage``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonStorage + ) + _failure_type: typing.ClassVar[str] = "StreamStorageError" + _retryable: typing.ClassVar[bool] = True + + +class StreamUnsupportedError(StreamError): + """The error of ``StreamErrorReason.StreamErrorReasonUnsupported``.""" + + _reason: typing.ClassVar[StreamErrorReason] = ( + StreamErrorReason.StreamErrorReasonUnsupported + ) + _failure_type: typing.ClassVar[str] = "StreamUnsupportedError" + + +_STREAM_ERRORS: dict[StreamErrorReason, type[StreamError]] = { + StreamErrorReason.StreamErrorReasonNotFound: StreamNotFoundError, + StreamErrorReason.StreamErrorReasonExpired: StreamExpiredError, + StreamErrorReason.StreamErrorReasonClosed: StreamClosedError, + StreamErrorReason.StreamErrorReasonFenced: StreamFencedError, + StreamErrorReason.StreamErrorReasonSeqDivergent: StreamSeqDivergentError, + StreamErrorReason.StreamErrorReasonNotOwner: StreamNotOwnerError, + StreamErrorReason.StreamErrorReasonOffset: StreamOffsetError, + StreamErrorReason.StreamErrorReasonTooLarge: StreamTooLargeError, + StreamErrorReason.StreamErrorReasonStorage: StreamStorageError, + StreamErrorReason.StreamErrorReasonUnsupported: StreamUnsupportedError, +} + + +def stream_error(details: StreamErrorDetails, message: str) -> StreamError: + """The error of the reason in the details of a failed call. A reason with no class of its + own, or one that this code doesn't know, gives ``StreamError``. + """ + reason = details.reason or StreamErrorReason.StreamErrorReasonUnspecified + return _STREAM_ERRORS.get(reason, StreamError)(message, reason=reason) diff --git a/temporalio/common/_streams/system/models.py b/temporalio/common/_streams/system/models.py index b937cdfd28..8358e32a5c 100644 --- a/temporalio/common/_streams/system/models.py +++ b/temporalio/common/_streams/system/models.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -28,6 +28,14 @@ class StreamProducerKind(enum.IntEnum): StreamProducerKindStandaloneActivity = 5 +def _stream_producer_kind_from_proto(value: int) -> StreamProducerKind: + try: + return StreamProducerKind(value) + except ValueError: + # A newer peer can send a case that this enum doesn't know. + return StreamProducerKind.StreamProducerKindUnspecified + + class ExecutionType(enum.IntEnum): ExecutionTypeUnspecified = 0 ExecutionTypeWorkflow = 1 @@ -35,11 +43,18 @@ class ExecutionType(enum.IntEnum): ExecutionTypeNexusOperation = 3 -class StreamStatus(enum.IntEnum): - StreamStatusUnspecified = 0 - StreamStatusOpen = 1 - StreamStatusClosed = 2 - StreamStatusExpired = 3 +def _execution_type_from_proto(value: int) -> ExecutionType: + try: + return ExecutionType(value) + except ValueError: + # A newer peer can send a case that this enum doesn't know. + return ExecutionType.ExecutionTypeUnspecified + + +class StreamReadStart(enum.IntEnum): + StreamReadStartUnspecified = 0 + StreamReadStartBeginning = 1 + StreamReadStartEnd = 2 class StreamCloseStatus(enum.IntEnum): @@ -51,21 +66,69 @@ class StreamCloseStatus(enum.IntEnum): StreamCloseStatusTimedOut = 5 -class StreamReadStart(enum.IntEnum): - StreamReadStartUnspecified = 0 - StreamReadStartBeginning = 1 - StreamReadStartEnd = 2 +def _stream_close_status_from_proto(value: int) -> StreamCloseStatus: + try: + return StreamCloseStatus(value) + except ValueError: + # A newer peer can send a case that this enum doesn't know. + return StreamCloseStatus.StreamCloseStatusUnspecified + + +class StreamStatus(enum.IntEnum): + StreamStatusUnspecified = 0 + StreamStatusOpen = 1 + StreamStatusClosed = 2 + StreamStatusExpired = 3 + + +def _stream_status_from_proto(value: int) -> StreamStatus: + try: + return StreamStatus(value) + except ValueError: + # A newer peer can send a case that this enum doesn't know. + return StreamStatus.StreamStatusUnspecified + + +class StreamErrorReason(enum.IntEnum): + StreamErrorReasonUnspecified = 0 + StreamErrorReasonNotFound = 1 + StreamErrorReasonExpired = 2 + StreamErrorReasonClosed = 3 + StreamErrorReasonFenced = 4 + StreamErrorReasonSequence = 5 + StreamErrorReasonSeqDivergent = 6 + StreamErrorReasonNotOwner = 7 + StreamErrorReasonFormatTooNew = 8 + StreamErrorReasonOffset = 9 + StreamErrorReasonTooLarge = 10 + StreamErrorReasonInvalid = 11 + StreamErrorReasonStorage = 12 + StreamErrorReasonUnknownWriter = 13 + StreamErrorReasonUnsupported = 14 + + +def _stream_error_reason_from_proto(value: int) -> StreamErrorReason: + try: + return StreamErrorReason(value) + except ValueError: + # A newer peer can send a case that this enum doesn't know. + return StreamErrorReason.StreamErrorReasonUnspecified + + +BodyT = typing.TypeVar("BodyT") @dataclasses.dataclass(slots=True) class CreateStreamRequest: - """Create a Stream, or get the one that exists, which keeps its Retention.""" + """Create a Stream, or get the open one. A closed, expired or another owner's Stream + fails. + """ - stream_id: str + id: str retention: datetime.timedelta | None = None withdraw_run: str | None = None owner: Producer | None = None - ref: StreamRef | None = None + stream_handle: StreamHandle | None = None class _CreateStreamRequestTransferTypeConverter( @@ -85,9 +148,9 @@ def from_transfer_type( value: temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest, type_hint: type[CreateStreamRequest], ) -> CreateStreamRequest: - stream_id = value.stream_id + id = value.stream_id return CreateStreamRequest( - stream_id=stream_id, + id=id, retention=_support.duration_from_proto(value.retention) if value.HasField("retention") else None, @@ -97,10 +160,10 @@ def from_transfer_type( ) if value.HasField("owner") else None, - ref=_StreamRefTransferTypeConverter().from_transfer_type( - value.ref, StreamRef + stream_handle=_StreamHandleTransferTypeConverter().from_transfer_type( + value.stream_handle, StreamHandle ) - if value.HasField("ref") + if value.HasField("stream_handle") else None, ) @@ -112,7 +175,7 @@ def to_transfer_type( message = ( temporalio.api.streamservice.v1.request_response_pb2.CreateStreamRequest() ) - message.stream_id = value.stream_id + message.stream_id = value.id if value.retention is not None: message.retention.CopyFrom(_support.duration_to_proto(value.retention)) if value.withdraw_run is not None: @@ -121,9 +184,11 @@ def to_transfer_type( message.owner.CopyFrom( _ProducerTransferTypeConverter().to_transfer_type(value.owner) ) - if value.ref is not None: - message.ref.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.ref) + if value.stream_handle is not None: + message.stream_handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type( + value.stream_handle + ) ) return message @@ -158,10 +223,14 @@ def from_transfer_type( value: temporalio.api.stream.v1.message_pb2.Producer, type_hint: type[Producer], ) -> Producer: + if not value.id: + raise ValueError("missing required field Producer.id") id = value.id return Producer( id=id, - kind=StreamProducerKind(value.kind) if value.kind != 0 else None, + kind=_stream_producer_kind_from_proto(value.kind) + if value.kind != 0 + else None, epoch=_ProducerEpochTransferTypeConverter().from_transfer_type( value.epoch, ProducerEpoch ) @@ -316,7 +385,7 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class StreamRef: +class StreamHandle: """Identifies one Stream. The id is opaque to callers.""" id: str @@ -324,23 +393,25 @@ class StreamRef: backend: str | None = None -class _StreamRefTransferTypeConverter( +class _StreamHandleTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - StreamRef, temporalio.api.stream.v1.message_pb2.StreamRef + StreamHandle, temporalio.api.stream.v1.message_pb2.StreamHandle ] ): - transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRef] | None = ( - temporalio.api.stream.v1.message_pb2.StreamRef + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamHandle] | None = ( + temporalio.api.stream.v1.message_pb2.StreamHandle ) @typing_extensions.override def from_transfer_type( self, - value: temporalio.api.stream.v1.message_pb2.StreamRef, - type_hint: type[StreamRef], - ) -> StreamRef: + value: temporalio.api.stream.v1.message_pb2.StreamHandle, + type_hint: type[StreamHandle], + ) -> StreamHandle: + if not value.id: + raise ValueError("missing required field StreamHandle.id") id = value.id - return StreamRef( + return StreamHandle( id=id, owner=_ExecutionTransferTypeConverter().from_transfer_type( value.owner, Execution @@ -353,9 +424,9 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: StreamRef, - ) -> temporalio.api.stream.v1.message_pb2.StreamRef: - message = temporalio.api.stream.v1.message_pb2.StreamRef() + value: StreamHandle, + ) -> temporalio.api.stream.v1.message_pb2.StreamHandle: + message = temporalio.api.stream.v1.message_pb2.StreamHandle() message.id = value.id if value.owner is not None: message.owner.CopyFrom( @@ -366,8 +437,8 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible(_StreamRefTransferTypeConverter)( - StreamRef +temporalio.converter.transfer_type_convertible(_StreamHandleTransferTypeConverter)( + StreamHandle ) # pyright: ignore[reportUnusedCallResult] @@ -396,7 +467,7 @@ def from_transfer_type( type_hint: type[Execution], ) -> Execution: return Execution( - type=ExecutionType(value.type) if value.type != 0 else None, + type=_execution_type_from_proto(value.type) if value.type != 0 else None, business_id=value.business_id if bool(value.business_id) else None, run_id=value.run_id if bool(value.run_id) else None, ) @@ -425,9 +496,8 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) class CreateStreamResult: - stream: StreamRef | None = None + handle: StreamHandle | None = None created: bool | None = None - status: StreamStatus | None = None class _CreateStreamResultTransferTypeConverter( @@ -448,13 +518,12 @@ def from_transfer_type( type_hint: type[CreateStreamResult], ) -> CreateStreamResult: return CreateStreamResult( - stream=_StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + handle=_StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) - if value.HasField("stream") + if value.HasField("handle") else None, created=value.created if bool(value.created) else None, - status=StreamStatus(value.status) if value.status != 0 else None, ) @typing_extensions.override @@ -465,16 +534,12 @@ def to_transfer_type( message = ( temporalio.api.streamservice.v1.request_response_pb2.CreateStreamResponse() ) - if value.stream is not None: - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + if value.handle is not None: + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) if value.created is not None: message.created = value.created - if value.status is not None: - message.status = temporalio.api.enums.v1.stream_pb2.StreamStatus.ValueType( - value.status - ) return message @@ -484,13 +549,13 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class AppendRecordsRequest: +class AppendRecordsRequest(typing.Generic[BodyT]): """Append. FAILED_PRECONDITION if closed, fenced or format too new. NOT_FOUND if absent. """ - stream: StreamRef - records: list[AppendRecord] | None = dataclasses.field(default_factory=list) + handle: StreamHandle + records: list[AppendRecord[BodyT]] | None = dataclasses.field(default_factory=list) failed_tasks: list[int] | None = dataclasses.field(default_factory=list) producer: Producer | None = None sequence: int | None = None @@ -500,7 +565,7 @@ class AppendRecordsRequest: class _AppendRecordsRequestTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - AppendRecordsRequest, + AppendRecordsRequest[BodyT], temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest, ] ): @@ -513,18 +578,22 @@ class _AppendRecordsRequestTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest, - type_hint: type[AppendRecordsRequest], - ) -> AppendRecordsRequest: - if not value.HasField("stream"): - raise ValueError("missing required field AppendRecordsRequest.stream") - stream = _StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + type_hint: type[AppendRecordsRequest[BodyT]], + ) -> AppendRecordsRequest[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) + if not value.HasField("handle"): + raise ValueError("missing required field AppendRecordsRequest.handle") + handle = _StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) return AppendRecordsRequest( - stream=stream, + handle=handle, records=[ - _AppendRecordTransferTypeConverter().from_transfer_type( - item, AppendRecord + typing.cast( + AppendRecord[BodyT], + _AppendRecordTransferTypeConverter().from_transfer_type( + item, AppendRecord[body_type] + ), ) for item in value.records ], @@ -546,13 +615,13 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: AppendRecordsRequest, + value: AppendRecordsRequest[BodyT], ) -> temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest: message = ( temporalio.api.streamservice.v1.request_response_pb2.AppendRecordsRequest() ) - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) if value.records: for element in value.records: @@ -580,20 +649,20 @@ def to_transfer_type( temporalio.converter.transfer_type_convertible( - _AppendRecordsRequestTransferTypeConverter + _AppendRecordsRequestTransferTypeConverter[typing.Any] )(AppendRecordsRequest) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) -class AppendRecord: +class AppendRecord(typing.Generic[BodyT]): """One Record to append.""" - body: typing.Any + body: BodyT class _AppendRecordTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - AppendRecord, temporalio.api.stream.v1.message_pb2.AppendRecord + AppendRecord[BodyT], temporalio.api.stream.v1.message_pb2.AppendRecord ] ): transfer_type: type[temporalio.api.stream.v1.message_pb2.AppendRecord] | None = ( @@ -604,11 +673,12 @@ class _AppendRecordTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.stream.v1.message_pb2.AppendRecord, - type_hint: type[AppendRecord], - ) -> AppendRecord: + type_hint: type[AppendRecord[BodyT]], + ) -> AppendRecord[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) if not value.HasField("body"): raise ValueError("missing required field AppendRecord.body") - body = _support.payload_from_proto(value.body) + body = typing.cast(BodyT, _support.payload_from_proto(value.body, body_type)) return AppendRecord( body=body, ) @@ -616,16 +686,16 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: AppendRecord, + value: AppendRecord[BodyT], ) -> temporalio.api.stream.v1.message_pb2.AppendRecord: message = temporalio.api.stream.v1.message_pb2.AppendRecord() message.body.CopyFrom(_support.payload_to_proto(value.body)) return message -temporalio.converter.transfer_type_convertible(_AppendRecordTransferTypeConverter)( - AppendRecord -) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible( + _AppendRecordTransferTypeConverter[typing.Any] +)(AppendRecord) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) @@ -738,13 +808,10 @@ def to_transfer_type( class CloseStreamRequest: """Close a Stream. Readers drain what is left, then end. Later appends are refused.""" - stream: StreamRef - failed_tasks: list[int] | None = dataclasses.field(default_factory=list) + handle: StreamHandle producer: Producer | None = None result: typing.Any | None = None failure: BaseException | None = None - status: StreamCloseStatus | None = None - previous_completed_task: WorkflowTaskEpoch | None = None class _CloseStreamRequestTransferTypeConverter( @@ -764,14 +831,13 @@ def from_transfer_type( value: temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest, type_hint: type[CloseStreamRequest], ) -> CloseStreamRequest: - if not value.HasField("stream"): - raise ValueError("missing required field CloseStreamRequest.stream") - stream = _StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + if not value.HasField("handle"): + raise ValueError("missing required field CloseStreamRequest.handle") + handle = _StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) return CloseStreamRequest( - stream=stream, - failed_tasks=list(value.failed_tasks), + handle=handle, producer=_ProducerTransferTypeConverter().from_transfer_type( value.producer, Producer ) @@ -783,12 +849,6 @@ def from_transfer_type( failure=_support.failure_from_proto(value.failure) if value.HasField("failure") else None, - status=StreamCloseStatus(value.status) if value.status != 0 else None, - previous_completed_task=_WorkflowTaskEpochTransferTypeConverter().from_transfer_type( - value.previous_completed_task, WorkflowTaskEpoch - ) - if value.HasField("previous_completed_task") - else None, ) @typing_extensions.override @@ -799,11 +859,9 @@ def to_transfer_type( message = ( temporalio.api.streamservice.v1.request_response_pb2.CloseStreamRequest() ) - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) - if value.failed_tasks: - message.failed_tasks.extend(value.failed_tasks) if value.producer is not None: message.producer.CopyFrom( _ProducerTransferTypeConverter().to_transfer_type(value.producer) @@ -812,18 +870,6 @@ def to_transfer_type( message.result.CopyFrom(_support.payload_to_proto(value.result)) if value.failure is not None: message.failure.CopyFrom(_support.failure_to_proto(value.failure)) - if value.status is not None: - message.status = ( - temporalio.api.enums.v1.stream_pb2.StreamCloseStatus.ValueType( - value.status - ) - ) - if value.previous_completed_task is not None: - message.previous_completed_task.CopyFrom( - _WorkflowTaskEpochTransferTypeConverter().to_transfer_type( - value.previous_completed_task - ) - ) return message @@ -876,7 +922,7 @@ def to_transfer_type( class ReadRecordsRequest: """Read Records from an offset.""" - stream: StreamRef + handle: StreamHandle from_offset: int | None = None start: StreamReadStart | None = None max_records: int | None = None @@ -903,13 +949,13 @@ def from_transfer_type( value: temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest, type_hint: type[ReadRecordsRequest], ) -> ReadRecordsRequest: - if not value.HasField("stream"): - raise ValueError("missing required field ReadRecordsRequest.stream") - stream = _StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + if not value.HasField("handle"): + raise ValueError("missing required field ReadRecordsRequest.handle") + handle = _StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) return ReadRecordsRequest( - stream=stream, + handle=handle, from_offset=value.from_offset if value.from_offset != 0 else None, start=StreamReadStart(value.start) if value.start != 0 else None, max_records=value.max_records if value.max_records != 0 else None, @@ -931,8 +977,8 @@ def to_transfer_type( message = ( temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsRequest() ) - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) if value.from_offset is not None: message.from_offset = value.from_offset @@ -961,15 +1007,15 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class ReadRecordsResponse: - records: list[StreamRecord] | None = dataclasses.field(default_factory=list) +class ReadRecordsResponse(typing.Generic[BodyT]): + records: list[StreamRecord[BodyT]] | None = dataclasses.field(default_factory=list) next_offset: int | None = None done: bool | None = None class _ReadRecordsResponseTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - ReadRecordsResponse, + ReadRecordsResponse[BodyT], temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse, ] ): @@ -982,12 +1028,16 @@ class _ReadRecordsResponseTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse, - type_hint: type[ReadRecordsResponse], - ) -> ReadRecordsResponse: + type_hint: type[ReadRecordsResponse[BodyT]], + ) -> ReadRecordsResponse[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) return ReadRecordsResponse( records=[ - _StreamRecordTransferTypeConverter().from_transfer_type( - item, StreamRecord + typing.cast( + StreamRecord[BodyT], + _StreamRecordTransferTypeConverter().from_transfer_type( + item, StreamRecord[body_type] + ), ) for item in value.records ], @@ -998,7 +1048,7 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: ReadRecordsResponse, + value: ReadRecordsResponse[BodyT], ) -> temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse: message = ( temporalio.api.streamservice.v1.request_response_pb2.ReadRecordsResponse() @@ -1017,21 +1067,21 @@ def to_transfer_type( temporalio.converter.transfer_type_convertible( - _ReadRecordsResponseTransferTypeConverter + _ReadRecordsResponseTransferTypeConverter[typing.Any] )(ReadRecordsResponse) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) -class StreamRecord: +class StreamRecord(typing.Generic[BodyT]): """One item that a read gives.""" - item: StreamItem + item: StreamItem[BodyT] | None = None offset: int | None = None class _StreamRecordTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - StreamRecord, temporalio.api.stream.v1.message_pb2.StreamRecord + StreamRecord[BodyT], temporalio.api.stream.v1.message_pb2.StreamRecord ] ): transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRecord] | None = ( @@ -1042,20 +1092,26 @@ class _StreamRecordTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.stream.v1.message_pb2.StreamRecord, - type_hint: type[StreamRecord], - ) -> StreamRecord: - _oneof_item: StreamItemData | StreamItemSuperseded | StreamItemClose + type_hint: type[StreamRecord[BodyT]], + ) -> StreamRecord[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) + _oneof_item: ( + StreamItemData[BodyT] | StreamItemSuperseded | StreamItemClose | None + ) match value.WhichOneof("item"): case "data": _oneof_item = StreamItemData( - _StreamDataTransferTypeConverter().from_transfer_type( - value.data, StreamData + typing.cast( + StreamData[BodyT], + _StreamDataTransferTypeConverter().from_transfer_type( + value.data, StreamData[body_type] + ), ) ) case "superseded": _oneof_item = StreamItemSuperseded( - _SupersessionTransferTypeConverter().from_transfer_type( - value.superseded, Supersession + _StreamSupersededTransferTypeConverter().from_transfer_type( + value.superseded, StreamSuperseded ) ) case "close": @@ -1065,7 +1121,7 @@ def from_transfer_type( ) ) case None: - raise ValueError("missing required field StreamRecord.item") + _oneof_item = None return StreamRecord( item=_oneof_item, offset=value.offset if value.offset != 0 else None, @@ -1074,49 +1130,50 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: StreamRecord, + value: StreamRecord[BodyT], ) -> temporalio.api.stream.v1.message_pb2.StreamRecord: message = temporalio.api.stream.v1.message_pb2.StreamRecord() - match value.item: - case StreamItemData(): - message.data.CopyFrom( - _StreamDataTransferTypeConverter().to_transfer_type( - value.item.value + if value.item is not None: + match value.item: + case StreamItemData(): + message.data.CopyFrom( + _StreamDataTransferTypeConverter().to_transfer_type( + value.item.value + ) ) - ) - case StreamItemSuperseded(): - message.superseded.CopyFrom( - _SupersessionTransferTypeConverter().to_transfer_type( - value.item.value + case StreamItemSuperseded(): + message.superseded.CopyFrom( + _StreamSupersededTransferTypeConverter().to_transfer_type( + value.item.value + ) ) - ) - case StreamItemClose(): - message.close.CopyFrom( - _StreamCloseTransferTypeConverter().to_transfer_type( - value.item.value + case StreamItemClose(): + message.close.CopyFrom( + _StreamCloseTransferTypeConverter().to_transfer_type( + value.item.value + ) ) - ) if value.offset is not None: message.offset = value.offset return message -temporalio.converter.transfer_type_convertible(_StreamRecordTransferTypeConverter)( - StreamRecord -) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible( + _StreamRecordTransferTypeConverter[typing.Any] +)(StreamRecord) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) -class StreamData: +class StreamData(typing.Generic[BodyT]): """One Record that a Producer appended.""" - body: typing.Any + body: BodyT producer: Producer | None = None class _StreamDataTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - StreamData, temporalio.api.stream.v1.message_pb2.StreamData + StreamData[BodyT], temporalio.api.stream.v1.message_pb2.StreamData ] ): transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamData] | None = ( @@ -1127,11 +1184,12 @@ class _StreamDataTransferTypeConverter( def from_transfer_type( self, value: temporalio.api.stream.v1.message_pb2.StreamData, - type_hint: type[StreamData], - ) -> StreamData: + type_hint: type[StreamData[BodyT]], + ) -> StreamData[BodyT]: + (body_type,) = typing.get_args(type_hint) or (typing.Any,) if not value.HasField("body"): raise ValueError("missing required field StreamData.body") - body = _support.payload_from_proto(value.body) + body = typing.cast(BodyT, _support.payload_from_proto(value.body, body_type)) return StreamData( body=body, producer=_ProducerTransferTypeConverter().from_transfer_type( @@ -1144,7 +1202,7 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: StreamData, + value: StreamData[BodyT], ) -> temporalio.api.stream.v1.message_pb2.StreamData: message = temporalio.api.stream.v1.message_pb2.StreamData() message.body.CopyFrom(_support.payload_to_proto(value.body)) @@ -1155,13 +1213,13 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible(_StreamDataTransferTypeConverter)( - StreamData -) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible( + _StreamDataTransferTypeConverter[typing.Any] +)(StreamData) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) -class Supersession: +class StreamSuperseded: """Which earlier Records a Superseded item withdraws.""" producer_id: str @@ -1173,26 +1231,28 @@ class Supersession: run_id: str | None = None -class _SupersessionTransferTypeConverter( +class _StreamSupersededTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - Supersession, temporalio.api.stream.v1.message_pb2.Supersession + StreamSuperseded, temporalio.api.stream.v1.message_pb2.StreamSuperseded ] ): - transfer_type: type[temporalio.api.stream.v1.message_pb2.Supersession] | None = ( - temporalio.api.stream.v1.message_pb2.Supersession - ) + transfer_type: ( + type[temporalio.api.stream.v1.message_pb2.StreamSuperseded] | None + ) = temporalio.api.stream.v1.message_pb2.StreamSuperseded @typing_extensions.override def from_transfer_type( self, - value: temporalio.api.stream.v1.message_pb2.Supersession, - type_hint: type[Supersession], - ) -> Supersession: + value: temporalio.api.stream.v1.message_pb2.StreamSuperseded, + type_hint: type[StreamSuperseded], + ) -> StreamSuperseded: + if not value.producer_id: + raise ValueError("missing required field StreamSuperseded.producer_id") producer_id = value.producer_id - return Supersession( + return StreamSuperseded( producer_id=producer_id, failed_tasks=list(value.failed_tasks), - producer_kind=StreamProducerKind(value.producer_kind) + producer_kind=_stream_producer_kind_from_proto(value.producer_kind) if value.producer_kind != 0 else None, epoch=_ProducerEpochTransferTypeConverter().from_transfer_type( @@ -1212,9 +1272,9 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: Supersession, - ) -> temporalio.api.stream.v1.message_pb2.Supersession: - message = temporalio.api.stream.v1.message_pb2.Supersession() + value: StreamSuperseded, + ) -> temporalio.api.stream.v1.message_pb2.StreamSuperseded: + message = temporalio.api.stream.v1.message_pb2.StreamSuperseded() message.producer_id = value.producer_id if value.failed_tasks: message.failed_tasks.extend(value.failed_tasks) @@ -1241,8 +1301,8 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible(_SupersessionTransferTypeConverter)( - Supersession +temporalio.converter.transfer_type_convertible(_StreamSupersededTransferTypeConverter)( + StreamSuperseded ) # pyright: ignore[reportUnusedCallResult] @@ -1283,7 +1343,9 @@ def from_transfer_type( failure=_support.failure_from_proto(value.failure) if value.HasField("failure") else None, - status=StreamCloseStatus(value.status) if value.status != 0 else None, + status=_stream_close_status_from_proto(value.status) + if value.status != 0 + else None, ) @typing_extensions.override @@ -1316,7 +1378,7 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) class DescribeStreamRequest: - stream: StreamRef + handle: StreamHandle class _DescribeStreamRequestTransferTypeConverter( @@ -1336,13 +1398,13 @@ def from_transfer_type( value: temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest, type_hint: type[DescribeStreamRequest], ) -> DescribeStreamRequest: - if not value.HasField("stream"): - raise ValueError("missing required field DescribeStreamRequest.stream") - stream = _StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + if not value.HasField("handle"): + raise ValueError("missing required field DescribeStreamRequest.handle") + handle = _StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) return DescribeStreamRequest( - stream=stream, + handle=handle, ) @typing_extensions.override @@ -1353,8 +1415,8 @@ def to_transfer_type( message = ( temporalio.api.streamservice.v1.request_response_pb2.DescribeStreamRequest() ) - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) return message @@ -1391,7 +1453,9 @@ def from_transfer_type( type_hint: type[DescribeStreamResponse], ) -> DescribeStreamResponse: return DescribeStreamResponse( - status=StreamStatus(value.status) if value.status != 0 else None, + status=_stream_status_from_proto(value.status) + if value.status != 0 + else None, last_offset=value.last_offset if value.last_offset != 0 else None, retention=_support.duration_from_proto(value.retention) if value.HasField("retention") @@ -1421,13 +1485,61 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class StreamItemData: - value: StreamData +class StreamErrorDetails: + """The details of a failed Stream call.""" + + reason: StreamErrorReason | None = None + + +class _StreamErrorDetailsTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + StreamErrorDetails, temporalio.api.stream.v1.message_pb2.StreamErrorDetails + ] +): + transfer_type: ( + type[temporalio.api.stream.v1.message_pb2.StreamErrorDetails] | None + ) = temporalio.api.stream.v1.message_pb2.StreamErrorDetails + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.stream.v1.message_pb2.StreamErrorDetails, + type_hint: type[StreamErrorDetails], + ) -> StreamErrorDetails: + return StreamErrorDetails( + reason=_stream_error_reason_from_proto(value.reason) + if value.reason != 0 + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: StreamErrorDetails, + ) -> temporalio.api.stream.v1.message_pb2.StreamErrorDetails: + message = temporalio.api.stream.v1.message_pb2.StreamErrorDetails() + if value.reason is not None: + message.reason = ( + temporalio.api.enums.v1.stream_pb2.StreamErrorReason.ValueType( + value.reason + ) + ) + return message + + +temporalio.converter.transfer_type_convertible( + _StreamErrorDetailsTransferTypeConverter +)(StreamErrorDetails) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class StreamItemData(typing.Generic[BodyT]): + value: StreamData[BodyT] @dataclasses.dataclass(slots=True) class StreamItemSuperseded: - value: Supersession + value: StreamSuperseded @dataclasses.dataclass(slots=True) @@ -1435,5 +1547,5 @@ class StreamItemClose: value: StreamClose -StreamItem = StreamItemData | StreamItemSuperseded | StreamItemClose +StreamItem = StreamItemData[BodyT] | StreamItemSuperseded | StreamItemClose """A read item: a Record, a Superseded item or the close.""" diff --git a/temporalio/common/_streams/system/operations/__init__.py b/temporalio/common/_streams/system/operations/__init__.py index 0dd909df4d..3bd09999cf 100644 --- a/temporalio/common/_streams/system/operations/__init__.py +++ b/temporalio/common/_streams/system/operations/__init__.py @@ -1,3 +1,3 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/common/_streams/system/operations/create_stream.py b/temporalio/common/_streams/system/operations/create_stream.py index 94fc4a4926..7da16a3403 100644 --- a/temporalio/common/_streams/system/operations/create_stream.py +++ b/temporalio/common/_streams/system/operations/create_stream.py @@ -1,64 +1,103 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations +import asyncio import datetime +import typing +import temporalio.nexus.system._support as _support import temporalio.workflow +from .._resource_call import resource_call from .._resources import Stream +from ..errors import StreamError, stream_error from ..models import ( CreateStreamRequest, CreateStreamResult, - Producer, - StreamRef, + StreamErrorDetails, + StreamHandle, ) +if typing.TYPE_CHECKING: + import temporalio.common + + +BodyT = typing.TypeVar("BodyT") + async def _create_stream( request: CreateStreamRequest, -) -> Stream: + body_type: type[BodyT], +) -> Stream[BodyT]: nexus_client = temporalio.workflow.create_nexus_client( service="temporal.api.streamservice.v1.StreamService", endpoint="__temporal_system", ) - handle = await nexus_client.start_operation( - operation="CreateStream", - input=request, - output_type=CreateStreamResult, - ) - result = await handle - if result.stream is None: - raise RuntimeError("CreateStream answered without `stream`") - return Stream( - stream=result.stream, - created=result.created or None, - status=result.status or None, - ) + try: + handle = await nexus_client.start_operation( + operation="CreateStream", + input=request, + output_type=CreateStreamResult, + ) + result = await handle + if result.handle is None: + raise RuntimeError("CreateStream answered without `handle`") + return Stream( + body_type=body_type, + handle=result.handle, + created=result.created or None, + ) + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err +@resource_call async def create_stream( + definition: type[temporalio.common.StreamDefinition[BodyT]], *, - stream_id: str, retention: datetime.timedelta | None = None, - withdraw_run: str | None = None, - owner: Producer | None = None, - ref: StreamRef | None = None, -) -> Stream: + backend: str | None = None, +) -> Stream[BodyT]: """ Args: - 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. - owner: The Workflow run that owns the Stream. The Workflow transport sets it. - ref: The Stream's ref, with the owner fields. + backend: Registered backend name, empty for the default. The Store never reads + it. """ + (body_type,) = _support.definition_type_arguments(definition) + id = _support.workflow_stream_id(definition) + stream_handle = StreamHandle( + backend=backend, + id=id, + owner=_support.workflow_handle_owner(), + ) request = CreateStreamRequest( - stream_id=stream_id, + id=id, retention=retention, - withdraw_run=withdraw_run, - owner=owner, - ref=ref, + withdraw_run=_support.withdraw_run(), + owner=_support.workflow_owner(), + stream_handle=stream_handle, + ) + return await _create_stream(request, body_type) + + +def start_create_stream( + definition: type[temporalio.common.StreamDefinition[BodyT]], + *, + retention: datetime.timedelta | None = None, + backend: str | None = None, +) -> asyncio.Task[Stream[BodyT]]: + """Start ``create_stream`` at the call, as ``workflow.start_activity`` places an + Activity, so sync Workflow code keeps the operation's place in the command order. + The task gives what ``create_stream`` returns. + """ + return _support.start_now( + create_stream(definition, retention=retention, backend=backend) ) - return await _create_stream(request) diff --git a/temporalio/common/_streams/system/services.py b/temporalio/common/_streams/system/services.py index bbb5f522cc..f4e1815134 100644 --- a/temporalio/common/_streams/system/services.py +++ b/temporalio/common/_streams/system/services.py @@ -1,7 +1,9 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations +import typing + from nexusrpc import Operation, service from .models import ( @@ -31,7 +33,7 @@ class StreamService: ] = Operation(name="CreateStream") append_records: Operation[ - AppendRecordsRequest, + AppendRecordsRequest[typing.Any], AppendRecordsResponse, ] = Operation(name="AppendRecords") @@ -42,7 +44,7 @@ class StreamService: read_records: Operation[ ReadRecordsRequest, - ReadRecordsResponse, + ReadRecordsResponse[typing.Any], ] = Operation(name="ReadRecords") describe_stream: Operation[ diff --git a/temporalio/common/_streams/writer_system/__init__.py b/temporalio/common/_streams/writer_system/__init__.py index d1849f5471..c7723818d4 100644 --- a/temporalio/common/_streams/writer_system/__init__.py +++ b/temporalio/common/_streams/writer_system/__init__.py @@ -1,4 +1,5 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. from __future__ import annotations @@ -15,8 +16,8 @@ OpenWriterRequest, Producer, ProducerEpoch, + StreamHandle, StreamProducerKind, - StreamRef, Timestamp, WorkflowTaskEpoch, WriteRecordsRequest, @@ -38,8 +39,8 @@ "OpenWriterRequest", "Producer", "ProducerEpoch", + "StreamHandle", "StreamProducerKind", - "StreamRef", "Timestamp", "WorkflowTaskEpoch", "WriteRecordsRequest", diff --git a/temporalio/common/_streams/writer_system/_resources/__init__.py b/temporalio/common/_streams/writer_system/_resources/__init__.py new file mode 100644 index 0000000000..bc59364ba6 --- /dev/null +++ b/temporalio/common/_streams/writer_system/_resources/__init__.py @@ -0,0 +1,10 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +from .writer import Writer + +__all__ = [ + "Writer", +] diff --git a/temporalio/common/_streams/writer_system/_resources/writer.py b/temporalio/common/_streams/writer_system/_resources/writer.py new file mode 100644 index 0000000000..2148260fb5 --- /dev/null +++ b/temporalio/common/_streams/writer_system/_resources/writer.py @@ -0,0 +1,13 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +import dataclasses + + +@dataclasses.dataclass +class Writer: + """An open writer of one Producer. The crate keeps its epoch and sequence.""" + + writer_id: int diff --git a/temporalio/common/_streams/writer_system/client.py b/temporalio/common/_streams/writer_system/client.py deleted file mode 100644 index 7046efe0c2..0000000000 --- a/temporalio/common/_streams/writer_system/client.py +++ /dev/null @@ -1,256 +0,0 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! - -from __future__ import annotations - -import abc -import collections.abc -import dataclasses -import datetime -import typing - -from .models import ( - AppendRecord, - CloseWriterRequest, - CloseWriterResponse, - CloseWritersRequest, - CloseWritersResponse, - FreeWriterRequest, - FreeWriterResponse, - OpenWriterRequest, - Producer, - StreamRef, - WorkflowTaskEpoch, - WriteRecordsRequest, - WriteRecordsResponse, - WriterResult, -) - -_OutputT = typing.TypeVar("_OutputT") - - -class _NexusClient(typing.Protocol): - """Starts the Nexus operations of one service from code outside a Workflow.""" - - @typing.overload - async def start_operation( - self, - operation: str, - input: typing.Any, - *, - output_type: type[_OutputT], - ) -> collections.abc.Awaitable[_OutputT]: ... - - @typing.overload - async def start_operation( - self, - operation: str, - input: typing.Any, - ) -> collections.abc.Awaitable[None]: ... - - -class StreamWriterClientExtension(abc.ABC): - """Starts the operations of ``StreamWriter`` from code outside a Workflow. - - A client inherits this class and implements ``_create_nexus_client``. - - .. warning:: - This API is experimental and subject to change. - """ - - @abc.abstractmethod - def _create_nexus_client( - self, - *, - service: str, - endpoint: str | None, - ) -> _NexusClient: - """Return the Nexus client that starts this service's operations. - - ``endpoint`` is ``None`` for a service with no endpoint, which the client serves. - """ - ... - - async def _open_writer( - self, - request: OpenWriterRequest, - ) -> Writer: - nexus_client = self._create_nexus_client( - service="temporal.sdk.streams.v1.StreamWriter", - endpoint=None, - ) - handle = await nexus_client.start_operation( - operation="OpenWriter", - input=request, - output_type=WriterResult, - ) - result = await handle - if result.writer_id is None: - raise RuntimeError("OpenWriter answered without `writer_id`") - return Writer( - _client=self, - writer_id=result.writer_id, - ) - - async def open_writer( - self, - *, - stream: StreamRef, - failed_tasks: list[int] | None = None, - producer: Producer | None = None, - max_batch_size: int | None = None, - max_retry_duration: datetime.timedelta | None = None, - previous_completed_task: WorkflowTaskEpoch | None = None, - ) -> Writer: - """ - Args: - failed_tasks: Local Activity only. The run's failed or timed out Workflow Tasks. - max_batch_size: A write sends at most this many Records in one batch. Zero is - 1000. - max_retry_duration: How long a write retries a Store failure. Empty is 10 s. - previous_completed_task: Local Activity only. The run's last completed Workflow - Task. - """ - request = OpenWriterRequest( - stream=stream, - failed_tasks=failed_tasks or [], - producer=producer, - max_batch_size=max_batch_size, - max_retry_duration=max_retry_duration, - previous_completed_task=previous_completed_task, - ) - return await self._open_writer(request) - - async def _write_records( - self, - request: WriteRecordsRequest, - ) -> collections.abc.Awaitable[WriteRecordsResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.sdk.streams.v1.StreamWriter", - endpoint=None, - ) - return await nexus_client.start_operation( - operation="WriteRecords", - input=request, - output_type=WriteRecordsResponse, - ) - - async def write_records( - self, - *, - records: list[AppendRecord] | None = None, - writer_id: int | None = None, - ) -> collections.abc.Awaitable[WriteRecordsResponse,]: - request = WriteRecordsRequest( - records=records or [], - writer_id=writer_id, - ) - return await self._write_records(request) - - async def _close_writer( - self, - request: CloseWriterRequest, - ) -> collections.abc.Awaitable[CloseWriterResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.sdk.streams.v1.StreamWriter", - endpoint=None, - ) - return await nexus_client.start_operation( - operation="CloseWriter", - input=request, - output_type=CloseWriterResponse, - ) - - async def close_writer( - self, - *, - writer_id: int | None = None, - ) -> collections.abc.Awaitable[CloseWriterResponse,]: - request = CloseWriterRequest( - writer_id=writer_id, - ) - return await self._close_writer(request) - - async def _free_writer( - self, - request: FreeWriterRequest, - ) -> collections.abc.Awaitable[FreeWriterResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.sdk.streams.v1.StreamWriter", - endpoint=None, - ) - return await nexus_client.start_operation( - operation="FreeWriter", - input=request, - output_type=FreeWriterResponse, - ) - - async def free_writer( - self, - *, - writer_id: int | None = None, - ) -> collections.abc.Awaitable[FreeWriterResponse,]: - request = FreeWriterRequest( - writer_id=writer_id, - ) - return await self._free_writer(request) - - async def _close_writers( - self, - request: CloseWritersRequest, - ) -> collections.abc.Awaitable[CloseWritersResponse,]: - nexus_client = self._create_nexus_client( - service="temporal.sdk.streams.v1.StreamWriter", - endpoint=None, - ) - return await nexus_client.start_operation( - operation="CloseWriters", - input=request, - output_type=CloseWritersResponse, - ) - - async def close_writers( - self, - ) -> collections.abc.Awaitable[CloseWritersResponse,]: - """Drop every open writer, for example at shutdown.""" - request = CloseWritersRequest() - return await self._close_writers(request) - - -@dataclasses.dataclass -class Writer: - """An open writer of one Producer. The crate keeps its epoch and sequence.""" - - _client: StreamWriterClientExtension = dataclasses.field(repr=False, compare=False) - writer_id: int - - async def write_records( - self, - records: list[AppendRecord] | None = None, - ) -> WriteRecordsResponse: - """Write Records at once. The answer has the first and the last offset.""" - request = WriteRecordsRequest( - writer_id=self.writer_id, - records=records or [], - ) - handle = await self._client._write_records(request) - return await handle - - async def close( - self, - ) -> CloseWriterResponse: - """Release the handle. The last one waits for a write, then drops the writer.""" - request = CloseWriterRequest( - writer_id=self.writer_id, - ) - handle = await self._client._close_writer(request) - return await handle - - async def free( - self, - ) -> FreeWriterResponse: - """Release the handle with no wait, for a handle that lang drops.""" - request = FreeWriterRequest( - writer_id=self.writer_id, - ) - handle = await self._client._free_writer(request) - return await handle diff --git a/temporalio/common/_streams/writer_system/client/__init__.py b/temporalio/common/_streams/writer_system/client/__init__.py new file mode 100644 index 0000000000..ec54ba2a09 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client/__init__.py @@ -0,0 +1,12 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +from ._extensions import StreamWriterClientExtension +from ._resources import Writer + +__all__ = [ + "StreamWriterClientExtension", + "Writer", +] diff --git a/temporalio/common/_streams/writer_system/client/_extensions.py b/temporalio/common/_streams/writer_system/client/_extensions.py new file mode 100644 index 0000000000..ca761b3bc6 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client/_extensions.py @@ -0,0 +1,338 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +import abc +import collections.abc +import datetime +import typing + +import temporalio.nexus.system._support as _support + +# The writer has no error enum of its own, so its calls raise the Stream service's errors. +from ...system.errors import StreamError, stream_error +from ...system.models import StreamErrorDetails +from ..models import ( + AppendRecord, + CloseWriterRequest, + CloseWriterResponse, + CloseWritersRequest, + CloseWritersResponse, + FreeWriterRequest, + FreeWriterResponse, + OpenWriterRequest, + Producer, + StreamHandle, + WorkflowTaskEpoch, + WriteRecordsRequest, + WriteRecordsResponse, + WriterResult, +) +from ._nexus_client import ( + _NexusClient, + _operation_id, +) +from ._resources import Writer + + +# The runtime serves the writer in process, so its calls name no endpoint. +class StreamWriterClientExtension(abc.ABC): + """Starts the operations of ``StreamWriter`` from code outside a Workflow. + + A client class inherits this class and implements ``_create_nexus_client``. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init_subclass__(cls, **kwargs: typing.Any) -> None: + super().__init_subclass__(**kwargs) + # The MRO picks one base for a name that two bases have, with no error. + for base in cls.__mro__: + if StreamWriterClientExtension in base.__mro__: + continue + for name in ( + "open_writer", + "write_records", + "close_writer", + "free_writer", + "close_writers", + "get_writer", + ): + if name in base.__dict__: + raise TypeError( + f"{cls.__qualname__} gets {name} from {base.__qualname__} and from " + "StreamWriterClientExtension" + ) + + @abc.abstractmethod + def _create_nexus_client( + self, + *, + service: str, + endpoint: str, + ) -> _NexusClient: + """Return the Nexus client that starts the operations of ``service``.""" + ... + + def get_writer( + self, + *, + writer_id: int, + ) -> Writer: + """A handle on an existing ``Writer``, with no call.""" + return Writer( + self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ), + writer_id=writer_id, + ) + + async def _stream_writer_open_writer( + self, + request: OpenWriterRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> Writer: + try: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ) + handle = await nexus_client.start_operation( + operation="OpenWriter", + arg=request, + id=_operation_id(), + result_type=WriterResult, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + result = await handle.result() + if result.writer_id is None: + raise RuntimeError("OpenWriter answered without `writer_id`") + resource = Writer( + nexus_client, + writer_id=result.writer_id, + ) + return resource + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def open_writer( + self, + *, + handle: StreamHandle, + failed_tasks: list[int] | None = None, + producer: Producer | None = None, + max_batch_size: int | None = None, + max_retry_duration: datetime.timedelta | None = None, + previous_completed_task: WorkflowTaskEpoch | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> Writer: + """ + Args: + failed_tasks: Local Activity only. The run's failed or timed out Workflow + Tasks. + max_batch_size: A write sends at most this many Records in one batch. Zero + is 1000. + max_retry_duration: How long a write retries a Store failure. Empty is 10 s. + previous_completed_task: Local Activity only. The run's last completed + Workflow Task. + """ + request = OpenWriterRequest( + handle=handle, + failed_tasks=failed_tasks or [], + producer=producer, + max_batch_size=max_batch_size, + max_retry_duration=max_retry_duration, + previous_completed_task=previous_completed_task, + ) + return await self._stream_writer_open_writer( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_writer_write_records( + self, + request: WriteRecordsRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> WriteRecordsResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ) + handle = await nexus_client.start_operation( + operation="WriteRecords", + arg=request, + id=_operation_id(), + result_type=WriteRecordsResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def write_records( + self, + *, + records: list[AppendRecord] | None = None, + writer_id: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> WriteRecordsResponse: + request = WriteRecordsRequest( + records=records or [], + writer_id=writer_id, + ) + return await self._stream_writer_write_records( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_writer_close_writer( + self, + request: CloseWriterRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseWriterResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ) + handle = await nexus_client.start_operation( + operation="CloseWriter", + arg=request, + id=_operation_id(), + result_type=CloseWriterResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def close_writer( + self, + *, + writer_id: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseWriterResponse: + request = CloseWriterRequest( + writer_id=writer_id, + ) + return await self._stream_writer_close_writer( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_writer_free_writer( + self, + request: FreeWriterRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> FreeWriterResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ) + handle = await nexus_client.start_operation( + operation="FreeWriter", + arg=request, + id=_operation_id(), + result_type=FreeWriterResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def free_writer( + self, + *, + writer_id: int | None = None, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> FreeWriterResponse: + request = FreeWriterRequest( + writer_id=writer_id, + ) + return await self._stream_writer_free_writer( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) + + async def _stream_writer_close_writers( + self, + request: CloseWritersRequest, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseWritersResponse: + try: + nexus_client = self._create_nexus_client( + service="temporal.sdk.streams.v1.StreamWriter", + endpoint="", + ) + handle = await nexus_client.start_operation( + operation="CloseWriters", + arg=request, + id=_operation_id(), + result_type=CloseWritersResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def close_writers( + self, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseWritersResponse: + """Drop every open writer, for example at shutdown.""" + request = CloseWritersRequest() + return await self._stream_writer_close_writers( + request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout + ) diff --git a/temporalio/common/_streams/writer_system/client/_nexus_client.py b/temporalio/common/_streams/writer_system/client/_nexus_client.py new file mode 100644 index 0000000000..d1e3ecdda0 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client/_nexus_client.py @@ -0,0 +1,51 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +import collections.abc +import datetime +import typing +import uuid + +_OutputT = typing.TypeVar("_OutputT") +_OutputT_co = typing.TypeVar("_OutputT_co", covariant=True) + + +class _NexusOperationHandle(typing.Protocol[_OutputT_co]): + """A started Nexus operation.""" + + async def result(self) -> _OutputT_co: ... + + +# The client modules import the names below, so they aren't unused. +class _NexusClient(typing.Protocol): # pyright: ignore[reportUnusedClass] + """Starts the operations of one Nexus service, like ``temporalio.client.NexusClient``.""" + + @typing.overload + async def start_operation( + self, + operation: str, + arg: typing.Any, + *, + id: str, + result_type: type[_OutputT], + rpc_metadata: collections.abc.Mapping[str, str | bytes] = ..., + rpc_timeout: datetime.timedelta | None = ..., + ) -> _NexusOperationHandle[_OutputT]: ... + + @typing.overload + async def start_operation( + self, + operation: str, + arg: typing.Any, + *, + id: str, + rpc_metadata: collections.abc.Mapping[str, str | bytes] = ..., + rpc_timeout: datetime.timedelta | None = ..., + ) -> _NexusOperationHandle[typing.Any]: ... + + +def _operation_id() -> str: # pyright: ignore[reportUnusedFunction] + """Return a new random operation ID.""" + return str(uuid.uuid4()) diff --git a/temporalio/common/_streams/writer_system/client/_resources/__init__.py b/temporalio/common/_streams/writer_system/client/_resources/__init__.py new file mode 100644 index 0000000000..bc59364ba6 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client/_resources/__init__.py @@ -0,0 +1,10 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +from .writer import Writer + +__all__ = [ + "Writer", +] diff --git a/temporalio/common/_streams/writer_system/client/_resources/writer.py b/temporalio/common/_streams/writer_system/client/_resources/writer.py new file mode 100644 index 0000000000..c25163adb3 --- /dev/null +++ b/temporalio/common/_streams/writer_system/client/_resources/writer.py @@ -0,0 +1,139 @@ +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. + +from __future__ import annotations + +import collections.abc +import dataclasses +import datetime + +import temporalio.nexus.system._support as _support + +# The writer has no error enum of its own, so its calls raise the Stream service's errors. +from ....system.errors import StreamError, stream_error +from ....system.models import StreamErrorDetails +from ...models import ( + AppendRecord, + CloseWriterRequest, + CloseWriterResponse, + FreeWriterRequest, + FreeWriterResponse, + WriteRecordsRequest, + WriteRecordsResponse, +) +from .._nexus_client import ( + _NexusClient, + _operation_id, +) + + +@dataclasses.dataclass(init=False) +class Writer: + """An open writer of one Producer. The crate keeps its epoch and sequence.""" + + writer_id: int + + def __init__( + self, + nexus_client: _NexusClient, + /, + *, + writer_id: int, + ) -> None: + self._nexus_client: _NexusClient = nexus_client + self.writer_id = writer_id + + async def write_records( + self, + records: list[AppendRecord] | None = None, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> WriteRecordsResponse: + """Write Records at once. The answer has the first and the last offset.""" + request = WriteRecordsRequest( + writer_id=self.writer_id, + records=records or [], + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="WriteRecords", + arg=request, + id=_operation_id(), + result_type=WriteRecordsResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def close( + self, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> CloseWriterResponse: + """Release the handle. The last one waits for a write, then drops the writer.""" + request = CloseWriterRequest( + writer_id=self.writer_id, + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="CloseWriter", + arg=request, + id=_operation_id(), + result_type=CloseWriterResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + async def free( + self, + *, + rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, + rpc_timeout: datetime.timedelta | None = None, + ) -> FreeWriterResponse: + """Release the handle with no wait, for a handle that lang drops.""" + request = FreeWriterRequest( + writer_id=self.writer_id, + ) + try: + nexus_client = self._nexus_client + handle = await nexus_client.start_operation( + operation="FreeWriter", + arg=request, + id=_operation_id(), + result_type=FreeWriterResponse, + rpc_metadata=rpc_metadata or {}, + rpc_timeout=rpc_timeout, + ) + return await handle.result() + except StreamError: + raise + except Exception as err: + found = _support.error_details(err, StreamErrorDetails) + if found is None: + raise + details, message = found + raise stream_error(details, message) from err + + +__all__ = ["Writer"] diff --git a/temporalio/common/_streams/writer_system/models.py b/temporalio/common/_streams/writer_system/models.py index 833405207f..06e23d5055 100644 --- a/temporalio/common/_streams/writer_system/models.py +++ b/temporalio/common/_streams/writer_system/models.py @@ -1,4 +1,5 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. from __future__ import annotations @@ -39,7 +40,7 @@ class StreamProducerKind(enum.IntEnum): class OpenWriterRequest: """Open a writer for one Producer of one Stream.""" - stream: StreamRef + handle: StreamHandle failed_tasks: list[int] | None = dataclasses.field(default_factory=list) producer: Producer | None = None max_batch_size: int | None = None @@ -63,13 +64,13 @@ def from_transfer_type( value: temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest, type_hint: type[OpenWriterRequest], ) -> OpenWriterRequest: - if not value.HasField("stream"): - raise ValueError("missing required field OpenWriterRequest.stream") - stream = _StreamRefTransferTypeConverter().from_transfer_type( - value.stream, StreamRef + if not value.HasField("handle"): + raise ValueError("missing required field OpenWriterRequest.handle") + handle = _StreamHandleTransferTypeConverter().from_transfer_type( + value.handle, StreamHandle ) return OpenWriterRequest( - stream=stream, + handle=handle, failed_tasks=list(value.failed_tasks), producer=_ProducerTransferTypeConverter().from_transfer_type( value.producer, Producer @@ -93,8 +94,8 @@ def to_transfer_type( value: OpenWriterRequest, ) -> temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest: message = temporalio.common._streams.proto.v1.runtime_pb2.OpenWriterRequest() - message.stream.CopyFrom( - _StreamRefTransferTypeConverter().to_transfer_type(value.stream) + message.handle.CopyFrom( + _StreamHandleTransferTypeConverter().to_transfer_type(value.handle) ) if value.failed_tasks: message.failed_tasks.extend(value.failed_tasks) @@ -123,7 +124,7 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) -class StreamRef: +class StreamHandle: """Identifies one Stream. The id is opaque to callers.""" id: str @@ -131,23 +132,25 @@ class StreamRef: backend: str | None = None -class _StreamRefTransferTypeConverter( +class _StreamHandleTransferTypeConverter( temporalio.converter.TransferTypeConverter[ - StreamRef, temporalio.api.stream.v1.message_pb2.StreamRef + StreamHandle, temporalio.api.stream.v1.message_pb2.StreamHandle ] ): - transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamRef] | None = ( - temporalio.api.stream.v1.message_pb2.StreamRef + transfer_type: type[temporalio.api.stream.v1.message_pb2.StreamHandle] | None = ( + temporalio.api.stream.v1.message_pb2.StreamHandle ) @typing_extensions.override def from_transfer_type( self, - value: temporalio.api.stream.v1.message_pb2.StreamRef, - type_hint: type[StreamRef], - ) -> StreamRef: + value: temporalio.api.stream.v1.message_pb2.StreamHandle, + type_hint: type[StreamHandle], + ) -> StreamHandle: + if not value.id: + raise ValueError("missing required field StreamHandle.id") id = value.id - return StreamRef( + return StreamHandle( id=id, owner=_ExecutionTransferTypeConverter().from_transfer_type( value.owner, Execution @@ -160,9 +163,9 @@ def from_transfer_type( @typing_extensions.override def to_transfer_type( self, - value: StreamRef, - ) -> temporalio.api.stream.v1.message_pb2.StreamRef: - message = temporalio.api.stream.v1.message_pb2.StreamRef() + value: StreamHandle, + ) -> temporalio.api.stream.v1.message_pb2.StreamHandle: + message = temporalio.api.stream.v1.message_pb2.StreamHandle() message.id = value.id if value.owner is not None: message.owner.CopyFrom( @@ -173,8 +176,8 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible(_StreamRefTransferTypeConverter)( - StreamRef +temporalio.converter.transfer_type_convertible(_StreamHandleTransferTypeConverter)( + StreamHandle ) # pyright: ignore[reportUnusedCallResult] @@ -255,6 +258,8 @@ def from_transfer_type( value: temporalio.api.stream.v1.message_pb2.Producer, type_hint: type[Producer], ) -> Producer: + if not value.id: + raise ValueError("missing required field Producer.id") id = value.id return Producer( id=id, @@ -289,7 +294,9 @@ def to_transfer_type( return message -temporalio.converter.transfer_type_convertible(_ProducerTransferTypeConverter)(Producer) # pyright: ignore[reportUnusedCallResult] +temporalio.converter.transfer_type_convertible(_ProducerTransferTypeConverter)( + Producer, +) # pyright: ignore[reportUnusedCallResult] @dataclasses.dataclass(slots=True) diff --git a/temporalio/common/_streams/writer_system/services.py b/temporalio/common/_streams/writer_system/services.py index 9f388ca0e3..37e13481f6 100644 --- a/temporalio/common/_streams/writer_system/services.py +++ b/temporalio/common/_streams/writer_system/services.py @@ -1,4 +1,5 @@ -# Generated by nexgen v0.2.10. DO NOT EDIT! +# Written by hand in the shape of nexgen's client output, since the writer has no WIT. Keep it +# close to that output, so a later switch to generated code changes no caller. from __future__ import annotations diff --git a/temporalio/nexus/notifications/__init__.py b/temporalio/nexus/notifications/__init__.py index 7a8bb24ead..6b8a5bdbba 100644 --- a/temporalio/nexus/notifications/__init__.py +++ b/temporalio/nexus/notifications/__init__.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/nexus/notifications/models.py b/temporalio/nexus/notifications/models.py index 64ad07fbae..ef5528a6f8 100644 --- a/temporalio/nexus/notifications/models.py +++ b/temporalio/nexus/notifications/models.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 0b3d0be51f..a6b621ae01 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -1,3 +1,4 @@ +import asyncio import collections.abc import dataclasses import typing @@ -12,7 +13,9 @@ import temporalio.api.workflow.v1 import temporalio.common import temporalio.converter +import temporalio.exceptions import temporalio.nexus.system +from temporalio.converter._payload_converter import _get_transfer_type_converter class SignalWithStartWorkflowModelRequest(typing.Protocol): @@ -343,3 +346,206 @@ def next_read_request( return dataclasses.replace( request, from_offset=response.next_offset, start=None, last=None ) + + +class _ReadItem(typing.Protocol): + @property + def offset(self) -> int | None: ... + + +def resume_read_request(request: _ReadRequestT, item: _ReadItem) -> _ReadRequestT: + """The request that resumes a Stream read after ``item``.""" + # The read resumes past the item, so the first page's start and last no longer apply. + return dataclasses.replace( + request, from_offset=(item.offset or 0) + 1, start=None, last=None + ) + + +def definition_type_arguments(definition: type[typing.Any]) -> tuple[typing.Any, ...]: + """The type arguments that a definition class gives its generic base, such as ``(Turn,)``.""" + for klass in definition.__mro__: + for base in klass.__dict__.get("__orig_bases__", ()): + arguments = typing.get_args(base) + if arguments and not any( + isinstance(argument, typing.TypeVar) for argument in arguments + ): + return arguments + raise TypeError(f"{definition.__qualname__} gives its base no type argument") + + +_DetailsT = typing.TypeVar("_DetailsT") + + +def error_details( + err: BaseException, details_type: type[_DetailsT] +) -> tuple[_DetailsT, str] | None: + """The details of a failed call, and the message of the error that holds them. + + A Nexus operation failure has the call's application error as its cause. Its details hold + the model, or the proto message that the model converts from. + """ + converter = _get_transfer_type_converter(details_type) + current: BaseException | None = err + while current is not None: + if isinstance(current, temporalio.exceptions.ApplicationError): + for detail in current.details: + if isinstance(detail, details_type): + return detail, current.message + transfer_type = converter.transfer_type if converter else None + if ( + converter is not None + and transfer_type is not None + and isinstance(detail, transfer_type) + ): + return ( + typing.cast( + _DetailsT, + converter.from_transfer_type(detail, details_type), + ), + current.message, + ) + current = current.__cause__ + return None + + +def exit_close( + resource: object, exc: BaseException | None +) -> tuple[bool, BaseException | None]: + """Whether the end of an ``async with`` block closes a Stream, and with which failure.""" + del resource + if exc is None: + return True, None + # The close status comes from the failure, and only Temporal's cancel failure cancels. + if isinstance(exc, asyncio.CancelledError): + return True, temporalio.exceptions.CancelledError(str(exc) or "Cancelled") + if isinstance(exc, Exception): + return True, exc + # An interrupt isn't a failure of the work, so the Stream stays open for a retry. + return False, None + + +_T = typing.TypeVar("_T") + + +def start_now(call: collections.abc.Awaitable[_T]) -> asyncio.Task[_T]: + """Run ``call`` up to its first wait now, in a task that gives its result. + + A Workflow call adds its command before it first waits, so the command keeps its place in + the order of the caller's commands, as ``workflow.start_activity`` does. + """ + + async def run() -> _T: + return await call + + # A task runs only when the loop gets to it, so step the call here and let a task drive the + # rest. Every Python version takes this path, so a replay on another version gives the same + # commands. + loop = asyncio.get_running_loop() + coroutine = run() + try: + waiting = coroutine.send(None) + except StopIteration as finished: + return loop.create_task(_finished(typing.cast(_T, finished.value))) + except BaseException as error: + return loop.create_task(_raised(error)) + return loop.create_task(_resume(coroutine, waiting)) + + +async def _finished(value: _T) -> _T: + return value + + +async def _raised(error: BaseException) -> typing.NoReturn: + raise error + + +async def _resume( + coroutine: collections.abc.Coroutine[object, None, _T], waiting: object +) -> _T: + # Drive a started coroutine as a task does. A cancel goes to the future that it waits on. + while True: + cancelled: asyncio.CancelledError | None = None + try: + if isinstance(waiting, asyncio.Future): + future = typing.cast("asyncio.Future[object]", waiting) + done = asyncio.get_running_loop().create_future() + + def wake(_: object, done: asyncio.Future[None] = done) -> None: + if not done.done(): + done.set_result(None) + + future.add_done_callback(wake) + await done + else: + await asyncio.sleep(0) + except asyncio.CancelledError as error: + if isinstance(waiting, asyncio.Future): + typing.cast("asyncio.Future[object]", waiting).cancel() + cancelled = error + try: + waiting = ( + coroutine.send(None) + if cancelled is None + else coroutine.throw(cancelled) + ) + except StopIteration as finished: + return typing.cast(_T, finished.value) + + +# The Stream sources. Each one checks where it runs and gives the value of that context, or +# raises where it has none to give. + + +def _no_workflow_streams(source: str) -> RuntimeError: + # This SDK doesn't run Workflow Stream calls yet, so a generated Workflow call fails here. + return RuntimeError(f"Workflow code can't call the Stream service yet ({source})") + + +def workflow_stream_id(definition: type[typing.Any]) -> str: + """The id of the Workflow's own Stream of ``definition``.""" + del definition + raise _no_workflow_streams("workflow_stream_id") + + +def workflow_handle_owner() -> typing.Any: + """The owner of the Workflow's own Stream in its handle.""" + raise _no_workflow_streams("workflow_handle_owner") + + +def workflow_owner() -> typing.Any: + """The Producer of the Workflow run that owns a Stream that it creates.""" + raise _no_workflow_streams("workflow_owner") + + +def workflow_producer() -> typing.Any: + """The Producer of the Workflow run's writes.""" + raise _no_workflow_streams("workflow_producer") + + +def withdraw_run() -> str | None: + """The previous run whose Records a retry run's first write withdraws.""" + from temporalio.workflow import in_workflow + + if in_workflow(): + raise _no_workflow_streams("withdraw_run") + # Code outside a Workflow isn't a retry run. + return None + + +def client_producer() -> typing.Any: + """The Producer of a write outside a Workflow.""" + # Each client has its own Producer, and a generated call doesn't say which client makes it. + # So the SDK's handles send their Producer, and a call without one fails. + raise RuntimeError("A Stream write outside a Workflow needs its Producer") + + +def local_activity_failed_tasks() -> list[int]: + """The run's failed Workflow Tasks, whose Records a Local Activity's write withdraws.""" + # Only a Local Activity has them, and this SDK's Local Activities don't write here yet. + return [] + + +def local_activity_previous_completed_task() -> typing.Any: + """The run's last completed Workflow Task, for a Local Activity's write.""" + # Only a Local Activity has it, and this SDK's Local Activities don't write here yet. + return None diff --git a/temporalio/nexus/system/workflow_service/__init__.py b/temporalio/nexus/system/workflow_service/__init__.py index f6a6f0c682..afb94f6059 100644 --- a/temporalio/nexus/system/workflow_service/__init__.py +++ b/temporalio/nexus/system/workflow_service/__init__.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -40,7 +40,9 @@ def __init__( ) -__nexus_operation_registry__ = { +__nexus_operation_registry__: dict[ + tuple[str, str], _NexusOperationInfo[typing.Any, typing.Any] +] = { ( "temporal.api.workflowservice.v1.WorkflowService", "SignalWithStartWorkflowExecution", diff --git a/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py index 5c62e00927..c1e992c8c7 100644 --- a/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py +++ b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -40,9 +40,7 @@ async def _start_system_nexus_operation( "temporalio.workflow.NexusOperationHandle[_OutputT]", await interceptor.start_signal_with_start_workflow(typed_input.input), ) - raise ValueError( - f"unsupported System Nexus operation: {input.service}/{input.operation_name}" - ) + return await interceptor.start_system_nexus_operation(input) class _SystemNexusWorkflowOutboundInterceptorBase(abc.ABC): @@ -51,6 +49,14 @@ def _next_system_nexus_interceptor( self, ) -> _SystemNexusWorkflowOutboundInterceptorBase: ... + @abc.abstractmethod + async def start_system_nexus_operation( + self, + input: StartNexusOperationInput[_InputT, _OutputT], + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + """Intercept a System Nexus operation through the generic hook.""" + ... + async def start_signal_with_start_workflow( self, request: models.SignalWithStartWorkflowRequest ) -> temporalio.workflow.NexusOperationHandle[ diff --git a/temporalio/nexus/system/workflow_service/models.py b/temporalio/nexus/system/workflow_service/models.py index dc7b5ce2e7..8fa35902fc 100644 --- a/temporalio/nexus/system/workflow_service/models.py +++ b/temporalio/nexus/system/workflow_service/models.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations @@ -18,7 +18,8 @@ @dataclasses.dataclass(slots=True, kw_only=True) class SignalWithStartWorkflowRequest: - """ + """Request fields for signaling a workflow, starting it first if needed. + .. warning:: This API is experimental and subject to change. """ @@ -223,7 +224,8 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) class UserMetadata: - """ + """Static metadata for a workflow execution. + .. warning:: This API is experimental and subject to change. """ @@ -276,7 +278,8 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) class SignalWithStartWorkflowResponse: - """ + """Result of signaling a workflow and starting it if needed. + .. warning:: This API is experimental and subject to change. """ diff --git a/temporalio/nexus/system/workflow_service/operations/__init__.py b/temporalio/nexus/system/workflow_service/operations/__init__.py index ff561b55c0..3bd09999cf 100644 --- a/temporalio/nexus/system/workflow_service/operations/__init__.py +++ b/temporalio/nexus/system/workflow_service/operations/__init__.py @@ -1,3 +1,3 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/nexus/system/workflow_service/operations/signal_with_start_workflow.py b/temporalio/nexus/system/workflow_service/operations/signal_with_start_workflow.py index 28b4735436..057c6078ed 100644 --- a/temporalio/nexus/system/workflow_service/operations/signal_with_start_workflow.py +++ b/temporalio/nexus/system/workflow_service/operations/signal_with_start_workflow.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/temporalio/nexus/system/workflow_service/services.py b/temporalio/nexus/system/workflow_service/services.py index c95fff35ad..61cf577bbe 100644 --- a/temporalio/nexus/system/workflow_service/services.py +++ b/temporalio/nexus/system/workflow_service/services.py @@ -1,4 +1,4 @@ -# Generated by nexgen v0.2.8. DO NOT EDIT! +# Generated by nexgen v0.2.11. DO NOT EDIT! from __future__ import annotations diff --git a/tests/contrib/streams/test_bindings.py b/tests/contrib/streams/test_bindings.py index 923800d361..017d285e16 100644 --- a/tests/contrib/streams/test_bindings.py +++ b/tests/contrib/streams/test_bindings.py @@ -1,35 +1,52 @@ -from collections.abc import Awaitable +import asyncio +from collections.abc import Mapping from datetime import timedelta -from typing import Any +from typing import Any, get_origin +import pytest from typing_extensions import override +import temporalio.nexus.system._support as _support +from temporalio.common import StreamDefinition from temporalio.common._streams.system import client, models -def _data(offset: int) -> models.StreamRecord: +class Tokens(StreamDefinition[str]): + pass + + +def _data(offset: int) -> models.StreamRecord[str]: return models.StreamRecord( - item=models.StreamItemData(models.StreamData(body=None)), offset=offset + item=models.StreamItemData(models.StreamData(body="token")), offset=offset ) class FakeNexusClient: - def __init__(self, *answers: models.ReadRecordsResponse) -> None: + def __init__(self, *answers: models.ReadRecordsResponse[str]) -> None: self.requests: list[models.ReadRecordsRequest] = [] self.answers = list(answers) async def start_operation( - self, operation: str, input: Any, *, output_type: type[Any] | None = None - ) -> Awaitable[Any]: - assert operation == "ReadRecords" - assert output_type is models.ReadRecordsResponse - self.requests.append(input) + self, + operation: str, + arg: Any, + *, + id: str, + result_type: type[Any] | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Any: + del rpc_metadata, rpc_timeout + assert operation == "ReadRecords" and id + assert get_origin(result_type) is models.ReadRecordsResponse + self.requests.append(arg) answer = self.answers.pop(0) - async def done() -> Any: - return answer + class Done: + async def result(self) -> Any: + return answer - return done() + return Done() class FakeService(client.StreamServiceClientExtension): @@ -37,9 +54,7 @@ def __init__(self, nexus_client: FakeNexusClient) -> None: self.nexus_client = nexus_client @override - def _create_nexus_client( - self, *, service: str, endpoint: str | None - ) -> FakeNexusClient: + def _create_nexus_client(self, *, service: str, endpoint: str) -> FakeNexusClient: assert service == "temporal.api.streamservice.v1.StreamService" assert endpoint == "__temporal_system" return self.nexus_client @@ -54,17 +69,12 @@ async def test_a_poll_carries_the_next_offset_until_the_close(): ), models.ReadRecordsResponse(next_offset=10, done=True), ) - stream = client.Stream( - _client=FakeService(nexus_client), - stream=models.StreamRef(id="wf:run"), - created=None, - status=None, - ) - pages = [ - page - async for page in stream.iter_read(from_offset=1, wait=timedelta(seconds=5)) + stream = FakeService(nexus_client).get_stream(Tokens, id="wf:run") + items = [ + item + async for item in stream.iter_read(from_offset=1, wait=timedelta(seconds=5)) ] - assert [page.next_offset for page in pages] == [8, 10, 10] + assert [item.offset for item in items] == [9] assert [request.from_offset for request in nexus_client.requests] == [1, 8, 10] assert all(request.wait for request in nexus_client.requests) @@ -77,15 +87,51 @@ async def test_a_read_with_no_wait_ends_at_an_empty_page_that_didnt_move(): ), models.ReadRecordsResponse(next_offset=5), ) - pages = [ - page - async for page in FakeService(nexus_client).iter_read_records( - stream=models.StreamRef(id="wf:run"), + items = [ + item + async for item in FakeService(nexus_client).iter_read_records( + str, + handle=models.StreamHandle(id="wf:run"), start=models.StreamReadStart.StreamReadStartBeginning, ) ] - assert len(pages) == 2 + assert [item.offset for item in items] == [4] first, second = nexus_client.requests assert first.start == models.StreamReadStart.StreamReadStartBeginning # The second page starts where the first ended, so the start no longer applies. assert (second.from_offset, second.start) == (5, None) + + +async def test_a_poll_resumes_after_the_last_item_it_gave(): + nexus_client = FakeNexusClient( + models.ReadRecordsResponse(records=[_data(4), _data(7)], next_offset=8), + ) + items = FakeService(nexus_client).iter_read_records( + str, handle=models.StreamHandle(id="wf:run"), last=2 + ) + first = await anext(items) + assert first.offset == 4 + # A read that resumes here starts past the item, and doesn't read the newest Records again. + assert items.next_request is not None + assert (items.next_request.from_offset, items.next_request.last) == (5, None) + + +async def test_start_now_runs_a_call_to_its_first_wait_at_once(): + steps: list[str] = [] + gate: asyncio.Future[str] = asyncio.get_running_loop().create_future() + + async def call() -> str: + steps.append("started") + return await gate + + task = _support.start_now(call()) + # The call's first step ran before start_now returned, as a command would. + assert steps == ["started"] + gate.set_result("done") + assert await task == "done" + + async def fails() -> None: + raise ValueError("refused") + + with pytest.raises(ValueError, match="refused"): + await _support.start_now(fails())