From 80a351fb6113a18005d7103c0e52ce7f8c571605 Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Wed, 25 Sep 2024 14:20:02 +0900
Subject: [PATCH] fix: Delete vfolder invitation and permission rows when
 deleting vfolders (#2780)

Backported-from: main (24.09)
Backported-to: 24.03
Backport-of: 2780
---
 changes/2780.fix.md                      |  1 +
 src/ai/backend/manager/api/vfolder.py    | 11 ++++-
 src/ai/backend/manager/models/vfolder.py | 57 +++++++++++++++++++++++-
 3 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 changes/2780.fix.md

diff --git a/changes/2780.fix.md b/changes/2780.fix.md
new file mode 100644
index 0000000000..71f7f7d98b
--- /dev/null
+++ b/changes/2780.fix.md
@@ -0,0 +1 @@
+Delete vfolder invitation and permission rows when deleting vfolders.
diff --git a/src/ai/backend/manager/api/vfolder.py b/src/ai/backend/manager/api/vfolder.py
index 7445eda8a8..438130656a 100644
--- a/src/ai/backend/manager/api/vfolder.py
+++ b/src/ai/backend/manager/api/vfolder.py
@@ -100,7 +100,11 @@
     vfolder_status_map,
     vfolders,
 )
-from ..models.utils import execute_with_retry
+from ..models.utils import execute_with_retry, execute_with_txn_retry
+from ..models.vfolder import (
+    VFolderPermissionRow,
+    delete_vfolder_relation_rows,
+)
 from .auth import admin_required, auth_required, superadmin_required
 from .exceptions import (
     BackendAgentError,
@@ -2254,9 +2258,12 @@ async def _delete(
             permission=VFolderHostPermission.DELETE,
         )
 
+    vfolder_row_ids = (entry["id"],)
+    async with root_ctx.db.connect() as db_conn:
+        await delete_vfolder_relation_rows(db_conn, root_ctx.db.begin_session, vfolder_row_ids)
     await update_vfolder_status(
         root_ctx.db,
-        (entry["id"],),
+        vfolder_row_ids,
         VFolderOperationStatus.DELETE_PENDING,
     )
 
diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py
index 7ec8731f26..686de10463 100644
--- a/src/ai/backend/manager/models/vfolder.py
+++ b/src/ai/backend/manager/models/vfolder.py
@@ -4,9 +4,23 @@
 import logging
 import os.path
 import uuid
+from collections.abc import Container, Iterable, Mapping
+from contextlib import AbstractAsyncContextManager as AbstractAsyncCtxMgr
+from dataclasses import dataclass
 from datetime import datetime
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING, Any, Final, List, Mapping, NamedTuple, Optional, Sequence
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Final,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    TypeAlias,
+    cast,
+)
 
 import aiohttp
 import aiotools
@@ -74,7 +88,7 @@
 from .minilang.ordering import OrderSpecItem, QueryOrderParser
 from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter
 from .user import UserRole, UserRow
-from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge
+from .utils import ExtendedAsyncSAEngine, execute_with_retry, execute_with_txn_retry, sql_json_merge
 
 if TYPE_CHECKING:
     from ..api.context import BackgroundTaskManager
@@ -398,6 +412,12 @@ class VFolderCloneInfo(NamedTuple):
 )
 
 
+class VFolderInvitationRow(Base):
+    __table__ = vfolder_invitations
+
+    vfolder_row = relationship("VFolderRow", back_populates="invitation_rows")
+
+
 vfolder_permissions = sa.Table(
     "vfolder_permissions",
     metadata,
@@ -426,6 +446,8 @@ class VFolderRow(Base):
         back_populates="vfolder_rows",
         primaryjoin="GroupRow.id == foreign(VFolderRow.group)",
     )
+    permission_rows = relationship(VFolderPermissionRow, back_populates="vfolder_row")
+    invitation_rows = relationship(VFolderInvitationRow, back_populates="vfolder_row")
 
     @classmethod
     async def get(
@@ -1144,6 +1166,34 @@ async def _update_source_vfolder() -> None:
     return task_id, target_folder_id.folder_id
 
 
+async def _delete_vfolder_permission_rows(
+    db_session: SASession,
+    vfolder_row_ids: Iterable[uuid.UUID],
+) -> None:
+    stmt = sa.delete(VFolderInvitationRow).where(VFolderInvitationRow.vfolder.in_(vfolder_row_ids))
+    await db_session.execute(stmt)
+
+
+async def _delete_vfolder_invitation_rows(
+    db_session: SASession,
+    vfolder_row_ids: Iterable[uuid.UUID],
+) -> None:
+    stmt = sa.delete(VFolderPermissionRow).where(VFolderPermissionRow.vfolder.in_(vfolder_row_ids))
+    await db_session.execute(stmt)
+
+
+async def delete_vfolder_relation_rows(
+    db_conn: SAConnection,
+    begin_session: Callable[..., AbstractAsyncCtxMgr[SASession]],
+    vfolder_row_ids: Iterable[uuid.UUID],
+) -> None:
+    async def _delete(db_session: SASession) -> None:
+        await _delete_vfolder_invitation_rows(db_session, vfolder_row_ids)
+        await _delete_vfolder_permission_rows(db_session, vfolder_row_ids)
+
+    await execute_with_txn_retry(_delete, begin_session, db_conn)
+
+
 async def initiate_vfolder_deletion(
     db_engine: ExtendedAsyncSAEngine,
     requested_vfolders: Sequence[VFolderDeletionInfo],
@@ -1158,6 +1208,9 @@ async def initiate_vfolder_deletion(
         return 0
     elif vfolder_info_len == 1:
         vfolders.c.id == vfolder_ids[0]
+
+    async with db_engine.connect() as db_conn:
+        await delete_vfolder_relation_rows(db_conn, db_engine.begin_session, vfolder_ids)
     await update_vfolder_status(
         db_engine, vfolder_ids, VFolderOperationStatus.DELETE_ONGOING, do_log=False
     )