✨ Add RabbitMQ lifespan management and validation in FastAPI

Author: pcrespov

packages/service-library/src/servicelib/fastapi/rabbitmq_lifespan.py

@@ -0,0 +1,48 @@
+import logging
+from collections.abc import AsyncIterator
+
+from fastapi import FastAPI
+from fastapi_lifespan_manager import State
+from pydantic import BaseModel, ValidationError
+from servicelib.logging_utils import log_context
+from servicelib.rabbitmq import wait_till_rabbitmq_responsive
+from settings_library.rabbit import RabbitSettings
+
+from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
+
+_logger = logging.getLogger(__name__)
+
+
+class RabbitMQConfigurationError(LifespanOnStartupError):
+    msg_template = "Invalid RabbitMQ config on startup : {validation_error}"
+
+
+class RabbitMQLifespanState(BaseModel):
+    RABBIT_SETTINGS: RabbitSettings
+
+
+async def rabbitmq_connectivity_lifespan(
+    _: FastAPI, state: State
+) -> AsyncIterator[State]:
+    """Ensures RabbitMQ connectivity during lifespan.
+
+    For creating clients, use additional lifespans like rabbitmq_rpc_client_context.
+    """
+    _lifespan_name = f"{__name__}.{rabbitmq_connectivity_lifespan.__name__}"
+
+    with log_context(_logger, logging.INFO, _lifespan_name):
+
+        # Check if lifespan has already been called
+        called_state = record_lifespan_called_once(state, _lifespan_name)
+
+        # Validate input state
+        try:
+            rabbit_state = RabbitMQLifespanState.model_validate(state)
+            rabbit_dsn_with_secrets = rabbit_state.RABBIT_SETTINGS.dsn
+        except ValidationError as exc:
+            raise RabbitMQConfigurationError(validation_error=exc, state=state) from exc
+
+        # Wait for RabbitMQ to be responsive
+        await wait_till_rabbitmq_responsive(rabbit_dsn_with_secrets)
+
+        yield {"RABBIT_CONNECTIVITY_LIFESPAN_NAME": _lifespan_name, **called_state}

packages/service-library/src/servicelib/rabbitmq/__init__.py

@@ -1,7 +1,7 @@
 from models_library.rabbitmq_basic_types import RPCNamespace
 
 from ._client import RabbitMQClient
-from ._client_rpc import RabbitMQRPCClient
+from ._client_rpc import RabbitMQRPCClient, rabbitmq_rpc_client_context
 from ._constants import BIND_TO_ALL_TOPICS, RPC_REQUEST_DEFAULT_TIMEOUT_S
 from ._errors import (
     RemoteMethodNotRegisteredError,
@@ -28,6 +28,7 @@
     "RabbitMQRPCClient",
     "RemoteMethodNotRegisteredError",
     "is_rabbitmq_responsive",
+    "rabbitmq_rpc_client_context",
     "wait_till_rabbitmq_responsive",
 )
 

packages/service-library/src/servicelib/rabbitmq/_client_rpc.py

@@ -1,7 +1,8 @@
 import asyncio
 import functools
 import logging
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Any
 
@@ -156,3 +157,19 @@ async def unregister_handler(self, handler: Callable[..., Any]) -> None:
             raise RPCNotInitializedError
 
         await self._rpc.unregister(handler)
+
+
+@asynccontextmanager
+async def rabbitmq_rpc_client_context(
+    rpc_client_name: str, settings: RabbitSettings, **kwargs
+) -> AsyncIterator[RabbitMQRPCClient]:
+    """
+    Adapter to create and close a RabbitMQRPCClient using an async context manager.
+    """
+    rpc_client = await RabbitMQRPCClient.create(
+        client_name=rpc_client_name, settings=settings, **kwargs
+    )
+    try:
+        yield rpc_client
+    finally:
+        await rpc_client.close()

packages/service-library/tests/fastapi/test_rabbitmq_lifespan.py

@@ -0,0 +1,142 @@
+# pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
+# pylint: disable=too-many-arguments
+# pylint: disable=unused-argument
+# pylint: disable=unused-variable
+
+from collections.abc import AsyncIterator
+
+import pytest
+import servicelib.fastapi.rabbitmq_lifespan
+import servicelib.rabbitmq
+from asgi_lifespan import LifespanManager as ASGILifespanManager
+from fastapi import FastAPI
+from fastapi_lifespan_manager import LifespanManager, State
+from pydantic import Field
+from pytest_mock import MockerFixture, MockType
+from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
+from pytest_simcore.helpers.typing_env import EnvVarsDict
+from servicelib.fastapi.rabbitmq_lifespan import (
+    RabbitMQConfigurationError,
+    RabbitMQLifespanState,
+    rabbitmq_connectivity_lifespan,
+)
+from servicelib.rabbitmq import rabbitmq_rpc_client_context
+from settings_library.application import BaseApplicationSettings
+from settings_library.rabbit import RabbitSettings
+
+
+@pytest.fixture
+def mock_rabbitmq_connection(mocker: MockerFixture) -> MockType:
+    return mocker.patch.object(
+        servicelib.fastapi.rabbitmq_lifespan,
+        "wait_till_rabbitmq_responsive",
+        return_value=mocker.AsyncMock(),
+    )
+
+
+@pytest.fixture
+def mock_rabbitmq_rpc_client_class(mocker: MockerFixture) -> MockType:
+    return mocker.patch.object(
+        servicelib.rabbitmq._client_rpc,
+        "RabbitMQRPCClient",
+        return_value=mocker.AsyncMock(),
+    )
+
+
+@pytest.fixture
+def app_environment(monkeypatch: pytest.MonkeyPatch) -> EnvVarsDict:
+    return setenvs_from_dict(
+        monkeypatch, RabbitSettings.model_json_schema()["examples"][0]
+    )
+
+
+@pytest.fixture
+def app_lifespan(
+    app_environment: EnvVarsDict,
+    mock_rabbitmq_connection: MockType,
+    mock_rabbitmq_rpc_client_class: MockType,
+) -> LifespanManager:
+    assert app_environment
+
+    class AppSettings(BaseApplicationSettings):
+        RABBITMQ: RabbitSettings = Field(
+            ..., json_schema_extra={"auto_default_from_env": True}
+        )
+
+    async def my_app_settings(app: FastAPI) -> AsyncIterator[State]:
+        app.state.settings = AppSettings.create_from_envs()
+
+        yield RabbitMQLifespanState(
+            RABBIT_SETTINGS=app.state.settings.RABBITMQ,
+        ).model_dump()
+
+    async def my_app_rpc_server(app: FastAPI, state: State) -> AsyncIterator[State]:
+
+        async with rabbitmq_rpc_client_context(
+            "rpc_server", app.state.settings.RABBITMQ
+        ) as rpc_server:
+            app.state.rpc_server = rpc_server
+            yield {}
+
+    app_lifespan = LifespanManager()
+    app_lifespan.add(my_app_settings)
+    app_lifespan.add(rabbitmq_connectivity_lifespan)
+    app_lifespan.add(my_app_rpc_server)
+
+    assert not mock_rabbitmq_connection.called
+    assert not mock_rabbitmq_rpc_client_class.called
+
+    return app_lifespan
+
+
+async def test_lifespan_rabbitmq_in_an_app(
+    is_pdb_enabled: bool,
+    app_environment: EnvVarsDict,
+    mock_rabbitmq_connection: MockType,
+    mock_rabbitmq_rpc_client_class: MockType,
+    app_lifespan: LifespanManager,
+):
+    app = FastAPI(lifespan=app_lifespan)
+
+    async with ASGILifespanManager(
+        app,
+        startup_timeout=None if is_pdb_enabled else 10,
+        shutdown_timeout=None if is_pdb_enabled else 10,
+    ) as asgi_manager:
+        # Verify that RabbitMQ responsiveness was checked
+        mock_rabbitmq_connection.assert_called_once_with(
+            app.state.settings.RABBITMQ.dsn
+        )
+
+        # Verify that RabbitMQ settings are in the lifespan manager state
+        assert app.state.settings.RABBITMQ
+
+    # No explicit shutdown logic for RabbitMQ in this case
+    assert mock_rabbitmq_rpc_client_class.called
+
+
+async def test_lifespan_rabbitmq_with_invalid_settings(
+    is_pdb_enabled: bool,
+):
+    async def my_app_settings(app: FastAPI) -> AsyncIterator[State]:
+        yield {"RABBIT_SETTINGS": None}
+
+    app_lifespan = LifespanManager()
+    app_lifespan.add(my_app_settings)
+    app_lifespan.add(rabbitmq_connectivity_lifespan)
+
+    app = FastAPI(lifespan=app_lifespan)
+
+    with pytest.raises(RabbitMQConfigurationError, match="Invalid RabbitMQ") as excinfo:
+        async with ASGILifespanManager(
+            app,
+            startup_timeout=None if is_pdb_enabled else 10,
+            shutdown_timeout=None if is_pdb_enabled else 10,
+        ):
+            ...
+
+    exception = excinfo.value
+    assert isinstance(exception, RabbitMQConfigurationError)
+    assert exception.validation_error
+    assert exception.state["RABBIT_SETTINGS"] is None