rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/serialization/json.py

@@ -0,0 +1,122 @@
+"""JSON serializer with Pydantic V2 support."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, fields, is_dataclass
+from typing import Any
+
+
+class JSONSerializer:
+    """Default JSON serializer.
+
+    Supports:
+    - dict/list → json.dumps/loads
+    - Pydantic V2 models → model_validate_json/model_dump_json
+    - dataclasses → asdict → json.dumps
+    - str → encode to UTF-8
+    - bytes → pass through
+
+    By default the serializer **raises** on objects ``json`` cannot represent
+    (e.g. ``datetime``, ``Decimal``) instead of silently coercing them via
+    ``str()``. Pass ``coerce_unknown_to_str=True`` to restore the legacy
+    ``default=str`` coercion behaviour.
+
+    H14 — decoding into a stdlib dataclass does **no type validation or
+    coercion**: a field declared ``qty: int`` silently receives whatever
+    JSON type was actually present if the producer sent the wrong type.
+    Unknown keys in the payload are dropped rather than raising. Use a
+    Pydantic model as ``target_type`` (or a msgspec-based serializer)
+    instead of a stdlib dataclass for untrusted input where wrong-typed
+    fields must be rejected.
+    """
+
+    #: M7: sane non-None default (64 MiB, matching
+    #: ``CompressionMiddleware``'s ``max_decompressed_size`` default) so an
+    #: uncompressed body is bounded out of the box. Pass ``max_parse_bytes=None``
+    #: to opt out (unbounded) if you've already sized this elsewhere.
+    _DEFAULT_MAX_PARSE_BYTES = 64 * 1024 * 1024
+
+    def __init__(
+        self,
+        *,
+        coerce_unknown_to_str: bool = False,
+        max_parse_bytes: int | None = _DEFAULT_MAX_PARSE_BYTES,
+    ) -> None:
+        self._coerce = coerce_unknown_to_str
+        # Defense-in-depth cap on the input size before json.loads (M7) — the
+        # compression middleware already caps decompressed output; this
+        # bounds the case where compression is off and a large body arrives
+        # directly. Defaults to a sane non-None value rather than "off".
+        self._max_parse_bytes = max_parse_bytes
+
+    def _check_size(self, data: bytes) -> None:
+        if self._max_parse_bytes is not None and len(data) > self._max_parse_bytes:
+            raise ValueError(f"JSON input size {len(data)} exceeds max_parse_bytes={self._max_parse_bytes}")
+
+    @property
+    def content_type(self) -> str:
+        return "application/json"
+
+    def _default(self, obj: Any) -> Any:
+        if self._coerce:
+            return str(obj)
+        raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+    def encode(self, data: Any) -> bytes:
+        """Serialize data to JSON bytes."""
+        if isinstance(data, bytes):
+            return data
+        if isinstance(data, str):
+            return data.encode("utf-8")
+
+        # Pydantic V2 model
+        if hasattr(data, "model_dump_json"):
+            result = data.model_dump_json()
+            return result.encode("utf-8") if isinstance(result, str) else result
+
+        # Dataclass
+        if is_dataclass(data) and not isinstance(data, type):
+            return json.dumps(asdict(data), default=self._default).encode("utf-8")
+
+        # Dict, list, or other JSON-serializable
+        return json.dumps(data, default=self._default).encode("utf-8")
+
+    def decode(self, data: bytes, target_type: type) -> Any:
+        """Deserialize JSON bytes to target type."""
+        self._check_size(data)
+        if target_type is bytes:
+            return data
+        if target_type is str:
+            return data.decode("utf-8")
+        if target_type is dict:
+            return json.loads(data)
+        if target_type is list:
+            return json.loads(data)
+
+        # Pydantic V2 model
+        if hasattr(target_type, "model_validate_json"):
+            return target_type.model_validate_json(data)
+
+        # Pydantic V2 model via model_validate (dict input)
+        if hasattr(target_type, "model_validate"):
+            parsed = json.loads(data)
+            return target_type.model_validate(parsed)
+
+        # Dataclass (H14: no type validation/coercion -- see DataclassDecoder
+        # in serialization/pipeline.py for the full contract this mirrors.
+        # Unknown keys are dropped rather than raising; a genuinely wrong
+        # shape raises TypeError naming the target dataclass.)
+        if is_dataclass(target_type):
+            parsed = json.loads(data)
+            if isinstance(parsed, dict):
+                known_fields = {f.name for f in fields(target_type)}
+                filtered = {k: v for k, v in parsed.items() if k in known_fields}
+                try:
+                    return target_type(**filtered)
+                except TypeError as exc:
+                    raise TypeError(f"Cannot decode into {target_type.__name__}: {exc}") from exc
+            return parsed
+
+        # Fallback: json.loads
+        return json.loads(data)

rabbitkit/serialization/msgspec.py

@@ -0,0 +1,136 @@
+"""msgspec serializer — optional high-performance serialization.
+
+Requires: pip install rabbitkit[msgspec]
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class MsgspecSerializer:
+    """High-performance serializer using msgspec.
+
+    Requires msgspec to be installed (optional dependency).
+    Falls back with clear error if not available.
+
+    M7: caps the input size before decoding (64 MiB by default, matching
+    ``CompressionMiddleware``'s ``max_decompressed_size`` default) — without
+    it, a large uncompressed body is fully materialized with no bound.
+    Pass ``max_parse_bytes=None`` to opt out.
+    """
+
+    #: M7: see JSONSerializer's identical default for the rationale.
+    _DEFAULT_MAX_PARSE_BYTES = 64 * 1024 * 1024
+
+    def __init__(self, *, max_parse_bytes: int | None = _DEFAULT_MAX_PARSE_BYTES) -> None:
+        try:
+            import msgspec
+
+            self._msgspec = msgspec
+            self._encoder = msgspec.json.Encoder()
+            self._decoder = msgspec.json.Decoder()
+            self._decoders: dict[Any, Any] = {}  # cached decoders per target_type (perf)
+        except ImportError as e:  # pragma: no cover
+            raise ImportError(  # pragma: no cover
+                "msgspec is required for MsgspecSerializer. Install it with: pip install rabbitkit[msgspec]"
+            ) from e
+        self._max_parse_bytes = max_parse_bytes
+
+    def _check_size(self, data: bytes) -> None:
+        if self._max_parse_bytes is not None and len(data) > self._max_parse_bytes:
+            raise ValueError(f"JSON input size {len(data)} exceeds max_parse_bytes={self._max_parse_bytes}")
+
+    @property
+    def content_type(self) -> str:
+        """Advisory only (M10): this is the ``content_type`` used when
+        *publishing* (set on the outgoing AMQP message property). It is
+        **not** verified against an incoming message's declared
+        ``content_type`` on :meth:`decode` — decode is driven solely by the
+        handler's declared parameter type, matching every other built-in
+        serializer in this codebase (none of them negotiate/verify
+        content-type on the consume side). If a message's actual body
+        doesn't match what this serializer expects (e.g. it's not JSON at
+        all), :meth:`decode` raises with a message naming both the target
+        type and a content-type-mismatch hint, rather than a raw
+        ``msgspec.DecodeError``.
+        """
+        return "application/json"
+
+    def encode(self, data: Any) -> bytes:
+        """Serialize data to JSON bytes using msgspec."""
+        if isinstance(data, bytes):
+            return data
+        if isinstance(data, str):
+            return data.encode("utf-8")
+
+        # msgspec Struct
+        if isinstance(data, self._msgspec.Struct):
+            return self._msgspec.json.encode(data)
+
+        # General encoding
+        return self._encoder.encode(data)
+
+    def decode(self, data: bytes, target_type: type) -> Any:
+        """Deserialize JSON bytes using msgspec.
+
+        M10: raises with a clearer message (naming the target type and
+        hinting at a content-type mismatch) instead of a raw
+        ``msgspec.DecodeError`` if the body isn't valid JSON for
+        *target_type* — most commonly caused by a message whose actual
+        content_type doesn't match what this serializer expects (see
+        :attr:`content_type`'s docstring — this is never verified upfront).
+        """
+        try:
+            return self._decode(data, target_type)
+        except self._msgspec.DecodeError as exc:
+            raise ValueError(
+                f"Failed to decode message body as JSON for target type {target_type!r}: {exc}. "
+                "If the message's actual content_type doesn't match what this serializer expects "
+                "(application/json), that's the likely cause — MsgspecSerializer does not verify "
+                "content_type before decoding."
+            ) from exc
+
+    def _decode(self, data: bytes, target_type: type) -> Any:
+        self._check_size(data)
+        if target_type is bytes:
+            return data
+        if target_type is str:
+            return data.decode("utf-8")
+        if target_type is dict or target_type is list:
+            return self._msgspec.json.decode(data)
+
+        # Generic alias (dict[str, Any], list[int], …) — origin is dict/list
+        origin = getattr(target_type, "__origin__", None)
+        if origin is dict or origin is list:
+            decoder = self._decoders.get(target_type)
+            if decoder is None:
+                try:
+                    decoder = self._msgspec.json.Decoder(type=target_type)
+                except Exception:
+                    return self._msgspec.json.decode(data)
+                self._decoders[target_type] = decoder
+            return decoder.decode(data)
+
+        # Pydantic V2 model — model_validate_json is faster than json.loads + model_validate
+        if hasattr(target_type, "model_validate_json"):
+            return target_type.model_validate_json(data)
+
+        # msgspec Struct — fastest path for msgspec-native types
+        try:
+            is_struct = issubclass(target_type, self._msgspec.Struct)
+        except TypeError:
+            is_struct = False
+        if is_struct:
+            return self._msgspec.json.decode(data, type=target_type)
+
+        # General typed decoder — cached per target_type (Decoder(type=T) codegens
+        # a converter; rebuilding per call defeats msgspec's performance advantage).
+        decoder = self._decoders.get(target_type)
+        if decoder is None:
+            try:
+                decoder = self._msgspec.json.Decoder(type=target_type)
+            except Exception:
+                return self._msgspec.json.decode(data)
+            self._decoders[target_type] = decoder
+        return decoder.decode(data)

rabbitkit/serialization/pipeline.py

@@ -0,0 +1,255 @@
+"""Two-stage serialization pipeline (parser + decoder).
+
+Splits serialization into two composable stages so you can mix and match
+wire formats with type-mapping strategies independently.
+
+  Stage 1 — **Parser** (``MessageParser`` protocol)
+    raw ``bytes``  →  intermediate form (usually ``dict`` or ``list``)
+    Examples: ``JsonParser``, ``MsgpackParser`` (bring your own)
+
+  Stage 2 — **Decoder** (``MessageDecoder`` protocol)
+    intermediate  →  typed Python object
+    Examples: ``PydanticDecoder``, ``DataclassDecoder``, ``RawDecoder``
+
+``SerializationPipeline`` composes them and exposes the same ``encode`` /
+``decode`` / ``content_type`` interface as built-in serializers, so it plugs
+directly into ``broker(serializer=...)`` or ``@broker.subscriber(serializer=...)``.
+
+Quick start — JSON + Pydantic
+------------------------------
+    from pydantic import BaseModel
+    from rabbitkit.serialization.pipeline import (
+        SerializationPipeline, JsonParser, PydanticDecoder,
+    )
+
+    class Order(BaseModel):
+        id: int
+        item: str
+        qty: int
+
+    pipeline = SerializationPipeline(JsonParser(), PydanticDecoder())
+
+    broker = AsyncBroker(config, serializer=pipeline)
+
+    @broker.subscriber(queue="orders")
+    async def handle_order(order: Order) -> None:
+        # `order` is already a validated Pydantic model
+        print(order.id, order.item)
+
+JSON + stdlib dataclass
+-----------------------
+    from dataclasses import dataclass
+    from rabbitkit.serialization.pipeline import (
+        SerializationPipeline, JsonParser, DataclassDecoder,
+    )
+
+    @dataclass
+    class Event:
+        type: str
+        payload: dict
+
+    pipeline = SerializationPipeline(JsonParser(), DataclassDecoder())
+
+    @broker.subscriber(queue="events", serializer=pipeline)
+    def handle(event: Event) -> None:
+        print(event.type)
+
+Pass-through (raw bytes, no decoding)
+--------------------------------------
+    from rabbitkit.serialization.pipeline import (
+        SerializationPipeline, JsonParser, RawDecoder,
+    )
+
+    # Still parses bytes → dict for intermediate, but decoder returns as-is
+    pipeline = SerializationPipeline(JsonParser(), RawDecoder())
+
+Custom parser (msgpack example)
+---------------------------------
+    import msgpack
+
+    class MsgpackParser:
+        def parse(self, data: bytes, content_type=None):
+            return msgpack.unpackb(data, raw=False)
+
+        def serialize(self, data) -> bytes:
+            return msgpack.packb(data, use_bin_type=True)
+
+        @property
+        def content_type(self) -> str:
+            return "application/msgpack"
+
+    pipeline = SerializationPipeline(MsgpackParser(), PydanticDecoder())
+
+Encoding (publish direction)
+-----------------------------
+``pipeline.encode(my_model)`` calls ``decoder.encode(model) → dict`` then
+``parser.serialize(dict) → bytes``.  This means the same pipeline handles
+both inbound *and* outbound serialization transparently.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Protocol
+
+
+class MessageParser(Protocol):
+    """Stage 1: raw bytes → intermediate form."""
+
+    def parse(self, data: bytes, content_type: str | None = None) -> Any: ...
+    def serialize(self, data: Any) -> bytes: ...
+    @property
+    def content_type(self) -> str: ...
+
+
+class MessageDecoder(Protocol):
+    """Stage 2: intermediate → typed Python object."""
+
+    def decode(self, data: Any, target_type: type) -> Any: ...
+    def encode(self, data: Any) -> Any: ...
+
+
+class SerializationPipeline:
+    """Two-stage serialization. Implements Serializer protocol."""
+
+    def __init__(self, parser: MessageParser, decoder: MessageDecoder) -> None:
+        self._parser = parser
+        self._decoder = decoder
+
+    def encode(self, data: Any) -> bytes:
+        intermediate = self._decoder.encode(data)
+        return self._parser.serialize(intermediate)
+
+    def decode(self, data: bytes, target_type: type) -> Any:
+        intermediate = self._parser.parse(data)
+        return self._decoder.decode(intermediate, target_type)
+
+    @property
+    def content_type(self) -> str:
+        return self._parser.content_type
+
+
+class JsonParser:
+    """Built-in JSON parser.
+
+    By default ``serialize`` **raises** on objects ``json`` cannot represent
+    (e.g. ``datetime``, ``Decimal``) rather than silently coercing them via
+    ``str()``. Pass ``coerce_unknown_to_str=True`` to restore the legacy
+    ``default=str`` coercion behaviour.
+
+    M7: caps the input size before ``json.loads`` (64 MiB by default,
+    matching ``CompressionMiddleware``'s ``max_decompressed_size`` default)
+    -- without it, a large uncompressed body is fully materialized with no
+    bound. Pass ``max_parse_bytes=None`` to opt out.
+    """
+
+    #: M7: see JSONSerializer's identical default for the rationale.
+    _DEFAULT_MAX_PARSE_BYTES = 64 * 1024 * 1024
+
+    def __init__(
+        self,
+        *,
+        coerce_unknown_to_str: bool = False,
+        max_parse_bytes: int | None = _DEFAULT_MAX_PARSE_BYTES,
+    ) -> None:
+        self._coerce = coerce_unknown_to_str
+        self._max_parse_bytes = max_parse_bytes
+
+    def _check_size(self, data: bytes) -> None:
+        if self._max_parse_bytes is not None and len(data) > self._max_parse_bytes:
+            raise ValueError(f"JSON input size {len(data)} exceeds max_parse_bytes={self._max_parse_bytes}")
+
+    def parse(self, data: bytes, content_type: str | None = None) -> Any:
+        self._check_size(data)
+        return json.loads(data)
+
+    def _default(self, data: Any) -> Any:
+        if self._coerce:
+            return str(data)
+        raise TypeError(f"Object of type {type(data).__name__} is not JSON serializable")
+
+    def serialize(self, data: Any) -> bytes:
+        return json.dumps(data, default=self._default).encode("utf-8")
+
+    @property
+    def content_type(self) -> str:
+        return "application/json"
+
+
+class PydanticDecoder:
+    """Decoder that uses Pydantic model_validate for decoding.
+
+    Companion to ``DataclassDecoder``'s H14 note: this is the decoder that
+    note recommends for untrusted input, on the premise that
+    ``model_validate`` rejects wrong-shaped data instead of silently
+    accepting it. That premise used to be undermined by an
+    ``isinstance(data, dict)`` guard here -- a non-dict top-level payload
+    (a producer sending a JSON array/string/number instead of an object,
+    whether by bug or by an attacker probing for a decoder that skips
+    validation on the untaken branch) bypassed ``model_validate`` entirely
+    and returned the raw, un-validated value straight to the handler.
+    Pydantic already raises a clean ``ValidationError`` for a non-dict
+    input (``model_type`` error), so there is no reason to special-case
+    it -- doing so only suppressed the exact validation this decoder
+    exists to provide.
+    """
+
+    def decode(self, data: Any, target_type: type) -> Any:
+        if hasattr(target_type, "model_validate"):
+            return target_type.model_validate(data)
+        return data
+
+    def encode(self, data: Any) -> Any:
+        if hasattr(data, "model_dump"):
+            return data.model_dump()
+        return data
+
+
+class DataclassDecoder:
+    """Decoder for stdlib dataclasses.
+
+    H14 — no type validation or coercion: unlike ``PydanticDecoder``, this
+    does NOT check field types. A field declared ``qty: int`` silently
+    receives whatever JSON type was actually present (e.g. the string
+    ``"3"``) if the producer sent the wrong type — stdlib dataclasses
+    perform no runtime type checking on construction. **Use
+    ``PydanticDecoder`` (or a msgspec-based decoder) instead for untrusted
+    input where wrong-typed fields must be rejected.**
+
+    Unknown keys in the incoming dict (not a declared field) are silently
+    dropped rather than raising ``TypeError`` — this keeps a producer that
+    adds a new field (forward-compatible) from turning every message into a
+    decode failure (misclassified PERMANENT, straight to the DLQ) until
+    every consumer is upgraded. A genuinely wrong shape (a missing required
+    field, etc.) still raises — as a ``TypeError`` naming the target
+    dataclass, not a bare constructor traceback.
+    """
+
+    def decode(self, data: Any, target_type: type) -> Any:
+        import dataclasses
+
+        if dataclasses.is_dataclass(target_type) and isinstance(data, dict):
+            known_fields = {f.name for f in dataclasses.fields(target_type)}
+            filtered = {k: v for k, v in data.items() if k in known_fields}
+            try:
+                return target_type(**filtered)
+            except TypeError as exc:
+                raise TypeError(f"Cannot decode into {target_type.__name__}: {exc}") from exc
+        return data
+
+    def encode(self, data: Any) -> Any:
+        import dataclasses
+
+        if dataclasses.is_dataclass(data) and not isinstance(data, type):
+            return dataclasses.asdict(data)
+        return data
+
+
+class RawDecoder:
+    """Pass-through decoder — no transformation."""
+
+    def decode(self, data: Any, target_type: type) -> Any:
+        return data
+
+    def encode(self, data: Any) -> Any:
+        return data

rabbitkit/streams.py

@@ -0,0 +1,139 @@
+"""Stream queue consumer offset tracking.
+
+RabbitMQ stream queues support offset-based consuming, allowing consumers
+to start reading from a specific position in the log.
+
+Offset types:
+- "first": Start from the beginning of the stream
+- "last": Start from the end (new messages only)
+- "next": Start from the next unconsumed message (default)
+- timestamp: Start from messages published after a given timestamp
+- numeric offset: Start from a specific offset value
+
+Usage::
+
+    from rabbitkit.streams import StreamOffset, StreamConsumerConfig
+
+    # Start from beginning
+    config = StreamConsumerConfig(offset=StreamOffset.first())
+
+    # Start from specific offset
+    config = StreamConsumerConfig(offset=StreamOffset.offset(42))
+
+    # Start from timestamp
+    config = StreamConsumerConfig(offset=StreamOffset.timestamp(datetime(2026, 1, 1)))
+
+    # Use with broker
+    @broker.subscriber(
+        queue=RabbitQueue("events", queue_type=QueueType.STREAM),
+        stream_config=config,
+    )
+    def handle(body: bytes) -> None: ...
+"""
+
+from __future__ import annotations
+
+import enum
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class StreamOffsetType(str, enum.Enum):
+    """Stream offset specification types."""
+
+    FIRST = "first"
+    LAST = "last"
+    NEXT = "next"
+    OFFSET = "offset"
+    TIMESTAMP = "timestamp"
+
+
+@dataclass(frozen=True, slots=True)
+class StreamOffset:
+    """Stream queue offset specification.
+
+    Use class methods to create:
+        StreamOffset.first()
+        StreamOffset.last()
+        StreamOffset.next()
+        StreamOffset.offset(42)
+        StreamOffset.timestamp(datetime(...))
+    """
+
+    type: StreamOffsetType = StreamOffsetType.NEXT
+    value: int | datetime | None = None
+
+    @classmethod
+    def first(cls) -> StreamOffset:
+        """Start from the beginning of the stream."""
+        return cls(type=StreamOffsetType.FIRST)
+
+    @classmethod
+    def last(cls) -> StreamOffset:
+        """Start from the end (new messages only)."""
+        return cls(type=StreamOffsetType.LAST)
+
+    @classmethod
+    def next(cls) -> StreamOffset:
+        """Start from the next unconsumed message (default)."""
+        return cls(type=StreamOffsetType.NEXT)
+
+    @classmethod
+    def offset(cls, value: int) -> StreamOffset:
+        """Start from a specific numeric offset."""
+        if value < 0:
+            msg = "Stream offset must be non-negative"
+            raise ValueError(msg)
+        return cls(type=StreamOffsetType.OFFSET, value=value)
+
+    @classmethod
+    def timestamp(cls, value: datetime) -> StreamOffset:
+        """Start from messages published after the given timestamp."""
+        return cls(type=StreamOffsetType.TIMESTAMP, value=value)
+
+    def to_consume_arguments(self) -> dict[str, Any]:
+        """Convert to RabbitMQ consume arguments (x-stream-offset).
+
+        Returns dict suitable for merging into basic_consume arguments.
+        """
+        if self.type == StreamOffsetType.FIRST:
+            return {"x-stream-offset": "first"}
+        if self.type == StreamOffsetType.LAST:
+            return {"x-stream-offset": "last"}
+        if self.type == StreamOffsetType.NEXT:
+            return {"x-stream-offset": "next"}
+        if self.type == StreamOffsetType.OFFSET:
+            return {"x-stream-offset": self.value}
+        if self.type == StreamOffsetType.TIMESTAMP:
+            if not isinstance(self.value, datetime):
+                raise TypeError("TIMESTAMP stream offset requires a datetime value")
+            # RabbitMQ x-stream-offset by time expects a Unix timestamp in seconds.
+            return {"x-stream-offset": int(self.value.timestamp())}
+        return {}  # pragma: no cover
+
+    def __repr__(self) -> str:
+        if self.value is not None:
+            return f"StreamOffset({self.type.value}={self.value})"
+        return f"StreamOffset({self.type.value})"
+
+
+@dataclass(frozen=True, slots=True)
+class StreamConsumerConfig:
+    """Configuration for stream queue consumers.
+
+    Extends basic consumer behavior with stream-specific options.
+    """
+
+    offset: StreamOffset = field(default_factory=StreamOffset.next)
+    consumer_name: str | None = None  # x-stream-consumer-name for single-active-consumer
+
+    def to_consume_arguments(self) -> dict[str, Any]:
+        """Build consume arguments for stream queue subscription."""
+        args = self.offset.to_consume_arguments()
+        if self.consumer_name is not None:
+            args["x-stream-consumer-name"] = self.consumer_name
+        return args

rabbitkit/sync/__init__.py

@@ -0,0 +1,11 @@
+"""Sync transport module — pika-based I/O adapter."""
+
+from rabbitkit.sync.batch import SyncBatchPublisher
+from rabbitkit.sync.broker import SyncBroker
+from rabbitkit.sync.transport import SyncTransport
+
+__all__ = [
+    "SyncBatchPublisher",
+    "SyncBroker",
+    "SyncTransport",
+]