✨ Refactor: Simplify application lifespan management by moving logic to events module and updating related handlers

Author: pcrespov

services/catalog/src/simcore_service_catalog/core/application.py

@@ -1,16 +1,10 @@
 import logging
-from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
 from fastapi.middleware.gzip import GZipMiddleware
-from fastapi_lifespan_manager import LifespanManager, State
 from models_library.basic_types import BootModeEnum
 from servicelib.fastapi import timing_middleware
 from servicelib.fastapi.openapi import override_fastapi_openapi_method
-from servicelib.fastapi.postgres_lifespan import (
-    PostgresLifespanStateKeys,
-    postgres_lifespan,
-)
 from servicelib.fastapi.profiler import initialize_profiler
 from servicelib.fastapi.prometheus_instrumentation import (
     setup_prometheus_instrumentation,
@@ -27,15 +21,12 @@
 )
 from ..api.rest.routes import setup_rest_api_routes
 from ..api.rpc.routes import setup_rpc_api_routes
-from ..db.events import setup_database
 from ..exceptions.handlers import setup_exception_handlers
 from ..services.function_services import setup_function_services
 from ..services.rabbitmq import setup_rabbitmq
+from . import events
 from .events import (
-    create_on_shutdown,
-    create_on_startup,
-    flush_finished_banner,
-    flush_started_banner,
+    _create_on_shutdown,
 )
 from .settings import ApplicationSettings
 
@@ -52,30 +43,6 @@
 )
 
 
-async def _main_setup(app: FastAPI) -> AsyncIterator[State]:
-    flush_started_banner()
-
-    settings: ApplicationSettings = app.state.settings
-
-    yield {
-        PostgresLifespanStateKeys.POSTGRES_SETTINGS: settings.CATALOG_POSTGRES,
-    }
-
-    flush_finished_banner()
-
-
-def _create_app_lifespan():
-    # app lifespan
-    app_lifespan = LifespanManager()
-    app_lifespan.add(_main_setup)
-
-    # - postgres lifespan
-    postgres_lifespan.add(setup_database)
-    app_lifespan.include(postgres_lifespan)
-
-    return app_lifespan
-
-
 def create_app() -> FastAPI:
     # keep mostly quiet noisy loggers
     quiet_level: int = max(
@@ -96,7 +63,7 @@ def create_app() -> FastAPI:
         openapi_url=f"/api/{API_VTAG}/openapi.json",
         docs_url="/dev/doc",
         redoc_url=None,  # default disabled
-        lifespan=_create_app_lifespan(),
+        lifespan=events.create_app_lifespan(),
     )
     override_fastapi_openapi_method(app)
 
@@ -106,9 +73,6 @@ def create_app() -> FastAPI:
     if settings.CATALOG_TRACING:
         initialize_tracing(app, settings.CATALOG_TRACING, APP_NAME)
 
-    # STARTUP-EVENT
-    app.add_event_handler("startup", create_on_startup(app))
-
     # PLUGIN SETUP
     setup_function_services(app)
     setup_rabbitmq(app)
@@ -133,7 +97,7 @@ def create_app() -> FastAPI:
     setup_rpc_api_routes(app)
 
     # SHUTDOWN-EVENT
-    app.add_event_handler("shutdown", create_on_shutdown(app))
+    app.add_event_handler("shutdown", _create_on_shutdown(app))
 
     # EXCEPTIONS
     setup_exception_handlers(app)

services/catalog/src/simcore_service_catalog/core/events.py

@@ -1,13 +1,21 @@
+import contextlib
 import logging
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from typing import TypeAlias
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import LifespanManager, State
+from servicelib.fastapi.postgres_lifespan import (
+    PostgresLifespanStateKeys,
+    postgres_lifespan,
+)
 from servicelib.logging_utils import log_context
 
 from .._meta import APP_FINISHED_BANNER_MSG, APP_STARTED_BANNER_MSG
+from ..db.events import setup_database
 from ..services.director import close_director, setup_director
 from .background_tasks import start_registry_sync_task, stop_registry_sync_task
+from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
 
@@ -24,7 +32,19 @@ def flush_finished_banner() -> None:
     print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201
 
 
-def create_on_startup(app: FastAPI) -> EventCallable:
+async def _main_setup(app: FastAPI) -> AsyncIterator[State]:
+    flush_started_banner()
+
+    settings: ApplicationSettings = app.state.settings
+
+    yield {
+        PostgresLifespanStateKeys.POSTGRES_SETTINGS: settings.CATALOG_POSTGRES,
+    }
+
+    flush_finished_banner()
+
+
+def _create_on_startup(app: FastAPI) -> EventCallable:
     async def _() -> None:
 
         if app.state.settings.CATALOG_DIRECTOR:
@@ -40,7 +60,7 @@ async def _() -> None:
     return _
 
 
-def create_on_shutdown(app: FastAPI) -> EventCallable:
+def _create_on_shutdown(app: FastAPI) -> EventCallable:
     async def _() -> None:
 
         with log_context(_logger, logging.INFO, "Application shutdown"):
@@ -52,3 +72,27 @@ async def _() -> None:
                     _logger.exception("Unexpected error while closing application")
 
     return _
+
+
+@contextlib.asynccontextmanager
+async def _other_setup(app: FastAPI) -> AsyncIterator[State]:
+
+    await _create_on_startup(app)()
+
+    yield {}
+
+    await _create_on_shutdown(app)()
+
+
+def create_app_lifespan():
+    # app lifespan
+    app_lifespan = LifespanManager()
+    app_lifespan.add(_main_setup)
+
+    # - postgres lifespan
+    postgres_lifespan.add(setup_database)
+    app_lifespan.include(postgres_lifespan)
+
+    app_lifespan.add(_other_setup)
+
+    return app_lifespan

services/catalog/tests/unit/conftest.py

@@ -117,11 +117,11 @@ def spy_app(mocker: MockerFixture) -> AppLifeSpanSpyTargets:
     # work as expected
     return AppLifeSpanSpyTargets(
         on_startup=mocker.spy(
-            simcore_service_catalog.core.application,
+            simcore_service_catalog.core.events,
             "flush_started_banner",
         ),
         on_shutdown=mocker.spy(
-            simcore_service_catalog.core.application,
+            simcore_service_catalog.core.events,
             "flush_finished_banner",
         ),
     )
@@ -217,9 +217,7 @@ def service_caching_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
 
 @pytest.fixture
 def postgres_setup_disabled(mocker: MockerFixture) -> MockType:
-    return mocker.patch.object(
-        simcore_service_catalog.core.application, "postgres_lifespan"
-    )
+    return mocker.patch.object(simcore_service_catalog.core.events, "postgres_lifespan")
 
 
 @pytest.fixture

services/catalog/tests/unit/test_services_director.py

@@ -61,6 +61,7 @@ async def test_director_client_high_level_api(
 
 
 async def test_director_client_low_level_api(
+    postgres_setup_disabled: MockType,
     background_tasks_setup_disabled: None,
     rabbitmq_and_rpc_setup_disabled: None,
     mocked_director_service_api: MockRouter,