✨ Refactor: Rename async engine creation function and update references

Author: pcrespov

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

@@ -4,7 +4,6 @@
 SEE migration aiopg->asyncpg https://github.com/ITISFoundation/osparc-simcore/issues/4529
 """
 
-
 import logging
 from typing import Final
 
@@ -16,7 +15,7 @@
 )
 from sqlalchemy.ext.asyncio import AsyncEngine
 
-from ..db_asyncpg_utils import create_async_engine_and_pg_database_ready
+from ..db_asyncpg_utils import create_async_engine_and_database_ready
 from ..logging_utils import log_context
 
 APP_DB_ASYNC_ENGINE_KEY: Final[str] = f"{__name__ }.AsyncEngine"
@@ -56,7 +55,7 @@ async def connect_to_db(app: web.Application, settings: PostgresSettings) -> Non
         "Connecting app[APP_DB_ASYNC_ENGINE_KEY] to postgres with %s",
         f"{settings=}",
     ):
-        engine = await create_async_engine_and_pg_database_ready(settings)
+        engine = await create_async_engine_and_database_ready(settings)
         _set_async_engine_to_app_state(app, engine)
 
     _logger.info(

packages/service-library/src/servicelib/db_asyncpg_utils.py

@@ -4,9 +4,6 @@
 
 from models_library.healthchecks import IsNonResponsive, IsResponsive, LivenessResult
 from settings_library.postgres import PostgresSettings
-from simcore_postgres_database.utils_aiosqlalchemy import (  # type: ignore[import-not-found] # this on is unclear
-    raise_if_migration_not_ready,
-)
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from tenacity import retry
@@ -17,7 +14,7 @@
 
 
 @retry(**PostgresRetryPolicyUponInitialization(_logger).kwargs)
-async def create_async_engine_and_pg_database_ready(
+async def create_async_engine_and_database_ready(
     settings: PostgresSettings,
 ) -> AsyncEngine:
     """
@@ -26,6 +23,10 @@ async def create_async_engine_and_pg_database_ready(
     - waits until db data is migrated (i.e. ready to use)
     - returns engine
     """
+    from simcore_postgres_database.utils_aiosqlalchemy import (  # type: ignore[import-not-found] # this on is unclear
+        raise_if_migration_not_ready,
+    )
+
     server_settings = None
     if settings.POSTGRES_CLIENT_NAME:
         server_settings = {
@@ -43,9 +44,10 @@ async def create_async_engine_and_pg_database_ready(
 
     try:
         await raise_if_migration_not_ready(engine)
-    except Exception:
+    except Exception as exc:
         # NOTE: engine must be closed because retry will create a new engine
         await engine.dispose()
+        exc.add_note("Failed during migration check. Created engine disposed.")
         raise
 
     return engine

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

@@ -7,7 +7,7 @@
 )
 from sqlalchemy.ext.asyncio import AsyncEngine
 
-from ..db_asyncpg_utils import create_async_engine_and_pg_database_ready
+from ..db_asyncpg_utils import create_async_engine_and_database_ready
 from ..logging_utils import log_context
 
 _logger = logging.getLogger(__name__)
@@ -19,7 +19,7 @@ async def connect_to_db(app: FastAPI, settings: PostgresSettings) -> None:
         logging.DEBUG,
         f"Connecting and migraging {settings.dsn_with_async_sqlalchemy}",
     ):
-        engine = await create_async_engine_and_pg_database_ready(settings)
+        engine = await create_async_engine_and_database_ready(settings)
 
     app.state.engine = engine
     _logger.debug(

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

@@ -0,0 +1,31 @@
+import logging
+from collections.abc import AsyncIterator
+
+from fastapi import FastAPI
+from fastapi_lifespan_manager import LifespanManager, State
+from servicelib.logging_utils import log_catch, log_context
+from settings_library.postgres import PostgresSettings
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from ..db_asyncpg_utils import create_async_engine_and_database_ready
+
+_logger = logging.getLogger(__name__)
+
+
+postgres_lifespan = LifespanManager()
+
+
+@postgres_lifespan.add
+async def setup_postgres_database(app: FastAPI) -> AsyncIterator[State]:
+    with log_context(_logger, logging.INFO, f"{__name__}"):
+
+        pg_settings: PostgresSettings = app.state.settings.CATALOG_POSTGRES
+
+        async_engine: AsyncEngine = await create_async_engine_and_database_ready(
+            pg_settings
+        )
+
+        yield {"postgres.async_engine": async_engine}
+
+        with log_catch(_logger, reraise=False):
+            await async_engine.dispose()

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

@@ -0,0 +1,91 @@
+# pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
+# pylint: disable=too-many-arguments
+# pylint: disable=unused-argument
+# pylint: disable=unused-variable
+
+from typing import Annotated, Any
+
+import pytest
+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.postgres_lifespan import postgres_lifespan
+from settings_library.application import BaseApplicationSettings
+from settings_library.postgres import PostgresSettings
+
+
+@pytest.fixture
+def mock_create_async_engine_and_database_ready(mocker: MockerFixture) -> MockType:
+    return mocker.patch(
+        "servicelib.fastapi.postgres_lifespan.create_async_engine_and_database_ready",
+        return_value=mocker.AsyncMock(),
+    )
+
+
+@pytest.fixture
+def app_environment(monkeypatch: pytest.MonkeyPatch) -> EnvVarsDict:
+    return setenvs_from_dict(
+        monkeypatch, PostgresSettings.model_json_schema()["examples"][0]
+    )
+
+
+async def test_setup_postgres_database_in_an_app(
+    is_pdb_enabled: bool,
+    app_environment: EnvVarsDict,
+    mock_create_async_engine_and_database_ready: MockType,
+):
+    assert app_environment
+
+    @postgres_lifespan.add
+    async def my_db_setup(app: FastAPI, state: State):
+        app.state.my_db_engine = state["postgres.async_engine"]
+
+        assert (
+            app.state.my_db_engine
+            == mock_create_async_engine_and_database_ready.return_value
+        )
+
+        yield
+
+    # compose lifespans
+    app_lifespan = LifespanManager()
+    app_lifespan.include(postgres_lifespan)
+
+    # define app
+    app = FastAPI(lifespan=app_lifespan)
+
+    # settings
+    class AppSettings(BaseApplicationSettings):
+        CATALOG_POSTGRES: Annotated[
+            PostgresSettings,
+            Field(json_schema_extra={"auto_default_from_env": True}),
+        ]
+
+    app.state.settings = AppSettings.create_from_envs()
+
+    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 the async engine was created
+        mock_create_async_engine_and_database_ready.assert_called_once_with(
+            app.state.settings.CATALOG_POSTGRES
+        )
+
+        # Verify that the async engine is in the lifespan manager state
+        assert "postgres.async_engine" in asgi_manager._state  # noqa: SLF001
+        assert app.state.my_db_engine
+        assert (
+            app.state.my_db_engine
+            == asgi_manager._state["postgres.async_engine"]  # noqa: SLF001
+        )
+
+    # Verify that the engine was disposed
+    async_engine: Any = mock_create_async_engine_and_database_ready.return_value
+    async_engine.dispose.assert_called_once()

services/director-v2/src/simcore_service_director_v2/modules/db/_asyncpg.py

@@ -1,7 +1,9 @@
 import logging
 
 from fastapi import FastAPI
-from servicelib.db_asyncpg_utils import create_async_engine_and_pg_database_ready
+from servicelib.db_asyncpg_utils import (
+    create_async_engine_and_database_ready,
+)
 from servicelib.logging_utils import log_context
 from settings_library.postgres import PostgresSettings
 from simcore_postgres_database.utils_aiosqlalchemy import get_pg_engine_stateinfo
@@ -15,7 +17,7 @@ async def asyncpg_connect_to_db(app: FastAPI, settings: PostgresSettings) -> Non
         logging.DEBUG,
         f"Connecting and migraging {settings.dsn_with_async_sqlalchemy}",
     ):
-        engine = await create_async_engine_and_pg_database_ready(settings)
+        engine = await create_async_engine_and_database_ready(settings)
 
     app.state.asyncpg_engine = engine
     _logger.debug(