acemq_fastapi.integration

  1# Copyright 2026 AceMQ.
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     https://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""The object an application holds, and the lifespan that owns the connection.
 16
 17One :class:`AceMQ` per application. It is built at import time, collects consumer
 18registrations from decorators, and does nothing at all until its lifespan opens —
 19which is the point: a module that connects to a broker when it is imported is a
 20module that cannot be imported by a test, a migration script or ``--help``.
 21"""
 22
 23from __future__ import annotations
 24
 25import asyncio
 26import logging
 27from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
 28from contextlib import asynccontextmanager, suppress
 29from dataclasses import dataclass, field
 30from typing import TYPE_CHECKING, Any, TypeVar
 31
 32import acemq_amqp
 33from acemq_amqp import (
 34    AceMQError,
 35    Ack,  # noqa: F401 -- see below
 36    Codec,
 37    Connection,
 38    ConsumeInterceptor,
 39    Consumer,
 40    Handler,
 41    Observer,
 42    Publisher,
 43    PublishInterceptor,
 44    RetryPolicy,
 45)
 46
 47from .settings import AceMQSettings
 48
 49if TYPE_CHECKING:  # pragma: no cover - typing only
 50    from fastapi import FastAPI
 51
 52    from .health import AceMQHealth
 53
 54__all__ = ["AceMQ", "ConsumerRegistration", "DrainReport"]
 55
 56log = logging.getLogger("acemq_fastapi")
 57
 58# ``Ack`` is imported but not used here, and that is on purpose. The library's
 59# ``Handler`` alias is ``Callable[[Message], "Ack | Awaitable[Ack]"]`` — a forward
 60# reference — and anything resolving the annotations of a function in this module
 61# evaluates that string against *this* module's globals. Without the name here,
 62# pdoc renders the registry's signatures as unresolved and a reader of the API
 63# reference sees a handler type that says nothing.
 64
 65#: What builds the connection. Replaceable so a test can hand over a connection
 66#: on a fake transport, and so an application that needs a connection this
 67#: package does not model — an unusual ``transport_options``, a codec built at
 68#: runtime — can build its own and keep everything else.
 69ConnectionFactory = Callable[[AceMQSettings], Awaitable[Connection]]
 70
 71#: Bound to :data:`acemq_amqp.Handler` so the decorator gives back the function it
 72#: was handed, with its own signature: a handler stays as callable, and as typed,
 73#: after decoration as it was before.
 74HandlerT = TypeVar("HandlerT", bound=Handler)
 75
 76
 77@dataclass(frozen=True, slots=True)
 78class ConsumerRegistration:
 79    """One ``@acemq.consumer`` declaration, before anything is running.
 80
 81    ``None`` on the per-consumer settings means "whatever ``acemq.consumer.*``
 82    says", resolved when the lifespan opens rather than when the decorator runs,
 83    so settings read from the environment still reach a consumer registered at
 84    import time.
 85    """
 86
 87    #: What this consumer is called to :meth:`AceMQ.consumers`. Defaults to the
 88    #: queue name, and has to be unique because that is how it is addressed.
 89    name: str
 90    queue: str
 91    handler: Handler
 92    codec: Codec | None = None
 93    retry: RetryPolicy | None = None
 94    prefetch: int | None = None
 95    concurrency: int | None = None
 96    declare: bool | None = None
 97    #: What to call this consumer to the broker. Empty lets the broker name it.
 98    tag: str = ""
 99    args: Mapping[str, Any] = field(default_factory=dict)
100
101
102@dataclass(frozen=True, slots=True)
103class DrainReport:
104    """What a shutdown got through.
105
106    Returned by :meth:`AceMQ.drain` and logged by the lifespan. ``finished``
107    false means the deadline expired: the handlers still running were cancelled
108    and their messages left unsettled, so the broker will hand them to whoever
109    replaces this process. Nothing is lost; some work may be done twice.
110    """
111
112    finished: bool
113    #: Seconds spent draining.
114    waited: float
115    #: How many consumers were asked to stop.
116    consumers: int
117    #: How many messages were still in a handler when the deadline expired.
118    abandoned: int = 0
119
120
121class AceMQ:
122    """The connection, the consumers and the settings, for one application.
123
124    ::
125
126        acemq = AceMQ()
127
128        @acemq.consumer("orders")
129        async def handle(message: Message) -> Ack:
130            ...
131            return accept()
132
133        app = FastAPI(lifespan=acemq.lifespan)
134
135    :param settings: what to connect to and how. Read from the environment when
136        not given
137    :param connection_factory: what opens the connection. Replaceable for tests,
138        and for an application that builds its own
139    :param on_publish: interceptors wrapping every publish, outermost first
140    :param on_consume: interceptors wrapping every handler
141    :param observer: where the library's numbers go
142    """
143
144    def __init__(
145        self,
146        settings: AceMQSettings | None = None,
147        *,
148        connection_factory: ConnectionFactory | None = None,
149        on_publish: tuple[PublishInterceptor, ...] = (),
150        on_consume: tuple[ConsumeInterceptor, ...] = (),
151        observer: Observer | None = None,
152    ) -> None:
153        self._settings = settings if settings is not None else AceMQSettings()
154        self._connect: ConnectionFactory = connection_factory or self._open
155        self._on_publish = tuple(on_publish)
156        self._on_consume = tuple(on_consume)
157        self._observer = observer
158        self._registrations: dict[str, ConsumerRegistration] = {}
159        self._consumers: dict[str, Consumer] = {}
160        self._publishers: dict[tuple[Any, ...], Publisher] = {}
161        self._connection: Connection | None = None
162        self._lock = asyncio.Lock()
163
164    # ------------------------------------------------------------------ state
165
166    @property
167    def settings(self) -> AceMQSettings:
168        """What this integration was configured with."""
169        return self._settings
170
171    @property
172    def started(self) -> bool:
173        """Whether the connection is open."""
174        return self._connection is not None
175
176    @property
177    def connection(self) -> Connection:
178        """The open connection.
179
180        :raises AceMQError: before the lifespan has opened it, because the
181            alternative is a ``None`` that turns into an ``AttributeError``
182            three frames away from the mistake
183        """
184        if self._connection is None:
185            raise AceMQError(
186                "acemq-fastapi: there is no connection yet. It is opened by the lifespan, "
187                "so this is either a call made before the application started or an "
188                "application built without lifespan=acemq.lifespan"
189            )
190        return self._connection
191
192    @property
193    def registrations(self) -> tuple[ConsumerRegistration, ...]:
194        """Every ``@acemq.consumer`` declaration, in the order they were made."""
195        return tuple(self._registrations.values())
196
197    @property
198    def consumers(self) -> Mapping[str, Consumer]:
199        """The running consumers by name, for scaling, pausing and counters."""
200        return dict(self._consumers)
201
202    # -------------------------------------------------------------- consumers
203
204    def consumer(
205        self,
206        queue: str,
207        *,
208        name: str | None = None,
209        codec: Codec | None = None,
210        retry: RetryPolicy | None = None,
211        prefetch: int | None = None,
212        concurrency: int | None = None,
213        declare: bool | None = None,
214        tag: str = "",
215        args: Mapping[str, Any] | None = None,
216    ) -> Callable[[HandlerT], HandlerT]:
217        """Registers a handler for a queue. The Python analogue of ``@AceListener``.
218
219        The decorated function is returned unchanged, so it stays an ordinary
220        function that a test can call with a :class:`~acemq_amqp.Message` and
221        assert on, without a broker anywhere near it.
222
223        Nothing subscribes here. The registration is held until the lifespan
224        opens, which is what puts the handler on the application's own event loop
225        rather than on whichever loop happened to import the module.
226
227        :param queue: what to read
228        :param name: what to call this consumer in :attr:`consumers`. Defaults to
229            the queue
230        :param codec: a codec other than the connection's
231        :param retry: a policy other than ``acemq.consumer.retry``
232        :param prefetch: how many unacknowledged messages to hold
233        :param concurrency: how many messages to work on at once
234        :param declare: declare the queues a failure needs before subscribing
235        :param tag: what to call this consumer to the broker
236        :param args: broker-specific consumer arguments
237        :returns: the decorator
238        :raises AceMQError: when the name is already taken, or the lifespan has
239            already opened
240        """
241
242        def register(handler: HandlerT) -> HandlerT:
243            self.add_consumer(
244                queue,
245                handler,
246                name=name,
247                codec=codec,
248                retry=retry,
249                prefetch=prefetch,
250                concurrency=concurrency,
251                declare=declare,
252                tag=tag,
253                args=args,
254            )
255            return handler
256
257        return register
258
259    def add_consumer(
260        self,
261        queue: str,
262        handler: Handler,
263        *,
264        name: str | None = None,
265        codec: Codec | None = None,
266        retry: RetryPolicy | None = None,
267        prefetch: int | None = None,
268        concurrency: int | None = None,
269        declare: bool | None = None,
270        tag: str = "",
271        args: Mapping[str, Any] | None = None,
272    ) -> ConsumerRegistration:
273        """Registers a handler without the decorator.
274
275        What a factory function or a plugin uses. Same arguments as
276        :meth:`consumer`, and the same restriction: before the lifespan opens.
277        """
278        if self._connection is not None:
279            raise AceMQError(
280                f"acemq-fastapi: cannot register a consumer for {queue!r} now — the "
281                "lifespan has already opened. Register consumers while the module is "
282                "being imported, which is what the decorator is for"
283            )
284        chosen = name or queue
285        if chosen in self._registrations:
286            raise AceMQError(
287                f"acemq-fastapi: there is already a consumer called {chosen!r}. Two "
288                "consumers on one queue is a reasonable thing to want; give one of them "
289                "a name=... so they can be told apart"
290            )
291        registration = ConsumerRegistration(
292            name=chosen,
293            queue=queue,
294            handler=handler,
295            codec=codec,
296            retry=retry,
297            prefetch=prefetch,
298            concurrency=concurrency,
299            declare=declare,
300            tag=tag,
301            args=dict(args or {}),
302        )
303        self._registrations[chosen] = registration
304        return registration
305
306    async def start_consumer(self, name: str) -> Consumer:
307        """Subscribes one registered consumer that is not running.
308
309        What :attr:`consumers` is for in the other direction: a consumer stopped
310        for back pressure, or one left registered by
311        ``acemq.consumer.auto_start=false``, is started again by name.
312        """
313        registration = self._registrations.get(name)
314        if registration is None:
315            known = ", ".join(sorted(self._registrations)) or "none"
316            raise AceMQError(
317                f"acemq-fastapi: no consumer is registered as {name!r}; registered: {known}"
318            )
319        if name in self._consumers:
320            return self._consumers[name]
321        consumer = await self._subscribe(registration)
322        self._consumers[name] = consumer
323        return consumer
324
325    async def stop_consumer(self, name: str) -> None:
326        """Drains and closes one consumer, leaving it registered.
327
328        It finishes what its handlers are holding, exactly as a shutdown does.
329        """
330        consumer = self._consumers.pop(name, None)
331        if consumer is not None:
332            await consumer.close()
333
334    # ------------------------------------------------------------- publishing
335
336    def publisher(
337        self,
338        exchange: str = "",
339        routing_key: str = "",
340        *,
341        codec: Codec | None = None,
342        persistent: bool = True,
343        mandatory: bool = False,
344    ) -> Publisher:
345        """A publisher for a destination, built once and reused.
346
347        Publishers are cheap to keep and safe to share, and building one per
348        request is a small waste repeated a lot. This caches by every argument
349        that changes what a publisher does, so two routes asking for the same
350        destination get the same object and a route asking for a different one
351        does not quietly get somebody else's.
352
353        :raises AceMQError: before the lifespan has opened the connection
354        """
355        key = (exchange, routing_key, id(codec) if codec else None, persistent, mandatory)
356        publisher = self._publishers.get(key)
357        if publisher is None:
358            publisher = self.connection.publisher(
359                exchange,
360                routing_key,
361                codec=codec,
362                persistent=persistent,
363                mandatory=mandatory,
364            )
365            self._publishers[key] = publisher
366        return publisher
367
368    # ----------------------------------------------------------------- health
369
370    def health_check(self) -> AceMQHealth:
371        """This connection as a :class:`~acemq_amqp.HealthCheck`.
372
373        Pass it to :func:`acemq_amqp.aggregate_health` beside an application's
374        own checks, or mount :meth:`health_router` and let it answer on its own.
375        """
376        from .health import AceMQHealth
377
378        return AceMQHealth(self)
379
380    async def health(self) -> acemq_amqp.HealthReport:
381        """What a readiness probe should see. See :class:`~acemq_fastapi.AceMQHealth`."""
382        return await self.health_check().check()
383
384    def health_router(self, *, path: str | None = None, tags: list[Any] | None = None) -> Any:
385        """An ``APIRouter`` with one route reporting on the connection.
386
387        ``app.include_router(acemq.health_router())`` and nothing else. Mounting
388        it is explicit rather than automatic because a route that appears in an
389        application's OpenAPI document without anybody adding it is a route
390        somebody has to go looking for.
391        """
392        from .health import health_router
393
394        return health_router(self, path=path, tags=tags)
395
396    # --------------------------------------------------------------- lifespan
397
398    @asynccontextmanager
399    async def lifespan(self, app: FastAPI | None = None) -> AsyncIterator[Mapping[str, Any]]:
400        """Opens the connection and the consumers, and drains them on the way out.
401
402        ``FastAPI(lifespan=acemq.lifespan)``. To keep an application's own
403        lifespan as well, see :func:`acemq_fastapi.compose_lifespans` — this one
404        belongs innermost, so the consumers start after whatever they depend on
405        and drain before it goes away.
406
407        The ``app`` is optional so the same context manager can be used in a test
408        or a worker with no ASGI application at all. When one is given it gets
409        ``app.state.acemq``, which is how the request dependencies find this
410        object without a module-level global.
411        """
412        started = await self.start(app)
413        try:
414            yield {"acemq": self}
415        finally:
416            report = await self.drain()
417            if not report.finished:
418                log.warning(
419                    "acemq-fastapi: the drain did not finish in %s; %d message(s) were "
420                    "left unsettled and the broker will redeliver them",
421                    self._settings.consumer.shutdown_timeout,
422                    report.abandoned,
423                )
424            await self.aclose()
425            log.info(
426                "acemq-fastapi: closed after %.3fs, %d consumer(s) drained",
427                report.waited,
428                started,
429            )
430
431    async def start(self, app: FastAPI | None = None) -> int:
432        """Connects, applies the topology and subscribes the consumers.
433
434        Called by :meth:`lifespan`. Public because a worker process that serves no
435        HTTP at all still wants the same start-up, and should not have to
436        construct an ASGI application to get it.
437
438        :returns: how many consumers were started
439        """
440        async with self._lock:
441            if self._connection is not None:
442                raise AceMQError("acemq-fastapi: this integration is already started")
443            timeout = self._settings.connection_timeout.total_seconds()
444            try:
445                connection = await asyncio.wait_for(self._connect(self._settings), timeout)
446            except asyncio.TimeoutError as expired:
447                raise AceMQError(
448                    f"acemq-fastapi: the broker at {self._settings.broker_url} did not "
449                    f"answer within {timeout:g}s"
450                ) from expired
451            self._connection = connection
452
453        try:
454            for publishing in self._on_publish:
455                connection.intercept_publish(publishing)
456            for consuming in self._on_consume:
457                connection.intercept_consume(consuming)
458
459            topology = self._settings.build_topology()
460            if topology is not None:
461                await connection.declare(topology)
462
463            if self._settings.consumer.auto_start:
464                for registration in self._registrations.values():
465                    self._consumers[registration.name] = await self._subscribe(registration)
466        except BaseException:
467            # A half-started application is worse than one that refused to start:
468            # the routes would be up, the consumers would not, and nothing outside
469            # could tell. Undo what was done and let the failure out.
470            await self.drain()
471            await self.aclose()
472            raise
473
474        if app is not None:
475            app.state.acemq = self
476        log.info(
477            "acemq-fastapi: connected to %s, %d consumer(s) running",
478            self._settings.broker_url,
479            len(self._consumers),
480        )
481        return len(self._consumers)
482
483    async def drain(self) -> DrainReport:
484        """Stops every consumer, finishing the handlers already running.
485
486        What this finishes, and what it gives back, is written down in the README
487        under *What a drain finishes*. In short: a handler in flight runs to
488        completion; a delivery fetched but not yet started is handed back to the
489        broker; a retry waiting out a short backoff in this process is waited out
490        in full; a publish waiting for its confirm is not waited for.
491
492        Bounded by ``acemq.consumer.shutdown_timeout``. When that expires the
493        handlers still running are cancelled and their messages left unsettled —
494        the broker redelivers them, so the work may be done twice and nothing is
495        dropped.
496        """
497        if self._connection is None or not self._consumers:
498            self._consumers.clear()
499            return DrainReport(finished=True, waited=0.0, consumers=0)
500
501        consumers = tuple(self._consumers.values())
502        self._consumers.clear()
503        deadline = self._settings.consumer.shutdown_timeout.total_seconds()
504        loop = asyncio.get_running_loop()
505        started = loop.time()
506
507        closing = [asyncio.create_task(consumer.close()) for consumer in consumers]
508        done, pending = await asyncio.wait(closing, timeout=deadline)
509        waited = loop.time() - started
510
511        if not pending:
512            for task in done:
513                if task.cancelled():
514                    # Its workers were already gone — cancelled from outside, or
515                    # killed by something the library logged. Nothing here can
516                    # finish a handler that is not running any more.
517                    log.warning(
518                        "acemq-fastapi: a consumer's workers had already stopped; "
519                        "whatever they were holding is back with the broker"
520                    )
521                elif task.exception() is not None:
522                    log.warning(
523                        "acemq-fastapi: a consumer failed to close: %s", task.exception()
524                    )
525            return DrainReport(finished=True, waited=waited, consumers=len(consumers))
526
527        # Out of time. Cancelling is the only lever there is, and it is pulled
528        # last rather than first: it aborts the handlers still running and leaves
529        # their deliveries unsettled for the broker to hand on. That is the price
530        # of a deadline, and it is cheaper than being SIGKILLed with the same
531        # messages unsettled and the socket gone.
532        abandoned = sum(consumer.in_flight for consumer in consumers)
533        for task in pending:
534            task.cancel()
535        await asyncio.gather(*pending, return_exceptions=True)
536        return DrainReport(
537            finished=False,
538            waited=loop.time() - started,
539            consumers=len(consumers),
540            abandoned=abandoned,
541        )
542
543    async def aclose(self) -> None:
544        """Releases the connection. Idempotent, and safe after a failed start."""
545        connection = self._connection
546        self._connection = None
547        self._publishers.clear()
548        self._consumers.clear()
549        if connection is not None:
550            with suppress(Exception):
551                await connection.close()
552
553    # ---------------------------------------------------------------- private
554
555    async def _subscribe(self, registration: ConsumerRegistration) -> Consumer:
556        defaults = self._settings.consumer
557        return await self.connection.consume(
558            registration.queue,
559            registration.handler,
560            codec=registration.codec,
561            retry=registration.retry or self._settings.build_retry(),
562            prefetch=registration.prefetch
563            if registration.prefetch is not None
564            else defaults.prefetch,
565            concurrency=registration.concurrency
566            if registration.concurrency is not None
567            else defaults.concurrency,
568            tag=registration.tag,
569            args=registration.args,
570            declare=(
571                registration.declare if registration.declare is not None else defaults.declare
572            ),
573        )
574
575    async def _open(self, settings: AceMQSettings) -> Connection:
576        """The default factory: :func:`acemq_amqp.connect` with the settings applied."""
577        return await acemq_amqp.connect(
578            settings.broker_url,
579            codec=settings.build_codec(),
580            origin=settings.client_name,
581            retry=settings.build_retry(),
582            prefetch=settings.prefetch,
583            max_outstanding_publishes=settings.max_outstanding_publishes,
584            confirm_timeout=settings.confirm_timeout,
585            security=settings.build_security(),
586            observer=self._observer,
587        )