From 4897f2038c395230beb1340e3a7c4e4305dce8eb Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Thu, 11 Jan 2024 23:49:10 +0900
Subject: [PATCH] fix: Regression of purging user due to fkey constraint of
 `main_access_key` (#1775)

Backported-from: main
Backported-to: 23.09
---
 ...change_main_access_key_ondelete_to_set_.py | 39 +++++++++++++++++++
 src/ai/backend/manager/models/keypair.py      | 10 +++++
 src/ai/backend/manager/models/user.py         |  2 +-
 .../backend/manager/scheduler/predicates.py   |  7 +++-
 4 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 src/ai/backend/manager/models/alembic/versions/c5ed277b7f7b_change_main_access_key_ondelete_to_set_.py

diff --git a/src/ai/backend/manager/models/alembic/versions/c5ed277b7f7b_change_main_access_key_ondelete_to_set_.py b/src/ai/backend/manager/models/alembic/versions/c5ed277b7f7b_change_main_access_key_ondelete_to_set_.py
new file mode 100644
index 0000000000..aedce4188d
--- /dev/null
+++ b/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",
+    )
diff --git a/src/ai/backend/manager/models/keypair.py b/src/ai/backend/manager/models/keypair.py
index edc8d9c4c0..b0e1bf2341 100644
--- a/src/ai/backend/manager/models/keypair.py
+++ b/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:
diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py
index 5f1a71238b..15a63485aa 100644
--- a/src/ai/backend/manager/models/user.py
+++ b/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
     ),
 )
diff --git a/src/ai/backend/manager/scheduler/predicates.py b/src/ai/backend/manager/scheduler/predicates.py
index e8a822add5..dc38c37bd3 100644
--- a/src/ai/backend/manager/scheduler/predicates.py
+++ b/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,