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)