refactor: Remove the foreign key constraint of vfolders.{user,group}

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.

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)',
+    )

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:

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):

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

@@ -75,7 +75,7 @@
 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 .user import UserRole, UserRow
 from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge
 
 if TYPE_CHECKING:
@@ -326,8 +326,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",
@@ -346,15 +346,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",
-    ),
 )
 
 
@@ -426,17 +417,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))
@@ -1222,7 +1221,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
@@ -2088,8 +2086,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: