♻️Director v2: remove aiopg usage (#7576)

Author: sanderegg

.github/copilot-instructions.md

@@ -22,6 +22,7 @@ This document provides guidelines and best practices for using GitHub Copilot in
 - ensure we use `pydantic` >2 compatible code.
 - ensure we use `fastapi` >0.100 compatible code
 - use f-string formatting
+- Only add comments in function if strictly necessary
 
 
 ### Json serialization

packages/common-library/src/common_library/async_tools.py

@@ -2,7 +2,8 @@
 import functools
 from collections.abc import Awaitable, Callable
 from concurrent.futures import Executor
-from typing import ParamSpec, TypeVar
+from inspect import isawaitable
+from typing import ParamSpec, TypeVar, overload
 
 R = TypeVar("R")
 P = ParamSpec("P")
@@ -22,3 +23,42 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         return wrapper
 
     return decorator
+
+
+_AwaitableResult = TypeVar("_AwaitableResult")
+
+
+@overload
+async def maybe_await(obj: Awaitable[_AwaitableResult]) -> _AwaitableResult: ...
+
+
+@overload
+async def maybe_await(obj: _AwaitableResult) -> _AwaitableResult: ...
+
+
+async def maybe_await(
+    obj: Awaitable[_AwaitableResult] | _AwaitableResult,
+) -> _AwaitableResult:
+    """Helper function to handle both async and sync database results.
+
+    This function allows code to work with both aiopg (async) and asyncpg (sync) result methods
+    by automatically detecting and handling both cases.
+
+    Args:
+        obj: Either an awaitable coroutine or direct result value
+
+    Returns:
+        The result value, after awaiting if necessary
+
+    Example:
+        ```python
+        result = await conn.execute(query)
+        # Works with both aiopg and asyncpg
+        row = await maybe_await(result.fetchone())
+        ```
+    """
+    if isawaitable(obj):
+        assert isawaitable(obj)  # nosec
+        return await obj
+    assert not isawaitable(obj)  # nosec
+    return obj

packages/common-library/tests/test_async_tools.py

@@ -1,8 +1,9 @@
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any
 
 import pytest
-from common_library.async_tools import make_async
+from common_library.async_tools import make_async, maybe_await
 
 
 @make_async()
@@ -43,3 +44,52 @@ def heavy_computation(x: int) -> int:
 
     result = await heavy_computation(4)
     assert result == 16, "Function should return 16"
+
+
+@pytest.mark.asyncio
+async def test_maybe_await_with_coroutine():
+    """Test maybe_await with an async function"""
+
+    async def async_value():
+        return 42
+
+    result = await maybe_await(async_value())
+    assert result == 42
+
+
+@pytest.mark.asyncio
+async def test_maybe_await_with_direct_value():
+    """Test maybe_await with a direct value"""
+    value = 42
+    result = await maybe_await(value)
+    assert result == 42
+
+
+@pytest.mark.asyncio
+async def test_maybe_await_with_none():
+    """Test maybe_await with None value"""
+    result = await maybe_await(None)
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_maybe_await_with_result_proxy():
+    """Test maybe_await with both async and sync ResultProxy implementations"""
+
+    class AsyncResultProxy:
+        """Mock async result proxy (aiopg style)"""
+
+        async def fetchone(self) -> Any:  # pylint: disable=no-self-use
+            return {"id": 1, "name": "test"}
+
+    class SyncResultProxy:
+        """Mock sync result proxy (asyncpg style)"""
+
+        def fetchone(self) -> Any:  # pylint: disable=no-self-use
+            return {"id": 2, "name": "test2"}
+
+    async_result = await maybe_await(AsyncResultProxy().fetchone())
+    assert async_result == {"id": 1, "name": "test"}
+
+    sync_result = await maybe_await(SyncResultProxy().fetchone())
+    assert sync_result == {"id": 2, "name": "test2"}

packages/postgres-database/src/simcore_postgres_database/_protocols.py

@@ -1,35 +1,49 @@
 """Common protocols to annotate equivalent connections:
-        - sqlalchemy.ext.asyncio.AsyncConnection
-        - aiopg.sa.connection.SAConnection
+    - sqlalchemy.ext.asyncio.AsyncConnection
+    - aiopg.sa.connection.SAConnection
 
-
-    Purpose: to reduce dependency wit aiopg (expected full migration to asyncpg)
+Purpose: to reduce dependency with aiopg (expected full migration to asyncpg)
 """
 
-from typing import Protocol
+from collections.abc import Awaitable
+from typing import Any, Protocol, TypeAlias, TypeVar
 
+from sqlalchemy.sql import Executable
 
-class DBConnection(Protocol):
-    # Prototype to account for aiopg and asyncio connection classes, i.e.
-    #   from aiopg.sa.connection import SAConnection
-    #   from sqlalchemy.ext.asyncio import AsyncConnection
-    async def scalar(self, *args, **kwargs):
-        ...
+# Type for query results
+Result = TypeVar("Result")
 
-    async def execute(self, *args, **kwargs):
-        ...
+# Type alias for methods that can be either async or sync
+MaybeCoro: TypeAlias = Awaitable[Result] | Result
 
-    async def begin(self):
-        ...
 
+class ResultProxy(Protocol):
+    """Protocol for query result objects from both engines
 
-class AiopgConnection(Protocol):
-    # Prototype to account for aiopg-only (this protocol avoids import <-> installation)
-    async def scalar(self, *args, **kwargs):
-        ...
+    Handles both aiopg's async methods and SQLAlchemy asyncpg's sync methods.
+    This is temporary until we fully migrate to asyncpg.
+    """
 
-    async def execute(self, *args, **kwargs):
-        ...
+    def fetchall(self) -> MaybeCoro[list[Any]]: ...
+    def fetchone(self) -> MaybeCoro[Any | None]: ...
+    def first(self) -> MaybeCoro[Any | None]: ...
 
-    async def begin(self):
-        ...
+
+class DBConnection(Protocol):
+    """Protocol to account for both aiopg and SQLAlchemy async connections"""
+
+    async def scalar(
+        self,
+        statement: Executable,
+        parameters: dict[str, Any] | None = None,
+        *,
+        execution_options: dict[str, Any] | None = None,
+    ) -> Any: ...
+
+    async def execute(
+        self,
+        statement: Executable,
+        parameters: dict[str, Any] | None = None,
+        *,
+        execution_options: dict[str, Any] | None = None,
+    ) -> ResultProxy: ...

packages/postgres-database/src/simcore_postgres_database/utils_aiosqlalchemy.py

@@ -1,4 +1,9 @@
+from typing import Any, TypeAlias, TypeVar
+
 import sqlalchemy as sa
+import sqlalchemy.exc as sql_exc
+from common_library.errors_classes import OsparcErrorMixin
+from sqlalchemy.dialects.postgresql.asyncpg import AsyncAdapt_asyncpg_dbapi
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 from .utils_migration import get_current_head
@@ -29,3 +34,49 @@ async def raise_if_migration_not_ready(engine: AsyncEngine) -> None:
         if version_num != head_version_num:
             msg = f"Migration is incomplete, expected {head_version_num} but got {version_num}"
             raise DBMigrationError(msg)
+
+
+AsyncpgSQLState: TypeAlias = str
+ErrorT = TypeVar("ErrorT", bound=OsparcErrorMixin)
+ErrorKwars: TypeAlias = dict[str, Any]
+
+
+def map_db_exception(
+    exception: Exception,
+    exception_map: dict[AsyncpgSQLState, tuple[type[ErrorT], ErrorKwars]],
+    default_exception: type[ErrorT] | None = None,
+) -> ErrorT | Exception:
+    """Maps SQLAlchemy database exceptions to domain-specific exceptions.
+
+    This function inspects SQLAlchemy and asyncpg exceptions to identify the error type
+    by checking pgcodes or error messages, and converts them to appropriate domain exceptions.
+
+    Args:
+        exception: The original exception from SQLAlchemy or the database driver
+        exception_map: Dictionary mapping pgcode
+        default_exception: Exception class to use if no matching error is found
+
+    Returns:
+        Domain-specific exception instance or the original exception if no mapping found
+        and no default_exception provided
+    """
+    pgcode = None
+
+    # Handle SQLAlchemy wrapped exceptions
+    if isinstance(exception, sql_exc.IntegrityError) and hasattr(exception, "orig"):
+        orig_error = exception.orig
+        # Handle asyncpg adapter exceptions
+        if isinstance(orig_error, AsyncAdapt_asyncpg_dbapi.IntegrityError) and hasattr(
+            orig_error, "pgcode"
+        ):
+            assert hasattr(orig_error, "pgcode")  # nosec
+            pgcode = orig_error.pgcode
+
+    # Match by pgcode if available
+    if pgcode:
+        for key, (exc_class, params) in exception_map.items():
+            if key == pgcode:
+                return exc_class(**params)
+
+    # If no match found, return default exception or original
+    return default_exception() if default_exception else exception

packages/postgres-database/src/simcore_postgres_database/utils_groups_extra_properties.py

@@ -1,14 +1,15 @@
 import datetime
 import logging
 import warnings
+from collections.abc import Callable
 from dataclasses import dataclass, fields
-from typing import Any, Callable
+from typing import Any
 
 import sqlalchemy as sa
-from aiopg.sa.connection import SAConnection
-from aiopg.sa.result import RowProxy
+from common_library.async_tools import maybe_await
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
+from ._protocols import DBConnection
 from .models.groups import GroupType, groups, user_to_groups
 from .models.groups_extra_properties import groups_extra_properties
 from .utils_models import FromRowMixin
@@ -22,12 +23,10 @@
 )
 
 
-class GroupExtraPropertiesError(Exception):
-    ...
+class GroupExtraPropertiesError(Exception): ...
 
 
-class GroupExtraPropertiesNotFoundError(GroupExtraPropertiesError):
-    ...
+class GroupExtraPropertiesNotFoundError(GroupExtraPropertiesError): ...
 
 
 @dataclass(frozen=True, slots=True, kw_only=True)
@@ -99,24 +98,6 @@ def _get_stmt(gid: int, product_name: str):
             & (groups_extra_properties.c.product_name == product_name)
         )
 
-    @staticmethod
-    async def get(
-        connection: SAConnection, *, gid: int, product_name: str
-    ) -> GroupExtraProperties:
-        warnings.warn(
-            _WARNING_FMSG.format("get", "get_v2"),
-            DeprecationWarning,
-            stacklevel=1,
-        )
-
-        query = GroupExtraPropertiesRepo._get_stmt(gid, product_name)
-        result = await connection.execute(query)
-        assert result  # nosec
-        if row := await result.first():
-            return GroupExtraProperties.from_row_proxy(row)
-        msg = f"Properties for group {gid} not found"
-        raise GroupExtraPropertiesNotFoundError(msg)
-
     @staticmethod
     async def get_v2(
         engine: AsyncEngine,
@@ -174,7 +155,7 @@ def _aggregate(
 
     @staticmethod
     async def get_aggregated_properties_for_user(
-        connection: SAConnection,
+        connection: DBConnection,
         *,
         user_id: int,
         product_name: str,
@@ -197,30 +178,9 @@ async def get_aggregated_properties_for_user(
         )
         assert result  # nosec
 
-        rows: list[RowProxy] | None = await result.fetchall()
-        assert rows is not None  # nosec
+        rows = await maybe_await(result.fetchall())
+        assert isinstance(rows, list)  # nosec
 
         return GroupExtraPropertiesRepo._aggregate(
-            rows, user_id, product_name, GroupExtraProperties.from_row_proxy
+            rows, user_id, product_name, GroupExtraProperties.from_row
         )
-
-    @staticmethod
-    async def get_aggregated_properties_for_user_v2(
-        engine: AsyncEngine,
-        connection: AsyncConnection | None = None,
-        *,
-        user_id: int,
-        product_name: str,
-    ) -> GroupExtraProperties:
-        async with pass_or_acquire_connection(engine, connection) as conn:
-
-            list_stmt = _list_table_entries_ordered_by_group_type_stmt(
-                user_id=user_id, product_name=product_name
-            )
-            result = await conn.stream(
-                sa.select(list_stmt).order_by(list_stmt.c.type_order)
-            )
-            rows = [row async for row in result]
-            return GroupExtraPropertiesRepo._aggregate(
-                rows, user_id, product_name, GroupExtraProperties.from_row
-            )