UsersRepo.new_user

Author: pcrespov

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

@@ -11,11 +11,13 @@
 import sqlalchemy as sa
 from common_library.async_tools import maybe_await
 from sqlalchemy import Column
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio.engine import AsyncConnection, AsyncEngine
 
 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
+from .utils_repos import pass_or_acquire_connection, transaction_context
 
 
 class BaseUserRepoError(Exception):
@@ -53,9 +55,12 @@ def generate_alternative_username(username: str) -> str:
 
 
 class UsersRepo:
+
     @staticmethod
     async def new_user(
-        conn: DBConnection,
+        engine: AsyncEngine,
+        connection: AsyncConnection | None = None,
+        *,
         email: str,
         password_hash: str,
         status: UserStatus,
@@ -73,26 +78,28 @@ async def new_user(
         user_id = None
         while user_id is None:
             try:
-                user_id = await conn.scalar(
-                    users.insert().values(**data).returning(users.c.id)
-                )
-            except UniqueViolation:
+                async with transaction_context(engine, connection) as conn:
+                    user_id = await conn.scalar(
+                        users.insert().values(**data).returning(users.c.id)
+                    )
+            except IntegrityError:
                 data["name"] = generate_alternative_username(data["name"])
 
-        result = await conn.execute(
-            sa.select(
-                users.c.id,
-                users.c.name,
-                users.c.email,
-                users.c.role,
-                users.c.status,
-            ).where(users.c.id == user_id)
-        )
-        return await maybe_await(result.first())
+        async with pass_or_acquire_connection(engine, connection) as conn:
+            result = await conn.execute(
+                sa.select(
+                    users.c.id,
+                    users.c.name,
+                    users.c.email,
+                    users.c.role,
+                    users.c.status,
+                ).where(users.c.id == user_id)
+            )
+            return result.one()
 
     @staticmethod
     async def link_and_update_user_from_pre_registration(
-        conn: DBConnection,
+        connection: AsyncConnection,
         *,
         new_user_id: int,
         new_user_email: str,
@@ -107,7 +114,7 @@ async def link_and_update_user_from_pre_registration(
         assert new_user_id > 0  # nosec
 
         # link both tables first
-        result = await conn.execute(
+        result = await connection.execute(
             users_pre_registration_details.update()
             .where(users_pre_registration_details.c.pre_email == new_user_email)
             .values(user_id=new_user_id)
@@ -133,14 +140,14 @@ async def link_and_update_user_from_pre_registration(
                 and c.name.startswith("pre_")
             }, "Different pre-cols detected. This code might need an update update"
 
-            result = await conn.execute(
+            result = await connection.execute(
                 sa.select(*pre_columns).where(
                     users_pre_registration_details.c.pre_email == new_user_email
                 )
             )
             if pre_registration_details_data := result.first():
                 # NOTE: could have many products! which to use?
-                await conn.execute(
+                await connection.execute(
                     users.update()
                     .where(users.c.id == new_user_id)
                     .values(

packages/postgres-database/tests/test_users.py

@@ -11,14 +11,15 @@
 from pytest_simcore.helpers.faker_factories import random_user
 from simcore_postgres_database.models.users import UserRole, UserStatus, users
 from simcore_postgres_database.utils_repos import (
+    pass_or_acquire_connection,
     transaction_context,
 )
 from simcore_postgres_database.utils_users import (
     UsersRepo,
     _generate_username_from_email,
     generate_alternative_username,
 )
-from sqlalchemy.exc import DataError, IntegrityError
+from sqlalchemy.exc import DBAPIError, IntegrityError
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.sql import func
 
@@ -52,7 +53,7 @@ async def test_user_status_as_pending(
     data = random_user(faker, status="PENDING")
     assert data["status"] == "PENDING"
     async with transaction_context(asyncpg_engine) as connection:
-        with pytest.raises(DataError) as err_info:
+        with pytest.raises(DBAPIError) as err_info:
             await connection.execute(users.insert().values(data))
 
         assert (
@@ -113,16 +114,20 @@ async def test_unique_username(
         assert user.id == user_id
         assert user.name == "pcrespov"
 
+    async with transaction_context(asyncpg_engine) as connection:
         # same name fails
         data["email"] = faker.email()
         with pytest.raises(IntegrityError):
             await connection.scalar(users.insert().values(data).returning(users.c.id))
 
+    async with transaction_context(asyncpg_engine) as connection:
         # generate new name
         data["name"] = _generate_username_from_email(user.email)
         data["email"] = faker.email()
         await connection.scalar(users.insert().values(data).returning(users.c.id))
 
+    async with transaction_context(asyncpg_engine) as connection:
+
         # and another one
         data["name"] = generate_alternative_username(data["name"])
         data["email"] = faker.email()
@@ -138,21 +143,21 @@ async def test_new_user(
         "status": UserStatus.ACTIVE,
         "expires_at": datetime.utcnow(),
     }
-    async with transaction_context(asyncpg_engine) as connection:
-        new_user = await UsersRepo.new_user(connection, **data)
+    new_user = await UsersRepo.new_user(asyncpg_engine, **data)
 
-        assert new_user.email == data["email"]
-        assert new_user.status == data["status"]
-        assert new_user.role == UserRole.USER
+    assert new_user.email == data["email"]
+    assert new_user.status == data["status"]
+    assert new_user.role == UserRole.USER
 
-        other_email = f"{new_user.name}@other-domain.com"
-        assert _generate_username_from_email(other_email) == new_user.name
-        other_data = {**data, "email": other_email}
+    other_email = f"{new_user.name}@other-domain.com"
+    assert _generate_username_from_email(other_email) == new_user.name
+    other_data = {**data, "email": other_email}
 
-        other_user = await UsersRepo.new_user(connection, **other_data)
-        assert other_user.email != new_user.email
-        assert other_user.name != new_user.name
+    other_user = await UsersRepo.new_user(asyncpg_engine, **other_data)
+    assert other_user.email != new_user.email
+    assert other_user.name != new_user.name
 
+    async with pass_or_acquire_connection(asyncpg_engine) as connection:
         assert await UsersRepo.get_email(connection, other_user.id) == other_user.email
         assert await UsersRepo.get_role(connection, other_user.id) == other_user.role
         assert (

packages/postgres-database/tests/test_users_details.py

@@ -258,6 +258,7 @@ async def test_create_and_link_user_from_pre_registration(
     async with transaction_context(asyncpg_engine) as connection:
         # user gets created
         new_user = await UsersRepo.new_user(
+            asyncpg_engine,
             connection,
             email=pre_email,
             password_hash="123456",  # noqa: S106
@@ -292,6 +293,7 @@ async def test_get_billing_details_from_pre_registration(
     # Create the user
     async with transaction_context(asyncpg_engine) as connection:
         new_user = await UsersRepo.new_user(
+            asyncpg_engine,
             connection,
             email=pre_email,
             password_hash="123456",  # noqa: S106
@@ -332,6 +334,7 @@ async def test_update_user_from_pre_registration(
     # Create the user and link to pre-registration
     async with transaction_context(asyncpg_engine) as connection:
         new_user = await UsersRepo.new_user(
+            asyncpg_engine,
             connection,
             email=pre_email,
             password_hash="123456",  # noqa: S106
@@ -490,6 +493,7 @@ async def test_user_preregisters_for_multiple_products_with_different_outcomes(
     # 3.Now create a user account with the approved pre-registration
     async with transaction_context(asyncpg_engine) as connection:
         new_user = await UsersRepo.new_user(
+            asyncpg_engine,
             connection,
             email=user_email,
             password_hash="123456",  # noqa: S106

services/web/server/src/simcore_service_webserver/login/_auth_service.py

@@ -31,8 +31,10 @@ async def create_user(
     expires_at: datetime | None,
 ) -> dict[str, Any]:
 
-    async with transaction_context(get_asyncpg_engine(app)) as conn:
+    asyncpg_engine = get_asyncpg_engine(app)
+    async with transaction_context(asyncpg_engine) as conn:
         user = await UsersRepo.new_user(
+            asyncpg_engine,
             conn,
             email=email,
             password_hash=security_service.encrypt_password(password),