Added a client for the server-side stream service.

scripts/gen_stream_protos.py

@@ -0,0 +1,193 @@
+"""Regenerate the vendored server-side stream protos.
+
+The stream service is still defined in the server rather than in the public
+API repo, so its protos are copied in and compiled here. Run this whenever the
+server-side definitions change:
+
+    uv run python scripts/gen_stream_protos.py /path/to/temporal
+
+Run it with an interpreter whose ``grpcio-tools`` matches the oldest protobuf
+runtime this package has to load on. Generated code refuses to load on a
+runtime older than the one it was built against, and these modules end up in
+applications that pin protobuf themselves. The stubs come from
+``mypy-protobuf`` rather than protoc's own ``--pyi_out``, which that
+generation of protoc does not have.
+
+Two rewrites make the copies compile against this SDK. The server's own
+routing and API-category options are dropped, because they mean nothing to a
+client and their proto files are not vendored. Paths move under
+``temporalio/`` so the generated modules import each other and the public API
+messages by the names this package actually publishes.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+BASE = Path(__file__).parent.parent
+OUT = BASE / "temporalio" / "api" / "streamservice" / "v1"
+# The stream protos reference the public API messages. They normally come from
+# the sdk-core submodule; TEMPORAL_API_PROTOS points at a `temporalio/api`
+# checkout instead when the submodule is not initialized.
+UPSTREAM_API = (
+    BASE
+    / "temporalio"
+    / "bridge"
+    / "sdk-core"
+    / "crates"
+    / "protos"
+    / "protos"
+    / "api_upstream"
+)
+
+
+def api_proto_path() -> Path:
+    override = os.environ.get("TEMPORAL_API_PROTOS")
+    if override:
+        return Path(override).resolve()
+    if (
+        UPSTREAM_API / "temporal" / "api" / "common" / "v1" / "message.proto"
+    ).is_file():
+        return UPSTREAM_API
+    raise SystemExit(
+        "no api protos found: initialize the sdk-core submodule or set "
+        "TEMPORAL_API_PROTOS to a temporalio/api checkout"
+    )
+
+
+SERVER_PROTO_DIR = Path("chasm/lib/stream/proto/v1")
+STAGE_PROTO_DIR = Path("temporalio/api/stream/v1")
+FILES = [
+    "message.proto",
+    "stream_state.proto",
+    "request_response.proto",
+    "service.proto",
+]
+
+# Server-only options. They carry shard routing and rate-limit category, which
+# a client neither reads nor can resolve, since the files defining them are not
+# vendored.
+DROP_IMPORT = re.compile(r'^import "temporal/server/.*?";\n', re.M)
+DROP_OPTION = re.compile(r"^\s*option \(temporal\.server\.api\..*?\n", re.M)
+
+# protoc names a python module after the proto path. These put the generated
+# modules where this package publishes them.
+FIX_PY = [
+    (
+        re.compile(r"from temporalio\.api\.stream\.v1 import"),
+        "from temporalio.api.streamservice.v1 import",
+    ),
+    (re.compile(r"temporalio\.api\.stream\.v1\."), "temporalio.api.streamservice.v1."),
+    (re.compile(r"from temporal\.api\."), "from temporalio.api."),
+    (re.compile(r"import temporal\.api\."), "import temporalio.api."),
+    # mypy-protobuf also writes the public API types fully qualified in the
+    # stubs, `temporal.api.common.v1.message_pb2.Payload`, which the import
+    # rewrites above do not reach. The `_pb2` suffix keeps this off the proto
+    # package names inside the serialized descriptors, which stay `temporal.api`.
+    # After the streamservice rewrite on purpose, so the public
+    # `temporal.api.stream.v1` package is not sent to the vendored one.
+    (
+        re.compile(r"\btemporal\.api\.(\w+)\.v1\.(\w+_pb2)\b"),
+        r"temporalio.api.\1.v1.\2",
+    ),
+]
+
+
+def stage(server: Path, into: Path) -> list[str]:
+    dest = into / STAGE_PROTO_DIR
+    dest.mkdir(parents=True)
+    for name in FILES:
+        text = (server / SERVER_PROTO_DIR / name).read_text()
+        text = text.replace(str(SERVER_PROTO_DIR) + "/", str(STAGE_PROTO_DIR) + "/")
+        text = DROP_IMPORT.sub("", text)
+        text = DROP_OPTION.sub("", text)
+        (dest / name).write_text(text)
+    return [str(STAGE_PROTO_DIR / name) for name in FILES]
+
+
+def main() -> None:
+    if len(sys.argv) != 2:
+        sys.exit("usage: gen_stream_protos.py /path/to/temporal-server-checkout")
+    server = Path(sys.argv[1]).resolve()
+    if not (server / SERVER_PROTO_DIR).is_dir():
+        sys.exit(f"no stream protos under {server / SERVER_PROTO_DIR}")
+
+    with tempfile.TemporaryDirectory() as tmp:
+        work = Path(tmp)
+        targets = stage(server, work)
+        out = work / "out"
+        out.mkdir()
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "grpc_tools.protoc",
+                f"--proto_path={work}",
+                f"--proto_path={api_proto_path()}",
+                f"--python_out={out}",
+                f"--grpc_python_out={out}",
+                f"--mypy_out={out}",
+                f"--mypy_grpc_out={out}",
+                *targets,
+            ],
+            check=True,
+            cwd=work,
+        )
+
+        generated = out / STAGE_PROTO_DIR
+        for path in sorted(generated.iterdir()):
+            if path.name == "__init__.py":
+                continue
+            text = path.read_text()
+            for pattern, repl in FIX_PY:
+                text = pattern.sub(repl, text)
+            (OUT / path.name).write_text(text)
+            print(f"wrote {OUT / path.name}")
+
+    shutil.rmtree(OUT / "__pycache__", ignore_errors=True)
+    write_init()
+
+
+MODULES = ["message_pb2", "stream_state_pb2", "request_response_pb2"]
+
+INIT_HEADER = """\
+# Generated by scripts/gen_stream_protos.py. Do not edit.
+#
+# Vendored rather than taken from the api submodule, because the stream service
+# is still defined in the server. See temporalio/client_stream.py.
+"""
+
+
+def write_init() -> None:
+    """Re-export every generated message and enum, so the package surface
+    cannot drift from the protos it was built from."""
+    import importlib
+
+    sys.path.insert(0, str(BASE))
+    blocks, exported = [], []
+    for module in MODULES:
+        mod = importlib.import_module(f"temporalio.api.streamservice.v1.{module}")
+        names = list(mod.DESCRIPTOR.message_types_by_name)
+        for enum in mod.DESCRIPTOR.enum_types_by_name.values():
+            names.append(enum.name)
+            names.extend(value.name for value in enum.values)
+        names.sort()
+        exported.extend(names)
+        joined = "".join(f"    {name},\n" for name in names)
+        blocks.append(f"from .{module} import (\n{joined})\n")
+
+    listed = "".join(f'    "{name}",\n' for name in sorted(exported))
+    (OUT / "__init__.py").write_text(
+        INIT_HEADER + "\n" + "\n".join(blocks) + f"\n__all__ = [\n{listed}]\n"
+    )
+    print(f"wrote {OUT / '__init__.py'}")
+
+
+if __name__ == "__main__":
+    main()

temporalio/api/streamservice/v1/__init__.py

@@ -0,0 +1,158 @@
+# Generated by scripts/gen_stream_protos.py. Do not edit.
+#
+# Vendored rather than taken from the api submodule, because the stream service
+# is still defined in the server. See temporalio/client_stream.py.
+
+from .message_pb2 import (
+    StreamRecord,
+    StreamRecordBatch,
+)
+from .request_response_pb2 import (
+    STREAM_OWNER_KIND_ACTIVITY,
+    STREAM_OWNER_KIND_UNSPECIFIED,
+    STREAM_OWNER_KIND_WORKFLOW,
+    STREAM_OWNER_KIND_WORKFLOW_ACTIVITY,
+    AddMessagesInput,
+    AddMessagesOutput,
+    AddMessagesRequest,
+    AddMessagesResponse,
+    AddWorkflowMessagesInput,
+    AddWorkflowMessagesRequest,
+    AddWorkflowMessagesResponse,
+    AdvanceConsumerHeadInput,
+    AdvanceConsumerHeadOutput,
+    AdvanceConsumerHeadRequest,
+    AdvanceConsumerHeadResponse,
+    CloseStreamInput,
+    CloseStreamOutput,
+    CloseStreamRequest,
+    CloseStreamResponse,
+    CreateStreamInput,
+    CreateStreamOutput,
+    CreateStreamRequest,
+    CreateStreamResponse,
+    DeleteStreamInput,
+    DeleteStreamOutput,
+    DeleteStreamRequest,
+    DeleteStreamResponse,
+    DescribeStreamInput,
+    DescribeStreamOutput,
+    DescribeStreamRequest,
+    DescribeStreamResponse,
+    DescribeWorkflowStreamInput,
+    DescribeWorkflowStreamRequest,
+    DescribeWorkflowStreamResponse,
+    FinishWritingInput,
+    FinishWritingOutput,
+    FinishWritingRequest,
+    FinishWritingResponse,
+    ListStreamsInput,
+    ListStreamsOutput,
+    ListStreamsRequest,
+    ListStreamsResponse,
+    PollMessagesInput,
+    PollMessagesOutput,
+    PollMessagesRequest,
+    PollMessagesResponse,
+    PollWorkflowMessagesInput,
+    PollWorkflowMessagesRequest,
+    PollWorkflowMessagesResponse,
+    RegisterStreamConsumerInput,
+    RegisterStreamConsumerOutput,
+    RegisterStreamConsumerRequest,
+    RegisterStreamConsumerResponse,
+    StreamListEntry,
+    StreamOwner,
+    StreamOwnerKind,
+    SubscribeWorkflowInput,
+    SubscribeWorkflowOutput,
+    SubscribeWorkflowRequest,
+    SubscribeWorkflowResponse,
+    TruncateStreamInput,
+    TruncateStreamOutput,
+    TruncateStreamRequest,
+    TruncateStreamResponse,
+)
+from .stream_state_pb2 import (
+    ConsumerCursor,
+    ProducerCursor,
+    StreamBudget,
+    StreamLifecycle,
+    StreamState,
+    WorkflowStreamCursor,
+)
+
+__all__ = [
+    "AddMessagesInput",
+    "AddMessagesOutput",
+    "AddMessagesRequest",
+    "AddMessagesResponse",
+    "AddWorkflowMessagesInput",
+    "AddWorkflowMessagesRequest",
+    "AddWorkflowMessagesResponse",
+    "AdvanceConsumerHeadInput",
+    "AdvanceConsumerHeadOutput",
+    "AdvanceConsumerHeadRequest",
+    "AdvanceConsumerHeadResponse",
+    "CloseStreamInput",
+    "CloseStreamOutput",
+    "CloseStreamRequest",
+    "CloseStreamResponse",
+    "ConsumerCursor",
+    "CreateStreamInput",
+    "CreateStreamOutput",
+    "CreateStreamRequest",
+    "CreateStreamResponse",
+    "DeleteStreamInput",
+    "DeleteStreamOutput",
+    "DeleteStreamRequest",
+    "DeleteStreamResponse",
+    "DescribeStreamInput",
+    "DescribeStreamOutput",
+    "DescribeStreamRequest",
+    "DescribeStreamResponse",
+    "DescribeWorkflowStreamInput",
+    "DescribeWorkflowStreamRequest",
+    "DescribeWorkflowStreamResponse",
+    "FinishWritingInput",
+    "FinishWritingOutput",
+    "FinishWritingRequest",
+    "FinishWritingResponse",
+    "ListStreamsInput",
+    "ListStreamsOutput",
+    "ListStreamsRequest",
+    "ListStreamsResponse",
+    "PollMessagesInput",
+    "PollMessagesOutput",
+    "PollMessagesRequest",
+    "PollMessagesResponse",
+    "PollWorkflowMessagesInput",
+    "PollWorkflowMessagesRequest",
+    "PollWorkflowMessagesResponse",
+    "ProducerCursor",
+    "RegisterStreamConsumerInput",
+    "RegisterStreamConsumerOutput",
+    "RegisterStreamConsumerRequest",
+    "RegisterStreamConsumerResponse",
+    "STREAM_OWNER_KIND_ACTIVITY",
+    "STREAM_OWNER_KIND_UNSPECIFIED",
+    "STREAM_OWNER_KIND_WORKFLOW",
+    "STREAM_OWNER_KIND_WORKFLOW_ACTIVITY",
+    "StreamBudget",
+    "StreamLifecycle",
+    "StreamListEntry",
+    "StreamOwner",
+    "StreamOwnerKind",
+    "StreamRecord",
+    "StreamRecordBatch",
+    "StreamState",
+    "SubscribeWorkflowInput",
+    "SubscribeWorkflowOutput",
+    "SubscribeWorkflowRequest",
+    "SubscribeWorkflowResponse",
+    "TruncateStreamInput",
+    "TruncateStreamOutput",
+    "TruncateStreamRequest",
+    "TruncateStreamResponse",
+    "WorkflowStreamCursor",
+]