🔧 Refactor: Update setup functions to use AsyncIterator for improved lifecycle management

Author: pcrespov

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

@@ -13,13 +13,17 @@
 _logger = logging.getLogger(__name__)
 
 
-async def connect_to_db(app: FastAPI, settings: PostgresSettings) -> None:
+async def connect_to_postgres_until_ready(settings: PostgresSettings) -> AsyncEngine:
     with log_context(
         _logger,
         logging.DEBUG,
-        f"Connecting and migraging {settings.dsn_with_async_sqlalchemy}",
+        f"Connecting to {settings.dsn_with_async_sqlalchemy}",
     ):
-        engine = await create_async_engine_and_pg_database_ready(settings)
+        return await create_async_engine_and_pg_database_ready(settings)
+
+
+async def connect_to_db(app: FastAPI, settings: PostgresSettings) -> None:
+    engine = await connect_to_postgres_until_ready(settings)
 
     app.state.engine = engine
     _logger.debug(

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

@@ -1,7 +1,9 @@
 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
@@ -11,8 +13,6 @@
 )
 from servicelib.fastapi.tracing import initialize_tracing
 from simcore_service_catalog.core.background_tasks import setup_background_task
-from simcore_service_catalog.infrastructure.director import setup_director
-from simcore_service_catalog.infrastructure.postgres import setup_postgres_database
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from .._meta import (
@@ -27,7 +27,10 @@
 from ..api.rest.routes import setup_rest_api_routes
 from ..api.rpc.routes import setup_rpc_api_routes
 from ..exceptions.handlers import setup_exception_handlers
-from ..infrastructure.rabbitmq import setup_rabbitmq
+from ..infrastructure.director import director_lifespan
+from ..infrastructure.postgres import postgres_lifespan
+from ..infrastructure.rabbitmq import rabbitmq_lifespan
+from ..repository.setup import setup_repository
 from ..service.function_services import setup_function_services
 from .settings import ApplicationSettings
 
@@ -44,15 +47,33 @@
 )
 
 
-def _flush_started_banner() -> None:
+async def _setup_banner(app: FastAPI) -> AsyncIterator[State]:
     # WARNING: this function is spied in the tests
+    assert app
     print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
 
+    yield {}
 
-def _flush_finished_banner() -> None:
     print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201
 
 
+def _create_app_lifespan(settings: ApplicationSettings):
+    app_lifespan = LifespanManager()
+
+    app_lifespan.add(_setup_banner)
+
+    postgres_lifespan.add(setup_repository)
+    app_lifespan.include(postgres_lifespan)
+
+    app_lifespan.include(director_lifespan)
+    app_lifespan.add(rabbitmq_lifespan)
+
+    app_lifespan.add(setup_function_services)
+    app_lifespan.add(setup_background_task)
+
+    return app_lifespan
+
+
 def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
     # keep mostly quiet noisy loggers
     quiet_level: int = max(
@@ -76,6 +97,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(settings),
     )
     override_fastapi_openapi_method(app)
 
@@ -85,15 +107,6 @@ def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
     if settings.CATALOG_TRACING:
         initialize_tracing(app, settings.CATALOG_TRACING, APP_NAME)
 
-    # STARTUP-EVENT
-    app.add_event_handler("startup", _flush_started_banner)
-
-    setup_postgres_database(app)
-    setup_director(app)
-    setup_function_services(app)
-    setup_rabbitmq(app)
-    setup_background_task(app)
-
     if app.state.settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
         setup_prometheus_instrumentation(app)
 
@@ -113,9 +126,6 @@ def create_app(settings: ApplicationSettings | None = None) -> FastAPI:
     setup_rest_api_routes(app, vtag=API_VTAG)
     setup_rpc_api_routes(app)
 
-    # SHUTDOWN-EVENT
-    app.add_event_handler("shutdown", _flush_finished_banner)
-
     # EXCEPTIONS
     setup_exception_handlers(app)
 

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

@@ -11,11 +11,13 @@
 
 import asyncio
 import logging
+from collections.abc import AsyncIterator
 from contextlib import suppress
 from pprint import pformat
 from typing import Final
 
 from fastapi import FastAPI, HTTPException
+from fastapi_lifespan_manager import State
 from models_library.services import ServiceMetaDataPublished
 from models_library.services_types import ServiceKey, ServiceVersion
 from packaging.version import Version
@@ -213,29 +215,23 @@ async def _sync_services_task(app: FastAPI) -> None:
             )
 
 
-async def _start_registry_sync_task(app: FastAPI) -> None:
+async def setup_background_task(app: FastAPI) -> AsyncIterator[State]:
+    # FIXME: check director service is in place and ready. Hand-shake??
+    # SEE https://github.com/ITISFoundation/osparc-simcore/issues/1728
+
     # FIXME: added this variable to overcome the state in which the
     # task cancelation is ignored and the exceptions enter in a loop
     # that never stops the background task. This flag is an additional
     # mechanism to enforce stopping the background task
-    app.state.registry_syncer_running = True
     task = asyncio.create_task(_sync_services_task(app))
-    app.state.registry_sync_task = task
+
     _logger.info("registry syncing task started")
 
+    yield {"registry_syncer_running": True, "registry_sync_task": task}
 
-async def _stop_registry_sync_task(app: FastAPI) -> None:
-    if task := app.state.registry_sync_task:
-        with suppress(asyncio.CancelledError):
-            app.state.registry_syncer_running = False
-            task.cancel()
-            await task
-        app.state.registry_sync_task = None
-    _logger.info("registry syncing task stopped")
+    with suppress(asyncio.CancelledError):
+        app.state.registry_syncer_running = False
+        task.cancel()
+        await task
 
-
-def setup_background_task(app: FastAPI):
-    # FIXME: check director service is in place and ready. Hand-shake??
-    # SEE https://github.com/ITISFoundation/osparc-simcore/issues/1728
-    app.add_event_handler("startup", _start_registry_sync_task)
-    app.add_event_handler("shutdown", _stop_registry_sync_task)
+    _logger.info("registry syncing task stopped")

services/catalog/src/simcore_service_catalog/infrastructure/director.py

@@ -3,14 +3,15 @@
 import json
 import logging
 import urllib.parse
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import suppress
 from pprint import pformat
 from typing import Any, Final
 
 import httpx
 from common_library.json_serialization import json_dumps
 from fastapi import FastAPI, HTTPException
+from fastapi_lifespan_manager import LifespanManager, State
 from models_library.api_schemas_directorv2.services import ServiceExtras
 from models_library.services_metadata_published import ServiceMetaDataPublished
 from models_library.services_types import ServiceKey, ServiceVersion
@@ -288,38 +289,30 @@ async def get_service_extras(
         return TypeAdapter(ServiceExtras).validate_python(result)
 
 
-async def _initialize_director_client(app: FastAPI) -> None:
-    if settings := app.state.settings.CATALOG_DIRECTOR:
-        with log_context(
-            _logger, logging.DEBUG, "Setup director at %s", f"{settings.base_url=}"
-        ):
-            async for attempt in AsyncRetrying(**_director_startup_retry_policy):
-                client = DirectorApi(base_url=settings.base_url, app=app)
-                with attempt:
-                    client = DirectorApi(base_url=settings.base_url, app=app)
-                    if not await client.is_responsive():
-                        with suppress(Exception):
-                            await client.close()
-                        raise DirectorUnresponsiveError
-
-                _logger.info(
-                    "Connection to director-v0 succeded [%s]",
-                    json_dumps(attempt.retry_state.retry_object.statistics),
-                )
-
-            # set when connected
-            app.state.director_api = client
+director_lifespan = LifespanManager()
 
 
-async def _shutdown_director_client(app: FastAPI) -> None:
-    client: DirectorApi | None
-    if client := app.state.director_api:
-        await client.close()
-
-    _logger.debug("Director client closed successfully")
+@director_lifespan.add
+async def _setup_director(app: FastAPI) -> AsyncIterator[State]:
+    settings = app.state.settings.CATALOG_DIRECTOR
 
+    with log_context(
+        _logger, logging.DEBUG, "Setup director at %s", f"{settings.base_url=}"
+    ):
+        async for attempt in AsyncRetrying(**_director_startup_retry_policy):
+            client = DirectorApi(base_url=settings.base_url, app=app)
+            with attempt:
+                client = DirectorApi(base_url=settings.base_url, app=app)
+                if not await client.is_responsive():
+                    with suppress(Exception):
+                        await client.close()
+                    raise DirectorUnresponsiveError
+
+            _logger.info(
+                "Connection to director-v0 succeded [%s]",
+                json_dumps(attempt.retry_state.retry_object.statistics),
+            )
 
-def setup_director(app: FastAPI):
+        yield {"director_api": client}
 
-    app.add_event_handler("startup", _initialize_director_client)
-    app.add_event_handler("shutdown", _shutdown_director_client)
+        await client.close()

services/catalog/src/simcore_service_catalog/infrastructure/postgres.py

@@ -1,30 +1,29 @@
-import contextlib
 import logging
+from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
-from servicelib.fastapi.db_asyncpg_engine import close_db_connection, connect_to_db
+from fastapi_lifespan_manager import LifespanManager, State
+from servicelib.fastapi.db_asyncpg_engine import connect_to_postgres_until_ready
 from servicelib.logging_utils import log_catch, log_context
-
-from ..repository.products import setup_default_product
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 _logger = logging.getLogger(__name__)
 
 
-def setup_postgres_database(app: FastAPI):
+postgres_lifespan = LifespanManager()
 
-    async def _():
-        with log_context(_logger, logging.INFO, f"{__name__} startup ..."):
-            # connection
-            await connect_to_db(app, app.state.settings.CATALOG_POSTGRES)
 
-            # configuring default product
-            await setup_default_product(app)
+@postgres_lifespan.add
+async def setup_postgres_database(app: FastAPI) -> AsyncIterator[State]:
 
-        yield
+    with log_context(_logger, logging.INFO, f"{__name__} startup ..."):
+        engine: AsyncEngine = await connect_to_postgres_until_ready(
+            app.state.settings.CATALOG_POSTGRES
+        )
 
-        with log_context(
-            _logger, logging.INFO, f"{__name__} shutdown ..."
-        ), contextlib.suppress(Exception), log_catch(_logger):
-            await close_db_connection(app)
+    yield {"engine": engine}
 
-    return _
+    with log_context(_logger, logging.INFO, f"{__name__} shutdown ..."), log_catch(
+        _logger, reraise=False
+    ):
+        await engine.dispose()

services/catalog/src/simcore_service_catalog/infrastructure/rabbitmq.py

@@ -1,7 +1,9 @@
 import logging
+from collections.abc import AsyncIterator
 from typing import cast
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import LifespanManager, State
 from servicelib.rabbitmq import RabbitMQRPCClient, wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
@@ -15,24 +17,23 @@ def get_rabbitmq_settings(app: FastAPI) -> RabbitSettings:
     return settings
 
 
-def setup_rabbitmq(app: FastAPI) -> None:
+rabbitmq_lifespan = LifespanManager()
+
+
+@rabbitmq_lifespan.add
+async def setup_rabbitmq(app: FastAPI) -> AsyncIterator[State]:
     settings: RabbitSettings = get_rabbitmq_settings(app)
-    app.state.rabbitmq_rpc_server = None
 
-    async def _on_startup() -> None:
-        await wait_till_rabbitmq_responsive(settings.dsn)
+    await wait_till_rabbitmq_responsive(settings.dsn)
 
-        app.state.rabbitmq_rpc_server = await RabbitMQRPCClient.create(
-            client_name=f"{PROJECT_NAME}_rpc_server", settings=settings
-        )
+    rabbitmq_rpc_server = await RabbitMQRPCClient.create(
+        client_name=f"{PROJECT_NAME}_rpc_server",
+        settings=settings,
+    )
 
-    async def _on_shutdown() -> None:
-        if app.state.rabbitmq_rpc_server:
-            await app.state.rabbitmq_rpc_server.close()
-            app.state.rabbitmq_rpc_server = None
+    yield {"rabbitmq_rpc_server": rabbitmq_rpc_server}
 
-    app.add_event_handler("startup", _on_startup)
-    app.add_event_handler("shutdown", _on_shutdown)
+    await rabbitmq_rpc_server.close()
 
 
 def get_rabbitmq_rpc_server(app: FastAPI) -> RabbitMQRPCClient:

services/catalog/src/simcore_service_catalog/repository/setup.py

@@ -1,8 +1,12 @@
+from collections.abc import AsyncIterator
+
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 
 from .products import ProductsRepository
 
 
-async def setup_repository(app: FastAPI):
+async def setup_repository(app: FastAPI) -> AsyncIterator[State]:
     repo = ProductsRepository(db_engine=app.state.engine)
-    app.state.default_product_name = await repo.get_default_product_name()
+
+    yield {"default_product_name": await repo.get_default_product_name()}

services/catalog/src/simcore_service_catalog/service/function_services.py

@@ -1,9 +1,12 @@
+from collections.abc import AsyncIterator
+
 # mypy: disable-error-code=truthy-function
 from typing import Any
 
 from fastapi import status
 from fastapi.applications import FastAPI
 from fastapi.exceptions import HTTPException
+from fastapi_lifespan_manager import State
 from models_library.function_services_catalog import (
     is_function_service,
     iter_service_docker_data,
@@ -31,12 +34,13 @@ def get_function_service(key, version) -> ServiceMetaDataPublished:
         ) from err
 
 
-def setup_function_services(app: FastAPI):
-    def _on_startup() -> None:
-        catalog = [_as_dict(metadata) for metadata in iter_service_docker_data()]
-        app.state.frontend_services_catalog = catalog
+async def setup_function_services(app: FastAPI) -> AsyncIterator[State]:
+    assert app  # nosec
+    assert not hasattr(app.state, "frontend_services_catalog")  # nosec
+
+    catalog = [_as_dict(metadata) for metadata in iter_service_docker_data()]
 
-    app.add_event_handler("startup", _on_startup)
+    yield {"frontend_services_catalog": catalog}
 
 
 __all__: tuple[str, ...] = (