connect with asyncpg

Author: sanderegg

services/api-server/src/simcore_service_api_server/core/events.py

@@ -4,7 +4,13 @@
 from fastapi import FastAPI
 
 from .._meta import APP_FINISHED_BANNER_MSG, APP_STARTED_BANNER_MSG
-from ..db.events import close_db_connection, connect_to_db
+from ..db.events import (
+    asyncpg_close_db_connection,
+    asyncpg_connect_to_db,
+    close_db_connection,
+    connect_to_db,
+)
+from .settings import ApplicationSettings
 
 logger = logging.getLogger(__name__)
 
@@ -14,10 +20,12 @@ async def _on_startup() -> None:
         logger.info("Application starting ...")
         if app.state.settings.API_SERVER_POSTGRES:
             # database
+            assert isinstance(app.state.settings, ApplicationSettings)  # nosec
             await connect_to_db(app)
+            await asyncpg_connect_to_db(app, app.state.settings.API_SERVER_POSTGRES)
             assert app.state.engine  # nosec
 
-        print(APP_STARTED_BANNER_MSG, flush=True)
+        print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
 
     return _on_startup
 
@@ -27,7 +35,9 @@ async def _on_shutdown() -> None:
         logger.info("Application stopping, ...")
 
         if app.state.settings.API_SERVER_POSTGRES:
+            assert isinstance(app.state.settings, ApplicationSettings)  # nosec
             try:
+                await asyncpg_close_db_connection(app)
                 await close_db_connection(app)
 
             except Exception as err:  # pylint: disable=broad-except
@@ -38,6 +48,6 @@ async def _on_shutdown() -> None:
                     stack_info=app.state.settings.debug,
                 )
 
-        print(APP_FINISHED_BANNER_MSG, flush=True)
+        print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201
 
     return _on_shutdown

services/api-server/src/simcore_service_api_server/db/events.py

@@ -2,13 +2,16 @@
 
 from aiopg.sa import Engine, create_engine
 from fastapi import FastAPI
+from servicelib.db_asyncpg_utils import create_async_engine_and_pg_database_ready
+from servicelib.logging_utils import log_context
 from servicelib.retry_policies import PostgresRetryPolicyUponInitialization
 from settings_library.postgres import PostgresSettings
 from simcore_postgres_database.utils_aiopg import (
     close_engine,
     get_pg_engine_info,
     raise_if_migration_not_ready,
 )
+from simcore_postgres_database.utils_aiosqlalchemy import get_pg_engine_stateinfo
 from tenacity import retry
 
 from .._meta import PROJECT_NAME
@@ -53,3 +56,30 @@ async def close_db_connection(app: FastAPI) -> None:
         await close_engine(engine)
 
     logger.debug("Disconnected from %s", engine.dsn)
+
+
+async def asyncpg_connect_to_db(app: FastAPI, settings: PostgresSettings) -> None:
+    with log_context(
+        logger,
+        logging.DEBUG,
+        f"Connecting and migraging {settings.dsn_with_async_sqlalchemy}",
+    ):
+        engine = await create_async_engine_and_pg_database_ready(settings)
+
+    app.state.asyncpg_engine = engine
+    logger.debug(
+        "Setup engine: %s",
+        await get_pg_engine_stateinfo(engine),
+    )
+
+
+async def asyncpg_close_db_connection(app: FastAPI) -> None:
+    with log_context(
+        logger, logging.DEBUG, f"db disconnect of {app.state.asyncpg_engine}"
+    ):
+        if engine := app.state.asyncpg_engine:
+            await engine.dispose()
+
+
+def get_asyncpg_engine(app: FastAPI):
+    return app.state.asyncpg_engine