fix: Regression of purging user due to fkey constraint of `main_acces…

…s_key` (#1775)

Backported-from: main
Backported-to: 23.09

src/ai/backend/manager/models/alembic/versions/c5ed277b7f7b_change_main_access_key_ondelete_to_set_.py

@@ -0,0 +1,39 @@
+"""change_main_access_key_ondelete_to_set_null
+
+Revision ID: c5ed277b7f7b
+Revises: d3f8c74bf148
+Create Date: 2024-01-11 15:37:24.097596
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "c5ed277b7f7b"
+down_revision = "d3f8c74bf148"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.drop_constraint("fk_users_main_access_key_keypairs", "users", type_="foreignkey")
+    op.create_foreign_key(
+        op.f("fk_users_main_access_key_keypairs"),
+        "users",
+        "keypairs",
+        ["main_access_key"],
+        ["access_key"],
+        ondelete="SET NULL",
+    )
+
+
+def downgrade():
+    op.drop_constraint(op.f("fk_users_main_access_key_keypairs"), "users", type_="foreignkey")
+    op.create_foreign_key(
+        "fk_users_main_access_key_keypairs",
+        "users",
+        "keypairs",
+        ["main_access_key"],
+        ["access_key"],
+        ondelete="RESTRICT",
+    )

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

@@ -624,7 +624,17 @@ async def mutate(
         info: graphene.ResolveInfo,
         access_key: AccessKey,
     ) -> DeleteKeyPair:
+        from .user import UserRow
+
         ctx: GraphQueryContext = info.context
+        async with ctx.db.begin_readonly_session() as db_session:
+            user_query = (
+                sa.select([sa.func.count()])
+                .select_from(UserRow)
+                .where(UserRow.main_access_key == access_key)
+            )
+            if (await db_session.scalar(user_query)) > 0:
+                return DeleteKeyPair(False, "the keypair is used as main access key by any user")
         delete_query = sa.delete(keypairs).where(keypairs.c.access_key == access_key)
         result = await simple_db_mutate(cls, ctx, delete_query)
         if result.ok:

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

@@ -156,7 +156,7 @@ class UserStatus(str, enum.Enum):
     sa.Column(
         "main_access_key",
         sa.String(length=20),
-        sa.ForeignKey("keypairs.access_key", ondelete="RESTRICT"),
+        sa.ForeignKey("keypairs.access_key", ondelete="SET NULL"),
         nullable=True,  # keypairs.user is non-nullable
     ),
 )

src/ai/backend/manager/scheduler/predicates.py

@@ -196,7 +196,12 @@ async def check_user_resource_limit(
     select_query = sa.select(KeyPairResourcePolicyRow).where(
         KeyPairResourcePolicyRow.name == resouce_policy_q.scalar_subquery()
     )
-    resource_policy: KeyPairResourcePolicyRow = (await db_sess.scalars(select_query)).first()
+    resource_policy: KeyPairResourcePolicyRow | None = (await db_sess.scalars(select_query)).first()
+    if resource_policy is None:
+        return PredicateResult(
+            False,
+            f"User has no main-keypair or the main-keypair has no keypair resource policy (uid: {sess_ctx.user_uuid})",
+        )
 
     resource_policy_map = {
         "total_resource_slots": resource_policy.total_resource_slots,