migration

Author: sanderegg

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

@@ -1,8 +1,9 @@
 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
@@ -22,12 +23,10 @@
 )
 
 
-class GroupExtraPropertiesError(Exception):
-    ...
+class GroupExtraPropertiesError(Exception): ...
 
 
-class GroupExtraPropertiesNotFoundError(GroupExtraPropertiesError):
-    ...
+class GroupExtraPropertiesNotFoundError(GroupExtraPropertiesError): ...
 
 
 @dataclass(frozen=True, slots=True, kw_only=True)
@@ -213,7 +212,6 @@ async def get_aggregated_properties_for_user_v2(
         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
             )

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

@@ -9,10 +9,10 @@
 from typing import Any, Final
 
 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 import Column
 
+from ._protocols import DBConnection
 from .aiopg_errors import UniqueViolation
 from .models.users import UserRole, UserStatus, users
 from .models.users_details import users_pre_registration_details
@@ -55,12 +55,12 @@ def generate_alternative_username(username: str) -> str:
 class UsersRepo:
     @staticmethod
     async def new_user(
-        conn: SAConnection,
+        conn: DBConnection,
         email: str,
         password_hash: str,
         status: UserStatus,
         expires_at: datetime | None,
-    ) -> RowProxy:
+    ) -> Any:
         data: dict[str, Any] = {
             "name": _generate_username_from_email(email),
             "email": email,
@@ -88,13 +88,15 @@ async def new_user(
                 users.c.status,
             ).where(users.c.id == user_id)
         )
-        row = await result.first()
+        row = await maybe_await(result.first())
+        from aiopg.sa.result import RowProxy
+
         assert isinstance(row, RowProxy)  # nosec
         return row
 
     @staticmethod
     async def join_and_update_from_pre_registration_details(
-        conn: SAConnection, new_user_id: int, new_user_email: str
+        conn: DBConnection, new_user_id: int, new_user_email: str
     ) -> None:
         """After a user is created, it can be associated with information provided during invitation
 
@@ -111,6 +113,10 @@ async def join_and_update_from_pre_registration_details(
             .values(user_id=new_user_id)
         )
 
+        from aiopg.sa.result import ResultProxy
+
+        assert isinstance(result, ResultProxy)  # nosec
+
         if result.rowcount:
             pre_columns = (
                 users_pre_registration_details.c.pre_first_name,
@@ -135,7 +141,7 @@ async def join_and_update_from_pre_registration_details(
                     users_pre_registration_details.c.pre_email == new_user_email
                 )
             )
-            if details := await result.fetchone():
+            if details := await maybe_await(result.fetchone()):
                 await conn.execute(
                     users.update()
                     .where(users.c.id == new_user_id)
@@ -169,15 +175,14 @@ def get_billing_details_query(user_id: int):
         )
 
     @staticmethod
-    async def get_billing_details(conn: SAConnection, user_id: int) -> RowProxy | None:
+    async def get_billing_details(conn: DBConnection, user_id: int) -> Any | None:
         result = await conn.execute(
             UsersRepo.get_billing_details_query(user_id=user_id)
         )
-        value: RowProxy | None = await result.fetchone()
-        return value
+        return await maybe_await(result.fetchone())
 
     @staticmethod
-    async def get_role(conn: SAConnection, user_id: int) -> UserRole:
+    async def get_role(conn: DBConnection, user_id: int) -> UserRole:
         value: UserRole | None = await conn.scalar(
             sa.select(users.c.role).where(users.c.id == user_id)
         )
@@ -188,7 +193,7 @@ async def get_role(conn: SAConnection, user_id: int) -> UserRole:
         raise UserNotFoundInRepoError
 
     @staticmethod
-    async def get_email(conn: SAConnection, user_id: int) -> str:
+    async def get_email(conn: DBConnection, user_id: int) -> str:
         value: str | None = await conn.scalar(
             sa.select(users.c.email).where(users.c.id == user_id)
         )
@@ -199,7 +204,7 @@ async def get_email(conn: SAConnection, user_id: int) -> str:
         raise UserNotFoundInRepoError
 
     @staticmethod
-    async def get_active_user_email(conn: SAConnection, user_id: int) -> str:
+    async def get_active_user_email(conn: DBConnection, user_id: int) -> str:
         value: str | None = await conn.scalar(
             sa.select(users.c.email).where(
                 (users.c.status == UserStatus.ACTIVE) & (users.c.id == user_id)
@@ -212,7 +217,7 @@ async def get_active_user_email(conn: SAConnection, user_id: int) -> str:
         raise UserNotFoundInRepoError
 
     @staticmethod
-    async def is_email_used(conn: SAConnection, email: str) -> bool:
+    async def is_email_used(conn: DBConnection, email: str) -> bool:
         email = email.lower()
 
         registered = await conn.scalar(

packages/postgres-database/tests/test_utils_projects_nodes.py

@@ -28,6 +28,8 @@
 )
 from sqlalchemy.ext.asyncio import AsyncConnection
 
+# NOTE: Temporary usage of connection_factory until asyncpg is used
+
 
 async def _delete_project(
     connection_factory: SAConnection, project_uuid: uuid.UUID

packages/postgres-database/tests/test_utils_users.py

@@ -12,6 +12,7 @@
 from pytest_simcore.helpers.faker_factories import random_user
 from simcore_postgres_database.models.users import UserRole, users
 from simcore_postgres_database.utils_users import UserNotFoundInRepoError, UsersRepo
+from sqlalchemy.ext.asyncio import AsyncConnection
 
 
 @pytest.fixture
@@ -25,11 +26,13 @@ async def user(connection: SAConnection, faker: Faker) -> dict[str, Any]:
     return data
 
 
-async def test_users_repo_get(connection: SAConnection, user: dict[str, Any]):
+async def test_users_repo_get(
+    connection_factory: SAConnection | AsyncConnection, user: dict[str, Any]
+):
     repo = UsersRepo()
-
-    assert await repo.get_email(connection, user_id=user["id"]) == user["email"]
-    assert await repo.get_role(connection, user_id=user["id"]) == user["role"]
+    # NOTE: Temporary usage of connection_factory until asyncpg is used
+    assert await repo.get_email(connection_factory, user_id=user["id"]) == user["email"]
+    assert await repo.get_role(connection_factory, user_id=user["id"]) == user["role"]
 
     with pytest.raises(UserNotFoundInRepoError):
-        await repo.get_role(connection, user_id=55)
+        await repo.get_role(connection_factory, user_id=55)