Source code for sqlspec.extensions.events._payload

"""Shared payload encoding/decoding utilities for event backends."""

import contextlib
from datetime import datetime, timezone
from typing import Any

from sqlspec.exceptions import EventChannelError
from sqlspec.extensions.events._models import EventMessage
from sqlspec.utils.serializers import from_json, to_json
from sqlspec.utils.uuids import uuid4

__all__ = (
    "MAX_NOTIFY_BYTES",
    "coerce_dict",
    "coerce_optional_dict",
    "decode_notify_payload",
    "encode_notify_payload",
    "fits_notify_payload",
    "measure_notify_payload",
    "parse_event_timestamp",
)

MAX_NOTIFY_BYTES = 7999


def coerce_dict(value: Any) -> "dict[str, Any]":
    """Coerce a value to a dict, wrapping non-dict values as {'value': ...}."""
    return value if isinstance(value, dict) else {"value": value}


def coerce_optional_dict(value: Any) -> "dict[str, Any] | None":
    """Coerce a value to a dict or None, wrapping non-dict values as {'value': ...}."""
    return value if value is None or isinstance(value, dict) else {"value": value}


def _serialize_notify_envelope(
    event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None", published_at: "datetime"
) -> bytes:
    """Serialize a native notification envelope to UTF-8 JSON bytes.

    The publication timestamp is normalized to UTC with microsecond precision so
    the encoded envelope width is independent of the clock reading.
    """
    return to_json(
        {
            "event_id": event_id,
            "payload": payload,
            "metadata": metadata,
            "published_at": published_at.astimezone(timezone.utc).isoformat(timespec="microseconds"),
        },
        as_bytes=True,
    )


def encode_notify_payload(event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None") -> str:
    """Encode event data as JSON for NOTIFY payload.

    Raises:
        EventChannelError: If the encoded envelope exceeds the PostgreSQL notification budget.
    """
    encoded = _serialize_notify_envelope(event_id, payload, metadata, datetime.now(timezone.utc))
    encoded_bytes = len(encoded)
    if encoded_bytes > MAX_NOTIFY_BYTES:
        msg = (
            f"PostgreSQL NOTIFY payload is {encoded_bytes} encoded bytes and exceeds the "
            f"{MAX_NOTIFY_BYTES}-byte maximum. Use fits_notify_payload() or "
            "measure_notify_payload() to split the batch before publishing."
        )
        raise EventChannelError(msg)
    return encoded.decode("utf-8")
def measure_notify_payload( payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None, *, event_id: "str | None" = None ) -> int: """Return the encoded UTF-8 byte size of the native notification envelope. The measurement covers the complete envelope rather than only the payload mapping. Omitting ``event_id`` measures the canonical backend UUID-hex shape. """ resolved_event_id = uuid4().hex if event_id is None else event_id return len(_serialize_notify_envelope(resolved_event_id, payload, metadata, datetime.now(timezone.utc)))