dory-processor-sdk 0.0.1__py3-none-any.whl

dory/output/rabbitmq.py

@@ -0,0 +1,632 @@
+"""
+RabbitMQ publisher with circuit breaker, retry, and in-memory buffering.
+
+Provides a resilient async publisher for sending processor output to RabbitMQ.
+Requires the ``aio-pika`` optional dependency: ``pip install dory-processor-sdk[production]``
+"""
+
+import asyncio
+import logging
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Optional
+from urllib.parse import urlparse
+
+import aio_pika
+import yarl
+
+from dory.output.formatter import JSONFormatter, OutputFormatter
+from dory.resilience import CircuitBreaker, CircuitOpenError, RetryExhaustedError
+
+logger = logging.getLogger(__name__)
+
+
+class PublishError(Exception):
+    """Raised when a publish operation fails."""
+
+
+def _mask_url(url: str) -> str:
+    """Mask credentials in an AMQP URL for safe logging.
+
+    Returns only the host (and port if present), stripping user/pass/vhost.
+    """
+    try:
+        parsed = urlparse(url)
+        host = parsed.hostname or "unknown"
+        if parsed.port:
+            return f"{host}:{parsed.port}"
+        return host
+    except Exception:
+        return "<masked>"
+
+
+@dataclass
+class PublisherConfig:
+    """Configuration for RabbitMQ publisher.
+
+    Attributes:
+        url: AMQP connection URL.
+        exchange: Default exchange name.
+        exchange_type: Exchange type (topic, direct, fanout, headers).
+        durable: Whether the exchange is durable.
+        connection_timeout: Connection timeout in seconds.
+        heartbeat: AMQP heartbeat interval in seconds.
+        buffer_enabled: Enable in-memory message buffering on failure.
+        buffer_max_size: Maximum number of buffered messages.
+        buffer_max_bytes: Maximum total size of buffered messages in bytes.
+        retry_max_attempts: Max retry attempts per publish.
+        retry_initial_delay: Initial retry delay in seconds.
+        retry_max_delay: Maximum retry delay in seconds.
+        circuit_breaker_failure_threshold: Failures before circuit opens.
+        circuit_breaker_timeout: Seconds before circuit half-opens.
+    """
+
+    url: str = "amqp://guest:guest@localhost:5672/"
+    exchange: str = "dory.output"
+    exchange_type: str = "topic"
+    durable: bool = True
+    connection_timeout: float = 10.0
+    heartbeat: int = 60
+    buffer_enabled: bool = True
+    buffer_max_size: int = 10000
+    buffer_max_bytes: int = 100 * 1024 * 1024  # 100MB
+    retry_max_attempts: int = 3
+    retry_initial_delay: float = 1.0
+    retry_max_delay: float = 30.0
+    circuit_breaker_failure_threshold: int = 5
+    circuit_breaker_timeout: float = 60.0
+
+
+@dataclass
+class BufferedMessage:
+    """A message buffered for later delivery.
+
+    Attributes:
+        exchange: Target exchange name.
+        routing_key: Message routing key.
+        body: Serialized message body.
+        headers: Optional message headers.
+        timestamp: Time the message was originally published.
+    """
+
+    exchange: str
+    routing_key: str
+    body: bytes
+    headers: dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+
+    def size_bytes(self) -> int:
+        """Return approximate size of this message in bytes."""
+        header_size = sum(
+            len(str(k)) + len(str(v)) for k, v in self.headers.items()
+        ) if self.headers else 0
+        return len(self.body) + len(self.routing_key) + len(self.exchange) + header_size
+
+
+class RabbitMQPublisher:
+    """Async RabbitMQ publisher with circuit breaker, retry, and buffering.
+
+    Publishes formatted messages to RabbitMQ with full resilience:
+    - Circuit breaker wraps retry to prevent cascading failures
+    - Retry with exponential backoff for transient errors
+    - In-memory buffer for messages that cannot be delivered
+    - Automatic buffer flush on reconnection
+    - MessageEnvelope wrapping for unified output
+
+    Args:
+        config: Publisher configuration.
+        formatter: Inner serializer (defaults to JSONFormatter). Wrapped in
+            an EnvelopeFormatter automatically.
+        metrics: Optional MetricsCollector for recording publish metrics.
+    """
+
+    def __init__(
+        self,
+        config: PublisherConfig | None = None,
+        formatter: OutputFormatter | None = None,
+        metrics: Any | None = None,
+        url_provider: Callable[[], Any] | None = None,
+        **kwargs: Any,
+    ):
+        self._config = config or PublisherConfig()
+        self._url_provider = url_provider
+        self._formatter = formatter or JSONFormatter()
+        self._metrics = metrics
+
+        # Envelope wrapping (always on)
+        from dory.output.envelope import EnvelopeFormatter, ENVELOPE_SCHEMA_VERSION
+
+        self._envelope_formatter = EnvelopeFormatter(
+            formatter=self._formatter,
+            schema_version=ENVELOPE_SCHEMA_VERSION,
+        )
+
+        # Connection state
+        self._connection: Optional[aio_pika.abc.AbstractRobustConnection] = None
+        self._channel: Optional[aio_pika.abc.AbstractChannel] = None
+        self._exchange: Optional[aio_pika.abc.AbstractExchange] = None
+        self._connected = False
+        self._closing = False
+
+        # Buffer
+        self._buffer: deque[BufferedMessage] = deque()
+        self._buffer_bytes: int = 0
+        self._buffer_lock = asyncio.Lock()
+
+        # Circuit breaker (wraps retry internally)
+        self._circuit_breaker = CircuitBreaker(
+            name="rabbitmq_publisher",
+            failure_threshold=self._config.circuit_breaker_failure_threshold,
+            timeout_seconds=self._config.circuit_breaker_timeout,
+        )
+
+        # Metrics counters
+        self._publish_total = 0
+        self._publish_success = 0
+        self._publish_failed = 0
+        self._publish_buffered = 0
+        self._buffer_dropped = 0
+
+        masked = _mask_url(self._config.url)
+        logger.info(f"RabbitMQ publisher created for {masked}")
+
+    async def connect(self) -> None:
+        """Establish connection to RabbitMQ.
+
+        Creates a robust connection, opens a channel, and declares the exchange.
+        If a ``url_provider`` was given, it is called to obtain the current URL
+        before connecting (e.g., to refresh an OAuth2 token).  A close callback
+        is also registered so that ``aio_pika``'s automatic reconnection uses a
+        freshly-fetched URL (with a new token) instead of the stale original.
+        Automatically flushes any buffered messages after reconnection.
+        """
+        if self._url_provider is not None:
+            self._config.url = await self._url_provider()
+
+        masked = _mask_url(self._config.url)
+        logger.info(f"Connecting to RabbitMQ at {masked}")
+
+        self._connection = await aio_pika.connect_robust(
+            self._config.url,
+            timeout=self._config.connection_timeout,
+            heartbeat=self._config.heartbeat,
+        )
+
+        # When url_provider is set (e.g. OAuth2), register a close callback
+        # that refreshes the connection URL before aio-pika's automatic
+        # reconnection factory retries.  Execution order on disconnect:
+        #   1. close_callbacks fire (async, awaited)
+        #   2. __connection_close_event wakes the reconnection loop
+        #   3. reconnection loop calls Connection.connect() using self.url
+        # So updating connection.url here ensures a fresh token is used.
+        if self._url_provider is not None:
+            self._connection.close_callbacks.add(self._refresh_url_on_close)
+
+        self._channel = await self._connection.channel()
+        try:
+            self._exchange = await self._channel.declare_exchange(
+                self._config.exchange,
+                type=aio_pika.ExchangeType(self._config.exchange_type),
+                durable=self._config.durable,
+            )
+        except Exception as e:
+            # If declare fails (e.g. OAuth2 user lacks configure permission),
+            # fall back to using the exchange without declaring it.
+            logger.info(
+                f"Exchange declare skipped (using existing): {e}"
+            )
+            self._exchange = await self._channel.get_exchange(
+                self._config.exchange, ensure=False,
+            )
+        self._connected = True
+
+        self._record_metric("connection_status", gauge=1.0)
+        logger.info(f"Connected to RabbitMQ at {masked}")
+
+        # Flush buffered messages
+        if self._buffer:
+            logger.info(f"Flushing {len(self._buffer)} buffered messages")
+            await self.flush_buffer()
+
+    async def _refresh_url_on_close(
+        self,
+        _sender: Any,
+        _exc: BaseException | None = None,
+    ) -> None:
+        """Close callback that refreshes the AMQP URL for reconnection.
+
+        Called by ``aio_pika`` when the connection drops, *before* the
+        reconnection factory loop retries.  Fetches a new token via the
+        ``url_provider`` and updates ``connection.url`` so the next
+        reconnection attempt uses valid credentials.
+        """
+        if self._url_provider is None or self._closing:
+            return
+        try:
+            fresh_url = await self._url_provider()
+            self._config.url = fresh_url
+            if self._connection is not None:
+                self._connection.url = yarl.URL(fresh_url)
+            logger.info(
+                f"Refreshed AMQP URL for reconnection to "
+                f"{_mask_url(fresh_url)}"
+            )
+        except Exception as e:
+            logger.warning(f"Failed to refresh AMQP URL on disconnect: {e}")
+
+    async def publish(
+        self,
+        routing_key: str,
+        data: Any,
+        exchange: str | None = None,
+        headers: dict[str, Any] | None = None,
+        raw: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Publish data to RabbitMQ.
+
+        When ``raw=False`` (default), data is wrapped in a MessageEnvelope
+        before publishing.  When ``raw=True``, data is serialized as-is
+        (caller is responsible for providing a properly formatted payload,
+        e.g. an envelope dict).
+
+        Args:
+            routing_key: Message routing key.
+            data: Data to publish.
+            exchange: Optional exchange override.
+            headers: Optional message headers.
+            raw: If True, skip envelope wrapping and serialize data directly.
+
+        Raises:
+            PublishError: If publish fails and buffering is disabled.
+        """
+        if self._closing:
+            logger.warning("Publisher is closing, cannot accept new messages")
+            return
+
+        self._publish_total += 1
+        self._record_metric("publish_total", counter=1)
+
+        target_exchange = exchange or self._config.exchange
+
+        # Format data (envelope wrapping unless raw=True)
+        try:
+            if raw:
+                body = self._formatter.format(data)
+            else:
+                body = self._envelope_formatter.format(data)
+        except Exception as e:
+            self._publish_failed += 1
+            self._record_metric("publish_failed_total", counter=1)
+            raise PublishError(f"Failed to format message: {e}") from e
+
+        # Try publishing through circuit breaker (which wraps retry)
+        try:
+            await self._circuit_breaker.call(
+                self._publish_with_retry,
+                target_exchange,
+                routing_key,
+                body,
+                headers,
+            )
+            self._publish_success += 1
+            self._record_metric("publish_success_total", counter=1)
+        except CircuitOpenError:
+            logger.warning(
+                f"Circuit breaker open, buffering message for {routing_key}"
+            )
+            await self._buffer_message(
+                target_exchange, routing_key, body, headers or {}
+            )
+        except (PublishError, RetryExhaustedError) as e:
+            logger.warning(f"Publish failed after retries: {e}")
+            self._publish_failed += 1
+            self._record_metric("publish_failed_total", counter=1)
+            if self._config.buffer_enabled:
+                await self._buffer_message(
+                    target_exchange, routing_key, body, headers or {}
+                )
+            else:
+                raise PublishError(f"Publish failed: {e}") from e
+        except Exception as e:
+            logger.warning(f"Publish failed: {e}")
+            self._publish_failed += 1
+            self._record_metric("publish_failed_total", counter=1)
+            if self._config.buffer_enabled:
+                await self._buffer_message(
+                    target_exchange, routing_key, body, headers or {}
+                )
+            else:
+                raise PublishError(f"Publish failed: {e}") from e
+
+    @staticmethod
+    def _is_token_expiry_error(error: Exception) -> bool:
+        """Check if an error indicates an expired OAuth2/JWT token."""
+        msg = str(error).lower()
+        return "access_refused" in msg and (
+            "token has expired" in msg
+            or "jwt" in msg and "expired" in msg
+        )
+
+    async def _reconnect_with_fresh_token(self) -> None:
+        """Close the current connection and reconnect with a fresh token.
+
+        Called when a publish fails due to an expired JWT while the AMQP
+        connection is still alive.  Fetches a new token via ``url_provider``,
+        then re-establishes the channel and exchange.
+        """
+        if self._url_provider is None:
+            return
+
+        logger.info("Token expired — refreshing token and reconnecting")
+
+        try:
+            # Fetch fresh URL (triggers OAuth2TokenProvider.get_token())
+            fresh_url = await self._url_provider()
+            self._config.url = fresh_url
+        except Exception as e:
+            logger.warning(f"Failed to refresh token: {e}")
+            raise
+
+        # Close existing connection (suppress errors on stale state)
+        if self._connection is not None:
+            try:
+                await self._connection.close()
+            except Exception:
+                pass
+            self._connection = None
+            self._channel = None
+            self._exchange = None
+            self._connected = False
+
+        # Reconnect with the fresh token
+        await self.connect()
+        logger.info("Reconnected with fresh token")
+
+    async def _publish_with_retry(
+        self,
+        exchange_name: str,
+        routing_key: str,
+        body: bytes,
+        headers: dict[str, Any] | None,
+    ) -> None:
+        """Publish with retry logic. Called inside circuit breaker.
+
+        If a publish fails due to an expired JWT token and a ``url_provider``
+        is configured, the connection is torn down and re-established with a
+        fresh token before retrying.
+        """
+        last_error: Exception | None = None
+        delay = self._config.retry_initial_delay
+        token_refreshed = False
+
+        for attempt in range(self._config.retry_max_attempts):
+            try:
+                await self._publish_internal(
+                    exchange_name, routing_key, body, headers
+                )
+                return
+            except Exception as e:
+                last_error = e
+
+                # On token expiry, reconnect with a fresh token once
+                if (
+                    not token_refreshed
+                    and self._url_provider is not None
+                    and self._is_token_expiry_error(e)
+                ):
+                    try:
+                        await self._reconnect_with_fresh_token()
+                        token_refreshed = True
+                        # Retry immediately after reconnect (no backoff)
+                        continue
+                    except Exception as refresh_err:
+                        logger.warning(f"Token refresh failed: {refresh_err}")
+
+                if attempt < self._config.retry_max_attempts - 1:
+                    logger.warning(
+                        f"Publish attempt {attempt + 1}/{self._config.retry_max_attempts} "
+                        f"failed: {e}. Retrying in {delay:.1f}s"
+                    )
+                    await asyncio.sleep(delay)
+                    delay = min(delay * 2, self._config.retry_max_delay)
+
+        raise PublishError(
+            f"Publish failed after {self._config.retry_max_attempts} attempts: {last_error}"
+        )
+
+    async def _publish_internal(
+        self,
+        exchange_name: str,
+        routing_key: str,
+        body: bytes,
+        headers: dict[str, Any] | None,
+    ) -> None:
+        """Perform the actual AMQP publish."""
+        if not self._connected or self._exchange is None:
+            raise PublishError("Not connected to RabbitMQ")
+
+        # Use declared exchange if name matches, otherwise get exchange by name
+        if exchange_name == self._config.exchange and self._exchange is not None:
+            exchange = self._exchange
+        else:
+            if self._channel is None:
+                raise PublishError("Channel not available")
+            exchange = await self._channel.declare_exchange(
+                exchange_name,
+                type=aio_pika.ExchangeType(self._config.exchange_type),
+                durable=self._config.durable,
+            )
+
+        message = aio_pika.Message(
+            body=body,
+            content_type=self._formatter.content_type,
+            headers=headers or {},
+            timestamp=time.time(),
+        )
+
+        logger.info(
+            "AMQP publish exchange=%s routing_key=%s size=%d bytes",
+            exchange_name,
+            routing_key,
+            len(body),
+        )
+        await exchange.publish(message, routing_key=routing_key)
+
+    async def _buffer_message(
+        self,
+        exchange: str,
+        routing_key: str,
+        body: bytes,
+        headers: dict[str, Any],
+    ) -> None:
+        """Buffer a message for later delivery."""
+        if not self._config.buffer_enabled:
+            return
+
+        msg = BufferedMessage(
+            exchange=exchange,
+            routing_key=routing_key,
+            body=body,
+            headers=headers,
+        )
+        msg_size = msg.size_bytes()
+
+        async with self._buffer_lock:
+            # Evict oldest messages if buffer is full (by count)
+            while (
+                len(self._buffer) >= self._config.buffer_max_size
+                and self._buffer
+            ):
+                evicted = self._buffer.popleft()
+                self._buffer_bytes -= evicted.size_bytes()
+                self._buffer_dropped += 1
+                logger.warning(
+                    f"Buffer full (max {self._config.buffer_max_size} messages), "
+                    f"dropping oldest message for {evicted.routing_key}"
+                )
+
+            # Evict oldest messages if buffer is full (by bytes)
+            while (
+                self._config.buffer_max_bytes > 0
+                and self._buffer_bytes + msg_size > self._config.buffer_max_bytes
+                and self._buffer
+            ):
+                evicted = self._buffer.popleft()
+                self._buffer_bytes -= evicted.size_bytes()
+                self._buffer_dropped += 1
+                logger.warning(
+                    f"Buffer full (max {self._config.buffer_max_bytes} bytes), "
+                    f"dropping oldest message for {evicted.routing_key}"
+                )
+
+            self._buffer.append(msg)
+            self._buffer_bytes += msg_size
+
+        self._publish_buffered += 1
+        self._record_metric("publish_buffered_total", counter=1)
+        self._record_metric("buffer_size_messages", gauge=float(len(self._buffer)))
+        self._record_metric("buffer_size_bytes", gauge=float(self._buffer_bytes))
+
+    async def flush_buffer(self) -> None:
+        """Flush all buffered messages to RabbitMQ.
+
+        Messages that fail to publish remain in the buffer.
+        """
+        if not self._buffer:
+            return
+
+        async with self._buffer_lock:
+            remaining: deque[BufferedMessage] = deque()
+            remaining_bytes = 0
+
+            while self._buffer:
+                msg = self._buffer.popleft()
+                self._buffer_bytes -= msg.size_bytes()
+
+                try:
+                    await self._publish_internal(
+                        msg.exchange, msg.routing_key, msg.body, msg.headers
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to flush buffered message for {msg.routing_key}: {e}"
+                    )
+                    remaining.append(msg)
+                    remaining_bytes += msg.size_bytes()
+
+            # Put failed messages back
+            self._buffer = remaining
+            self._buffer_bytes = remaining_bytes
+
+        self._record_metric("buffer_size_messages", gauge=float(len(self._buffer)))
+        self._record_metric("buffer_size_bytes", gauge=float(self._buffer_bytes))
+
+    async def close(self) -> None:
+        """Close the publisher, flushing the buffer first."""
+        if self._closing:
+            return
+        self._closing = True
+
+        logger.info("Closing RabbitMQ publisher")
+
+        # Flush buffer before closing
+        if self._buffer and self._connected:
+            try:
+                await self.flush_buffer()
+            except Exception as e:
+                logger.warning(f"Error flushing buffer during close: {e}")
+
+        # Close channel and connection
+        if self._channel is not None:
+            try:
+                await self._channel.close()
+            except Exception:
+                pass
+            self._channel = None
+
+        if self._connection is not None:
+            try:
+                await self._connection.close()
+            except Exception:
+                pass
+            self._connection = None
+
+        self._connected = False
+        self._exchange = None
+        self._record_metric("connection_status", gauge=0.0)
+        logger.info("RabbitMQ publisher closed")
+
+    def is_connected(self) -> bool:
+        """Check if the publisher is currently connected."""
+        return self._connected
+
+    def get_buffer_size(self) -> dict[str, int]:
+        """Get current buffer size.
+
+        Returns:
+            Dictionary with 'messages' count and 'bytes' total size.
+        """
+        return {
+            "messages": len(self._buffer),
+            "bytes": self._buffer_bytes,
+        }
+
+    def _record_metric(
+        self,
+        name: str,
+        counter: float = 0,
+        gauge: float | None = None,
+    ) -> None:
+        """Record a metric if a metrics collector is available."""
+        if self._metrics is None:
+            return
+        try:
+            full_name = f"publisher_{name}"
+            if gauge is not None:
+                self._metrics.set_gauge(full_name, gauge)
+            elif counter > 0:
+                self._metrics.inc_counter(full_name, counter)
+        except Exception:
+            pass  # Metrics should never break publishing