feat: Add new vfolder API to update sharing status (#2740)

Backported-from: main (24.09)
Backported-to: 23.09
Backport-of: 2740

changes/2740.feature.md

@@ -0,0 +1 @@
+Add new vfolder API to update sharing status.

src/ai/backend/manager/api/vfolder.py

@@ -6,13 +6,15 @@
 import logging
 import math
 import stat
+import textwrap
 import uuid
 from datetime import datetime
 from enum import StrEnum
 from pathlib import Path
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     Awaitable,
     Callable,
@@ -34,7 +36,13 @@
 import sqlalchemy as sa
 import trafaret as t
 from aiohttp import web
-from sqlalchemy.orm import selectinload
+from pydantic import (
+    AliasChoices,
+    BaseModel,
+    Field,
+)
+from sqlalchemy.ext.asyncio import AsyncSession as SASession
+from sqlalchemy.orm import load_only, selectinload
 
 from ai.backend.common import msgpack, redis_helper
 from ai.backend.common import validators as tx
@@ -86,7 +94,8 @@
     vfolder_permissions,
     vfolders,
 )
-from ..models.utils import execute_with_retry
+from ..models.utils import execute_with_retry, execute_with_txn_retry
+from ..models.vfolder import VFolderPermissionRow
 from .auth import admin_required, auth_required, superadmin_required
 from .exceptions import (
     BackendAgentError,
@@ -2553,6 +2562,100 @@ async def update_shared_vfolder(request: web.Request, params: Any) -> web.Respon
     return web.json_response(resp, status=200)
 
 
+class UserPermMapping(BaseModel):
+    user_id: Annotated[
+        uuid.UUID,
+        Field(
+            validation_alias=AliasChoices("user", "user_id", "userID"),
+            description="Target user id to update sharing status.",
+        ),
+    ]
+    perm: Annotated[
+        VFolderPermission | None,
+        Field(
+            validation_alias=AliasChoices("perm", "permission"),
+            default=None,
+            description=textwrap.dedent(
+                "Permission to update. Delete the sharing between vfolder and user if this value is null. "
+                f"Should be one of {[p.value for p in VFolderPermission]}. "
+                "Default value is null."
+            ),
+        ),
+    ]
+
+
+class UpdateSharedRequestModel(BaseModel):
+    vfolder_id: Annotated[
+        uuid.UUID,
+        Field(
+            validation_alias=AliasChoices("vfolder_id", "vfolderId", "vfolder"),
+            description="Target vfolder id to update sharing status.",
+        ),
+    ]
+    user_perm_list: Annotated[
+        list[UserPermMapping],
+        Field(
+            validation_alias=AliasChoices("user_perm", "user_perm_list", "userPermList"),
+            description="A list of user and permission mappings.",
+        ),
+    ]
+
+
+@auth_required
+@server_status_required(ALL_ALLOWED)
+@pydantic_params_api_handler(UpdateSharedRequestModel)
+async def update_vfolder_sharing_status(
+    request: web.Request, params: UpdateSharedRequestModel
+) -> web.Response:
+    """
+    Update permission for shared vfolders.
+    """
+    root_ctx: RootContext = request.app["_root.context"]
+    access_key = request["keypair"]["access_key"]
+    vfolder_id = params.vfolder_id
+    user_perm_list = params.user_perm_list
+    log.info(
+        "VFOLDER.UPDATE_VFOLDER_SHARING_STATUS(email:{}, ak:{}, vfid:{}, data:{})",
+        request["user"]["email"],
+        access_key,
+        vfolder_id,
+        user_perm_list,
+    )
+
+    to_delete: list[uuid.UUID] = []
+    to_update: list[Mapping[str, Any]] = []
+    for mapping in user_perm_list:
+        if mapping.perm is None:
+            to_delete.append(mapping.user_id)
+        else:
+            to_update.append({
+                "user_id": mapping.user_id,
+                "perm": mapping.perm,
+            })
+
+    async def _update_or_delete(db_session: SASession) -> None:
+        if to_delete:
+            stmt = (
+                sa.delete(VFolderPermissionRow)
+                .where(VFolderPermissionRow.vfolder == vfolder_id)
+                .where(VFolderPermissionRow.user.in_(to_delete))
+            )
+            await db_session.execute(stmt)
+
+        if to_update:
+            stmt = (
+                sa.update(VFolderPermissionRow)
+                .values(permission=sa.bindparam("perm"))
+                .where(VFolderPermissionRow.vfolder == vfolder_id)
+                .where(VFolderPermissionRow.user == sa.bindparam("user_id"))
+            )
+            await db_session.execute(stmt, to_update)
+
+    async with root_ctx.db.connect() as db_conn:
+        await execute_with_txn_retry(_update_or_delete, root_ctx.db.begin_session, db_conn)
+    return web.Response(status=201)
+
+
 @superadmin_required
 @server_status_required(READ_ALLOWED)
 @check_api_params(
@@ -3136,6 +3239,7 @@ def create_app(default_cors_options):
     cors.add(add_route("DELETE", r"/invitations/delete", delete_invitation))
     cors.add(add_route("GET", r"/_/shared", list_shared_vfolders))
     cors.add(add_route("POST", r"/_/shared", update_shared_vfolder))
+    cors.add(add_route("POST", r"/_/sharing", update_vfolder_sharing_status))
     cors.add(add_route("GET", r"/_/fstab", get_fstab_contents))
     cors.add(add_route("GET", r"/_/mounts", list_mounts))
     cors.add(add_route("POST", r"/_/mounts", mount_host))