diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 578d714c6..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, @@ -32,6 +33,100 @@ import temporalio.api.workflow.v1 import temporalio.types +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 class RetryPolicy: @@ -1482,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 3888aaa78..517ce4888 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -1 +1,52 @@ """The Streams API, which :py:mod:`temporalio.common` exports.""" + +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 new file mode 100644 index 000000000..48d235233 --- /dev/null +++ b/temporalio/common/_streams/_defn.py @@ -0,0 +1,107 @@ +"""Stream definitions, which give the Stream name and its one payload type.""" + +from __future__ import annotations + +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..e9f751aaa --- /dev/null +++ b/temporalio/common/_streams/_errors.py @@ -0,0 +1,186 @@ +"""Errors of the Streams API in temporalio.common.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import temporalio.exceptions +import temporalio.service +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.api.stream.v1 import StreamErrorDetails + + +class StreamError(temporalio.exceptions.ApplicationError): + """Base class of every Stream error. + + Each Stream error is a non-retryable application error, except + :py:class:`StreamStorageError`. So an uncaught Stream error in Workflow code fails the + Workflow, and doesn't fail the Workflow Task again and again. A reason with no class of its + own, such as ``STREAM_ERROR_REASON_SEQUENCE``, gives this base class. + + .. warning:: + This API is experimental. + """ + + _retryable = False + + def __init__( + self, + message: str, + *, + reason: StreamErrorReason.ValueType = StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED, + code: temporalio.service.RPCStatusCode | None = None, + ) -> None: + """Create a Stream error.""" + # The reason goes in the details, so it comes back from a Local Activity. + super().__init__( + message, + *([int(reason)] if reason else []), + type=type(self).__name__, + non_retryable=not self._retryable, + ) + self._reason = reason + self._code = code + + @property + def reason(self) -> StreamErrorReason.ValueType: + """The ``StreamErrorReason`` of the cause. ``STREAM_ERROR_REASON_UNSPECIFIED`` if none.""" + return self._reason + + @property + def code(self) -> temporalio.service.RPCStatusCode | None: + """Status code of the failed call, as in gRPC. + + None if no call failed, or at a Workflow await, where the error comes back from a Local + Activity. + """ + return self._code + + +class StreamNotFoundError(StreamError): + """The Stream does not exist, or it was not created in time. + + A publish needs a Stream that a create made first. + """ + + +class StreamClosedError(StreamError): + """The Stream is closed, so it takes no more Records, and a create of it fails.""" + + +class StreamExpiredError(StreamError): + """Retention removed the Stream, or the Records at the offset that a read asked for.""" + + +class StreamFencedError(StreamError): + """A newer attempt of the same Producer wrote to the Stream, so this attempt can't.""" + + +class StreamStorageError(StreamError): + """The Store failed. A retry of the Activity may succeed.""" + + _retryable = True + + +class StreamOffsetError(StreamError): + """The offset is below zero or past the newest one plus one.""" + + +class StreamTooLargeError(StreamError): + """A Record or a batch is over the Store's limit.""" + + +class StreamUnsupportedError(StreamError): + """The Store lacks a capability that the call needs. + + The message ends with the name of the capability. + """ + + +class StreamNotOwnerError(StreamError): + """The caller isn't allowed to close the Stream or write as that Producer. + + Only the Workflow's own run chain can close a Workflow's Stream. Only a client can close a + client's Stream. + """ + + +class StreamSeqDivergentError(StreamError): + """A resent batch has the sequence number of the Producer's last batch, but other content. + + A writer that restarts with the same Producer id and epoch sends its first sequence number + again. Give it a higher epoch instead. + """ + + +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}." + ) + + +_BY_REASON: dict[StreamErrorReason.ValueType, type[StreamError]] = { + StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND: StreamNotFoundError, + StreamErrorReason.STREAM_ERROR_REASON_EXPIRED: StreamExpiredError, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED: StreamClosedError, + StreamErrorReason.STREAM_ERROR_REASON_FENCED: StreamFencedError, + StreamErrorReason.STREAM_ERROR_REASON_SEQ_DIVERGENT: StreamSeqDivergentError, + StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER: StreamNotOwnerError, + StreamErrorReason.STREAM_ERROR_REASON_OFFSET: StreamOffsetError, + StreamErrorReason.STREAM_ERROR_REASON_TOO_LARGE: StreamTooLargeError, + StreamErrorReason.STREAM_ERROR_REASON_STORAGE: StreamStorageError, + StreamErrorReason.STREAM_ERROR_REASON_UNSUPPORTED: StreamUnsupportedError, +} + + +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``. + """ + code, message, details = err.args + reason = StreamErrorDetails.FromString(details).reason if details else 0 + return from_reason(reason, message, temporalio.service.RPCStatusCode(code)) + + +def from_reason( + reason: int, + message: str, + code: temporalio.service.RPCStatusCode | None = None, +) -> StreamError: + """The Stream error for a ``StreamErrorReason``.""" + value = StreamErrorReason.ValueType(reason) + return _BY_REASON.get(value, StreamError)(message, reason=value, code=code) + + +_BY_TYPE: dict[str, type[StreamError]] = { + "StreamError": StreamError, + **{cls.__name__: cls for cls in _BY_REASON.values()}, +} + + +def from_activity_error(err: Exception) -> Exception: + """The Stream error that a write Local Activity raised, else ``err``.""" + cause = err.cause if isinstance(err, temporalio.exceptions.ActivityError) else err + if ( + isinstance(cause, temporalio.exceptions.ApplicationError) + and cause.type in _BY_TYPE + ): + reason = StreamErrorReason.ValueType( + int(cause.details[0]) if cause.details else 0 + ) + return _BY_TYPE[cause.type](cause.message, reason=reason) + return err diff --git a/temporalio/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py new file mode 100644 index 000000000..0c383442a --- /dev/null +++ b/temporalio/common/_streams/_handle.py @@ -0,0 +1,315 @@ +"""The shared part of the Stream handles of each context.""" + +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) + + +H = TypeVar("H", bound="StreamHandle[Any]") + + +# 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 ref. 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.StreamRef, as MessageToJson writes it, so every + # SDK reads a handle that another wrote. + ref = models.StreamRef( + 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._StreamRefTransferTypeConverter().to_transfer_type(ref) + ) + + +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)}") + if origin is StreamHandle: + raise TypeError( + "Give a Stream parameter the handle type of its context, such as " + "ActivityStream[Turn]" + ) + 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 ActivityStream[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]] +): + # Every handle type has the same JSON form, and the parameter's type hint picks the class. + # 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]): + """Base class of the Stream handle of each context. ``T`` is the Stream's payload type. + + A handle is data. It's passed as a Workflow or Activity argument in the proto JSON form of + the ``StreamRef`` message. The parameter's type, such as ``ActivityStream[Turn]``, sets the + handle class and 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 a context's create and get functions 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 + + @classmethod + def from_handle(cls: type[H], handle: StreamHandle[Any]) -> H: + """Return the same Stream as a handle of this class. + + For example, a Workflow converts a ``WorkflowStream`` to an ``ActivityStream`` for an + Activity's input, so the input has the type that the Activity declares. + """ + return cls( + id=handle.id, + workflow_id=handle.workflow_id, + first_execution_run_id=handle.first_execution_run_id, + backend=handle.backend, + payload_type=handle._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..3403dfc8b --- /dev/null +++ b/temporalio/common/_streams/_items.py @@ -0,0 +1,215 @@ +"""The items that a Stream read gives.""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +from temporalio.api.enums.v1 import StreamCloseStatus, StreamProducerKind +from temporalio.common._streams.system import models +from temporalio.exceptions import is_cancelled_exception + +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.""" + + supersession: models.Supersession + """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.supersession, 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_status(failure: BaseException | None) -> StreamCloseStatus.ValueType | None: + """The status a caller's close asks for. None lets the backend pick completed or failed.""" + if failure is not None and is_cancelled_exception(failure): + return StreamCloseStatus.STREAM_CLOSE_STATUS_CANCELED + return None + + +@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, supersession: models.Supersession) -> bool: + if producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_WORKFLOW: + return producer.id == supersession.producer_id + if producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_LOCAL_ACTIVITY: + return (producer.run_id or "") == (supersession.run_id or "") + return False + + +def withdraws( + supersession: models.Supersession, 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 supersession.withdraw_run and producer.run_id == supersession.withdraw_run: + return True + if supersession.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 supersession.epoch is None or producer.epoch is None: + return False + if supersession.producer_kind != StreamProducerKind.STREAM_PRODUCER_KIND_WORKFLOW: + return ( + (producer.kind or 0) == (supersession.producer_kind or 0) + and producer.id == supersession.producer_id + and _epoch(producer.epoch) < _epoch(supersession.epoch) + ) + if not _of_run(producer, supersession): + return False + task = _task(supersession.epoch) + old = _task(producer.epoch) + started = producer.epoch.task_started_event_id or 0 + if started in (supersession.failed_tasks or []): + return True + completed = supersession.previous_completed_task + after = (completed.started_event_id or 0) if completed is not None else 0 + if after < started < (supersession.epoch.task_started_event_id or 0): + return True + if started == (supersession.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..5a05d301c --- /dev/null +++ b/tests/contrib/streams/test_handles.py @@ -0,0 +1,183 @@ +import json +import subprocess +import sys +import types +from dataclasses import dataclass +from typing import Any, TypeVar + +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 + +T = TypeVar("T") + + +@dataclass +class Turn: + kind: str + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn +class Raw(common.StreamDefinition[RawValue]): ... + + +class PlainStream(common.StreamHandle[T]): + pass + + +OWNED: PlainStream[Turn] = PlainStream( + 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_ref(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + ref = temporalio.api.stream.v1.StreamRef( + 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(ref) + assert json.loads(converter.to_payload(PlainStream(id="s")).data) == {"id": "s"} + + +def test_type_hint_picks_the_class_and_the_payload_type(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + decoded = converter.from_payload(payload, PlainStream[Turn]) + assert type(decoded) is PlainStream + assert decoded == OWNED + assert decoded._payload_type is Turn + listed = converter.from_payload(payload, PlainStream[list[str]]) + assert listed._payload_type == list[str] + bare = converter.from_payload(payload, PlainStream) + assert bare._payload_type is None and bare.backend == "redis" + with pytest.raises(Exception, match="handle type of its context"): + converter.from_payload(payload, common.StreamHandle) + + +@dataclass +class CallModelInput: + prompt: str + stream: PlainStream[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: PlainStream[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), PlainStream[Any]) + assert type(decoded) is PlainStream 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..f1979ba74 --- /dev/null +++ b/tests/contrib/streams/test_vectors.py @@ -0,0 +1,124 @@ +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["supersession"] + completed = given.get("previous_completed_task") + supersession = models.Supersession( + 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(supersession, producer) == case["withdraws"], case[ + "name" + ]