removed confusion

Author: sanderegg

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

@@ -1,5 +1,7 @@
+import contextlib
 import logging
 import time
+from collections.abc import AsyncIterator
 from datetime import timedelta
 
 from models_library.healthchecks import IsNonResponsive, IsResponsive, LivenessResult
@@ -8,6 +10,7 @@
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from tenacity import retry
 
+from .logging_utils import log_context
 from .retry_policies import PostgresRetryPolicyUponInitialization
 
 _logger = logging.getLogger(__name__)
@@ -64,3 +67,34 @@ async def check_postgres_liveness(engine: AsyncEngine) -> LivenessResult:
         return IsResponsive(elapsed=timedelta(seconds=elapsed_time))
     except SQLAlchemyError as err:
         return IsNonResponsive(reason=f"{err}")
+
+
+@contextlib.asynccontextmanager
+async def with_async_pg_engine(
+    settings: PostgresSettings,
+) -> AsyncIterator[AsyncEngine]:
+    """
+    Creates an asyncpg engine and ensures it is properly closed after use.
+    """
+    try:
+        with log_context(
+            _logger,
+            logging.DEBUG,
+            f"connection to db {settings.dsn_with_async_sqlalchemy}",
+        ):
+            server_settings = None
+            if settings.POSTGRES_CLIENT_NAME:
+                assert isinstance(settings.POSTGRES_CLIENT_NAME, str)
+
+            engine = create_async_engine(
+                settings.dsn_with_async_sqlalchemy,
+                pool_size=settings.POSTGRES_MINSIZE,
+                max_overflow=settings.POSTGRES_MAXSIZE - settings.POSTGRES_MINSIZE,
+                connect_args={"server_settings": server_settings},
+                pool_pre_ping=True,  # https://docs.sqlalchemy.org/en/14/core/pooling.html#dealing-with-disconnects
+                future=True,  # this uses sqlalchemy 2.0 API, shall be removed when sqlalchemy 2.0 is released
+            )
+        yield engine
+    finally:
+        with log_context(_logger, logging.DEBUG, f"db disconnect of {engine}"):
+            await engine.dispose()

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

@@ -1,8 +1,7 @@
 import logging
 
 from fastapi import FastAPI
-from servicelib.db_async_engine import close_db_connection
-from servicelib.fastapi.db_asyncpg_engine import connect_to_db
+from servicelib.fastapi.db_asyncpg_engine import close_db_connection, connect_to_db
 from servicelib.retry_policies import PostgresRetryPolicyUponInitialization
 from sqlalchemy.ext.asyncio import AsyncEngine
 from tenacity import retry