Flask-RMQ 0.1.0__py3-none-any.whl

examples/basic_app/basic_app/__init__.py

@@ -0,0 +1,3 @@
+from .app import create_app
+
+__all__ = ['create_app']

examples/basic_app/basic_app/app.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import os
+
+from flask import Flask
+
+from flask_rmq import FlaskRMQ
+
+from .messaging import register_messaging
+from .views import api
+
+rmq = FlaskRMQ()
+
+
+def create_app() -> Flask:
+    app = Flask(__name__)
+    app.config['RABBITMQ_CONNECTIONS'] = {
+        'default': {
+            'HOST': os.getenv('RABBITMQ_HOST', 'localhost'),
+            'PORT': int(os.getenv('RABBITMQ_PORT', '5672')),
+            'VIRTUAL_HOST': os.getenv('RABBITMQ_VHOST', '/'),
+            'USER': os.getenv('RABBITMQ_USER', 'guest'),
+            'PASSWORD': os.getenv('RABBITMQ_PASSWORD', 'guest'),
+            'HEARTBEAT': 60,
+            'BLOCKED_CONNECTION_TIMEOUT': 30,
+            'RECONNECT_INITIAL_BACKOFF': 1,
+            'RECONNECT_MAX_BACKOFF': 15,
+        }
+    }
+    app.config['CONSUMED_FILE'] = os.getenv('CONSUMED_FILE', 'consumed.jsonl')
+    rmq.init_app(app)
+    app.register_blueprint(api)
+    with app.app_context():
+        register_messaging()
+    return app

examples/basic_app/basic_app/messaging.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from flask import current_app
+from pika.adapters.blocking_connection import BlockingChannel
+
+from flask_rmq import Consumer, Producer, QueueConfig, get_consumers_registry, get_setup_registry
+
+EVENTS_QUEUE = QueueConfig(
+    name='flask_rmq.events',
+    durable=True,
+    dead_letter_exchange='flask_rmq.dlx',
+    dead_letter_routing_key='flask_rmq.events.failed',
+)
+producer = Producer(exchange='flask_rmq.events', queue=EVENTS_QUEUE)
+consumer = Consumer(queue=EVENTS_QUEUE)
+
+
+def setup_topology(channel: BlockingChannel) -> None:
+    channel.exchange_declare(exchange='flask_rmq.events', exchange_type='direct', durable=True)
+    channel.exchange_declare(exchange='flask_rmq.dlx', exchange_type='direct', durable=True)
+    channel.queue_declare(queue=EVENTS_QUEUE.name, durable=True, arguments=EVENTS_QUEUE.arguments)
+    channel.queue_bind(
+        exchange='flask_rmq.events',
+        queue=EVENTS_QUEUE.name,
+        routing_key=EVENTS_QUEUE.name,
+    )
+    channel.queue_declare(queue='flask_rmq.events.failed', durable=True)
+    channel.queue_bind(
+        exchange='flask_rmq.dlx',
+        queue='flask_rmq.events.failed',
+        routing_key='flask_rmq.events.failed',
+    )
+
+
+@consumer
+def save_event(channel, method, properties, body: bytes) -> None:
+    output = Path(current_app.config['CONSUMED_FILE'])
+    with output.open('ab') as stream:
+        stream.write(body + b'\n')
+    current_app.logger.info('Consumed: %s', body.decode('utf-8'))
+    channel.basic_ack(delivery_tag=method.delivery_tag)
+
+
+def register_messaging() -> None:
+    get_setup_registry().register(setup_topology)
+    get_consumers_registry().register(consumer)

examples/basic_app/basic_app/views.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import json
+from datetime import UTC, datetime
+from uuid import uuid4
+
+from flask import Blueprint, jsonify, request
+
+from .messaging import producer
+
+api = Blueprint('api', __name__)
+
+
+@api.get('/health')
+def health():
+    return {'status': 'ok'}
+
+
+@api.post('/events')
+def publish_event():
+    incoming = request.get_json(silent=True) or {}
+    event = {
+        'id': str(uuid4()),
+        'created_at': datetime.now(UTC).isoformat(),
+        'payload': incoming,
+    }
+    producer.publish(json.dumps(event))
+    return jsonify(event), 202

flask_rmq/__init__.py

@@ -0,0 +1,28 @@
+from flask_rmq.consumer import Consumer, MessageCallback
+from flask_rmq.exceptions import ConfigurationError, FlaskRMQError, NotInitializedError
+from flask_rmq.extension import FlaskRMQ
+from flask_rmq.producer import Producer
+from flask_rmq.queues import QueueConfig
+from flask_rmq.registries import (
+    ConsumersRegistry,
+    SetupRegistry,
+    get_consumers_registry,
+    get_setup_registry,
+)
+
+__all__ = [
+    'ConfigurationError',
+    'Consumer',
+    'ConsumersRegistry',
+    'FlaskRMQ',
+    'FlaskRMQError',
+    'MessageCallback',
+    'NotInitializedError',
+    'Producer',
+    'QueueConfig',
+    'SetupRegistry',
+    'get_consumers_registry',
+    'get_setup_registry',
+]
+
+__version__ = '0.1.0'

flask_rmq/cli.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import signal
+import threading
+from types import FrameType
+
+import click
+from flask import Flask, current_app
+from flask.cli import with_appcontext
+
+from flask_rmq.consumer import Consumer
+from flask_rmq.state import get_state
+
+
+def _aliases(using: str | None) -> list[str]:
+    state = get_state()
+    if using is not None:
+        if using not in state.connection_managers:
+            raise click.ClickException(f'Unknown RabbitMQ alias {using!r}')
+        return [using]
+    return list(state.connection_managers)
+
+
+@click.group('rmq')
+def rmq_cli() -> None:
+    """Set up RabbitMQ topology and run registered consumers."""
+
+
+@rmq_cli.command('check')
+@with_appcontext
+def check_command() -> None:
+    """Validate configuration and show registered objects without connecting."""
+
+    state = get_state()
+    for alias in state.connection_managers:
+        setups = len(state.setup_registries[alias].all())
+        consumers = len(state.consumers_registries[alias].all())
+        click.echo(f'{alias}: configuration OK, {setups} setup callback(s), {consumers} consumer(s)')
+
+
+@rmq_cli.command('setup')
+@click.option('--using', help='Configure only this RABBITMQ_CONNECTIONS alias.')
+@with_appcontext
+def setup_command(using: str | None) -> None:
+    """Run all registered idempotent topology callbacks."""
+
+    state = get_state()
+    for alias in _aliases(using):
+        callbacks = state.setup_registries[alias].all()
+        click.echo(f'[{alias}] applying {len(callbacks)} topology callback(s)')
+        channel = state.connection_managers[alias].get_producer_channel()
+        for callback in callbacks:
+            callback(channel)
+            click.echo(f'  ✓ {callback.__module__}.{callback.__name__}')
+
+
+def _run_consumer(app: Flask, consumer: Consumer, stop_event: threading.Event) -> None:
+    with app.app_context():
+        consumer.consume(stop_event=stop_event)
+
+
+@rmq_cli.command('consume')
+@click.option('--using', help='Start only consumers registered for this alias.')
+@with_appcontext
+def consume_command(using: str | None) -> None:
+    """Run every registered consumer in its own thread until SIGINT/SIGTERM."""
+
+    state = get_state()
+    app = current_app
+    pairs = [(alias, consumer) for alias in _aliases(using) for consumer in state.consumers_registries[alias].all()]
+    if not pairs:
+        click.echo('No consumers registered.')
+        return
+
+    stop_event = threading.Event()
+
+    def request_stop(_signum: int, _frame: FrameType | None) -> None:
+        stop_event.set()
+
+    previous_sigint = signal.signal(signal.SIGINT, request_stop)
+    previous_sigterm = signal.signal(signal.SIGTERM, request_stop)
+
+    threads = []
+    try:
+        for alias, consumer in pairs:
+            thread = threading.Thread(
+                target=_run_consumer,
+                args=(app, consumer, stop_event),
+                name=f'flask-rmq-{alias}-{consumer.queue}',
+                daemon=False,
+            )
+            thread.start()
+            threads.append(thread)
+            click.echo(f'[{alias}] started {consumer.queue} → {consumer.handler_name}')
+        for thread in threads:
+            while thread.is_alive():
+                thread.join(timeout=0.5)
+    except KeyboardInterrupt:
+        stop_event.set()
+    finally:
+        stop_event.set()
+        for thread in threads:
+            thread.join(timeout=5)
+        signal.signal(signal.SIGINT, previous_sigint)
+        signal.signal(signal.SIGTERM, previous_sigterm)
+
+
+def main() -> None:
+    """Explain the app-bound entry point when invoked as ``flask-rmq``."""
+
+    raise click.ClickException('Flask-RMQ commands are app-bound. Run: flask --app your_module rmq --help')

flask_rmq/config.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from flask_rmq.exceptions import ConfigurationError
+
+
+@dataclass(frozen=True, slots=True)
+class RabbitMQConfig:
+    """Resolved settings for one independent RabbitMQ broker."""
+
+    host: str
+    port: int = 5672
+    virtual_host: str = '/'
+    user: str = 'guest'
+    password: str = 'guest'
+    heartbeat: int = 600
+    blocked_connection_timeout: float = 300.0
+    reconnect_initial_backoff: float = 1.0
+    reconnect_max_backoff: float = 30.0
+
+    @classmethod
+    def from_mapping(cls, alias: str, values: Mapping[str, Any]) -> RabbitMQConfig:
+        """Build and validate a config, accepting Flask-style uppercase keys."""
+
+        normalized = {str(key).lower(): value for key, value in values.items()}
+        if not normalized.get('host'):
+            raise ConfigurationError(f'RABBITMQ_CONNECTIONS[{alias!r}] requires HOST')
+        try:
+            config = cls(
+                host=str(normalized['host']),
+                port=int(normalized.get('port', 5672)),
+                virtual_host=str(normalized.get('virtual_host', '/')),
+                user=str(normalized.get('user', 'guest')),
+                password=str(normalized.get('password', 'guest')),
+                heartbeat=int(normalized.get('heartbeat', 600)),
+                blocked_connection_timeout=float(normalized.get('blocked_connection_timeout', 300.0)),
+                reconnect_initial_backoff=float(normalized.get('reconnect_initial_backoff', 1.0)),
+                reconnect_max_backoff=float(normalized.get('reconnect_max_backoff', 30.0)),
+            )
+        except (TypeError, ValueError) as exc:
+            raise ConfigurationError(f'Invalid RabbitMQ values for alias {alias!r}: {exc}') from exc
+        if not 1 <= config.port <= 65535:
+            raise ConfigurationError(f'RabbitMQ PORT for {alias!r} must be between 1 and 65535')
+        if config.heartbeat < 0 or config.blocked_connection_timeout <= 0:
+            raise ConfigurationError(f'RabbitMQ timeouts for {alias!r} must be positive')
+        if config.reconnect_initial_backoff <= 0 or config.reconnect_max_backoff <= 0:
+            raise ConfigurationError(f'RabbitMQ reconnect backoffs for {alias!r} must be positive')
+        if config.reconnect_initial_backoff > config.reconnect_max_backoff:
+            raise ConfigurationError(
+                f'RabbitMQ RECONNECT_INITIAL_BACKOFF for {alias!r} cannot exceed RECONNECT_MAX_BACKOFF'
+            )
+        return config

flask_rmq/connections.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import logging
+import threading
+
+from pika import BlockingConnection, ConnectionParameters, PlainCredentials
+from pika.adapters.blocking_connection import BlockingChannel
+from pika.exceptions import AMQPError
+
+from flask_rmq.config import RabbitMQConfig
+from flask_rmq.state import get_state, resolve_alias
+
+logger = logging.getLogger('rabbitmq')
+
+
+def get_connection_manager(using: str | None = None) -> RabbitMQConnectionManager:
+    return resolve_alias(get_state().connection_managers, using)
+
+
+class RabbitMQConnectionManager:
+    """Own thread-local, role-separated Pika connections for one broker alias."""
+
+    def __init__(self, config: RabbitMQConfig) -> None:
+        self.config = config
+        self._parameters = ConnectionParameters(
+            host=config.host,
+            port=config.port,
+            virtual_host=config.virtual_host,
+            credentials=PlainCredentials(config.user, config.password),
+            heartbeat=config.heartbeat,
+            blocked_connection_timeout=config.blocked_connection_timeout,
+        )
+        self._local = threading.local()
+
+    def _connection(self, attr: str) -> BlockingConnection:
+        connection = getattr(self._local, attr, None)
+        if connection is None or not connection.is_open:
+            connection = BlockingConnection(self._parameters)
+            setattr(self._local, attr, connection)
+        return connection
+
+    def get_producer_connection(self) -> BlockingConnection:
+        return self._connection('producer_connection')
+
+    def get_consumer_connection(self) -> BlockingConnection:
+        return self._connection('consumer_connection')
+
+    def get_producer_channel(self) -> BlockingChannel:
+        connection = self.get_producer_connection()
+        channel = getattr(self._local, 'producer_channel', None)
+        if channel is None or not channel.is_open:
+            channel = connection.channel()
+            channel.confirm_delivery()
+            self._local.producer_channel = channel
+        return channel
+
+    def reset_producer_channel(self) -> None:
+        """Close and forget producer resources after a transport failure."""
+
+        for attr in ('producer_channel', 'producer_connection'):
+            resource = getattr(self._local, attr, None)
+            if resource is None:
+                continue
+            try:
+                if resource.is_open:
+                    resource.close()
+            except (AMQPError, OSError):
+                logger.warning('Failed to close stale RabbitMQ %s', attr, exc_info=True)
+            finally:
+                delattr(self._local, attr)
+
+    def close_thread_connections(self) -> None:
+        """Close all resources owned by the calling thread."""
+
+        self.reset_producer_channel()
+        connection = getattr(self._local, 'consumer_connection', None)
+        if connection is not None:
+            try:
+                if connection.is_open:
+                    connection.close()
+            except (AMQPError, OSError):
+                logger.warning('Failed to close consumer connection', exc_info=True)
+            finally:
+                delattr(self._local, 'consumer_connection')

flask_rmq/consumer.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import functools
+import logging
+import threading
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from pika.exceptions import (
+    AMQPConnectionError,
+    ChannelClosed,
+    ChannelClosedByBroker,
+    ConnectionClosed,
+    StreamLostError,
+)
+
+from flask_rmq.connections import get_connection_manager
+from flask_rmq.queues import QueueConfig
+
+if TYPE_CHECKING:
+    from pika.adapters.blocking_connection import BlockingChannel, BlockingConnection
+    from pika.spec import Basic, BasicProperties
+
+MessageCallback = Callable[['BlockingChannel', 'Basic.Deliver', 'BasicProperties', bytes], None]
+logger = logging.getLogger('rabbitmq')
+_RECONNECTABLE_ERRORS = (
+    AMQPConnectionError,
+    ConnectionClosed,
+    ChannelClosed,
+    ChannelClosedByBroker,
+    StreamLostError,
+    ConnectionResetError,
+)
+
+
+class Consumer:
+    """Consume one queue with reconnect backoff and graceful shutdown."""
+
+    def __init__(
+        self,
+        queue: QueueConfig | str,
+        prefetch_count: int = 1,
+        reconnect_initial_backoff: float | None = None,
+        reconnect_max_backoff: float | None = None,
+        using: str | None = None,
+    ) -> None:
+        if prefetch_count < 1:
+            raise ValueError('prefetch_count must be at least 1')
+        self.queue = str(queue)
+        self._queue_config = queue
+        self._prefetch_count = prefetch_count
+        self._using = using
+        self._initial_override = reconnect_initial_backoff
+        self._max_override = reconnect_max_backoff
+        self._handler: MessageCallback | None = None
+
+    @property
+    def prefetch_count(self) -> int:
+        return self._prefetch_count
+
+    @property
+    def using(self) -> str | None:
+        return self._using
+
+    @property
+    def handler_name(self) -> str:
+        return self._handler.__name__ if self._handler else 'unregistered'
+
+    def handler(self, func: MessageCallback) -> MessageCallback:
+        if self._handler is not None:
+            raise RuntimeError(f'Consumer for queue {self.queue!r} already has a handler')
+        self._handler = func
+        return func
+
+    def __call__(self, func: MessageCallback) -> MessageCallback:
+        return self.handler(func)
+
+    def consume(self, stop_event: threading.Event | None = None) -> None:
+        if self._handler is None:
+            raise RuntimeError(f'Consumer for queue {self.queue!r} has no handler')
+        stop_event = stop_event or threading.Event()
+        config = get_connection_manager(self._using).config
+        initial = self._initial_override if self._initial_override is not None else config.reconnect_initial_backoff
+        maximum = self._max_override if self._max_override is not None else config.reconnect_max_backoff
+        backoff = initial
+        while not stop_event.is_set():
+            try:
+                self._run_session(self._handler, stop_event)
+                return
+            except _RECONNECTABLE_ERRORS:
+                logger.warning(
+                    'RabbitMQ consumer %s disconnected; retrying in %.2fs',
+                    self.queue,
+                    backoff,
+                    exc_info=True,
+                )
+                stop_event.wait(backoff)
+                backoff = min(backoff * 2, maximum)
+
+    def _run_session(self, handler: MessageCallback, stop_event: threading.Event) -> None:
+        connection: BlockingConnection = get_connection_manager(self._using).get_consumer_connection()
+        channel: BlockingChannel = connection.channel()
+        try:
+            self._declare_queue(channel)
+            channel.basic_qos(prefetch_count=self._prefetch_count)
+            channel.basic_consume(queue=self.queue, on_message_callback=functools.partial(self._dispatch, handler))
+            logger.info('Consuming queue %s with %s', self.queue, handler.__name__)
+            while not stop_event.is_set():
+                connection.process_data_events(time_limit=1)
+        finally:
+            if channel.is_open:
+                channel.stop_consuming()
+                channel.close()
+
+    def _declare_queue(self, channel: BlockingChannel) -> None:
+        if isinstance(self._queue_config, QueueConfig):
+            channel.queue_declare(
+                queue=self.queue,
+                durable=self._queue_config.durable,
+                arguments=self._queue_config.arguments,
+            )
+        else:
+            channel.queue_declare(queue=self.queue, durable=True)
+
+    def _dispatch(
+        self,
+        handler: MessageCallback,
+        ch: BlockingChannel,
+        method: Basic.Deliver,
+        props: BasicProperties,
+        body: bytes,
+    ) -> None:
+        try:
+            handler(ch, method, props, body)
+        except Exception:
+            logger.exception('Handler %s failed for queue %s; nacking', handler.__name__, self.queue)
+            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)

flask_rmq/exceptions.py

@@ -0,0 +1,10 @@
+class FlaskRMQError(Exception):
+    """Base exception raised by Flask-RMQ."""
+
+
+class ConfigurationError(FlaskRMQError):
+    """Raised when Flask-RMQ configuration is missing or invalid."""
+
+
+class NotInitializedError(FlaskRMQError):
+    """Raised when an API is used before :class:`FlaskRMQ` initialization."""

flask_rmq/extension.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, cast
+
+from flask import Flask
+
+from flask_rmq.config import RabbitMQConfig
+from flask_rmq.connections import RabbitMQConnectionManager
+from flask_rmq.exceptions import ConfigurationError
+from flask_rmq.registries import ConsumersRegistry, SetupRegistry
+from flask_rmq.state import RMQState, register_state
+
+
+class FlaskRMQ:
+    """Flask extension that configures RabbitMQ aliases and Click commands."""
+
+    def __init__(self, app: Flask | None = None) -> None:
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app: Flask) -> None:
+        if 'flask-rmq' in app.extensions:
+            raise ConfigurationError('Flask-RMQ is already initialized for this app')
+        raw_connections: Any = app.config.get('RABBITMQ_CONNECTIONS')
+        if not isinstance(raw_connections, Mapping) or not raw_connections:
+            raise ConfigurationError('Flask-RMQ requires a non-empty RABBITMQ_CONNECTIONS mapping in app.config')
+
+        managers: dict[str, RabbitMQConnectionManager] = {}
+        setups: dict[str, SetupRegistry] = {}
+        consumers: dict[str, ConsumersRegistry] = {}
+        for raw_alias, values in raw_connections.items():
+            alias = str(raw_alias)
+            if not isinstance(values, Mapping):
+                raise ConfigurationError(f'RABBITMQ_CONNECTIONS[{alias!r}] must be a mapping')
+            managers[alias] = RabbitMQConnectionManager(RabbitMQConfig.from_mapping(alias, values))
+            setups[alias] = SetupRegistry()
+            consumers[alias] = ConsumersRegistry()
+
+        state = RMQState(managers, setups, consumers)
+        app.extensions['flask-rmq'] = state
+        register_state(app, state)
+
+        from .cli import rmq_cli
+
+        app.cli.add_command(rmq_cli)
+
+    @staticmethod
+    def get_state(app: Flask) -> RMQState:
+        try:
+            return cast(RMQState, app.extensions['flask-rmq'])
+        except KeyError as exc:
+            raise ConfigurationError('Flask-RMQ is not initialized for this app') from exc