From a77a4ccd7e89f3c4cfcd54396f06547628035d7e Mon Sep 17 00:00:00 2001
From: octodog <mu001@lablup.com>
Date: Wed, 25 Sep 2024 23:07:57 +0900
Subject: [PATCH] fix: Delete vfolder invitation and permission rows when
 deleting vfolders (#2780) (#2860)

Co-authored-by: Sanghun Lee <sanghun@lablup.com>
Co-authored-by: Kyujin Cho <kyujin.cho@lablup.com>
---
 changes/2780.fix.md                      |  1 +
 src/ai/backend/manager/api/vfolder.py    |  7 +++-
 src/ai/backend/manager/models/vfolder.py | 48 +++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 2 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..21b53182c1 100644
--- a/src/ai/backend/manager/api/vfolder.py
+++ b/src/ai/backend/manager/api/vfolder.py
@@ -101,6 +101,9 @@
     vfolders,
 )
 from ..models.utils import execute_with_retry
+from ..models.vfolder import (
+    delete_vfolder_relation_rows,
+)
 from .auth import admin_required, auth_required, superadmin_required
 from .exceptions import (
     BackendAgentError,
@@ -2254,9 +2257,11 @@ async def _delete(
             permission=VFolderHostPermission.DELETE,
         )
 
+    vfolder_row_ids = (entry["id"],)
+    await delete_vfolder_relation_rows(root_ctx.db, 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..2b5a4648be 100644
--- a/src/ai/backend/manager/models/vfolder.py
+++ b/src/ai/backend/manager/models/vfolder.py
@@ -4,9 +4,18 @@
 import logging
 import os.path
 import uuid
+from collections.abc import Iterable, Mapping
 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,
+    Final,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+)
 
 import aiohttp
 import aiotools
@@ -398,6 +407,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 +441,7 @@ class VFolderRow(Base):
         back_populates="vfolder_rows",
         primaryjoin="GroupRow.id == foreign(VFolderRow.group)",
     )
+    invitation_rows = relationship(VFolderInvitationRow, back_populates="vfolder_row")
 
     @classmethod
     async def get(
@@ -1144,6 +1160,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(vfolder_permissions).where(vfolder_permissions.c.vfolder.in_(vfolder_row_ids))
+    await db_session.execute(stmt)
+
+
+async def delete_vfolder_relation_rows(
+    db_engine: ExtendedAsyncSAEngine,
+    vfolder_row_ids: Iterable[uuid.UUID],
+) -> None:
+    async def _delete() -> None:
+        async with db_engine.begin_session() as db_session:
+            await _delete_vfolder_invitation_rows(db_session, vfolder_row_ids)
+            await _delete_vfolder_permission_rows(db_session, vfolder_row_ids)
+
+    await execute_with_retry(_delete)
+
+
 async def initiate_vfolder_deletion(
     db_engine: ExtendedAsyncSAEngine,
     requested_vfolders: Sequence[VFolderDeletionInfo],
@@ -1158,6 +1202,8 @@ async def initiate_vfolder_deletion(
         return 0
     elif vfolder_info_len == 1:
         vfolders.c.id == vfolder_ids[0]
+
+    await delete_vfolder_relation_rows(db_engine, vfolder_ids)
     await update_vfolder_status(
         db_engine, vfolder_ids, VFolderOperationStatus.DELETE_ONGOING, do_log=False
     )