feature: Prevent deleting mounted folders (#2036)

resolves #776

- The mount lookup happens only when we **"delete"** vfolders, not when we **"delete-forever"** them since it is already disallowed to delete-forever the READY vfolders
- Apply [GIN index](https://www.postgresql.org/docs/current/datatype-json.html#JSON-INDEXING) to `sessions.vfolder_mounts` column

**Checklist:** (if applicable)

- [x] Milestone metadata specifying the target backport version
- [x] Mention to the original issue
- [x] API server-client counterparts (e.g., manager API -> client SDK)
- [ ] Test case(s) to:
  - Demonstrate the difference of before/after
  - Demonstrate the flow of abstract/conceptual models with a concrete implementation
- [ ] Documentation
  - Contents in the `docs` directory
  - docstrings in public interfaces and type annotations

changes/2036.feature.md

@@ -0,0 +1 @@
+Prevent deleting mounted folders.

src/ai/backend/manager/models/alembic/versions/31463788c713_add_index_sessions_vfolder_mounts.py

@@ -0,0 +1,29 @@
+"""add_index_sessions_vfolder_mounts
+
+Revision ID: 31463788c713
+Revises: dddf9be580f5
+Create Date: 2024-04-19 18:53:29.903113
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "31463788c713"
+down_revision = "dddf9be580f5"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_index(
+        "ix_sessions_vfolder_mounts",
+        "sessions",
+        ["vfolder_mounts"],
+        unique=False,
+        postgresql_using="gin",
+    )
+
+
+def downgrade():
+    op.drop_index("ix_sessions_vfolder_mounts", table_name="sessions", postgresql_using="gin")

src/ai/backend/manager/models/base.py

@@ -417,6 +417,9 @@ def __init__(self, schema: Type[JSONSerializableMixin]) -> None:
         super().__init__()
         self._schema = schema
 
+    def coerce_compared_value(self, op, value):
+        return JSONB()
+
     def process_bind_param(self, value, dialect):
         return [self._schema.to_json(item) for item in value]
 

src/ai/backend/manager/models/session.py

@@ -676,6 +676,7 @@ class SessionRow(Base):
             ),
             unique=False,
         ),
+        sa.Index("ix_sessions_vfolder_mounts", vfolder_mounts, postgresql_using="gin"),
     )
 
     @property

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

@@ -6,7 +6,7 @@
 import uuid
 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, Mapping, NamedTuple, Optional, Sequence, cast
 
 import aiohttp
 import aiotools
@@ -23,7 +23,7 @@
 from sqlalchemy.engine.row import Row
 from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
-from sqlalchemy.orm import selectinload
+from sqlalchemy.orm import load_only, selectinload
 
 from ai.backend.common.bgtask import ProgressReporter
 from ai.backend.common.config import model_definition_iv
@@ -32,6 +32,7 @@
     MountPermission,
     QuotaScopeID,
     QuotaScopeType,
+    SessionId,
     VFolderHostPermission,
     VFolderHostPermissionMap,
     VFolderID,
@@ -72,6 +73,7 @@
 from .group import GroupRow, ProjectType
 from .minilang.ordering import OrderSpecItem, QueryOrderParser
 from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter
+from .session import DEAD_SESSION_STATUSES, SessionRow
 from .user import UserRole
 from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge
 
@@ -935,6 +937,20 @@ async def update_vfolder_status(
 
     now = datetime.now(tzutc())
 
+    if update_status == VFolderOperationStatus.DELETE_PENDING:
+        select_stmt = sa.select(VFolderRow).where(VFolderRow.id.in_(vfolder_ids))
+        async with engine.begin_readonly_session() as db_session:
+            for vf_row in await db_session.scalars(select_stmt):
+                vf_row = cast(VFolderRow, vf_row)
+                mount_sessions = await get_sessions_by_mounted_folder(
+                    db_session, VFolderID.from_row(vf_row)
+                )
+                if mount_sessions:
+                    session_ids = [str(s) for s in mount_sessions]
+                    raise InvalidAPIParameters(
+                        f"Cannot delete the vfolder. The vfolder(id: {vf_row.id}) is mounted on sessions(ids: {session_ids})"
+                    )
+
     async def _update() -> None:
         async with engine.begin_session() as db_session:
             query = (
@@ -1213,6 +1229,26 @@ async def ensure_quota_scope_accessible_by_user(
     raise InvalidAPIParameters
 
 
+async def get_sessions_by_mounted_folder(
+    db_session: SASession, vfolder_id: VFolderID
+) -> tuple[SessionId]:
+    """
+    Return a tuple of sessions.id that the give folder is mounted on.
+    """
+
+    select_stmt = (
+        sa.select(SessionRow)
+        .where(
+            (SessionRow.status.not_in(DEAD_SESSION_STATUSES))
+            & SessionRow.vfolder_mounts.contains([{"vfid": str(vfolder_id)}])
+        )
+        .options(load_only(SessionRow.id))
+    )
+
+    session_rows = (await db_session.scalars(select_stmt)).all()
+    return tuple([session.id for session in session_rows])
+
+
 class VirtualFolder(graphene.ObjectType):
     class Meta:
         interfaces = (Item,)