✨ Refactor catalog: Implement application lifespan management with Postgres integration and cleanup old connection logic

Author: pcrespov

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

@@ -1,24 +1,42 @@
 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,
 )
 from servicelib.fastapi.tracing import initialize_tracing
 from starlette.middleware.base import BaseHTTPMiddleware
 
-from .._meta import API_VERSION, API_VTAG, APP_NAME, PROJECT_NAME, SUMMARY
+from .._meta import (
+    API_VERSION,
+    API_VTAG,
+    APP_NAME,
+    PROJECT_NAME,
+    SUMMARY,
+)
 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 .events import create_on_shutdown, create_on_startup
+from .events import (
+    create_on_shutdown,
+    create_on_startup,
+    flush_finished_banner,
+    flush_started_banner,
+)
 from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
@@ -34,18 +52,39 @@
 )
 
 
-def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
+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(
         min(logging.root.level + _LOG_LEVEL_STEP, logging.CRITICAL), logging.WARNING
     )
     for name in _NOISY_LOGGERS:
         logging.getLogger(name).setLevel(quiet_level)
 
-    if settings is None:
-        settings = ApplicationSettings.create_from_envs()
-
-    assert settings  # nosec
+    settings = ApplicationSettings.create_from_envs()
     _logger.debug(settings.model_dump_json(indent=2))
 
     app = FastAPI(
@@ -57,6 +96,7 @@ def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
         openapi_url=f"/api/{API_VTAG}/openapi.json",
         docs_url="/dev/doc",
         redoc_url=None,  # default disabled
+        lifespan=_create_app_lifespan(),
     )
     override_fastapi_openapi_method(app)
 
@@ -73,11 +113,11 @@ def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
     setup_function_services(app)
     setup_rabbitmq(app)
 
-    if app.state.settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
+    if settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
         setup_prometheus_instrumentation(app)
 
     # MIDDLEWARES
-    if app.state.settings.CATALOG_PROFILING:
+    if settings.CATALOG_PROFILING:
         initialize_profiler(app)
 
     if settings.SC_BOOT_MODE != BootModeEnum.PRODUCTION:

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

@@ -3,11 +3,9 @@
 from typing import TypeAlias
 
 from fastapi import FastAPI
-from servicelib.fastapi.db_asyncpg_engine import close_db_connection, connect_to_db
 from servicelib.logging_utils import log_context
 
 from .._meta import APP_FINISHED_BANNER_MSG, APP_STARTED_BANNER_MSG
-from ..db.events import setup_default_product
 from ..services.director import close_director, setup_director
 from .background_tasks import start_registry_sync_task, stop_registry_sync_task
 
@@ -17,23 +15,17 @@
 EventCallable: TypeAlias = Callable[[], Awaitable[None]]
 
 
-def _flush_started_banner() -> None:
+def flush_started_banner() -> None:
     # WARNING: this function is spied in the tests
     print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
 
 
-def _flush_finished_banner() -> None:
+def flush_finished_banner() -> None:
     print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201
 
 
 def create_on_startup(app: FastAPI) -> EventCallable:
     async def _() -> None:
-        _flush_started_banner()
-
-        # setup connection to pg db
-        if app.state.settings.CATALOG_POSTGRES:
-            await connect_to_db(app, app.state.settings.CATALOG_POSTGRES)
-            await setup_default_product(app)
 
         if app.state.settings.CATALOG_DIRECTOR:
             # setup connection to director
@@ -56,10 +48,7 @@ async def _() -> None:
                 try:
                     await stop_registry_sync_task(app)
                     await close_director(app)
-                    await close_db_connection(app)
                 except Exception:  # pylint: disable=broad-except
                     _logger.exception("Unexpected error while closing application")
 
-            _flush_finished_banner()
-
     return _

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

@@ -1,12 +1,20 @@
 import logging
+from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
+from servicelib.fastapi.postgres_lifespan import PostgresLifespanStateKeys
 
 from .repositories.products import ProductsRepository
 
 _logger = logging.getLogger(__name__)
 
 
-async def setup_default_product(app: FastAPI):
+async def setup_database(app: FastAPI, state: State) -> AsyncIterator[State]:
+    app.state.engine = state[PostgresLifespanStateKeys.POSTGRES_ASYNC_ENGINE]
+
     repo = ProductsRepository(db_engine=app.state.engine)
+
     app.state.default_product_name = await repo.get_default_product_name()
+
+    yield {}

services/catalog/tests/unit/conftest.py

@@ -15,6 +15,7 @@
 import pytest
 import respx
 import simcore_service_catalog
+import simcore_service_catalog.core.application
 import simcore_service_catalog.core.events
 import yaml
 from asgi_lifespan import LifespanManager
@@ -116,12 +117,12 @@ def spy_app(mocker: MockerFixture) -> AppLifeSpanSpyTargets:
     # work as expected
     return AppLifeSpanSpyTargets(
         on_startup=mocker.spy(
-            simcore_service_catalog.core.events,
-            "_flush_started_banner",
+            simcore_service_catalog.core.application,
+            "flush_started_banner",
         ),
         on_shutdown=mocker.spy(
-            simcore_service_catalog.core.events,
-            "_flush_finished_banner",
+            simcore_service_catalog.core.application,
+            "flush_finished_banner",
         ),
     )
 
@@ -139,7 +140,7 @@ async def app(
 
     # create instance
     assert app_environment
-    app_under_test = create_app(settings=app_settings)
+    app_under_test = create_app()
 
     assert spy_app.on_startup.call_count == 0
     assert spy_app.on_shutdown.call_count == 0
@@ -166,7 +167,7 @@ def client(
 
     # create instance
     assert app_environment
-    app_under_test = create_app(settings=app_settings)
+    app_under_test = create_app()
 
     assert (
         spy_app.on_startup.call_count == 0