From 247c20fbeb6e815b16d3af1c1c66c880c8db3c4d Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Thu, 5 Nov 2026 01:35:11 -0700 Subject: [PATCH] Added Stream refs, Records and errors in `temporalio.streams`. --- temporalio/streams/__init__.py | 87 ++++++++++- temporalio/streams/_errors.py | 94 +++++++++++ temporalio/streams/_ids.py | 44 ++++++ temporalio/streams/_items.py | 217 ++++++++++++++++++++++++++ temporalio/streams/_ref.py | 188 ++++++++++++++++++++++ tests/contrib/streams/_support.py | 41 +++++ tests/contrib/streams/test_refs.py | 211 +++++++++++++++++++++++++ tests/contrib/streams/test_vectors.py | 124 +++++++++++++++ 8 files changed, 1005 insertions(+), 1 deletion(-) create mode 100644 temporalio/streams/_errors.py create mode 100644 temporalio/streams/_ids.py create mode 100644 temporalio/streams/_items.py create mode 100644 temporalio/streams/_ref.py create mode 100644 tests/contrib/streams/_support.py create mode 100644 tests/contrib/streams/test_refs.py create mode 100644 tests/contrib/streams/test_vectors.py diff --git a/temporalio/streams/__init__.py b/temporalio/streams/__init__.py index 23e86ba97..e817a1bf8 100644 --- a/temporalio/streams/__init__.py +++ b/temporalio/streams/__init__.py @@ -1 +1,86 @@ -"""The Streams API, which Workflows, Activities and clients share.""" +"""The Streams API, which Workflows, Activities and clients share. + +.. warning:: + The Streams API is experimental. +""" + +from temporalio.streams._errors import ( + StreamAlreadyExistsError, + StreamClosedError, + StreamConfigError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetTooHighError, + StreamSeqDivergentError, + StreamStorageUnavailableError, + StreamTooLargeError, + StreamUnsupportedError, +) +from temporalio.streams._items import END, End, ItemKind, Record, Superseded +from temporalio.streams._ref import StreamOptions, StreamRef + +# nexgen generates these wire models. Refs, Records, Superseded items and errors carry them. +from temporalio.streams.system.models import ( + ActivityEpoch, + ClientEpoch, + Execution, + ExecutionType, + LocalActivityEpoch, + Producer, + ProducerEpoch, + ProducerEpochKind, + ProducerEpochKindActivity, + ProducerEpochKindClient, + ProducerEpochKindLocalActivity, + ProducerEpochKindStandaloneActivity, + ProducerEpochKindWorkflow, + StreamErrorReason, + StreamSuperseded, + StreamWriter, + Timestamp, + WorkflowTaskEpoch, +) + +__all__ = [ + "END", + "ActivityEpoch", + "ClientEpoch", + "End", + "Execution", + "ExecutionType", + "ItemKind", + "LocalActivityEpoch", + "Producer", + "ProducerEpoch", + "ProducerEpochKind", + "ProducerEpochKindActivity", + "ProducerEpochKindClient", + "ProducerEpochKindLocalActivity", + "ProducerEpochKindStandaloneActivity", + "ProducerEpochKindWorkflow", + "Record", + "StreamAlreadyExistsError", + "StreamClosedError", + "StreamConfigError", + "StreamError", + "StreamErrorReason", + "StreamExpiredError", + "StreamFencedError", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetTooHighError", + "StreamOptions", + "StreamRef", + "StreamSeqDivergentError", + "StreamStorageUnavailableError", + "StreamSuperseded", + "StreamTooLargeError", + "StreamUnsupportedError", + "StreamWriter", + "Superseded", + "Timestamp", + "WorkflowTaskEpoch", +] diff --git a/temporalio/streams/_errors.py b/temporalio/streams/_errors.py new file mode 100644 index 000000000..ecebc63f5 --- /dev/null +++ b/temporalio/streams/_errors.py @@ -0,0 +1,94 @@ +"""Errors of the Streams API.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import temporalio.api.stream.v1 +import temporalio.nexus.system._support as _support +from temporalio.streams.system import models + +# The generated classes are the Stream errors, one for each reason that has its own. A reason +# with no class of its own, such as STORAGE_CORRUPT, raises the base class. +from temporalio.streams.system.errors import ( + StreamAlreadyExistsError, + StreamClosedError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetTooHighError, + StreamSeqDivergentError, + StreamStorageUnavailableError, + StreamTooLargeError, + StreamUnsupportedError, + stream_error, +) + +__all__ = [ + "StreamAlreadyExistsError", + "StreamClosedError", + "StreamConfigError", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetTooHighError", + "StreamSeqDivergentError", + "StreamStorageUnavailableError", + "StreamTooLargeError", + "StreamUnsupportedError", +] + + +class StreamConfigError(RuntimeError): + """The Stream backend config is wrong, such as when no backend has a given name. + + It isn't a :py:class:`StreamError`, so ``except StreamError`` doesn't catch it. In Workflow + code it fails the Workflow Task. An Activity that raises it retries. + + .. warning:: + This API is experimental. + """ + + +def no_default_backend( + owner: str, names: Sequence[str], hint: str +) -> StreamConfigError: + """The error of a call with no backend name when no backend has ``default=True``.""" + return StreamConfigError( + f"No Stream backend of the {owner} ({', '.join(names)}) has default=True, so the " + "default is native, which this SDK does not have. Mark a backend with " + f"default=True, or give its name in {hint}." + ) + + +def from_reason(reason: int, message: str) -> StreamError: + """The Stream error of a ``StreamErrorReason`` number.""" + try: + known = models.StreamErrorReason(reason) + except ValueError: + # A reason that this SDK doesn't know gives the base class. + known = models.StreamErrorReason.StreamErrorReasonUnspecified + return stream_error(models.StreamErrorDetails(reason=known or None), message) + + +def from_call_error(err: Exception) -> StreamError: + """Turn a failed call of the crate's runtime into the Stream error for its reason. + + ``err.args`` are the gRPC code, the message and the encoded ``StreamErrorDetails``. + """ + _, message, details = err.args + proto = temporalio.api.stream.v1.StreamErrorDetails.FromString(details) + return from_reason(proto.reason, message) + + +def from_activity_error(err: Exception) -> Exception: + """The Stream error that a write Local Activity raised, else ``err``.""" + found = _support.error_details(err, models.StreamErrorDetails) + if found is None: + return err + details, message = found + return stream_error(details, message) diff --git a/temporalio/streams/_ids.py b/temporalio/streams/_ids.py new file mode 100644 index 000000000..1924e9b4c --- /dev/null +++ b/temporalio/streams/_ids.py @@ -0,0 +1,44 @@ +"""Producer ids. They must match the temporalio-streams crate byte for byte.""" + +from __future__ import annotations + +import urllib.parse + + +def encode_part(text: str) -> str: + """Percent-encode every byte but the unreserved ones.""" + return urllib.parse.quote(text, safe="") + + +def _join(*parts: str) -> str: + return ":".join(encode_part(part) for part in parts) + + +def workflow_producer_id(run_id: str) -> str: + """Producer id of a Workflow run.""" + return f"wf:{_join(run_id)}" + + +def local_activity_producer_id(run_id: str, call_key: str) -> str: + """Producer id of a Local Activity call.""" + return f"la:{_join(run_id, call_key)}" + + +def activity_producer_id(run_id: str, activity_id: str, schedule_key: str) -> str: + """Producer id of an Activity that a Workflow scheduled.""" + return f"act:{_join(run_id, activity_id, schedule_key)}" + + +def standalone_activity_producer_id(activity_id: str, activity_run_id: str) -> str: + """Producer id of a Standalone Activity.""" + return f"sa:{_join(activity_id, activity_run_id)}" + + +def client_producer_id(identity: str, process_key: str) -> str: + """Producer id of a client process.""" + return f"client:{_join(identity, process_key)}" + + +def named_client_producer_id(name: str) -> str: + """Producer id of a client with a name of its own, which it keeps across restarts.""" + return f"client:{_join(name)}" diff --git a/temporalio/streams/_items.py b/temporalio/streams/_items.py new file mode 100644 index 000000000..eddbc4553 --- /dev/null +++ b/temporalio/streams/_items.py @@ -0,0 +1,217 @@ +"""The items that a Stream read gives.""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import Any, Generic, Literal, TypeVar + +from temporalio.streams.system import models + +T = TypeVar("T") + + +class End(enum.Enum): + """An ``offset`` past the newest Record, to read only new Records. Use ``END``.""" + + END = "END" + + +END = End.END + + +class ItemKind(enum.Enum): + """The kind of a read item. Code matches on an item's ``kind``, which narrows its type. + + .. warning:: + This API is experimental. + """ + + RECORD = "record" + """A :py:class:`Record`.""" + + SUPERSEDED = "superseded" + """A :py:class:`Superseded` item.""" + + +@dataclass(frozen=True) +class Record(Generic[T]): + """A data Record of a Stream. + + .. warning:: + This API is experimental. + """ + + kind: Literal[ItemKind.RECORD] = field(default=ItemKind.RECORD, init=False) + """Always :py:attr:`ItemKind.RECORD`.""" + + data: T + """Value of the Record, decoded with the Stream's payload type. A list is one value.""" + + offset: int + """Position of the Record in the Stream. Offsets are ordered and unique, with gaps.""" + + producer: models.Producer | None + """Producer that wrote the Record.""" + + +@dataclass(frozen=True) +class Superseded: + """An item that withdraws the Records of a failed attempt. + + Every reader gets it at the same point in the Stream. A live view removes the Records in + :py:attr:`withdrawn_offsets`. A view that keeps Records from an earlier read, such as from + before a reconnect, checks them with :py:meth:`withdraws`. + + The next write of a run to a Stream withdraws the Records that the run's Local Activities + wrote there in Workflow Tasks that never completed. A Stream that the run never writes + again keeps them. + + .. warning:: + This API is experimental. + """ + + kind: Literal[ItemKind.SUPERSEDED] = field(default=ItemKind.SUPERSEDED, init=False) + """Always :py:attr:`ItemKind.SUPERSEDED`.""" + + offset: int + """Position of the item in the Stream.""" + + facts: models.StreamSuperseded + """The withdrawal facts from the Store. Its ``writer`` gives the Records that it withdraws.""" + + withdrawn_offsets: tuple[int, ...] = () + """The offsets of the Records that this read gave and this item withdraws, in order. + + A read gives each offset once. It doesn't include Records from before the read's start. + """ + + def withdraws(self, record: Record[Any] | models.Producer) -> bool: + """Whether this item withdraws an earlier Record, or the Records of a Producer. + + It answers for any earlier Record, such as one that an app showed before a resume, + from the Producer and epoch that each Record carries. So an app can keep only the + Producer of each Record that it shows. + """ + producer = record.producer if isinstance(record, Record) else record + return withdraws(self.facts, producer) + + +# The epoch rules match the crate's epoch keys. A field compares in the order of the key, and an +# empty or negative field compares as zero. +def _n(value: int | None) -> int: + return max(value or 0, 0) + + +def _time(value: models.Timestamp | None) -> tuple[int, int]: + if value is None: + return (0, 0) + return (_n(value.seconds), _n(value.nanos)) + + +_TaskKey = tuple[int, int, tuple[int, int]] + + +def _task_key(task: models.WorkflowTaskEpoch | None) -> _TaskKey: + if task is None: + return (0, 0, (0, 0)) + return (_n(task.started_event_id), _n(task.attempt), _time(task.started_time)) + + +def _task_of(epoch: models.ProducerEpoch) -> models.WorkflowTaskEpoch | None: + # A Workflow's own task, or the task that dispatched a Local Activity. The other kinds have + # none. + match epoch.kind: + case models.ProducerEpochKindWorkflow(value=task): + return task + case models.ProducerEpochKindLocalActivity(value=local): + return local.task + case _: + return None + + +def _epoch_key(epoch: models.ProducerEpoch) -> tuple[Any, ...]: + scheduled: models.Timestamp | None = None + attempt: int | None = None + match epoch.kind: + case models.ProducerEpochKindLocalActivity(value=local): + attempt = local.attempt + case ( + models.ProducerEpochKindActivity(value=activity) + | models.ProducerEpochKindStandaloneActivity(value=activity) + ): + scheduled, attempt = activity.attempt_scheduled_time, activity.attempt + case models.ProducerEpochKindClient(value=client): + attempt = client.attempt + case _: + pass + return (*_task_key(_task_of(epoch)), _time(scheduled), _n(attempt)) + + +def _kind(epoch: models.ProducerEpoch | None) -> type[Any] | None: + return type(epoch.kind) if epoch is not None and epoch.kind is not None else None + + +def is_client(producer: models.Producer) -> bool: + """Whether the Producer is a client, whose newer epoch withdraws nothing.""" + return _kind(producer.epoch) is models.ProducerEpochKindClient + + +def producer_key(producer: models.Producer) -> tuple[Any, ...]: + """The Producer fields that :py:func:`withdraws` reads, so equal keys are withdrawn together.""" + epoch = _epoch_key(producer.epoch) if producer.epoch is not None else None + return (_kind(producer.epoch), producer.id, producer.run_id or "", epoch) + + +def _of_run(producer: models.Producer, by: models.Producer) -> bool: + # The run's Workflow Producer, or one of its Local Activities. + kind = _kind(producer.epoch) + if kind is models.ProducerEpochKindWorkflow: + return producer.id == by.id + if kind is models.ProducerEpochKindLocalActivity: + return (producer.run_id or "") == (by.run_id or "") + return False + + +def withdraws(facts: models.StreamSuperseded, producer: models.Producer | None) -> bool: + """Whether a Superseded item withdraws a Record of ``producer``. + + The rule matches the crate's, so every SDK withdraws the same Records. + """ + writer = facts.writer + if producer is None or writer is None: + return False + if writer.withdraw_run and producer.run_id == writer.withdraw_run: + return True + by, old = writer.producer, producer.epoch + epoch = by.epoch if by is not None else None + if by is None or epoch is None or old is None: + return False + kind = _kind(epoch) + if kind is None or kind is models.ProducerEpochKindClient: + # A client's newer epoch only fences. Its earlier Records are real messages, not a + # failed attempt's. + return False + if kind is not models.ProducerEpochKindWorkflow: + return ( + _kind(old) is kind + and producer.id == by.id + and _epoch_key(old) < _epoch_key(epoch) + ) + if not _of_run(producer, by): + return False + task, old_task = _task_of(epoch), _task_of(old) + started = (old_task.started_event_id if old_task else None) or 0 + if started in (writer.failed_tasks or []): + return True + completed = writer.previous_completed_task + after = (completed.started_event_id if completed else None) or 0 + if after < started < ((task.started_event_id if task else None) or 0): + return True + # No task in the range can have completed, so a lower attempt at either end goes too. + ends = [task] if completed is None else [task, completed] + return any( + started == ((end.started_event_id if end else None) or 0) + and _task_key(old_task) < _task_key(end) + for end in ends + ) diff --git a/temporalio/streams/_ref.py b/temporalio/streams/_ref.py new file mode 100644 index 000000000..ebbdf14cc --- /dev/null +++ b/temporalio/streams/_ref.py @@ -0,0 +1,188 @@ +"""The Stream ref, the data that identifies one Stream.""" + +from __future__ import annotations + +import typing +from datetime import timedelta +from typing import Any, Generic, TypeVar + +import google.protobuf.json_format +from typing_extensions import Self, TypedDict + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.converter +from temporalio.streams.system import models + +T = TypeVar("T") + + +class StreamOptions(TypedDict, total=False): + """The options of a new Stream. ``create_or_get_stream`` ignores them on a Stream that exists. + + .. warning:: + This API is experimental. + """ + + retention: timedelta + """How long a Record stays readable, even while the Stream is open. + + A closed Stream is kept for one Retention after its close. Defaults to the backend's + Retention. + """ + + +def _model(ref: StreamRef[Any]) -> models.StreamRef: + return models.StreamRef(id=ref.id, owner=ref.owner, backend=ref.backend) + + +def to_proto(ref: StreamRef[Any]) -> stream_v1.StreamRef: + """The ``temporal.api.stream.v1.StreamRef`` message of a ref.""" + return models._StreamRefTransferTypeConverter().to_transfer_type(_model(ref)) + + +def _to_json(ref: StreamRef[Any]) -> dict[str, Any]: + # The proto JSON of temporal.api.stream.v1.StreamRef, as MessageToJson writes it, so every SDK + # reads a ref that another wrote. + return google.protobuf.json_format.MessageToDict(to_proto(ref)) + + +def _from_json(value: Any, type_hint: Any) -> StreamRef[Any]: + origin = typing.get_origin(type_hint) or type_hint + if not isinstance(origin, type) or not issubclass(origin, StreamRef): + raise TypeError(f"{type_hint!r} is not a Stream ref type") + if isinstance(value, origin): + return value + if not isinstance(value, dict) or not isinstance(value.get("id"), str): + raise TypeError(f"Expected the JSON object of a Stream ref, got {type(value)}") + message = google.protobuf.json_format.ParseDict( + value, stream_v1.StreamRef(), ignore_unknown_fields=True + ) + model = models._StreamRefTransferTypeConverter().from_transfer_type( + message, models.StreamRef + ) + # A ref is data, so it keeps any owner, such as a Standalone Activity's. A close checks it. + owner = model.owner if model.owner != models.Execution() else None + return origin(id=model.id, owner=owner, backend=model.backend) + + +class _StreamRefJson( + temporalio.converter.TransferTypeConverter["StreamRef[Any]", dict[str, Any]] +): + # The JSON converter writes a nested dataclass as its fields, so the ref isn't the generated + # dataclass. It goes through the generated converter, as a dict so a payload stays json/plain. + transfer_type: type[dict[str, Any]] | None = dict + + def to_transfer_type(self, value: StreamRef[Any]) -> dict[str, Any]: + return _to_json(value) + + def from_transfer_type( + self, value: dict[str, Any], type_hint: type[StreamRef[Any]] + ) -> StreamRef[Any]: + return _from_json(value, type_hint) + + +class StreamRef(Generic[T]): + """The data that identifies one Stream. ``T`` is the payload type, for type checking only. + + A ref has the fields of the wire model of the ``temporal.api.stream.v1.StreamRef`` message. + Code passes it as a Workflow or Activity argument, a query result, a Record, or a field of a + dataclass or a pydantic model, in the proto JSON form of that message. + + A ref has no calls. To make calls, Workflow code gets a Stream with + ``workflow.to_stream(ref)``, and other code with ``client.to_stream(ref)``. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + id: str, + owner: models.Execution | None = None, + backend: str | None = None, + ) -> None: + """Create a ref. Code gets one from a Stream's ``ref`` instead.""" + self._id = id + self._owner = owner + self._backend = backend + + @property + def id(self) -> str: + """The Stream id. It is opaque to callers.""" + return self._id + + @property + def owner(self) -> models.Execution | None: + """The owner of a bound Stream, such as a Workflow run chain. None for a standalone one.""" + return self._owner + + @property + def backend(self) -> str | None: + """Name of the backend that stores the Stream. None means the default one.""" + return self._backend + + def __eq__(self, other: object) -> bool: + """Refs of one class are equal when their fields are equal.""" + if type(other) is not type(self): + return NotImplemented + assert isinstance(other, StreamRef) + return (self._id, self._owner, self._backend) == ( + other._id, + other._owner, + other._backend, + ) + + def __hash__(self) -> int: + """Hash of the id and the backend, which identify the Stream.""" + return hash((type(self), self._id, self._backend)) + + def __repr__(self) -> str: + """Show the fields.""" + return ( + f"{type(self).__name__}(id={self._id!r}, owner={self._owner!r}, " + f"backend={self._backend!r})" + ) + + def __copy__(self) -> Self: + """A ref doesn't change, so a copy is the same ref.""" + return self + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + """Return the same ref, since dataclasses.asdict deep-copies a ref field.""" + return self + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: Any) -> Any: + """Give pydantic models the proto JSON form, as the payload converter does. + + Pydantic is optional, so it is imported only when a model uses a ref. + """ + from pydantic_core import core_schema + + return core_schema.no_info_plain_validator_function( + lambda value: _from_json(value, source), + json_schema_input_schema=core_schema.typed_dict_schema( + { + "id": core_schema.typed_dict_field(core_schema.str_schema()), + "owner": core_schema.typed_dict_field( + core_schema.typed_dict_schema( + { + key: core_schema.typed_dict_field( + core_schema.str_schema(), required=False + ) + for key in ("type", "businessId", "runId") + } + ), + required=False, + ), + "backend": core_schema.typed_dict_field( + core_schema.str_schema(), required=False + ), + } + ), + serialization=core_schema.plain_serializer_function_ser_schema(_to_json), + ) + + +# Registered as a call, not a decorator, so type checkers keep the class generic. +temporalio.converter.transfer_type_convertible(_StreamRefJson)(StreamRef) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py new file mode 100644 index 000000000..9f1a05f10 --- /dev/null +++ b/tests/contrib/streams/_support.py @@ -0,0 +1,41 @@ +"""Finds the Stream crates that the bridge links, for the shared vectors and protos.""" + +from __future__ import annotations + +import functools +import json +import subprocess +from pathlib import Path + +import pytest + +_BRIDGE = Path(__file__).parents[3] / "temporalio" / "bridge" + + +@functools.cache +def crate_dirs() -> dict[str, Path]: + try: + output = subprocess.run( + ["cargo", "metadata", "--offline", "--format-version", "1"], + cwd=_BRIDGE, + capture_output=True, + check=True, + timeout=120, + ).stdout + except (OSError, subprocess.SubprocessError): + return {} + return { + package["name"]: Path(package["manifest_path"]).parent + for package in json.loads(output)["packages"] + } + + +def crate_file(relative: str, crate: str = "temporalio-streams") -> Path: + directory = crate_dirs().get(crate) + if directory is None: + pytest.skip(f"The {crate} crate checkout is not available") + return directory / relative + + +def vectors(name: str) -> object: + return json.loads(crate_file(f"testdata/{name}").read_text(encoding="utf-8")) diff --git a/tests/contrib/streams/test_refs.py b/tests/contrib/streams/test_refs.py new file mode 100644 index 000000000..0ce3fe40d --- /dev/null +++ b/tests/contrib/streams/test_refs.py @@ -0,0 +1,211 @@ +import json +import subprocess +import sys +from dataclasses import dataclass +from typing import Any + +import pydantic +from google.protobuf import json_format + +import temporalio.api.common.v1 +import temporalio.api.enums.v1 +import temporalio.api.stream.v1 +import temporalio.common +from temporalio import streams +from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.converter import DataConverter +from temporalio.streams import _items + + +@dataclass +class Turn: + kind: str + + +OWNER = streams.Execution( + type=streams.ExecutionType.ExecutionTypeWorkflow, business_id="wf", run_id="run" +) +BOUND: streams.StreamRef[Turn] = streams.StreamRef( + id="session-7-turns", owner=OWNER, backend="redis" +) + + +def test_ref_is_the_proto_json_of_the_ref_message(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(BOUND) + message = temporalio.api.stream.v1.StreamRef( + id="session-7-turns", + owner=temporalio.api.common.v1.Execution( + type=temporalio.api.enums.v1.ExecutionType.EXECUTION_TYPE_WORKFLOW, + business_id="wf", + run_id="run", + ), + backend="redis", + ) + assert json.loads(payload.data) == json_format.MessageToDict(message) + assert payload.metadata["encoding"] == b"json/plain" + assert payload.data == ( + b'{"backend":"redis","id":"session-7-turns",' + b'"owner":{"businessId":"wf","runId":"run","type":"EXECUTION_TYPE_WORKFLOW"}}' + ) + assert json.loads(converter.to_payload(streams.StreamRef(id="s")).data) == { + "id": "s" + } + + +def test_ref_has_the_fields_of_the_wire_model(): + assert BOUND.id == "session-7-turns" + assert BOUND.owner is not None and BOUND.owner.business_id == "wf" + assert BOUND.owner.run_id == "run" and BOUND.backend == "redis" + assert streams.StreamRef(id="s") == streams.StreamRef("s") + assert streams.StreamRef(id="s") != streams.StreamRef(id="s", backend="redis") + same: streams.StreamRef[Turn] = streams.StreamRef("session-7-turns", OWNER, "redis") + assert same == BOUND and len({BOUND, same}) == 1 + + +def test_any_type_hint_decodes_the_ref(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(BOUND) + for hint in ( + streams.StreamRef[Turn], + streams.StreamRef[list[str]], + streams.StreamRef[Any], + streams.StreamRef, + ): + decoded = converter.from_payload(payload, hint) + assert type(decoded) is streams.StreamRef and decoded == BOUND + + +@dataclass +class CallModelInput: + prompt: str + stream: streams.StreamRef[Turn] + + +def test_ref_in_a_dataclass_field(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(CallModelInput("hi", BOUND)) + assert json.loads(payload.data)["stream"]["owner"]["businessId"] == "wf" + assert payload.data == ( + b'{"prompt":"hi","stream":{"backend":"redis","id":"session-7-turns",' + b'"owner":{"businessId":"wf","runId":"run","type":"EXECUTION_TYPE_WORKFLOW"}}}' + ) + decoded = converter.from_payload(payload, CallModelInput) + assert decoded.stream == BOUND + + +class CallModelModel(pydantic.BaseModel): + prompt: str + stream: streams.StreamRef[Turn] + + +def test_ref_in_a_pydantic_field(): + model = CallModelModel(prompt="hi", stream=BOUND) + dumped = json.loads(model.model_dump_json()) + assert dumped["stream"] == { + "id": "session-7-turns", + "owner": { + "type": "EXECUTION_TYPE_WORKFLOW", + "businessId": "wf", + "runId": "run", + }, + "backend": "redis", + } + loaded = CallModelModel.model_validate_json(model.model_dump_json()) + assert loaded.stream == BOUND + converter = pydantic_data_converter.payload_converter + payload = converter.to_payload(model) + decoded = converter.from_payload(payload, CallModelModel) + assert decoded.stream == BOUND + + +def test_a_ref_keeps_an_owner_of_any_type(): + # Data from a later stage or another SDK can carry a Standalone Activity's run. + owner = streams.Execution( + type=streams.ExecutionType.ExecutionTypeActivity, business_id="a", run_id="r" + ) + ref: streams.StreamRef[str] = streams.StreamRef(id="s", owner=owner) + converter = DataConverter.default.payload_converter + decoded = converter.from_payload(converter.to_payload(ref), streams.StreamRef[str]) + assert decoded == ref and decoded.owner == owner + + +def _producer(id: str, attempt: int) -> streams.Producer: + return streams.Producer( + id=id, + epoch=streams.ProducerEpoch( + kind=streams.ProducerEpochKindActivity( + streams.ActivityEpoch(attempt=attempt) + ) + ), + run_id="run", + ) + + +def test_read_items_have_a_kind_that_narrows_their_type(): + record = streams.Record(data=Turn("model"), offset=3, producer=_producer("a", 1)) + superseded = streams.Superseded( + offset=4, + facts=streams.StreamSuperseded( + writer=streams.StreamWriter(producer=_producer("a", 2), failed_tasks=[]) + ), + ) + items: list[streams.Record[Turn] | streams.Superseded] = [record, superseded] + shown: list[Turn] = [] + withdrawn: list[bool] = [] + for item in items: + match item.kind: + case streams.ItemKind.SUPERSEDED: + withdrawn.append(item.withdraws(record)) + case _: + shown.append(item.data) + assert shown == [Turn("model")] and withdrawn == [True] + assert record.kind is streams.ItemKind.RECORD + + +def test_superseded_withdraws_by_the_producer_of_an_earlier_record(): + superseded = streams.Superseded( + offset=9, + facts=streams.StreamSuperseded( + writer=streams.StreamWriter(producer=_producer("a", 2), failed_tasks=[]) + ), + ) + # An app that resumes keeps only the Producer of each Record that it showed. + assert superseded.withdraws(_producer("a", 1)) + assert not superseded.withdraws(_producer("a", 2)) + assert not superseded.withdraws(_producer("b", 1)) + + +def test_a_client_producer_is_never_withdrawn(): + client = streams.Producer( + id="client:me", + epoch=streams.ProducerEpoch( + kind=streams.ProducerEpochKindClient(streams.ClientEpoch(attempt=1)) + ), + ) + newer = streams.Producer( + id="client:me", + epoch=streams.ProducerEpoch( + kind=streams.ProducerEpochKindClient(streams.ClientEpoch(attempt=2)) + ), + ) + facts = streams.StreamSuperseded( + writer=streams.StreamWriter(producer=newer, failed_tasks=[]) + ) + assert _items.is_client(client) + assert not _items.withdraws(facts, client) + + +def test_temporalio_common_has_no_stream_api(): + assert not any(name.startswith("Stream") for name in vars(temporalio.common)) + + +def test_import_of_the_sdk_loads_no_stream_code(): + code = ( + "import sys, temporalio.client, temporalio.common, temporalio.worker, " + "temporalio.workflow; print('temporalio.streams' in sys.modules)" + ) + loaded = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ).stdout + assert loaded.strip() == "False" diff --git a/tests/contrib/streams/test_vectors.py b/tests/contrib/streams/test_vectors.py new file mode 100644 index 000000000..e73253f2c --- /dev/null +++ b/tests/contrib/streams/test_vectors.py @@ -0,0 +1,124 @@ +from pathlib import Path +from typing import Any + +from temporalio.streams import _ids, _items +from temporalio.streams.system import models +from tests.contrib.streams._support import crate_file, vectors + +_SCRIPTS = Path(__file__).parents[3] / "scripts" / "_proto" + + +def test_vendored_redis_proto_matches_the_crate(): + # The other Stream protos come from the Core submodule. The Redis config comes from the + # Redis crate, so the SDK keeps a copy. + relative = "temporal/sdk/streams/v1/redis.proto" + crate = crate_file(f"protos/{relative}", "temporalio-streams-redis") + expected = crate.read_text(encoding="utf-8") + assert (_SCRIPTS / relative).read_text(encoding="utf-8") == expected + + +def test_producer_ids_match_the_crate(): + producers: Any = vectors("ids.json") + producers = producers["producers"] + assert ( + _ids.workflow_producer_id(producers["workflow"]["run_id"]) + == (producers["workflow"]["id"]) + ) + la = producers["local_activity"] + assert _ids.local_activity_producer_id(la["run_id"], la["call_key"]) == la["id"] + act = producers["activity"] + assert ( + _ids.activity_producer_id( + act["run_id"], act["activity_id"], act["schedule_key"] + ) + == act["id"] + ) + sa = producers["standalone_activity"] + assert ( + _ids.standalone_activity_producer_id(sa["activity_id"], sa["activity_run_id"]) + == sa["id"] + ) + client = producers["client"] + assert ( + _ids.client_producer_id(client["identity"], client["process_key"]) + == client["id"] + ) + named = producers["named_client"] + assert _ids.named_client_producer_id(named["name"]) == named["id"] + + +def _time(value: list[int] | None) -> models.Timestamp | None: + return models.Timestamp(seconds=value[0], nanos=value[1]) if value else None + + +def _task(value: dict[str, Any], prefix: str = "") -> models.WorkflowTaskEpoch: + return models.WorkflowTaskEpoch( + started_event_id=value.get(f"{prefix}started_event_id"), + attempt=value.get(f"{prefix}attempt"), + started_time=_time(value.get(f"{prefix}started_time")), + ) + + +def _epoch(kind: str, value: dict[str, Any] | None) -> models.ProducerEpoch | None: + # The vectors keep the flat fields of every kind, so the case comes from the kind. + if value is None: + return None + activity = models.ActivityEpoch( + attempt_scheduled_time=_time(value.get("attempt_scheduled_time")), + attempt=value.get("attempt"), + ) + cases: dict[str, models.ProducerEpochKind] = { + "workflow": models.ProducerEpochKindWorkflow(_task(value, "task_")), + "local_activity": models.ProducerEpochKindLocalActivity( + models.LocalActivityEpoch( + task=_task(value, "task_"), attempt=value.get("attempt") + ) + ), + "activity": models.ProducerEpochKindActivity(activity), + "standalone_activity": models.ProducerEpochKindStandaloneActivity(activity), + "client": models.ProducerEpochKindClient( + models.ClientEpoch(attempt=value.get("attempt")) + ), + } + return models.ProducerEpoch(kind=cases[kind]) + + +def test_withdrawals_match_the_crate(): + cases: Any = vectors("withdrawals.json") + assert cases + for case in cases: + record = case["record"]["producer"] + producer = models.Producer( + id=record["id"], + run_id=record.get("run_id"), + epoch=_epoch(record["kind"], record.get("epoch")), + ) + given = case["superseded"] + completed = given.get("previous_completed_task") + writer = models.StreamWriter( + producer=models.Producer( + id=given["producer_id"], + epoch=_epoch(given["producer_kind"], given.get("epoch")), + run_id=given.get("run_id"), + ), + previous_completed_task=_task(completed) if completed else None, + failed_tasks=given.get("failed_tasks", []), + withdraw_run=given.get("withdraw_run"), + ) + facts = models.StreamSuperseded(writer=writer) + assert _items.withdraws(facts, producer) == case["withdraws"], case["name"] + + +def test_epoch_order_matches_the_crate(): + cases: Any = vectors("epochs.json") + keys: list[tuple[tuple[Any, ...], bytes]] = [] + for case in cases: + epoch = _epoch(case["kind"], case["epoch"]) if case["kind"] != "empty" else None + python = _items._epoch_key(epoch or models.ProducerEpoch()) + keys.append((python, bytes.fromhex(case["key"]))) + assert len(keys) > 1 + # A byte compare of the crate's keys orders the epochs, so the Python keys compare the same. + for python_a, crate_a in keys: + for python_b, crate_b in keys: + assert (python_a < python_b) == (crate_a < crate_b) + assert (python_a == python_b) == (crate_a == crate_b)