From 54c72f850735f4f52925f290d382952e09759270 Mon Sep 17 00:00:00 2001 From: octodog Date: Sat, 13 Jul 2024 23:05:02 +0900 Subject: [PATCH] refactor: Remove the foreign key constraint of vfolders.{user,group} (#2404) (#2406) refs #2376 Co-authored-by: achimnol Remove the foreign/check constraints preventing decoupling of user/vfolder deletions. Backported-from: main (24.09) Backported-to: 24.03 Backport-of: 2404 --- changes/2404.enhance.md | 1 + ...oreign_key_in_vfolders_user_and_project.py | 38 +++++++++++++++++++ src/ai/backend/manager/models/group.py | 6 ++- src/ai/backend/manager/models/user.py | 6 ++- src/ai/backend/manager/models/vfolder.py | 36 ++++++++---------- 5 files changed, 65 insertions(+), 22 deletions(-) create mode 100644 changes/2404.enhance.md create mode 100644 src/ai/backend/manager/models/alembic/versions/59a622c31820_remove_foreign_key_in_vfolders_user_and_project.py diff --git a/changes/2404.enhance.md b/changes/2404.enhance.md new file mode 100644 index 0000000000..10c83081a6 --- /dev/null +++ b/changes/2404.enhance.md @@ -0,0 +1 @@ +Remove database-level foreign key constraints in `vfolders.{user,group}` columns to decouple the timing of vfolder deletion and user/group deletion. diff --git a/src/ai/backend/manager/models/alembic/versions/59a622c31820_remove_foreign_key_in_vfolders_user_and_project.py b/src/ai/backend/manager/models/alembic/versions/59a622c31820_remove_foreign_key_in_vfolders_user_and_project.py new file mode 100644 index 0000000000..a3e659e45e --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/59a622c31820_remove_foreign_key_in_vfolders_user_and_project.py @@ -0,0 +1,38 @@ +"""Remove foreign key constraint from vfolders to users and projects + +Revision ID: 59a622c31820 +Revises: fdb2dcdb8811 +Create Date: 2024-07-08 22:54:20.762521 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "59a622c31820" +down_revision = "fdb2dcdb8811" +branch_labels = None +depends_on = None + + +def upgrade(): + op.drop_constraint("fk_vfolders_user_users", "vfolders", type_="foreignkey") + op.drop_constraint("fk_vfolders_group_groups", "vfolders", type_="foreignkey") + op.drop_constraint("ck_vfolders_ownership_type_match_with_user_or_group", "vfolders") + op.drop_constraint("ck_vfolders_either_one_of_user_or_group", "vfolders") + + +def downgrade(): + op.create_foreign_key("fk_vfolders_group_groups", "vfolders", "groups", ["group"], ["id"]) + op.create_foreign_key("fk_vfolders_user_users", "vfolders", "users", ["user"], ["uuid"]) + op.create_check_constraint( + "ck_vfolders_ownership_type_match_with_user_or_group", + "vfolders", + "(ownership_type = 'user' AND \"user\" IS NOT NULL) OR " + "(ownership_type = 'group' AND \"group\" IS NOT NULL)", + ) + op.create_check_constraint( + "ck_vfolders_either_one_of_user_or_group", + "vfolders", + '("user" IS NULL AND "group" IS NOT NULL) OR ("user" IS NOT NULL AND "group" IS NULL)', + ) diff --git a/src/ai/backend/manager/models/group.py b/src/ai/backend/manager/models/group.py index e902759cfa..a55991fd7a 100644 --- a/src/ai/backend/manager/models/group.py +++ b/src/ai/backend/manager/models/group.py @@ -197,7 +197,11 @@ class GroupRow(Base): users = relationship("AssocGroupUserRow", back_populates="group") resource_policy_row = relationship("ProjectResourcePolicyRow", back_populates="projects") kernels = relationship("KernelRow", back_populates="group_row") - vfolder_row = relationship("VFolderRow", back_populates="group_row") + vfolder_rows = relationship( + "VFolderRow", + back_populates="group_row", + primaryjoin="GroupRow.id == foreign(VFolderRow.group)", + ) def _build_group_query(cond: sa.sql.BinaryExpression, domain_name: str) -> sa.sql.Select: diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py index ca35baddef..1929f39cf5 100644 --- a/src/ai/backend/manager/models/user.py +++ b/src/ai/backend/manager/models/user.py @@ -186,7 +186,11 @@ class UserRow(Base): main_keypair = relationship("KeyPairRow", foreign_keys=users.c.main_access_key) - vfolder_row = relationship("VFolderRow", back_populates="user_row") + vfolder_rows = relationship( + "VFolderRow", + back_populates="user_row", + primaryjoin="UserRow.uuid == foreign(VFolderRow.user)", + ) class UserGroup(graphene.ObjectType): diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py index 50a0535d3a..fb810fc5d9 100644 --- a/src/ai/backend/manager/models/vfolder.py +++ b/src/ai/backend/manager/models/vfolder.py @@ -73,7 +73,7 @@ from .group import GroupRow, ProjectType from .minilang.ordering import OrderSpecItem, QueryOrderParser from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter -from .user import UserRole +from .user import UserRole, UserRow from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge if TYPE_CHECKING: @@ -323,8 +323,8 @@ class VFolderCloneInfo(NamedTuple): nullable=False, index=True, ), - sa.Column("user", GUID, sa.ForeignKey("users.uuid"), nullable=True), # owner if user vfolder - sa.Column("group", GUID, sa.ForeignKey("groups.id"), nullable=True), # owner if project vfolder + sa.Column("user", GUID, nullable=True), # owner if user vfolder + sa.Column("group", GUID, nullable=True), # owner if project vfolder sa.Column("cloneable", sa.Boolean, default=False, nullable=False), sa.Column( "status", @@ -343,15 +343,6 @@ class VFolderCloneInfo(NamedTuple): # } sa.Column("status_history", pgsql.JSONB(), nullable=True, default=sa.null()), sa.Column("status_changed", sa.DateTime(timezone=True), nullable=True, index=True), - sa.CheckConstraint( - "(ownership_type = 'user' AND \"user\" IS NOT NULL) OR " - "(ownership_type = 'group' AND \"group\" IS NOT NULL)", - name="ownership_type_match_with_user_or_group", - ), - sa.CheckConstraint( - '("user" IS NULL AND "group" IS NOT NULL) OR ("user" IS NOT NULL AND "group" IS NULL)', - name="either_one_of_user_or_group", - ), ) @@ -418,17 +409,25 @@ class VFolderRow(Base): __table__ = vfolders endpoints = relationship("EndpointRow", back_populates="model_row") - user_row = relationship("UserRow", back_populates="vfolder_row") - group_row = relationship("GroupRow", back_populates="vfolder_row") + user_row = relationship( + "UserRow", + back_populates="vfolder_rows", + primaryjoin="UserRow.uuid == foreign(VFolderRow.user)", + ) + group_row = relationship( + "GroupRow", + back_populates="vfolder_rows", + primaryjoin="GroupRow.id == foreign(VFolderRow.group)", + ) @classmethod async def get( cls, session: SASession, id: uuid.UUID, - load_user=False, - load_group=False, - ) -> "VFolderRow": + load_user: bool = False, + load_group: bool = False, + ) -> VFolderRow: query = sa.select(VFolderRow).where(VFolderRow.id == id) if load_user: query = query.options(selectinload(VFolderRow.user_row)) @@ -1211,7 +1210,6 @@ async def ensure_quota_scope_accessible_by_user( quota_scope: QuotaScopeID, user: Mapping[str, Any], ) -> None: - from ai.backend.manager.models import GroupRow, UserRow from ai.backend.manager.models import association_groups_users as agus # Lookup user table to match if quota is scoped to the user @@ -2059,8 +2057,6 @@ def resolve_id(self, info: graphene.ResolveInfo) -> str: return f"QuotaScope:{self.storage_host_name}/{self.quota_scope_id}" async def resolve_details(self, info: graphene.ResolveInfo) -> Optional[int]: - from ai.backend.manager.models import GroupRow, UserRow - graph_ctx: GraphQueryContext = info.context proxy_name, volume_name = graph_ctx.storage_manager.split_host(self.storage_host_name) try: