vercel-queue 0.7.0__py3-none-any.whl

vercel/queue/_internal/transports.py

@@ -0,0 +1,363 @@
+from __future__ import annotations
+
+from typing import Any, ForwardRef, Generic, Literal, TypeAlias, TypeVar, cast
+
+import inspect
+import json
+from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterable
+from importlib import import_module
+
+from .constants import CONTENT_TYPE_JSON, CONTENT_TYPE_OCTET_STREAM, CONTENT_TYPE_TEXT
+from .errors import SubscriptionError
+from .streams import (
+    AsyncStreamPayload,
+    AsyncTextStreamPayload,
+    SyncStreamPayload,
+    SyncTextStreamPayload,
+)
+from .types import (
+    RequestContent,
+    Topic,
+    Transport,
+)
+from .typeutils import (
+    annotation_needs_resolution,
+    args,
+    is_classvar,
+    is_final,
+    is_type_var,
+    is_union_type,
+    origin_is,
+    strip_annotated,
+)
+
+T = TypeVar("T")
+
+TransportKind: TypeAlias = Literal[
+    "json",
+    "byte-buffer",
+    "byte-stream",
+    "text-buffer",
+    "text-stream",
+]
+
+
+class ByteBufferTransport:
+    content_type = CONTENT_TYPE_OCTET_STREAM
+
+    def serialize(self, value: bytes | bytearray | memoryview) -> bytes:
+        if isinstance(value, bytes):
+            return value
+        return bytes(value)
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> bytes:
+        return await _collect_bytes_async(payload)
+
+
+class TextBufferTransport:
+    content_type = CONTENT_TYPE_TEXT
+
+    def serialize(self, value: str) -> bytes:
+        return value.encode("utf-8")
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> str:
+        return (await _collect_bytes_async(payload)).decode("utf-8")
+
+
+class ByteStreamTransport:
+    """Transport that preserves receive payloads as one-shot byte streams.
+
+    This transport avoids buffering message bytes on receive and passes stream
+    payloads through to the HTTP client on send.
+    """
+
+    content_type = CONTENT_TYPE_OCTET_STREAM
+
+    def serialize(
+        self,
+        value: bytes
+        | bytearray
+        | memoryview
+        | Iterable[bytes]
+        | AsyncIterable[bytes]
+        | SyncStreamPayload
+        | AsyncStreamPayload,
+    ) -> RequestContent:
+        if isinstance(value, bytes):
+            return value
+        if isinstance(value, (bytearray, memoryview)):
+            return bytes(value)
+        if isinstance(value, AsyncIterable):
+            return value
+        return value
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> AsyncStreamPayload:
+        return AsyncStreamPayload(payload)
+
+
+class RawJsonTransport(Generic[T]):
+    content_type = CONTENT_TYPE_JSON
+
+    def __init__(
+        self,
+        *,
+        json_encoder: type[json.JSONEncoder] | None = None,
+        json_decoder: type[json.JSONDecoder] | None = None,
+    ) -> None:
+        self.json_encoder = json_encoder
+        self.json_decoder = json_decoder
+
+    def serialize(self, value: T) -> bytes:
+        if self.json_encoder is None:
+            return json.dumps(value).encode("utf-8")
+        return json.dumps(value, cls=self.json_encoder).encode("utf-8")
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> T:
+        text = (await _collect_bytes_async(payload)).decode("utf-8")
+        if self.json_decoder is None:
+            return cast("T", json.loads(text))
+        return cast("T", json.loads(text, cls=self.json_decoder))
+
+
+class TypedJsonTransport(Generic[T]):
+    content_type = CONTENT_TYPE_JSON
+
+    def __init__(self, model: type[T]) -> None:
+        self.model = model
+        try:
+            type_adapter = import_module("pydantic").TypeAdapter
+        except ImportError as exc:  # pragma: no cover - depends on optional extra
+            raise RuntimeError("Install vercel-queue[typed] to use TypedJsonTransport") from exc
+        self._adapter = type_adapter(model)
+
+    def serialize(self, value: T) -> bytes:
+        if hasattr(value, "model_dump_json"):
+            return cast("Any", value).model_dump_json().encode("utf-8")
+        return self._adapter.dump_json(value)
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> T:
+        return self._adapter.validate_json(await _collect_bytes_async(payload))
+
+
+class _ModelDumpJsonTransport:
+    content_type = CONTENT_TYPE_JSON
+
+    def __init__(self, model_dump_json: Callable[[], str]) -> None:
+        self._model_dump_json = model_dump_json
+
+    def serialize(self, value: object) -> bytes:
+        del value
+        return self._model_dump_json().encode("utf-8")
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> object:
+        return await RawJsonTransport[Any]().deserialize(payload, content_type=content_type)
+
+
+class TextStreamTransport:
+    """Transport that sends and receives one-shot UTF-8 text streams."""
+
+    content_type = CONTENT_TYPE_TEXT
+
+    def serialize(
+        self,
+        value: str
+        | Iterable[str]
+        | AsyncIterable[str]
+        | SyncTextStreamPayload
+        | AsyncTextStreamPayload,
+    ) -> RequestContent:
+        if isinstance(value, str):
+            return value.encode("utf-8")
+        if isinstance(value, AsyncIterable):
+            return _encode_text_async(cast("AsyncIterable[str]", value))
+        return _encode_text_sync(value)
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> AsyncTextStreamPayload:
+        return AsyncTextStreamPayload(payload)
+
+
+def infer_send_transport(payload: object) -> object:
+    if model_dump_json := _model_dump_json(payload):
+        return _ModelDumpJsonTransport(model_dump_json)
+    if isinstance(payload, (bytes, bytearray, memoryview)):
+        return ByteBufferTransport()
+    if isinstance(payload, str):
+        return TextBufferTransport()
+    if isinstance(payload, (SyncTextStreamPayload, AsyncTextStreamPayload)):
+        return TextStreamTransport()
+    if isinstance(payload, (SyncStreamPayload, AsyncStreamPayload)):
+        return ByteStreamTransport()
+    return RawJsonTransport[Any]()
+
+
+def is_untyped_payload_annotation(annotation: Any) -> bool:
+    annotation = strip_annotated(annotation)
+    return annotation in {inspect.Signature.empty, Any, object}
+
+
+def reject_invalid_payload_annotation(annotation: Any) -> None:
+    annotation = strip_annotated(annotation)
+    if is_untyped_payload_annotation(annotation):
+        return
+    if is_type_var(annotation) or is_classvar(annotation) or is_final(annotation):
+        raise SubscriptionError(f"unsupported queue subscriber payload annotation: {annotation!r}")
+    if annotation in {list, dict, tuple, set, frozenset}:
+        raise SubscriptionError(
+            f"unsupported bare queue subscriber payload annotation: {annotation!r}"
+        )
+    if isinstance(annotation, str | ForwardRef) or origin_is(annotation, Literal):
+        return
+    if is_union_type(annotation):
+        for item in args(annotation):
+            reject_invalid_payload_annotation(item)
+        return
+    if origin_is(annotation, Iterable, AsyncIterable):
+        stream_args = args(annotation)
+        if len(stream_args) != 1 or stream_args[0] not in {bytes, str}:
+            raise SubscriptionError(
+                "queue stream subscriber annotation must be Iterable[bytes], "
+                "Iterable[str], AsyncIterable[bytes], or AsyncIterable[str]"
+            )
+        return
+    for item in args(annotation):
+        if item is not Ellipsis:
+            reject_invalid_payload_annotation(item)
+
+
+def payload_transport_kind(annotation: Any) -> TransportKind:
+    annotation = strip_annotated(annotation)
+    if annotation is bytes:
+        return "byte-buffer"
+    if annotation is str:
+        return "text-buffer"
+    if annotation in {SyncStreamPayload, AsyncStreamPayload}:
+        return "byte-stream"
+    if annotation in {SyncTextStreamPayload, AsyncTextStreamPayload}:
+        return "text-stream"
+    if origin_is(annotation, Iterable, AsyncIterable):
+        stream_args = args(annotation)
+        if stream_args == (bytes,):
+            return "byte-stream"
+        if stream_args == (str,):
+            return "text-stream"
+    return "json"
+
+
+def transport_for_kind(kind: TransportKind) -> Transport[Any]:
+    if kind == "byte-buffer":
+        return ByteBufferTransport()
+    if kind == "byte-stream":
+        return ByteStreamTransport()
+    if kind == "text-buffer":
+        return TextBufferTransport()
+    if kind == "text-stream":
+        return TextStreamTransport()
+    return RawJsonTransport[Any]()
+
+
+def receive_transport_for_annotation(annotation: Any) -> Transport[Any]:
+    annotation = strip_annotated(annotation)
+    if is_untyped_payload_annotation(annotation):
+        return RawJsonTransport[Any]()
+    if annotation_needs_resolution(annotation):
+        raise SubscriptionError(f"unsupported queue subscriber payload annotation: {annotation!r}")
+    reject_invalid_payload_annotation(annotation)
+    kind = payload_transport_kind(annotation)
+    if kind == "json":
+        return TypedJsonTransport[Any](annotation)
+    return transport_for_kind(kind)
+
+
+def receive_transport_for_topic(topic: object) -> Transport[Any]:
+    if not isinstance(topic, Topic):
+        return RawJsonTransport[Any]()
+    if topic.transport is not None:
+        return topic.transport
+    if getattr(type(topic), "__topic_origin__", None) is not Topic:
+        return RawJsonTransport[Any]()
+    return receive_transport_for_annotation(type(topic).__topic_payload_type__)
+
+
+def send_transport_for_topic(topic: object) -> Transport[Any] | None:
+    if isinstance(topic, Topic):
+        return topic.transport
+    return None
+
+
+def _model_dump_json(value: object) -> Callable[[], str] | None:
+    model_dump_json = getattr(value, "model_dump_json", None)
+    if callable(model_dump_json):
+        return cast("Callable[[], str]", model_dump_json)
+    return None
+
+
+async def _collect_bytes_async(payload: AsyncIterator[bytes]) -> bytes:
+    chunks = bytearray()
+    async for chunk in payload:
+        chunks.extend(chunk)
+    return bytes(chunks)
+
+
+def _encode_text_sync(payload: Iterable[str]) -> Iterable[bytes]:
+    for chunk in payload:
+        yield chunk.encode("utf-8")
+
+
+async def _encode_text_async(payload: AsyncIterable[str]) -> AsyncIterator[bytes]:
+    async for chunk in payload:
+        yield chunk.encode("utf-8")
+
+
+# Only add public symbols to __all__; internal helpers must stay unexported.
+__all__ = (
+    "ByteBufferTransport",
+    "ByteStreamTransport",
+    "RawJsonTransport",
+    "TextBufferTransport",
+    "TextStreamTransport",
+    "TransportKind",
+    "TypedJsonTransport",
+    "infer_send_transport",
+    "is_untyped_payload_annotation",
+    "payload_transport_kind",
+    "receive_transport_for_annotation",
+    "receive_transport_for_topic",
+    "reject_invalid_payload_annotation",
+    "send_transport_for_topic",
+    "transport_for_kind",
+)

vercel/queue/_internal/types.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeGuard, TypeVar, cast
+
+from collections.abc import (
+    AsyncIterable,
+    AsyncIterator,
+    Iterable,
+    Iterator,
+    Mapping,
+)
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+
+from .constants import DEFAULT_RETRY_AFTER_SECONDS
+from .names import SanitizedName, validate_name, validate_topic_name
+
+T = TypeVar("T")
+_TYPE_VAR_TYPE = type(T)
+
+Duration: TypeAlias = int | float | timedelta
+MessageID: TypeAlias = str
+ReceiptHandle: TypeAlias = str
+RequestContent: TypeAlias = bytes | Iterable[bytes] | AsyncIterable[bytes]
+Headers: TypeAlias = Mapping[str, str]
+RawHeaders: TypeAlias = Mapping[str, str]
+
+
+class StrContainer(Protocol):
+    """Container of strings that excludes bare strings structurally."""
+
+    # Bare str has __contains__(str), while normal containers accept object.
+    # Requiring the wider signature lets type checkers reject str here.
+
+    def __iter__(self) -> Iterator[str]: ...
+
+    def __contains__(self, item: object, /) -> bool: ...
+
+
+def _is_duration(value: object) -> TypeGuard[Duration]:
+    return isinstance(value, (int, float, timedelta)) and not isinstance(value, bool)
+
+
+def duration_to_seconds(duration: Duration) -> int:
+    if isinstance(duration, timedelta):
+        return int(duration.total_seconds())
+    if isinstance(duration, int) and not isinstance(duration, bool):
+        return duration
+    if isinstance(duration, float):
+        return int(duration)
+    raise TypeError("duration must be an int or float number of seconds or datetime.timedelta")
+
+
+def duration_to_float_seconds(duration: Duration) -> float:
+    if isinstance(duration, timedelta):
+        return duration.total_seconds()
+    if isinstance(duration, int) and not isinstance(duration, bool):
+        return float(duration)
+    if isinstance(duration, float):
+        return duration
+    raise TypeError("duration must be an int or float number of seconds or datetime.timedelta")
+
+
+@dataclass(frozen=True, kw_only=True, eq=False)
+class Topic(Generic[T]):
+    """A named Vercel Queues topic.
+
+    Topics identify the stream that messages are sent to and received from.
+    """
+
+    name: SanitizedName
+    """Topic name to send to or receive from."""
+
+    transport: Transport[Any] | None = None
+    """Optional transport used when sending to or polling this topic."""
+
+    __topic_origin__: ClassVar[type[Topic[Any]] | None] = None
+    __topic_payload_type__: ClassVar[Any] = None
+    _specializations: ClassVar[dict[Any, type[Topic[Any]]]] = {}
+
+    def __class_getitem__(cls, params: Any) -> type[Topic[Any]]:
+        if isinstance(params, tuple):
+            if len(params) != 1:
+                raise TypeError("Topic expects exactly one type argument")
+            params = params[0]
+        if isinstance(params, _TYPE_VAR_TYPE):
+            return cast("type[Topic[Any]]", cls)
+
+        try:
+            return cls._specializations[params]
+        except KeyError:
+            pass
+
+        payload_repr = _topic_payload_type_repr(params)
+        specialization = type(
+            f"Topic[{payload_repr}]",
+            (cls,),
+            {
+                "__module__": cls.__module__,
+                "__topic_origin__": Topic,
+                "__topic_payload_type__": params,
+            },
+        )
+        cls._specializations[params] = specialization
+        return cast("type[Topic[Any]]", specialization)
+
+    def __init__(
+        self,
+        name: str | SanitizedName,
+        *,
+        transport: Transport[Any] | None = None,
+    ) -> None:
+        object.__setattr__(self, "name", SanitizedName(validate_topic_name(name)))
+        object.__setattr__(self, "transport", transport)
+
+    def __repr__(self) -> str:
+        return f"Topic(name={self.name!r})"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Topic):
+            return NotImplemented
+        return self.name == other.name
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+
+def _topic_payload_type_repr(payload_type: object) -> str:
+    name = getattr(payload_type, "__qualname__", None)
+    if isinstance(name, str):
+        return name
+    name = getattr(payload_type, "__name__", None)
+    if isinstance(name, str):
+        return name
+    return repr(payload_type)
+
+
+@dataclass(frozen=True, kw_only=True)
+class MessageMetadata:
+    """Metadata attached to a queue message delivery.
+
+    Identifiers and delivery tokens are opaque service values and should be
+    passed back unchanged when acknowledging or extending a message.
+    """
+
+    message_id: MessageID
+    """Opaque message ID assigned by the service."""
+
+    delivery_count: int
+    """Number of delivery attempts observed for the message."""
+
+    created_at: datetime
+    """Message creation timestamp."""
+
+    topic: str
+    """Topic name the message belongs to."""
+
+    consumer_group: SanitizedName
+    """Consumer group that owns this delivery."""
+
+    receipt_handle: ReceiptHandle | None = None
+    """Opaque delivery token used for follow-up operations."""
+
+    content_type: str | None = None
+    """Stored message content type."""
+
+    region: str | None = None
+    """Queue region for follow-up operations."""
+
+    expires_at: datetime | None = None
+    """Message expiration timestamp, when supplied by the service."""
+
+    visibility_deadline: datetime | None = None
+    """Current processing deadline, when supplied by the service."""
+
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self,
+            "consumer_group",
+            SanitizedName(validate_name(self.consumer_group, field="consumer_group")),
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class Message(Generic[T]):
+    """A deserialized payload plus metadata for this delivery."""
+
+    payload: T
+    """Message payload returned by the configured transport."""
+
+    metadata: MessageMetadata
+    """Metadata for this delivery."""
+
+    @property
+    def message_id(self) -> MessageID:
+        return self.metadata.message_id
+
+
+class QueueDirective(Exception):  # noqa: N818
+    """Directive raised by a queue subscriber.
+
+    Queue delivery is at least once. Subscriber directives let a handler choose
+    how the SDK resolves the delivery after user code runs.
+    """
+
+    reason: object | None
+
+    def __init__(self, reason: object | None = None) -> None:
+        self.reason = reason
+        super().__init__(str(reason) if reason is not None else "")
+
+
+class Handoff(QueueDirective):
+    """Directive that tells the SDK a delivery was handed off.
+
+    ``accept_and_handle`` normally owns the full push-delivery lifecycle: it
+    accepts the delivery, keeps the processing lease alive while subscribers
+    run, and acknowledges the message when all matching subscribers return
+    successfully. ``Handoff`` tells the SDK that subscriber code successfully
+    handed the delivery to another system with its own processing lifecycle.
+
+    Raising this directive stops subscriber dispatch and suppresses the SDK's
+    automatic follow-up action. The SDK also stops its lease renewal and leaves
+    the current queue lease open. External code must take over from that point:
+    it is responsible for any further lease renewal needed while processing
+    continues, and for eventually acknowledging the message or changing its
+    visibility using the original message metadata. If external code does none
+    of those things, the message becomes eligible for redelivery when the lease
+    expires.
+
+    This is distinct from ``RetryAfter``, which asks the SDK to update
+    visibility immediately, and from raising a normal exception, which reports
+    handler failure. For example, the Celery broker uses ``Handoff`` after a
+    Vercel Queue push delivery has been accepted into Kombu, because Celery
+    must remain responsible for the eventual ACK or reject by delivery tag.
+    """
+
+
+class RetryAfter(QueueDirective):
+    """Directive that retries a message after a delay.
+
+    Raising ``RetryAfter`` stops subscriber dispatch and asks the SDK to make
+    the message visible again after ``delay``. The SDK keeps lifecycle ownership
+    through that follow-up action: it stops lease renewal, updates the delivery
+    visibility using the current receipt handle, and reports success to the
+    push-delivery caller once that update succeeds. Use this when the subscriber
+    cannot process the message now, but wants the SDK to schedule the next
+    delivery attempt instead of treating the handler as failed.
+
+    A delay of zero makes the message eligible for immediate redelivery.
+
+    Args:
+        delay: Retry delay in seconds or as a ``datetime.timedelta``. Defaults
+            to 60 seconds.
+        reason: Optional diagnostic value stored on the directive.
+
+    """
+
+    timeout_seconds: int
+
+    def __init__(
+        self,
+        delay: Duration = DEFAULT_RETRY_AFTER_SECONDS,
+        reason: object | None = None,
+    ) -> None:
+        self.timeout_seconds = duration_to_seconds(delay)
+        if self.timeout_seconds < 0:
+            raise ValueError("delay must be non-negative")
+        super().__init__(reason)
+
+    def __repr__(self) -> str:
+        return f"RetryAfter(timeout_seconds={self.timeout_seconds!r})"
+
+
+class Transport(Protocol[T]):
+    content_type: str
+
+    def serialize(self, value: T) -> RequestContent: ...
+
+    async def deserialize(
+        self,
+        payload: AsyncIterator[bytes],
+        *,
+        content_type: str,
+    ) -> T: ...
+
+
+# Only add public symbols to __all__; internal helpers must stay unexported.
+__all__: tuple[str, ...] = (
+    "Duration",
+    "Handoff",
+    "Message",
+    "MessageID",
+    "MessageMetadata",
+    "QueueDirective",
+    "ReceiptHandle",
+    "RetryAfter",
+    "StrContainer",
+    "Topic",
+)