From fcc99d9378190c93b978b97068443704ebeb5b9b Mon Sep 17 00:00:00 2001
From: Ryan <8829107+ryan-pip@users.noreply.github.com>
Date: Fri, 20 Mar 2026 15:34:13 +1300
Subject: [PATCH] fix: catch edge case for no pks and check on downgrade

---
 .../versions/c3d4e5f6a7b8_add_missing_primary_keys.py    | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/backend/open_webui/migrations/versions/c3d4e5f6a7b8_add_missing_primary_keys.py b/backend/open_webui/migrations/versions/c3d4e5f6a7b8_add_missing_primary_keys.py
index a9e8c22e9d..16d82dd046 100644
--- a/backend/open_webui/migrations/versions/c3d4e5f6a7b8_add_missing_primary_keys.py
+++ b/backend/open_webui/migrations/versions/c3d4e5f6a7b8_add_missing_primary_keys.py
@@ -33,7 +33,7 @@ def upgrade() -> None:
     inspector = Inspector.from_engine(conn)
 
     for table, column in TABLES:
-        pk = inspector.get_pk_constraint(table)
+        pk = inspector.get_pk_constraint(table) or {}
         if column in (pk.get("constrained_columns") or []):
             continue
 
@@ -46,6 +46,13 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)
+
     for table, column in TABLES:
+        pk = inspector.get_pk_constraint(table) or {}
+        if column not in (pk.get("constrained_columns") or []):
+            continue
+
         with op.batch_alter_table(table, schema=None) as batch_op:
             batch_op.drop_constraint(f"pk_{column}", type_="primary")