diff --git a/src/ai/backend/manager/api/vfolder.py b/src/ai/backend/manager/api/vfolder.py index 438130656a..21b53182c1 100644 --- a/src/ai/backend/manager/api/vfolder.py +++ b/src/ai/backend/manager/api/vfolder.py @@ -100,9 +100,8 @@ vfolder_status_map, vfolders, ) -from ..models.utils import execute_with_retry, execute_with_txn_retry +from ..models.utils import execute_with_retry from ..models.vfolder import ( - VFolderPermissionRow, delete_vfolder_relation_rows, ) from .auth import admin_required, auth_required, superadmin_required @@ -2259,8 +2258,7 @@ async def _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 delete_vfolder_relation_rows(root_ctx.db, vfolder_row_ids) await update_vfolder_status( root_ctx.db, vfolder_row_ids, diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py index 686de10463..2b5a4648be 100644 --- a/src/ai/backend/manager/models/vfolder.py +++ b/src/ai/backend/manager/models/vfolder.py @@ -4,22 +4,17 @@ 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 collections.abc import Iterable, Mapping from datetime import datetime from pathlib import PurePosixPath from typing import ( TYPE_CHECKING, Any, - Callable, Final, List, NamedTuple, Optional, Sequence, - TypeAlias, - cast, ) import aiohttp @@ -88,7 +83,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, execute_with_txn_retry, sql_json_merge +from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge if TYPE_CHECKING: from ..api.context import BackgroundTaskManager @@ -446,7 +441,6 @@ 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 @@ -1178,20 +1172,20 @@ 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)) + 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_conn: SAConnection, - begin_session: Callable[..., AbstractAsyncCtxMgr[SASession]], + db_engine: ExtendedAsyncSAEngine, 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) + 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_txn_retry(_delete, begin_session, db_conn) + await execute_with_retry(_delete) async def initiate_vfolder_deletion( @@ -1209,8 +1203,7 @@ async def initiate_vfolder_deletion( 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 delete_vfolder_relation_rows(db_engine, vfolder_ids) await update_vfolder_status( db_engine, vfolder_ids, VFolderOperationStatus.DELETE_ONGOING, do_log=False )