moves workspaces

Author: pcrespov

services/web/server/src/simcore_service_webserver/workspaces/_workspaces_db.py

@@ -22,11 +22,16 @@
 from simcore_postgres_database.models.workspaces_access_rights import (
     workspaces_access_rights,
 )
+from simcore_postgres_database.utils_repos import (
+    pass_or_acquire_connection,
+    transaction_context,
+)
 from sqlalchemy import asc, desc, func
 from sqlalchemy.dialects.postgresql import BOOLEAN, INTEGER
+from sqlalchemy.ext.asyncio import AsyncConnection
 from sqlalchemy.sql import Subquery, select
 
-from ..db.plugin import get_database_engine
+from ..db.plugin import get_asyncpg_engine
 from .errors import WorkspaceAccessForbiddenError, WorkspaceNotFoundError
 
 _logger = logging.getLogger(__name__)
@@ -45,13 +50,15 @@
 
 async def create_workspace(
     app: web.Application,
+    connection: AsyncConnection | None = None,
+    *,
     product_name: ProductName,
     owner_primary_gid: GroupID,
     name: str,
     description: str | None,
     thumbnail: str | None,
 ) -> WorkspaceDB:
-    async with get_database_engine(app).acquire() as conn:
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
         result = await conn.execute(
             workspaces.insert()
             .values(
@@ -65,11 +72,11 @@ async def create_workspace(
             )
             .returning(*_SELECTION_ARGS)
         )
-        row = await result.first()
+        row = result.first()
         return WorkspaceDB.from_orm(row)
 
 
-access_rights_subquery = (
+_access_rights_subquery = (
     select(
         workspaces_access_rights.c.workspace_id,
         func.jsonb_object_agg(
@@ -116,6 +123,7 @@ def _create_my_access_rights_subquery(user_id: UserID) -> Subquery:
 
 async def list_workspaces_for_user(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     user_id: UserID,
     product_name: ProductName,
@@ -128,11 +136,11 @@ async def list_workspaces_for_user(
     base_query = (
         select(
             *_SELECTION_ARGS,
-            access_rights_subquery.c.access_rights,
+            _access_rights_subquery.c.access_rights,
             my_access_rights_subquery.c.my_access_rights,
         )
         .select_from(
-            workspaces.join(access_rights_subquery).join(my_access_rights_subquery)
+            workspaces.join(_access_rights_subquery).join(my_access_rights_subquery)
         )
         .where(workspaces.c.product_name == product_name)
     )
@@ -148,12 +156,12 @@ async def list_workspaces_for_user(
         list_query = base_query.order_by(desc(getattr(workspaces.c, order_by.field)))
     list_query = list_query.offset(offset).limit(limit)
 
-    async with get_database_engine(app).acquire() as conn:
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
         count_result = await conn.execute(count_query)
-        total_count = await count_result.scalar()
+        total_count = count_result.scalar()
 
         result = await conn.execute(list_query)
-        rows = await result.fetchall() or []
+        rows = result.fetchall() or []
         results: list[UserWorkspaceAccessRightsDB] = [
             UserWorkspaceAccessRightsDB.from_orm(row) for row in rows
         ]
@@ -163,6 +171,8 @@ async def list_workspaces_for_user(
 
 async def get_workspace_for_user(
     app: web.Application,
+    connection: AsyncConnection | None = None,
+    *,
     user_id: UserID,
     workspace_id: WorkspaceID,
     product_name: ProductName,
@@ -172,21 +182,21 @@ async def get_workspace_for_user(
     base_query = (
         select(
             *_SELECTION_ARGS,
-            access_rights_subquery.c.access_rights,
+            _access_rights_subquery.c.access_rights,
             my_access_rights_subquery.c.my_access_rights,
         )
         .select_from(
-            workspaces.join(access_rights_subquery).join(my_access_rights_subquery)
+            workspaces.join(_access_rights_subquery).join(my_access_rights_subquery)
         )
         .where(
             (workspaces.c.workspace_id == workspace_id)
             & (workspaces.c.product_name == product_name)
         )
     )
 
-    async with get_database_engine(app).acquire() as conn:
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
         result = await conn.execute(base_query)
-        row = await result.first()
+        row = result.first()
         if row is None:
             raise WorkspaceAccessForbiddenError(
                 reason=f"User {user_id} does not have access to the workspace {workspace_id}. Or workspace does not exist.",
@@ -196,13 +206,15 @@ async def get_workspace_for_user(
 
 async def update_workspace(
     app: web.Application,
+    connection: AsyncConnection | None = None,
+    *,
     workspace_id: WorkspaceID,
     name: str,
     description: str | None,
     thumbnail: str | None,
     product_name: ProductName,
 ) -> WorkspaceDB:
-    async with get_database_engine(app).acquire() as conn:
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
         result = await conn.execute(
             workspaces.update()
             .values(
@@ -217,18 +229,20 @@ async def update_workspace(
             )
             .returning(*_SELECTION_ARGS)
         )
-        row = await result.first()
+        row = result.first()
         if row is None:
             raise WorkspaceNotFoundError(reason=f"Workspace {workspace_id} not found.")
         return WorkspaceDB.from_orm(row)
 
 
 async def delete_workspace(
     app: web.Application,
+    connection: AsyncConnection | None = None,
+    *,
     workspace_id: WorkspaceID,
     product_name: ProductName,
 ) -> None:
-    async with get_database_engine(app).acquire() as conn:
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
         await conn.execute(
             workspaces.delete().where(
                 (workspaces.c.workspace_id == workspace_id)