diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index e6260ee00..d4ae79334 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -13,6 +13,7 @@ from datetime import datetime, timedelta from enum import IntEnum from typing import ( + TYPE_CHECKING, Any, ClassVar, Generic, @@ -31,7 +32,100 @@ import temporalio.api.enums.v1 import temporalio.api.workflow.v1 import temporalio.types -from temporalio.common._streams._defn import StreamDefinition as StreamDefinition + +if TYPE_CHECKING: + from temporalio.common._streams._defn import StreamDefinition, defn + from temporalio.common._streams._errors import ( + StreamClosedError, + StreamConfigError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetError, + StreamSeqDivergentError, + StreamStorageError, + StreamTooLargeError, + StreamUnsupportedError, + ) + from temporalio.common._streams._handle import ( + Creating, + StreamHandle, + StreamOptions, + ) + from temporalio.common._streams._items import ( + END, + Closed, + CloseStatus, + End, + Record, + Superseded, + ) + +__all__ = [ + "ActivityIDConflictPolicy", + "ActivityIDReusePolicy", + "AutoUpgradeVersioningOverride", + "CloseStatus", + "Closed", + "Creating", + "END", + "End", + "HeaderCodecBehavior", + "MetricAttributes", + "MetricCommon", + "MetricCounter", + "MetricGauge", + "MetricGaugeFloat", + "MetricHistogram", + "MetricHistogramFloat", + "MetricHistogramTimedelta", + "MetricMeter", + "NexusOperationCancellationState", + "NexusOperationExecutionStatus", + "NexusOperationIDConflictPolicy", + "NexusOperationIDReusePolicy", + "OneTimeVersioningOverride", + "PendingNexusOperationExecutionState", + "PinnedVersioningOverride", + "Priority", + "QueryRejectCondition", + "RawValue", + "Record", + "RetryPolicy", + "SearchAttributeIndexedValueType", + "SearchAttributeKey", + "SearchAttributePair", + "SearchAttributeUpdate", + "SearchAttributeValue", + "SearchAttributeValueType", + "SearchAttributeValues", + "SearchAttributes", + "StreamClosedError", + "StreamConfigError", + "StreamDefinition", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamOptions", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Superseded", + "TypedSearchAttributes", + "VersioningBehavior", + "VersioningOverride", + "WorkerDeploymentVersion", + "WorkflowIDConflictPolicy", + "WorkflowIDReusePolicy", + "defn", +] @dataclass @@ -1483,3 +1577,42 @@ def wait_sync(self, timeout: float | None = None) -> None: if not self.thread_event: raise RuntimeError("Missing event") self.thread_event.wait(timeout) + + +# The Stream modules import temporalio.client and temporalio.workflow, which import this module, +# so this module imports the Stream names on first use. +_STREAM_NAMES = frozenset( + [ + "CloseStatus", + "Closed", + "Creating", + "END", + "End", + "Record", + "StreamClosedError", + "StreamConfigError", + "StreamDefinition", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamOptions", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Superseded", + "defn", + ] +) + + +def __getattr__(name: str) -> Any: + if name in _STREAM_NAMES: + import temporalio.common._streams + + return getattr(temporalio.common._streams, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 39ddb83f5..517ce4888 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1,7 +1,52 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" -from temporalio.common._streams._defn import StreamDefinition +from temporalio.common._streams._defn import StreamDefinition, defn +from temporalio.common._streams._errors import ( + StreamClosedError, + StreamConfigError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetError, + StreamSeqDivergentError, + StreamStorageError, + StreamTooLargeError, + StreamUnsupportedError, +) +from temporalio.common._streams._handle import Creating, StreamHandle, StreamOptions +from temporalio.common._streams._items import ( + END, + Closed, + CloseStatus, + End, + Record, + Superseded, +) __all__ = [ + "END", + "CloseStatus", + "Closed", + "Creating", + "End", + "Record", + "StreamClosedError", + "StreamConfigError", "StreamDefinition", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamHandle", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamOptions", + "StreamSeqDivergentError", + "StreamStorageError", + "StreamTooLargeError", + "StreamUnsupportedError", + "Superseded", + "defn", ] diff --git a/temporalio/common/_streams/_defn.py b/temporalio/common/_streams/_defn.py index 1141e9dd4..48d235233 100644 --- a/temporalio/common/_streams/_defn.py +++ b/temporalio/common/_streams/_defn.py @@ -2,14 +2,106 @@ from __future__ import annotations -from typing import Generic, TypeVar +import typing +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Generic, TypeVar, overload T = TypeVar("T") +ClassT = TypeVar("ClassT", bound=type) + +_DEFINITION_ATTR = "__temporal_stream_definition" class StreamDefinition(Generic[T]): """Base class of a Stream definition. Its type argument is the Stream's payload type. + Mark a subclass with :py:func:`defn`:: + + @common.defn(name="turns") + class Turns(common.StreamDefinition[Turn]): ... + + A definition doesn't name a backend, so Workflows, Activities and clients can share it. + Separate channels, such as turns and token deltas, are separate Streams. Use ``RawValue`` as + the payload type to read each Record's Payload with no conversion. + + .. warning:: + This API is experimental. + """ + + +@dataclass(frozen=True) +class DefinitionInfo: + """What :py:func:`defn` takes from a definition class.""" + + name: str + """Name of a Stream of this definition.""" + + type: Any + """The payload type, which a read decodes each Record's value with.""" + + +@overload +def defn(cls: ClassT) -> ClassT: ... + + +@overload +def defn(*, name: str | None = None) -> Callable[[ClassT], ClassT]: ... + + +def defn( + cls: ClassT | None = None, *, name: str | None = None +) -> ClassT | Callable[[ClassT], ClassT]: + """Mark a subclass of :py:class:`StreamDefinition` as a Stream definition. + + Args: + cls: The class. + name: The name of a Stream of this definition. Defaults to the class name. + + Raises: + TypeError: The class doesn't subclass ``StreamDefinition`` with a payload type. + .. warning:: This API is experimental. """ + + def decorate(cls: ClassT) -> ClassT: + info = DefinitionInfo(name=name or cls.__name__, type=_payload_type(cls)) + setattr(cls, _DEFINITION_ATTR, info) + return cls + + if cls is not None: + return decorate(cls) + return decorate + + +def _payload_type(cls: type) -> Any: + if not isinstance(cls, type) or not issubclass(cls, StreamDefinition): + raise TypeError( + f"{cls!r} must subclass StreamDefinition with its payload type, such as " + "StreamDefinition[Turn]" + ) + for klass in cls.__mro__: + for base in klass.__dict__.get("__orig_bases__", ()): + if typing.get_origin(base) is StreamDefinition: + [payload] = typing.get_args(base) + if not isinstance(payload, TypeVar): + return payload + raise TypeError( + f"{cls.__qualname__} gives StreamDefinition no payload type, such as " + "StreamDefinition[Turn]" + ) + + +def definition_of(cls: type) -> DefinitionInfo: + """The definition that :py:func:`defn` gave a class. + + Raises: + TypeError: The class has no ``@common.defn``. + """ + definition = cls.__dict__.get(_DEFINITION_ATTR) if isinstance(cls, type) else None + if not isinstance(definition, DefinitionInfo): + raise TypeError( + f"{cls!r} is not a Stream definition. Add @temporalio.common.defn." + ) + return definition diff --git a/temporalio/common/_streams/_errors.py b/temporalio/common/_streams/_errors.py new file mode 100644 index 000000000..9b3cf2338 --- /dev/null +++ b/temporalio/common/_streams/_errors.py @@ -0,0 +1,88 @@ +"""Errors of the Streams API in temporalio.common.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import temporalio.api.stream.v1 +import temporalio.nexus.system._support as _support +from temporalio.common._streams.system import models + +# The generated classes are the Stream errors, one for each reason that has its own. +from temporalio.common._streams.system.errors import ( + StreamClosedError, + StreamError, + StreamExpiredError, + StreamFencedError, + StreamNotFoundError, + StreamNotOwnerError, + StreamOffsetError, + StreamSeqDivergentError, + StreamStorageError, + StreamTooLargeError, + StreamUnsupportedError, + stream_error, +) + +__all__ = [ + "StreamClosedError", + "StreamConfigError", + "StreamError", + "StreamExpiredError", + "StreamFencedError", + "StreamNotFoundError", + "StreamNotOwnerError", + "StreamOffsetError", + "StreamSeqDivergentError", + "StreamStorageError", + "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. + """ + + +def no_default_backend( + owner: str, names: Sequence[str], hint: str +) -> StreamConfigError: + """The error of a call with no backend name when no backend has ``default=True``.""" + return StreamConfigError( + f"No Stream backend of the {owner} ({', '.join(names)}) has default=True, so the " + "default is native, which this SDK does not have. Mark a backend with " + f"default=True, or name one, such as with {hint}." + ) + + +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/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py new file mode 100644 index 000000000..1a12d1ec3 --- /dev/null +++ b/temporalio/common/_streams/_handle.py @@ -0,0 +1,292 @@ +"""The Stream handle, which Workflow, Activity and client code share.""" + +from __future__ import annotations + +import asyncio +import typing +from collections.abc import Callable, Coroutine, Generator +from datetime import timedelta +from types import TracebackType +from typing import Any, Generic, Protocol, TypeVar + +import google.protobuf.json_format +from typing_extensions import Self, TypedDict + +import temporalio.converter +import temporalio.exceptions +from temporalio.common._streams.system import models + +T = TypeVar("T") + + +def exit_failure(exc: BaseException | None) -> tuple[bool, BaseException | None]: + """Whether the end of an ``async with`` block closes the Stream, and with which failure. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open, since the caller's code didn't end the block. + """ + if isinstance(exc, asyncio.CancelledError): + return True, temporalio.exceptions.CancelledError(str(exc) or "Cancelled") + if exc is None or isinstance(exc, Exception): + return True, exc + return False, None + + +class StreamOptions(TypedDict, total=False): + """The options of a Stream create. A create of a Stream that exists ignores them. + + .. warning:: + This API is experimental. + """ + + retention: timedelta + """The age after which a Record is deleted, even while the Stream is open. + + A closed Stream is kept for one Retention after its close. Defaults to the backend's + Retention. + """ + + +class _ClosesOnExit(Protocol): + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: ... + + +S = TypeVar("S", bound=_ClosesOnExit) + + +class Creating(Generic[S]): + """A pending Stream create. + + Await it to get the handle, or use it in ``async with``, which closes the Stream at the end + of the block. The create starts when it's awaited or entered, not before. + + .. warning:: + This API is experimental. + """ + + def __init__(self, create: Callable[[], Coroutine[Any, Any, S]]) -> None: + """Create the object. Use a context's create function instead.""" + self._create = create + self._stream: S | None = None + + def __await__(self) -> Generator[Any, None, S]: + """Create the Stream and return its handle.""" + return self._create().__await__() + + async def __aenter__(self) -> S: + """Create the Stream and return its handle for the block.""" + self._stream = await self._create() + return self._stream + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream as the handle's own ``async with`` does.""" + assert self._stream is not None + await self._stream.__aexit__(exc_type, exc, tb) + + +# The owner type of a Workflow's Stream in proto JSON. A client's Stream has no owner. +_WORKFLOW_OWNER = "EXECUTION_TYPE_WORKFLOW" + + +def workflow_owner( + workflow_id: str | None, first_execution_run_id: str | None +) -> models.Execution | None: + """The owner of a Workflow's Stream in a handle. A client's Stream has no owner.""" + if not workflow_id and not first_execution_run_id: + return None + return models.Execution( + type=models.ExecutionType.ExecutionTypeWorkflow, + business_id=workflow_id or None, + run_id=first_execution_run_id or None, + ) + + +def _to_json(handle: StreamHandle[Any]) -> dict[str, Any]: + # The proto JSON of temporal.api.stream.v1.StreamHandle, as MessageToJson writes it, so every + # SDK reads a handle that another wrote. + message = models.StreamHandle( + id=handle.id, + owner=workflow_owner(handle.workflow_id, handle.first_execution_run_id), + backend=handle.backend or None, + ) + return google.protobuf.json_format.MessageToDict( + models._StreamHandleTransferTypeConverter().to_transfer_type(message) + ) + + +def _from_json(value: Any, type_hint: Any) -> StreamHandle[Any]: + origin = typing.get_origin(type_hint) or type_hint + if not isinstance(origin, type) or not issubclass(origin, StreamHandle): + raise TypeError(f"{type_hint!r} is not a Stream handle type") + if isinstance(value, origin): + return value + if isinstance(value, StreamHandle): + value = _to_json(value) + if not isinstance(value, dict) or not isinstance(value.get("id"), str): + raise TypeError(f"Expected the JSON object of a Stream, got {type(value)}") + owner = value.get("owner") or {} + if not isinstance(owner, dict) or ( + owner and owner.get("type") not in (_WORKFLOW_OWNER, 1) + ): + raise TypeError("A Stream handle's owner must be a Workflow run chain") + args = typing.get_args(type_hint) + # A hint such as StreamHandle[Any] gives no payload type, so values decode with no hint. + payload_type = ( + args[0] + if args and args[0] is not Any and not isinstance(args[0], TypeVar) + else None + ) + return origin( + id=value["id"], + workflow_id=owner.get("businessId") or None, + first_execution_run_id=owner.get("runId") or None, + backend=value.get("backend") or None, + payload_type=payload_type, + ) + + +class _StreamHandleJson( + temporalio.converter.TransferTypeConverter["StreamHandle[Any]", dict[str, str]] +): + # A handle has the same JSON form in every context, and the parameter's type hint gives the + # payload type. So a Workflow can pass a Stream to an Activity as data. + transfer_type: type[dict[str, str]] | None = dict + + def to_transfer_type(self, value: StreamHandle[Any]) -> dict[str, str]: + return _to_json(value) + + def from_transfer_type( + self, value: dict[str, str], type_hint: type[StreamHandle[Any]] + ) -> StreamHandle[Any]: + return _from_json(value, type_hint) + + +class StreamHandle(Generic[T]): + """A handle on one Stream in Workflow, Activity and client code. ``T`` is the payload type. + + A handle is data. It's passed as a Workflow or Activity argument in the proto JSON form of + the ``temporal.api.stream.v1.StreamHandle`` message. The parameter's type, such as + ``StreamHandle[Turn]``, gives the payload type. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + id: str, + workflow_id: str | None = None, + first_execution_run_id: str | None = None, + backend: str | None = None, + payload_type: Any = None, + ) -> None: + """Create a handle. Use ``create_stream`` or ``get_stream_handle`` instead.""" + self._id = id + self._workflow_id = workflow_id + self._first_execution_run_id = first_execution_run_id + self._backend = backend + # The type that a read decodes each value with. None decodes with no type hint. + self._payload_type = payload_type + + @property + def id(self) -> str: + """The Stream id. It is opaque to callers.""" + return self._id + + @property + def workflow_id(self) -> str | None: + """For a Workflow's Stream, the id of the Workflow that owns it.""" + return self._workflow_id + + @property + def first_execution_run_id(self) -> str | None: + """For a Workflow's Stream, the first run id of the owning run chain.""" + return self._first_execution_run_id + + @property + def backend(self) -> str | None: + """Name of the backend that stores the Stream. None means the default one.""" + return self._backend + + def _fields(self) -> dict[str, str]: + # The fields that name the Stream. They need no call, so a handle whose id is not known + # yet still compares, hashes and prints. + fields = {"id": self._id} + if self._workflow_id: + fields["workflowId"] = self._workflow_id + if self._first_execution_run_id: + fields["firstExecutionRunId"] = self._first_execution_run_id + if self._backend: + fields["backend"] = self._backend + return fields + + def __eq__(self, other: object) -> bool: + """Handles of one class are equal when they name the same Stream.""" + if type(other) is not type(self): + return NotImplemented + assert isinstance(other, StreamHandle) + return self._fields() == other._fields() + + def __hash__(self) -> int: + """Hash of the fields that never change, so it's the same after a client learns the id.""" + return hash((type(self), self._backend, self._workflow_id or self._id)) + + def __repr__(self) -> str: + """Show the Stream fields.""" + fields = ", ".join(f"{key}={value!r}" for key, value in self._fields().items()) + return f"{type(self).__name__}({fields})" + + def __copy__(self) -> Self: + """A copy is a handle of the same Stream in the same context.""" + return self + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + """Return the same handle, since dataclasses.asdict deep-copies a handle 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 handle. + """ + 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(_StreamHandleJson)(StreamHandle) diff --git a/temporalio/common/_streams/_ids.py b/temporalio/common/_streams/_ids.py new file mode 100644 index 000000000..8311dc8d5 --- /dev/null +++ b/temporalio/common/_streams/_ids.py @@ -0,0 +1,49 @@ +"""Stream and 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 scoped_stream_id(workflow_id: str, first_run_id: str, name: str) -> str: + """Id of a Stream that a Workflow creates with a name.""" + return _join(workflow_id, first_run_id, name) + + +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/common/_streams/_items.py b/temporalio/common/_streams/_items.py new file mode 100644 index 000000000..9f7df6695 --- /dev/null +++ b/temporalio/common/_streams/_items.py @@ -0,0 +1,215 @@ +"""The items that a Stream read gives.""" + +from __future__ import annotations + +import asyncio +import enum +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +import temporalio.exceptions +from temporalio.api.enums.v1 import StreamCloseStatus, StreamProducerKind +from temporalio.common._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 + + +@dataclass(frozen=True) +class Record(Generic[T]): + """A data Record of a Stream. + + .. warning:: + This API is experimental. + """ + + 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. + """ + + offset: int + """Position of the item in the Stream.""" + + facts: models.StreamSuperseded + """Which earlier Records the item 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]) -> bool: + """Whether this item withdraws an earlier Record.""" + return withdraws(self.facts, record.producer) + + +class CloseStatus(enum.IntEnum): + """How a Stream closed. + + .. warning:: + This API is experimental. + """ + + COMPLETED = int(StreamCloseStatus.STREAM_CLOSE_STATUS_COMPLETED) + """Closed with no failure.""" + + FAILED = int(StreamCloseStatus.STREAM_CLOSE_STATUS_FAILED) + """Closed with a failure.""" + + CANCELED = int(StreamCloseStatus.STREAM_CLOSE_STATUS_CANCELED) + """Closed with a cancel failure, such as by a cancelled owner run.""" + + TERMINATED = int(StreamCloseStatus.STREAM_CLOSE_STATUS_TERMINATED) + """The server closed it, since the owner run was terminated.""" + + TIMED_OUT = int(StreamCloseStatus.STREAM_CLOSE_STATUS_TIMED_OUT) + """The server closed it, since the owner run timed out.""" + + +def close_failure(failure: BaseException | None) -> BaseException | None: + """The failure that a close sends. The backend takes the close status from it.""" + # Only Temporal's cancel converts to a cancel failure, which closes the Stream as canceled. + if isinstance(failure, asyncio.CancelledError): + return temporalio.exceptions.CancelledError(str(failure) or "Cancelled") + return failure + + +@dataclass(frozen=True) +class Closed: + """The last item of a closed Stream. + + .. warning:: + This API is experimental. + """ + + offset: int + """Position of the close in the Stream.""" + + result: Any | None + """Result of the close, decoded with no type hint. None if it has none.""" + + failure: BaseException | None + """Failure of the close. None if it has none.""" + + status: CloseStatus = CloseStatus.COMPLETED + """How the Stream closed. + + A reader can tell a cancelled or terminated owner from a failed one without checking the + failure's type. + """ + + +# Fields compare in the order of the proto. An empty or negative field compares as zero, as in +# the crate's epoch keys. +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)) + + +def _task(epoch: models.ProducerEpoch) -> tuple[int, int, tuple[int, int]]: + return ( + _n(epoch.task_started_event_id), + _n(epoch.task_attempt), + _time(epoch.task_started_time), + ) + + +def _epoch(epoch: models.ProducerEpoch) -> tuple[Any, ...]: + return (*_task(epoch), _time(epoch.attempt_scheduled_time), _n(epoch.attempt)) + + +def _completed(task: models.WorkflowTaskEpoch) -> tuple[int, int, tuple[int, int]]: + return (_n(task.started_event_id), _n(task.attempt), _time(task.started_time)) + + +def producer_key(producer: models.Producer) -> tuple[Any, ...]: + """The Producer fields that :py:func:`withdraws` reads, so equal keys are withdrawn together.""" + epoch = _epoch(producer.epoch) if producer.epoch is not None else None + return (producer.kind or 0, producer.id, producer.run_id or "", epoch) + + +def _of_run(producer: models.Producer, facts: models.StreamSuperseded) -> bool: + if producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_WORKFLOW: + return producer.id == facts.producer_id + if producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_LOCAL_ACTIVITY: + return (producer.run_id or "") == (facts.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. + """ + if producer is None: + return False + if facts.withdraw_run and producer.run_id == facts.withdraw_run: + return True + if facts.producer_kind == StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT: + # A client's newer epoch only fences. Its earlier Records are real messages, not a + # failed attempt's. + return False + if facts.epoch is None or producer.epoch is None: + return False + if facts.producer_kind != StreamProducerKind.STREAM_PRODUCER_KIND_WORKFLOW: + return ( + (producer.kind or 0) == (facts.producer_kind or 0) + and producer.id == facts.producer_id + and _epoch(producer.epoch) < _epoch(facts.epoch) + ) + if not _of_run(producer, facts): + return False + task = _task(facts.epoch) + old = _task(producer.epoch) + started = producer.epoch.task_started_event_id or 0 + if started in (facts.failed_tasks or []): + return True + completed = facts.previous_completed_task + after = (completed.started_event_id or 0) if completed is not None else 0 + if after < started < (facts.epoch.task_started_event_id or 0): + return True + if started == (facts.epoch.task_started_event_id or 0) and old < task: + return True + if completed is None or started != (completed.started_event_id or 0): + return False + if (completed.attempt or 0) > 0: + return old < _completed(completed) + # Core sends the completed task with attempt 0 when it no longer knows the attempt. A task's + # attempts start in time order, so the started time still orders them. + return old[2] < _time(completed.started_time) diff --git a/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_handles.py b/tests/contrib/streams/test_handles.py new file mode 100644 index 000000000..1736cf844 --- /dev/null +++ b/tests/contrib/streams/test_handles.py @@ -0,0 +1,179 @@ +import json +import subprocess +import sys +import types +from dataclasses import dataclass +from typing import Any + +import pydantic +import pytest +from google.protobuf import json_format + +import temporalio.api.common.v1 +import temporalio.api.enums.v1 +import temporalio.api.stream.v1 +from temporalio import common +from temporalio.common import RawValue +from temporalio.common._streams._defn import DefinitionInfo, definition_of +from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.converter import DataConverter + + +@dataclass +class Turn: + kind: str + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn +class Raw(common.StreamDefinition[RawValue]): ... + + +OWNED: common.StreamHandle[Turn] = common.StreamHandle( + id="wf:run:turns", + workflow_id="wf", + first_execution_run_id="run", + backend="redis", + payload_type=Turn, +) + + +def test_definition_has_its_name_and_payload_type(): + assert definition_of(Turns) == DefinitionInfo(name="turns", type=Turn) + assert definition_of(Raw) == DefinitionInfo(name="Raw", type=RawValue) + + +def test_subclass_of_a_definition_keeps_the_payload_type(): + class Base(common.StreamDefinition[int]): ... + + @common.defn + class Sub(Base): ... + + assert definition_of(Sub).type is int + + +def test_definition_needs_the_base_class_and_a_payload_type(): + with pytest.raises(TypeError, match="must subclass StreamDefinition"): + + @common.defn + class Plain: # pyright: ignore[reportUnusedClass] + pass + + with pytest.raises(TypeError, match="no payload type"): + + @common.defn + class Bare(common.StreamDefinition): # pyright: ignore[reportUnusedClass] + pass + + +def test_class_without_defn_is_not_a_definition(): + class Unmarked(common.StreamDefinition[int]): ... + + with pytest.raises(TypeError, match="not a Stream definition"): + definition_of(Unmarked) + + +def test_handle_is_the_proto_json_of_the_handle_message(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + message = temporalio.api.stream.v1.StreamHandle( + id="wf:run: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 json.loads(converter.to_payload(common.StreamHandle(id="s")).data) == { + "id": "s" + } + + +def test_type_hint_gives_the_payload_type(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + decoded = converter.from_payload(payload, common.StreamHandle[Turn]) + assert type(decoded) is common.StreamHandle + assert decoded == OWNED + assert decoded._payload_type is Turn + listed = converter.from_payload(payload, common.StreamHandle[list[str]]) + assert listed._payload_type == list[str] + bare = converter.from_payload(payload, common.StreamHandle) + assert bare._payload_type is None and bare.backend == "redis" + + +@dataclass +class CallModelInput: + prompt: str + stream: common.StreamHandle[Turn] + + +def test_handle_in_a_dataclass_field(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(CallModelInput("hi", OWNED)) + assert json.loads(payload.data)["stream"]["owner"]["businessId"] == "wf" + decoded = converter.from_payload(payload, CallModelInput) + assert decoded.stream == OWNED and decoded.stream._payload_type is Turn + + +class CallModelModel(pydantic.BaseModel): + prompt: str + stream: common.StreamHandle[Turn] + + +def test_handle_in_a_pydantic_field(): + model = CallModelModel(prompt="hi", stream=OWNED) + dumped = json.loads(model.model_dump_json()) + assert dumped["stream"] == { + "id": "wf:run:turns", + "owner": { + "type": "EXECUTION_TYPE_WORKFLOW", + "businessId": "wf", + "runId": "run", + }, + "backend": "redis", + } + loaded = CallModelModel.model_validate_json(model.model_dump_json()) + assert loaded.stream == OWNED and loaded.stream._payload_type is Turn + converter = pydantic_data_converter.payload_converter + payload = converter.to_payload(model) + decoded = converter.from_payload(payload, CallModelModel) + assert decoded.stream == OWNED + + +def test_any_type_argument_has_no_payload_type(): + converter = DataConverter.default.payload_converter + decoded = converter.from_payload( + converter.to_payload(OWNED), common.StreamHandle[Any] + ) + assert type(decoded) is common.StreamHandle and decoded._payload_type is None + assert decoded.id == OWNED.id + + +def test_temporalio_common_exports_the_streams_api(): + # temporalio.common imports the Stream names on first use. Each is the private package's. + import temporalio.common._streams as streams_api + + names = { + name + for name, value in vars(streams_api).items() + if not name.startswith("_") and not isinstance(value, types.ModuleType) + } + assert names == common._STREAM_NAMES and names <= set(common.__all__) + for name in common._STREAM_NAMES: + assert getattr(common, name) is getattr(streams_api, name) + + +def test_import_of_temporalio_common_loads_no_stream_code(): + # The Stream modules import temporalio.client and temporalio.workflow, which import + # temporalio.common. + code = "import sys, temporalio.common; print('temporalio.common._streams' in sys.modules)" + loaded = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ).stdout + assert loaded.strip() == "False" diff --git a/tests/contrib/streams/test_vectors.py b/tests/contrib/streams/test_vectors.py new file mode 100644 index 000000000..056ffcb89 --- /dev/null +++ b/tests/contrib/streams/test_vectors.py @@ -0,0 +1,122 @@ +from pathlib import Path +from typing import Any + +import pytest + +from temporalio.common._streams import _ids, _items +from temporalio.common._streams.system import models +from tests.contrib.streams._support import crate_file, vectors + +_KINDS = { + "workflow": models.StreamProducerKind.StreamProducerKindWorkflow, + "local_activity": models.StreamProducerKind.StreamProducerKindLocalActivity, + "activity": models.StreamProducerKind.StreamProducerKindActivity, + "client": models.StreamProducerKind.StreamProducerKindClient, + "standalone_activity": models.StreamProducerKind.StreamProducerKindStandaloneActivity, +} + +_SCRIPTS = Path(__file__).parents[3] / "scripts" / "_proto" + + +@pytest.mark.parametrize( + ("crate", "relative"), + [ + ("temporalio-streams", "temporal/api/enums/v1/stream.proto"), + ("temporalio-streams", "temporal/api/stream/v1/message.proto"), + ("temporalio-streams", "temporal/api/streamservice/v1/request_response.proto"), + ("temporalio-streams", "temporal/sdk/streams/v1/runtime.proto"), + ("temporalio-streams-redis", "temporal/sdk/streams/v1/redis.proto"), + ], +) +def test_vendored_protos_match_the_crates(crate: str, relative: str): + expected = crate_file(f"protos/{relative}", crate).read_text(encoding="utf-8") + assert (_SCRIPTS / relative).read_text(encoding="utf-8") == expected + + +def test_stream_ids_match_the_crate(): + data: Any = vectors("ids.json") + for case in data["streams"]: + # Every Stream has a name, so the crate's cases with no name don't apply. + if case["name"] is None: + continue + built = _ids.scoped_stream_id( + case["workflow_id"], case["first_run_id"], case["name"] + ) + assert built == case["stream_id"] + + +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 _epoch(value: dict[str, Any] | None) -> models.ProducerEpoch | None: + if value is None: + return None + return models.ProducerEpoch( + task_started_event_id=value.get("task_started_event_id"), + task_attempt=value.get("task_attempt"), + task_started_time=_time(value.get("task_started_time")), + attempt_scheduled_time=_time(value.get("attempt_scheduled_time")), + attempt=value.get("attempt"), + ) + + +def test_withdrawals_match_the_crate(): + cases: Any = vectors("withdrawals.json") + assert cases + for case in cases: + record = case["record"]["producer"] + producer = models.Producer( + kind=_KINDS[record["kind"]], + id=record["id"], + run_id=record.get("run_id"), + epoch=_epoch(record.get("epoch")), + ) + given = case["superseded"] + completed = given.get("previous_completed_task") + facts = models.StreamSuperseded( + producer_kind=_KINDS[given["producer_kind"]], + producer_id=given["producer_id"], + epoch=_epoch(given.get("epoch")), + previous_completed_task=models.WorkflowTaskEpoch( + started_event_id=completed.get("started_event_id"), + attempt=completed.get("attempt"), + started_time=_time(completed.get("started_time")), + ) + if completed + else None, + failed_tasks=given.get("failed_tasks", []), + withdraw_run=given.get("withdraw_run"), + run_id=given.get("run_id"), + ) + assert _items.withdraws(facts, producer) == case["withdraws"], case["name"]