review @pcrespov

Author: matusdrobuliak66

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

@@ -81,11 +81,3 @@ def as_postgres_sql_query_str(statement) -> str:
         dialect=postgresql.dialect(),  # type: ignore[misc]
     )
     return f"{compiled}"
-
-
-def assemble_array_groups(user_group_ids: list[int]) -> str:
-    return (
-        "array[]::text[]"
-        if len(user_group_ids) == 0
-        else f"""array[{', '.join(f"'{group_id}'" for group_id in user_group_ids)}]"""
-    )

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

@@ -0,0 +1,6 @@
+def assemble_array_groups(user_group_ids: list[int]) -> str:
+    return (
+        "array[]::text[]"
+        if len(user_group_ids) == 0
+        else f"""array[{', '.join(f"'{group_id}'" for group_id in user_group_ids)}]"""
+    )

services/storage/src/simcore_service_storage/db_access_layer.py

@@ -51,7 +51,7 @@
     workspaces_access_rights,
 )
 from simcore_postgres_database.storage_models import file_meta_data, user_to_groups
-from simcore_postgres_database.utils import assemble_array_groups
+from simcore_postgres_database.utils_sql import assemble_array_groups
 
 logger = logging.getLogger(__name__)
 

services/web/server/src/simcore_service_webserver/folders/_folders_api.py

@@ -272,7 +272,7 @@ async def update_folder(
 
     folder_db = await folders_db.update(
         app,
-        folder_id=folder_id,
+        folders_id_or_ids=folder_id,
         name=name,
         parent_folder_id=parent_folder_id,
         product_name=product_name,

services/web/server/src/simcore_service_webserver/folders/_folders_db.py

@@ -23,12 +23,17 @@
 from simcore_postgres_database.models.workspaces_access_rights import (
     workspaces_access_rights,
 )
-from simcore_postgres_database.utils import assemble_array_groups
+from simcore_postgres_database.utils_repos import (
+    pass_or_acquire_connection,
+    transaction_context,
+)
+from simcore_postgres_database.utils_sql import assemble_array_groups
 from sqlalchemy import func
+from sqlalchemy.ext.asyncio import AsyncConnection
 from sqlalchemy.orm import aliased
 from sqlalchemy.sql import ColumnElement, CompoundSelect, Select, asc, desc, select
 
-from ..db.plugin import get_database_engine
+from ..db.plugin import get_asyncpg_engine
 from ..groups.api import list_all_user_groups
 from .errors import FolderAccessForbiddenError, FolderNotFoundError
 
@@ -61,6 +66,7 @@ def as_dict_exclude_unset(**params) -> dict[str, Any]:
 
 async def create(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     created_by_gid: GroupID,
     folder_name: str,
@@ -73,8 +79,8 @@ async def create(
         user_id is not None and workspace_id is not None
     ), "Both user_id and workspace_id cannot be provided at the same time. Please provide only one."
 
-    async with get_database_engine(app).acquire() as conn:
-        result = await conn.execute(
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
+        result = await conn.stream(
             folders_v2.insert()
             .values(
                 name=folder_name,
@@ -94,6 +100,7 @@ async def create(
 
 async def list_(  # pylint: disable=too-many-arguments,too-many-branches
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     product_name: ProductName,
     user_id: UserID,
@@ -234,18 +241,17 @@ async def list_(  # pylint: disable=too-many-arguments,too-many-branches
         )
     list_query = list_query.offset(offset).limit(limit)
 
-    async with get_database_engine(app).acquire() as conn:
-        count_result = await conn.execute(count_query)
-        total_count = await count_result.scalar()
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
+        total_count = await conn.scalar(count_query)
 
-        result = await conn.execute(list_query)
-        rows = await result.fetchall() or []
-        results: list[FolderDB] = [FolderDB.from_orm(row) for row in rows]
-        return cast(int, total_count), results
+        result = await conn.stream(list_query)
+        folders: list[FolderDB] = [FolderDB.from_orm(row) async for row in result]
+        return cast(int, total_count), folders
 
 
 async def get(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     folder_id: FolderID,
     product_name: ProductName,
@@ -259,8 +265,8 @@ async def get(
         )
     )
 
-    async with get_database_engine(app).acquire() as conn:
-        result = await conn.execute(query)
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
+        result = await conn.stream(query)
         row = await result.first()
         if row is None:
             raise FolderAccessForbiddenError(
@@ -271,6 +277,7 @@ async def get(
 
 async def get_for_user_or_workspace(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     folder_id: FolderID,
     product_name: ProductName,
@@ -295,8 +302,8 @@ async def get_for_user_or_workspace(
     else:
         query = query.where(folders_v2.c.workspace_id == workspace_id)
 
-    async with get_database_engine(app).acquire() as conn:
-        result = await conn.execute(query)
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
+        result = await conn.stream(query)
         row = await result.first()
         if row is None:
             raise FolderAccessForbiddenError(
@@ -305,8 +312,10 @@ async def get_for_user_or_workspace(
         return FolderDB.from_orm(row)
 
 
-async def _update_impl(
+async def update(
     app: web.Application,
+    connection: AsyncConnection | None = None,
+    *,
     folders_id_or_ids: FolderID | set[FolderID],
     product_name: ProductName,
     # updatable columns
@@ -339,64 +348,22 @@ async def _update_impl(
         # single-update
         query = query.where(folders_v2.c.folder_id == folders_id_or_ids)
 
-    async with get_database_engine(app).acquire() as conn:
-        result = await conn.execute(query)
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
+        result = await conn.stream(query)
         row = await result.first()
         if row is None:
             raise FolderNotFoundError(reason=f"Folder {folders_id_or_ids} not found.")
         return FolderDB.from_orm(row)
 
 
-async def update_batch(
-    app: web.Application,
-    *folder_id: FolderID,
-    product_name: ProductName,
-    # updatable columns
-    name: str | UnSet = _unset,
-    parent_folder_id: FolderID | None | UnSet = _unset,
-    trashed_at: datetime | None | UnSet = _unset,
-    trashed_explicitly: bool | UnSet = _unset,
-) -> FolderDB:
-    return await _update_impl(
-        app=app,
-        folders_id_or_ids=set(folder_id),
-        product_name=product_name,
-        name=name,
-        parent_folder_id=parent_folder_id,
-        trashed_at=trashed_at,
-        trashed_explicitly=trashed_explicitly,
-    )
-
-
-async def update(
-    app: web.Application,
-    *,
-    folder_id: FolderID,
-    product_name: ProductName,
-    # updatable columns
-    name: str | UnSet = _unset,
-    parent_folder_id: FolderID | None | UnSet = _unset,
-    trashed_at: datetime | None | UnSet = _unset,
-    trashed_explicitly: bool | UnSet = _unset,
-) -> FolderDB:
-    return await _update_impl(
-        app=app,
-        folders_id_or_ids=folder_id,
-        product_name=product_name,
-        name=name,
-        parent_folder_id=parent_folder_id,
-        trashed_at=trashed_at,
-        trashed_explicitly=trashed_explicitly,
-    )
-
-
 async def delete_recursively(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     folder_id: FolderID,
     product_name: ProductName,
 ) -> None:
-    async with get_database_engine(app).acquire() as conn, conn.begin():
+    async with transaction_context(get_asyncpg_engine(app), connection) as conn:
         # Step 1: Define the base case for the recursive CTE
         base_query = select(
             folders_v2.c.folder_id, folders_v2.c.parent_folder_id
@@ -422,10 +389,9 @@ async def delete_recursively(
 
         # Step 4: Execute the query to get all descendants
         final_query = select(folder_hierarchy_cte)
-        result = await conn.execute(final_query)
-        rows = (  # list of tuples [(folder_id, parent_folder_id), ...] ex. [(1, None), (2, 1)]
-            await result.fetchall() or []
-        )
+        result = await conn.stream(final_query)
+        # list of tuples [(folder_id, parent_folder_id), ...] ex. [(1, None), (2, 1)]
+        rows = [row async for row in result]
 
         # Sort folders so that child folders come first
         sorted_folders = sorted(
@@ -439,6 +405,7 @@ async def delete_recursively(
 
 async def get_projects_recursively_only_if_user_is_owner(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     folder_id: FolderID,
     private_workspace_user_id_or_none: UserID | None,
@@ -453,7 +420,8 @@ async def get_projects_recursively_only_if_user_is_owner(
     or the `users_to_groups` table for private workspace projects.
     """
 
-    async with get_database_engine(app).acquire() as conn, conn.begin():
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
+
         # Step 1: Define the base case for the recursive CTE
         base_query = select(
             folders_v2.c.folder_id, folders_v2.c.parent_folder_id
@@ -462,6 +430,7 @@ async def get_projects_recursively_only_if_user_is_owner(
             & (folders_v2.c.product_name == product_name)
         )
         folder_hierarchy_cte = base_query.cte(name="folder_hierarchy", recursive=True)
+
         # Step 2: Define the recursive case
         folder_alias = aliased(folders_v2)
         recursive_query = select(
@@ -472,16 +441,15 @@ async def get_projects_recursively_only_if_user_is_owner(
                 folder_alias.c.parent_folder_id == folder_hierarchy_cte.c.folder_id,
             )
         )
+
         # Step 3: Combine base and recursive cases into a CTE
         folder_hierarchy_cte = folder_hierarchy_cte.union_all(recursive_query)
+
         # Step 4: Execute the query to get all descendants
         final_query = select(folder_hierarchy_cte)
-        result = await conn.execute(final_query)
-        rows = (  # list of tuples [(folder_id, parent_folder_id), ...] ex. [(1, None), (2, 1)]
-            await result.fetchall() or []
-        )
-
-        folder_ids = [item[0] for item in rows]
+        result = await conn.stream(final_query)
+        # list of tuples [(folder_id, parent_folder_id), ...] ex. [(1, None), (2, 1)]
+        folder_ids = [item[0] async for item in result]
 
         query = (
             select(projects_to_folders.c.project_uuid)
@@ -494,19 +462,19 @@ async def get_projects_recursively_only_if_user_is_owner(
         if private_workspace_user_id_or_none is not None:
             query = query.where(projects.c.prj_owner == user_id)
 
-        result = await conn.execute(query)
-
-        rows = await result.fetchall() or []
-        return [ProjectID(row[0]) for row in rows]
+        result = await conn.stream(query)
+        return [ProjectID(row[0]) async for row in result]
 
 
 async def get_folders_recursively(
     app: web.Application,
+    connection: AsyncConnection | None = None,
     *,
     folder_id: FolderID,
     product_name: ProductName,
 ) -> list[FolderID]:
-    async with get_database_engine(app).acquire() as conn, conn.begin():
+    async with pass_or_acquire_connection(get_asyncpg_engine(app), connection) as conn:
+
         # Step 1: Define the base case for the recursive CTE
         base_query = select(
             folders_v2.c.folder_id, folders_v2.c.parent_folder_id
@@ -532,9 +500,5 @@ async def get_folders_recursively(
 
         # Step 4: Execute the query to get all descendants
         final_query = select(folder_hierarchy_cte)
-        result = await conn.execute(final_query)
-        rows = (  # list of tuples [(folder_id, parent_folder_id), ...] ex. [(1, None), (2, 1)]
-            await result.fetchall() or []
-        )
-
-        return [FolderID(row[0]) for row in rows]
+        result = await conn.stream(final_query)
+        return [FolderID(row[0]) async for row in result]