@pcrespov use fastapi lifespan

Author: bisgaard-itis

services/storage/src/simcore_service_storage/core/application.py

@@ -32,7 +32,7 @@
 from ..exceptions.handlers import set_exception_handlers
 from ..modules.db import setup_db
 from ..modules.long_running_tasks import setup_rest_api_long_running_tasks_for_uploads
-from ..modules.rabbitmq import setup as setup_rabbitmq
+from ..modules.rabbitmq import rabbitmq_lifespan
 from ..modules.redis import setup as setup_redis
 from ..modules.s3 import setup_s3
 from .settings import ApplicationSettings
@@ -69,6 +69,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
         openapi_url=f"/api/{API_VTAG}/openapi.json",
         docs_url="/dev/doc",
         redoc_url=None,  # default disabled
+        lifespan=rabbitmq_lifespan,
     )
     override_fastapi_openapi_method(app)
 
@@ -79,7 +80,6 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     setup_s3(app)
     setup_client_session(app)
 
-    setup_rabbitmq(app)
     setup_rpc_api_routes(app)
     setup_rest_api_long_running_tasks_for_uploads(app)
     setup_rest_api_routes(app, API_VTAG)

services/storage/src/simcore_service_storage/modules/rabbitmq.py

@@ -1,4 +1,5 @@
 import logging
+from contextlib import asynccontextmanager
 from typing import cast
 
 from fastapi import FastAPI
@@ -15,40 +16,38 @@
 _logger = logging.getLogger(__name__)
 
 
-def setup(app: FastAPI) -> None:
-    async def on_startup() -> None:
-        with log_context(
-            _logger,
-            logging.INFO,
-            msg="Storage startup Rabbitmq",
-        ):
-            app.state.rabbitmq_client = None
-            rabbit_settings: RabbitSettings | None = app.state.settings.STORAGE_RABBITMQ
-            if not rabbit_settings:
-                raise ConfigurationError(
-                    msg="RabbitMQ client is de-activated in the settings"
-                )
-            await wait_till_rabbitmq_responsive(rabbit_settings.dsn)
-            app.state.rabbitmq_client = RabbitMQClient(
-                client_name="storage", settings=rabbit_settings
+@asynccontextmanager
+async def rabbitmq_lifespan(app: FastAPI):
+    with log_context(
+        _logger,
+        logging.INFO,
+        msg="Storage startup Rabbitmq",
+    ):
+        app.state.rabbitmq_client = None
+        rabbit_settings: RabbitSettings | None = app.state.settings.STORAGE_RABBITMQ
+        if not rabbit_settings:
+            raise ConfigurationError(
+                msg="RabbitMQ client is de-activated in the settings"
             )
-            app.state.rabbitmq_rpc_server = await RabbitMQRPCClient.create(
-                client_name="storage_rpc_server", settings=rabbit_settings
-            )
-
-    async def on_shutdown() -> None:
-        with log_context(
-            _logger,
-            logging.INFO,
-            msg="Storage shutdown Rabbitmq",
-        ):
-            if app.state.rabbitmq_client:
-                await app.state.rabbitmq_client.close()
-            if app.state.rabbitmq_rpc_server:
-                await app.state.rabbitmq_rpc_server.close()
+        await wait_till_rabbitmq_responsive(rabbit_settings.dsn)
+        app.state.rabbitmq_client = RabbitMQClient(
+            client_name="storage", settings=rabbit_settings
+        )
+        app.state.rabbitmq_rpc_server = await RabbitMQRPCClient.create(
+            client_name="storage_rpc_server", settings=rabbit_settings
+        )
 
-    app.add_event_handler("startup", on_startup)
-    app.add_event_handler("shutdown", on_shutdown)
+    yield
+
+    with log_context(
+        _logger,
+        logging.INFO,
+        msg="Storage shutdown Rabbitmq",
+    ):
+        if app.state.rabbitmq_client:
+            await app.state.rabbitmq_client.close()
+        if app.state.rabbitmq_rpc_server:
+            await app.state.rabbitmq_rpc_server.close()
 
 
 def get_rabbitmq_client(app: FastAPI) -> RabbitMQClient: