vercel-queue 0.7.0__py3-none-any.whl

vercel/queue/_internal/polling.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from typing import Any
+
+import concurrent.futures
+import threading
+from collections.abc import Callable, Iterator
+
+import anyio
+
+from .errors import InvalidLimitError
+from .names import SanitizedName
+from .subscribers import call_subscriber, call_subscriber_sync, poll_targets_for_subscriber
+from .types import Duration, StrContainer, duration_to_float_seconds
+
+MAX_POLL_LIMIT = 10
+
+
+def _request_limit(limit: int | None) -> int:
+    if limit is None:
+        return MAX_POLL_LIMIT
+    if limit < 1 or limit > MAX_POLL_LIMIT:
+        raise InvalidLimitError(limit)
+    return limit
+
+
+def _idle_interval_seconds(interval: Duration) -> float:
+    seconds = duration_to_float_seconds(interval)
+    if seconds < 0:
+        raise ValueError("interval must be non-negative")
+    return seconds
+
+
+async def _poll_topic_batch_async(
+    *,
+    poll: Callable[..., Any],
+    subscriber: Callable[..., Any],
+    topic: str,
+    consumer_group: SanitizedName,
+    request_limit: int,
+    lease_duration: Duration | None,
+) -> int:
+    handled = 0
+    async for delivery in poll(
+        topic,
+        consumer_group,
+        limit=request_limit,
+        lease_duration=lease_duration,
+    ):
+        async with delivery as message:
+            await call_subscriber(subscriber, message)
+        handled += 1
+    return handled
+
+
+def _poll_topic_batch_sync(
+    *,
+    poll: Callable[..., Iterator[Any]],
+    subscriber: Callable[..., Any],
+    topic: str,
+    consumer_group: SanitizedName,
+    request_limit: int,
+    lease_duration: Duration | None,
+    stop: threading.Event,
+) -> int:
+    handled = 0
+    for delivery in poll(
+        topic,
+        consumer_group,
+        limit=request_limit,
+        lease_duration=lease_duration,
+    ):
+        if stop.is_set():
+            break
+        with delivery as message:
+            call_subscriber_sync(subscriber, message)
+        handled += 1
+    return handled
+
+
+async def poll_and_handle_async(
+    *,
+    poll: Callable[..., Any],
+    subscriber: Callable[..., Any],
+    topics: StrContainer | None,
+    interval: Duration,
+    limit: int | None,
+    lease_duration: Duration | None,
+) -> None:
+    targets = poll_targets_for_subscriber(subscriber, topics)
+    request_limit = _request_limit(limit)
+    idle_seconds = _idle_interval_seconds(interval)
+
+    while True:
+        handled_any = False
+        for topic, consumer_group in targets:
+            while True:
+                handled_batch = await _poll_topic_batch_async(
+                    poll=poll,
+                    subscriber=subscriber,
+                    topic=topic,
+                    consumer_group=consumer_group,
+                    request_limit=request_limit,
+                    lease_duration=lease_duration,
+                )
+                handled_any = handled_any or handled_batch > 0
+                if limit is not None or handled_batch == 0:
+                    break
+        if not handled_any:
+            await anyio.sleep(idle_seconds)
+
+
+def run_poll_and_handle_sync(
+    *,
+    poll: Callable[..., Iterator[Any]],
+    subscriber: Callable[..., Any],
+    topics: StrContainer | None,
+    interval: Duration,
+    limit: int | None,
+    lease_duration: Duration | None,
+    stop: threading.Event,
+) -> None:
+    targets = poll_targets_for_subscriber(subscriber, topics)
+    request_limit = _request_limit(limit)
+    idle_seconds = _idle_interval_seconds(interval)
+
+    while not stop.is_set():
+        handled_any = False
+        for topic, consumer_group in targets:
+            if stop.is_set():
+                break
+            while not stop.is_set():
+                handled_batch = _poll_topic_batch_sync(
+                    poll=poll,
+                    subscriber=subscriber,
+                    topic=topic,
+                    consumer_group=consumer_group,
+                    request_limit=request_limit,
+                    lease_duration=lease_duration,
+                    stop=stop,
+                )
+                handled_any = handled_any or handled_batch > 0
+                if limit is not None or handled_batch == 0:
+                    break
+        if not handled_any:
+            stop.wait(idle_seconds)
+
+
+class PollingFuture(concurrent.futures.Future[None]):
+    def __init__(self) -> None:
+        super().__init__()
+        self._stop = threading.Event()
+        self._cancel_requested = False
+
+    @property
+    def stop_event(self) -> threading.Event:
+        return self._stop
+
+    def cancel(self) -> bool:
+        if self.done():
+            return False
+        self._cancel_requested = True
+        self._stop.set()
+        return True
+
+    def cancelled(self) -> bool:
+        return self._cancel_requested and self.done()
+
+    def result(self, timeout: float | None = None) -> None:
+        result = super().result(timeout)
+        if self._cancel_requested:
+            raise concurrent.futures.CancelledError
+        return result
+
+    def exception(self, timeout: float | None = None) -> BaseException | None:
+        if self._cancel_requested and self.done():
+            raise concurrent.futures.CancelledError
+        return super().exception(timeout)
+
+
+def start_sync_polling_thread(
+    target: Callable[[threading.Event], None],
+    *,
+    name: str,
+) -> PollingFuture:
+    future = PollingFuture()
+
+    def run() -> None:
+        try:
+            target(future.stop_event)
+        except BaseException as exc:  # noqa: BLE001
+            if not future.done():
+                future.set_exception(exc)
+        else:
+            if not future.done():
+                future.set_result(None)
+
+    thread = threading.Thread(target=run, name=name, daemon=True)
+    thread.start()
+    return future

vercel/queue/_internal/push.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+from typing import Generic, TypeGuard, TypeVar, cast
+
+import json
+from collections.abc import AsyncIterable, AsyncIterator
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+from .asynctools import iter_bytes_async
+from .config import validate_region
+from .constants import (
+    CLOUD_EVENT_HEADER_TYPE,
+    CLOUD_EVENT_HEADER_VQS_CONSUMER_GROUP,
+    CLOUD_EVENT_HEADER_VQS_CREATED_AT,
+    CLOUD_EVENT_HEADER_VQS_DELIVERY_COUNT,
+    CLOUD_EVENT_HEADER_VQS_EXPIRES_AT,
+    CLOUD_EVENT_HEADER_VQS_MESSAGE_ID,
+    CLOUD_EVENT_HEADER_VQS_RECEIPT_HANDLE,
+    CLOUD_EVENT_HEADER_VQS_REGION,
+    CLOUD_EVENT_HEADER_VQS_TOPIC,
+    CLOUD_EVENT_HEADER_VQS_VISIBILITY_DEADLINE,
+    CLOUD_EVENT_TYPE_V2BETA,
+    CONTENT_TYPE_JSON,
+    CONTENT_TYPE_OCTET_STREAM,
+    HEADER_CONTENT_TYPE,
+)
+from .errors import ProtocolError
+from .http import (
+    AsyncHttpMessage,
+    AsyncPushDeliveryBody,
+    AsyncPushDeliveryInput,
+    HttpResponse,
+    PushDeliveryBody,
+    PushDeliveryInput,
+    headers_from_raw,
+    response_headers,
+)
+from .log import debug_log
+from .messages import parse_optional_datetime, parse_required_datetime
+from .names import SanitizedName
+from .types import (
+    Headers,
+    Message,
+    MessageMetadata,
+    RawHeaders,
+    ReceiptHandle,
+    Transport,
+)
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True, kw_only=True)
+class ParsedPushDelivery(Generic[T]):
+    metadata: MessageMetadata
+    message: Message[T] | None
+
+
+async def parse_push_delivery(
+    raw_body: PushDeliveryBody,
+    headers: RawHeaders | None = None,
+    *,
+    transport: Transport[T] | None = None,
+) -> ParsedPushDelivery[T]:
+    headers = headers_from_raw(headers or {})
+    metadata = _parse_push_delivery_metadata(headers)
+    debug_log(
+        "push.delivery_metadata",
+        metadata=_debug_metadata(metadata),
+    )
+    if not metadata.receipt_handle:
+        return ParsedPushDelivery(metadata=metadata, message=None)
+    # v2beta inline callbacks send the stored payload bytes directly; there is
+    # no JSON CloudEvent envelope to unwrap.
+    message = Message(
+        payload=await _deserialize_push_delivery(
+            raw_body,
+            metadata.content_type or CONTENT_TYPE_OCTET_STREAM,
+            transport,
+        ),
+        metadata=metadata,
+    )
+    return ParsedPushDelivery(
+        metadata=metadata,
+        message=message,
+    )
+
+
+def accept_input_sync(
+    raw_body_or_response: PushDeliveryBody | HttpResponse,
+    headers: RawHeaders | None,
+) -> tuple[PushDeliveryBody, Headers]:
+    return _accept_input_from_response(
+        raw_body_or_response,
+        headers,
+        allow_async_response=False,
+    )
+
+
+def accept_input(
+    raw_body_or_response: AsyncPushDeliveryInput,
+    headers: RawHeaders | None,
+) -> tuple[AsyncPushDeliveryBody, Headers]:
+    raw_body, resolved_headers = _accept_input_from_response(
+        raw_body_or_response,
+        headers,
+        allow_async_response=True,
+    )
+    return cast("AsyncPushDeliveryBody", raw_body), resolved_headers
+
+
+def _parse_push_delivery_metadata(headers: Headers) -> MessageMetadata:
+    event_type = headers.get(CLOUD_EVENT_HEADER_TYPE, "")
+    if event_type != CLOUD_EVENT_TYPE_V2BETA:
+        raise ValueError(
+            f"Invalid CloudEvent type: expected {CLOUD_EVENT_TYPE_V2BETA!r}, got {event_type!r}"
+        )
+
+    topic = headers.get(CLOUD_EVENT_HEADER_VQS_TOPIC, "")
+    consumer_group = headers.get(CLOUD_EVENT_HEADER_VQS_CONSUMER_GROUP, "")
+    message_id = headers.get(CLOUD_EVENT_HEADER_VQS_MESSAGE_ID, "")
+    receipt_handle = headers.get(CLOUD_EVENT_HEADER_VQS_RECEIPT_HANDLE, "")
+    if not consumer_group:
+        raise ProtocolError("push delivery metadata must include consumer_group")
+    if not topic or not message_id:
+        raise ProtocolError("missing required ce-vqs* headers")
+
+    region = validate_region(headers.get(CLOUD_EVENT_HEADER_VQS_REGION))
+    if not receipt_handle:
+        return MessageMetadata(
+            message_id=message_id,
+            delivery_count=1,
+            created_at=datetime.now(timezone.utc),
+            topic=topic,
+            consumer_group=SanitizedName(consumer_group),
+            receipt_handle=None,
+            content_type=headers.get(HEADER_CONTENT_TYPE),
+            region=region,
+        )
+
+    delivery_count_raw = headers.get(CLOUD_EVENT_HEADER_VQS_DELIVERY_COUNT, "1")
+    try:
+        delivery_count = int(delivery_count_raw)
+    except ValueError:
+        delivery_count = 1
+
+    content_type = headers.get(HEADER_CONTENT_TYPE, CONTENT_TYPE_OCTET_STREAM)
+    return MessageMetadata(
+        message_id=message_id,
+        delivery_count=delivery_count,
+        created_at=parse_required_datetime(
+            headers.get(CLOUD_EVENT_HEADER_VQS_CREATED_AT),
+            CLOUD_EVENT_HEADER_VQS_CREATED_AT,
+        ),
+        expires_at=parse_optional_datetime(
+            headers.get(CLOUD_EVENT_HEADER_VQS_EXPIRES_AT),
+            CLOUD_EVENT_HEADER_VQS_EXPIRES_AT,
+        ),
+        topic=topic,
+        consumer_group=SanitizedName(consumer_group),
+        receipt_handle=ReceiptHandle(receipt_handle),
+        content_type=content_type,
+        region=region,
+        visibility_deadline=parse_optional_datetime(
+            headers.get(CLOUD_EVENT_HEADER_VQS_VISIBILITY_DEADLINE),
+            CLOUD_EVENT_HEADER_VQS_VISIBILITY_DEADLINE,
+        ),
+    )
+
+
+def _debug_metadata(metadata: MessageMetadata) -> dict[str, object]:
+    return {
+        "message_id": metadata.message_id,
+        "delivery_count": metadata.delivery_count,
+        "created_at": metadata.created_at,
+        "expires_at": metadata.expires_at,
+        "topic": metadata.topic,
+        "consumer_group": str(metadata.consumer_group),
+        "receipt_handle": metadata.receipt_handle,
+        "content_type": metadata.content_type,
+        "region": metadata.region,
+        "visibility_deadline": metadata.visibility_deadline,
+    }
+
+
+def parse_push_delivery_metadata(headers: RawHeaders) -> MessageMetadata:
+    return _parse_push_delivery_metadata(headers_from_raw(headers))
+
+
+async def _deserialize_push_delivery(
+    payload: PushDeliveryBody,
+    content_type: str,
+    transport: Transport[T] | None,
+) -> T:
+    chunks = _push_delivery_chunks(payload)
+    if transport is not None:
+        return await transport.deserialize(chunks, content_type=content_type)
+    buffered = await _collect_push_delivery_body(chunks)
+    if CONTENT_TYPE_JSON in content_type.lower():
+        return json.loads(buffered.decode("utf-8"))
+    return cast("T", buffered)
+
+
+def _push_delivery_chunks(
+    payload: PushDeliveryBody,
+) -> AsyncIterator[bytes]:
+    if isinstance(payload, bytes):
+        return _one_chunk_async(payload)
+    if isinstance(payload, AsyncIterable):
+        return aiter(cast("AsyncIterable[bytes]", payload))
+    return iter_bytes_async(payload)
+
+
+async def _collect_push_delivery_body(payload: AsyncIterable[bytes]) -> bytes:
+    chunks = bytearray()
+    async for chunk in payload:
+        chunks.extend(chunk)
+    return bytes(chunks)
+
+
+async def _one_chunk_async(payload: bytes) -> AsyncIterator[bytes]:
+    yield payload
+
+
+def _accept_input_from_response(
+    raw_body_or_response: PushDeliveryInput | AsyncPushDeliveryInput,
+    headers: RawHeaders | None,
+    *,
+    allow_async_response: bool,
+) -> tuple[PushDeliveryBody, Headers]:
+    if headers is not None:
+        if allow_async_response and not _is_async_push_delivery_body(raw_body_or_response):
+            raise TypeError("async accept() requires bytes or an async byte iterable")
+        return cast("PushDeliveryBody", raw_body_or_response), headers_from_raw(headers)
+    if allow_async_response and isinstance(raw_body_or_response, AsyncHttpMessage):
+        return raw_body_or_response.aiter_bytes(), response_headers(raw_body_or_response)
+    if allow_async_response:
+        raise TypeError("async accept() requires headers or an async HTTP response")
+    if isinstance(raw_body_or_response, HttpResponse):
+        return raw_body_or_response.iter_bytes(), response_headers(raw_body_or_response)
+    raise TypeError("accept() missing required headers")
+
+
+def _is_async_push_delivery_body(value: object) -> TypeGuard[AsyncPushDeliveryBody]:
+    return isinstance(value, (bytes, AsyncIterable))

vercel/queue/_internal/response.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+
+from .errors import (
+    BadRequestError,
+    ConsumerDiscoveryError,
+    ConsumerRegistryNotConfiguredError,
+    DuplicateIdempotencyKeyError,
+    ForbiddenError,
+    InternalServerError,
+    MessageAlreadyProcessedError,
+    MessageLeaseExpiredError,
+    MessageLockedError,
+    MessageNotFoundError,
+    MessageNotInFlightError,
+    MessageUnavailableError,
+    ReceiptHandleMismatchError,
+    ServiceError,
+    ThrottledError,
+    UnauthorizedError,
+)
+from .http import AsyncHttpResponse, parse_retry_after, response_text
+from .types import MessageID
+
+_QUEUE_STATUS_ERRORS = {
+    401: UnauthorizedError,
+    403: ForbiddenError,
+}
+
+
+async def queue_response_error(
+    response: AsyncHttpResponse,
+    *,
+    message_id: MessageID | None = None,
+) -> Exception | None:
+    if response.status_code == 400:
+        return BadRequestError(await response_text(response) or "Invalid parameters")
+    if error_type := _QUEUE_STATUS_ERRORS.get(response.status_code):
+        return error_type()
+    if response.status_code == 404 and message_id:
+        return MessageNotFoundError(message_id)
+    if response.status_code == 409 and message_id:
+        return await _conflict_error(response, message_id)
+    if response.status_code == 410 and message_id:
+        return MessageAlreadyProcessedError(message_id)
+    if response.status_code == 429:
+        return ThrottledError(parse_retry_after(response))
+    if response.status_code >= 500:
+        return InternalServerError(
+            await response_text(response) or f"Server error: {response.status_code}"
+        )
+    if response.status_code >= 300:
+        return ServiceError(
+            response.status_code,
+            await response_text(response) or f"Unexpected queue response: {response.status_code}",
+        )
+    return None
+
+
+async def send_response_error(response: AsyncHttpResponse) -> Exception | None:
+    if response.status_code == 409:
+        return DuplicateIdempotencyKeyError("Duplicate idempotency key detected")
+    if response.status_code == 502:
+        return ConsumerDiscoveryError(await response_text(response) or None)
+    if response.status_code == 503:
+        return ConsumerRegistryNotConfiguredError(await response_text(response) or None)
+    return await queue_response_error(response)
+
+
+async def _conflict_error(response: AsyncHttpResponse, message_id: MessageID) -> Exception:
+    body = await _response_body(response)
+    retry_after = parse_retry_after(response)
+    if body.original_message_id:
+        return MessageUnavailableError(
+            message_id,
+            f"originalMessageId={body.original_message_id}",
+            original_message_id=body.original_message_id,
+        )
+    if body.error == "Message is not currently in-flight":
+        return MessageNotInFlightError(message_id, retry_after)
+    if body.error == "Message lease has expired":
+        return MessageLeaseExpiredError(message_id, retry_after)
+    if body.error == "Receipt handle does not match current lease holder":
+        return ReceiptHandleMismatchError(message_id, retry_after)
+    return MessageLockedError(message_id, retry_after, body.error)
+
+
+@dataclass(frozen=True)
+class _ResponseBody:
+    error: str | None = None
+    original_message_id: MessageID | None = None
+
+
+async def _response_body(response: AsyncHttpResponse) -> _ResponseBody:
+    text = await response_text(response)
+    if not text:
+        return _ResponseBody()
+    try:
+        data = json.loads(text)
+    except ValueError:
+        return _ResponseBody(error=text)
+    if not isinstance(data, dict):
+        return _ResponseBody(error=text)
+    original_message_id = data.get("originalMessageId")
+    error = data.get("error")
+    return _ResponseBody(
+        error=str(error) if error else None,
+        original_message_id=MessageID(str(original_message_id)) if original_message_id else None,
+    )

vercel/queue/_internal/retry.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import TypeVar
+
+import time
+from collections.abc import Awaitable, Callable
+
+import anyio
+
+from .errors import RetryableError, ThrottledError
+from .log import debug_log
+
+T = TypeVar("T")
+
+_TRANSIENT_FOLLOW_UP_RETRY_SECONDS = 0.1
+DIRECTIVE_FOLLOW_UP_ATTEMPTS = 3
+
+
+async def retry_async_follow_up(
+    operation: Callable[[], Awaitable[T]],
+    *,
+    event_prefix: str = "follow_up",
+    sleep: Callable[[float], Awaitable[None]] | None = None,
+) -> T:
+    sleep = sleep or anyio.sleep
+    last_error: BaseException | None = None
+    for attempt in range(DIRECTIVE_FOLLOW_UP_ATTEMPTS):
+        try:
+            debug_log(f"{event_prefix}.retry_attempt", attempt=attempt + 1)
+            return await operation()
+        except Exception as exc:
+            if not _is_retryable_follow_up_error(exc):
+                raise
+            last_error = exc
+            if attempt == DIRECTIVE_FOLLOW_UP_ATTEMPTS - 1:
+                break
+            await sleep(_follow_up_retry_delay(exc))
+    if last_error is not None:
+        debug_log(
+            f"{event_prefix}.retry_exhausted",
+            attempts=DIRECTIVE_FOLLOW_UP_ATTEMPTS,
+            exception_class=last_error.__class__.__name__,
+            exception_message=str(last_error),
+        )
+        raise last_error
+    raise RuntimeError("retry loop exited without result or error")
+
+
+def retry_sync_follow_up(
+    operation: Callable[[], T],
+    *,
+    event_prefix: str = "follow_up",
+    sleep: Callable[[float], None] | None = None,
+) -> T:
+    sleep = sleep or time.sleep
+    last_error: BaseException | None = None
+    for attempt in range(DIRECTIVE_FOLLOW_UP_ATTEMPTS):
+        try:
+            debug_log(f"{event_prefix}.retry_attempt", attempt=attempt + 1)
+            return operation()
+        except Exception as exc:
+            if not _is_retryable_follow_up_error(exc):
+                raise
+            last_error = exc
+            if attempt == DIRECTIVE_FOLLOW_UP_ATTEMPTS - 1:
+                break
+            sleep(_follow_up_retry_delay(exc))
+    if last_error is not None:
+        debug_log(
+            f"{event_prefix}.retry_exhausted",
+            attempts=DIRECTIVE_FOLLOW_UP_ATTEMPTS,
+            exception_class=last_error.__class__.__name__,
+            exception_message=str(last_error),
+        )
+        raise last_error
+    raise RuntimeError("retry loop exited without result or error")
+
+
+def _is_retryable_follow_up_error(exc: BaseException) -> bool:
+    return isinstance(exc, RetryableError) and exc.retryable
+
+
+def _follow_up_retry_delay(exc: BaseException) -> float:
+    if isinstance(exc, ThrottledError) and exc.retry_after is not None:
+        return max(1.0, float(exc.retry_after))
+    return _TRANSIENT_FOLLOW_UP_RETRY_SECONDS