migrated to asyncpg

Author: sanderegg

packages/simcore-sdk/src/simcore_sdk/node_ports_common/dbmanager.py

@@ -1,107 +1,77 @@
 import json
 import logging
-import os
-import socket
-from typing import Any
 
-import aiopg.sa
 import sqlalchemy as sa
-import tenacity
-from aiopg.sa.engine import Engine
-from aiopg.sa.result import RowProxy
 from models_library.projects import ProjectID
 from models_library.users import UserID
-from servicelib.common_aiopg_utils import DataSourceName, create_pg_engine
-from servicelib.retry_policies import PostgresRetryPolicyUponInitialization
+from pydantic import TypeAdapter
+from servicelib.db_asyncpg_utils import create_async_engine_and_pg_database_ready
 from settings_library.node_ports import NodePortsSettings
 from simcore_postgres_database.models.comp_tasks import comp_tasks
 from simcore_postgres_database.models.projects import projects
-from simcore_postgres_database.utils_aiopg import (
-    close_engine,
-    raise_if_migration_not_ready,
-)
-from sqlalchemy import and_
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
 from .exceptions import NodeNotFound, ProjectNotFoundError
 
-log = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 
 
 async def _get_node_from_db(
-    project_id: str, node_uuid: str, connection: aiopg.sa.SAConnection
-) -> RowProxy:
-    log.debug(
+    project_id: str, node_uuid: str, connection: AsyncConnection
+) -> sa.engine.Row:
+    _logger.debug(
         "Reading from comp_tasks table for node uuid %s, project %s",
         node_uuid,
         project_id,
     )
+    rows_count = await connection.scalar(
+        sa.select(sa.func.count()).select_from(
+            (comp_tasks.c.node_id == node_uuid)
+            & (comp_tasks.c.project_id == project_id),
+        )
+    )
+    if rows_count > 1:
+        _logger.error("the node id %s is not unique", node_uuid)
     result = await connection.execute(
         sa.select(comp_tasks).where(
-            and_(
-                comp_tasks.c.node_id == node_uuid,
-                comp_tasks.c.project_id == project_id,
-            )
+            (comp_tasks.c.node_id == node_uuid)
+            & (comp_tasks.c.project_id == project_id)
         )
     )
-    if result.rowcount > 1:
-        log.error("the node id %s is not unique", node_uuid)
-    node: RowProxy | None = await result.first()
+    node = result.one_or_none()
     if not node:
-        log.error("the node id %s was not found", node_uuid)
+        _logger.error("the node id %s was not found", node_uuid)
         raise NodeNotFound(node_uuid)
     return node
 
 
-@tenacity.retry(**PostgresRetryPolicyUponInitialization().kwargs)
-async def _ensure_postgres_ready(dsn: DataSourceName) -> Engine:
-    engine: aiopg.sa.Engine = await create_pg_engine(dsn, minsize=1, maxsize=4)
-    try:
-        await raise_if_migration_not_ready(engine)
-    except Exception:
-        await close_engine(engine)
-        raise
-    return engine
-
-
 class DBContextManager:
-    def __init__(self, db_engine: aiopg.sa.Engine | None = None):
-        self._db_engine: aiopg.sa.Engine | None = db_engine
+    def __init__(self, db_engine: AsyncEngine | None = None) -> None:
+        self._db_engine: AsyncEngine | None = db_engine
         self._db_engine_created: bool = False
 
     @staticmethod
-    async def _create_db_engine() -> aiopg.sa.Engine:
+    async def _create_db_engine() -> AsyncEngine:
         settings = NodePortsSettings.create_from_envs()
-        dsn = DataSourceName(
-            application_name=f"{__name__}_{socket.gethostname()}_{os.getpid()}",
-            database=settings.POSTGRES_SETTINGS.POSTGRES_DB,
-            user=settings.POSTGRES_SETTINGS.POSTGRES_USER,
-            password=settings.POSTGRES_SETTINGS.POSTGRES_PASSWORD.get_secret_value(),
-            host=settings.POSTGRES_SETTINGS.POSTGRES_HOST,
-            port=settings.POSTGRES_SETTINGS.POSTGRES_PORT,
+        engine = await create_async_engine_and_pg_database_ready(
+            settings.POSTGRES_SETTINGS
         )
-
-        engine: aiopg.sa.Engine = await _ensure_postgres_ready(dsn)
+        assert isinstance(engine, AsyncEngine)  # nosec
         return engine
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> AsyncEngine:
         if not self._db_engine:
             self._db_engine = await self._create_db_engine()
             self._db_engine_created = True
         return self._db_engine
 
-    async def __aexit__(self, exc_type, exc, tb):
+    async def __aexit__(self, exc_type, exc, tb) -> None:
         if self._db_engine and self._db_engine_created:
-            await close_engine(self._db_engine)
-            log.debug(
-                "engine '%s' after shutdown: closed=%s, size=%d",
-                self._db_engine.dsn,
-                self._db_engine.closed,
-                self._db_engine.size,
-            )
+            await self._db_engine.dispose()
 
 
 class DBManager:
-    def __init__(self, db_engine: aiopg.sa.Engine | None = None):
+    def __init__(self, db_engine: AsyncEngine | None = None):
         self._db_engine = db_engine
 
     async def write_ports_configuration(
@@ -111,20 +81,18 @@ async def write_ports_configuration(
             f"Writing port configuration to database for "
             f"project={project_id} node={node_uuid}: {json_configuration}"
         )
-        log.debug(message)
+        _logger.debug(message)
 
         node_configuration = json.loads(json_configuration)
         async with DBContextManager(
             self._db_engine
-        ) as engine, engine.acquire() as connection:
+        ) as engine, engine.begin() as connection:
             # update the necessary parts
             await connection.execute(
                 comp_tasks.update()
                 .where(
-                    and_(
-                        comp_tasks.c.node_id == node_uuid,
-                        comp_tasks.c.project_id == project_id,
-                    )
+                    (comp_tasks.c.node_id == node_uuid)
+                    & (comp_tasks.c.project_id == project_id),
                 )
                 .values(
                     schema=node_configuration["schema"],
@@ -137,13 +105,13 @@ async def write_ports_configuration(
     async def get_ports_configuration_from_node_uuid(
         self, project_id: str, node_uuid: str
     ) -> str:
-        log.debug(
+        _logger.debug(
             "Getting ports configuration of node %s from comp_tasks table", node_uuid
         )
         async with DBContextManager(
             self._db_engine
-        ) as engine, engine.acquire() as connection:
-            node: RowProxy = await _get_node_from_db(project_id, node_uuid, connection)
+        ) as engine, engine.connect() as connection:
+            node = await _get_node_from_db(project_id, node_uuid, connection)
             node_json_config = json.dumps(
                 {
                     "schema": node.schema,
@@ -152,18 +120,18 @@ async def get_ports_configuration_from_node_uuid(
                     "run_hash": node.run_hash,
                 }
             )
-        log.debug("Found and converted to json")
+        _logger.debug("Found and converted to json")
         return node_json_config
 
     async def get_project_owner_user_id(self, project_id: ProjectID) -> UserID:
         async with DBContextManager(
             self._db_engine
-        ) as engine, engine.acquire() as connection:
-            prj_owner: Any | None = await connection.scalar(
+        ) as engine, engine.connect() as connection:
+            prj_owner = await connection.scalar(
                 sa.select(projects.c.prj_owner).where(
                     projects.c.uuid == f"{project_id}"
                 )
             )
             if prj_owner is None:
                 raise ProjectNotFoundError(project_id)
-        return UserID(prj_owner)
+        return TypeAdapter(UserID).validate_python(prj_owner)