From e96b4bf42eb84c19925dc2f84aae1f195bfee49c Mon Sep 17 00:00:00 2001
From: Simon Young <40179067+Kastier1@users.noreply.github.com>
Date: Tue, 1 Oct 2024 14:32:05 -0700
Subject: [PATCH] a friendly little helper (#4021)

* a friendly little helper

* addressing comments

* update comment

---------

Co-authored-by: simon <simon@reflex.dev>
---
 reflex/model.py        |  5 ++---
 reflex/utils/compat.py | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/reflex/model.py b/reflex/model.py
index fefb1f9e9..0e8d62e90 100644
--- a/reflex/model.py
+++ b/reflex/model.py
@@ -22,7 +22,7 @@ from reflex import constants
 from reflex.base import Base
 from reflex.config import get_config
 from reflex.utils import console
-from reflex.utils.compat import sqlmodel
+from reflex.utils.compat import sqlmodel, sqlmodel_field_has_primary_key
 
 
 def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
@@ -166,8 +166,7 @@ class Model(Base, sqlmodel.SQLModel):  # pyright: ignore [reportGeneralTypeIssue
         non_default_primary_key_fields = [
             field_name
             for field_name, field in cls.__fields__.items()
-            if field_name != "id"
-            and getattr(field.field_info, "primary_key", None) is True
+            if field_name != "id" and sqlmodel_field_has_primary_key(field)
         ]
         if non_default_primary_key_fields:
             cls.__fields__.pop("id", None)
diff --git a/reflex/utils/compat.py b/reflex/utils/compat.py
index ef5fcd3e1..27c4753db 100644
--- a/reflex/utils/compat.py
+++ b/reflex/utils/compat.py
@@ -69,3 +69,21 @@ def pydantic_v1_patch():
 
 with pydantic_v1_patch():
     import sqlmodel as sqlmodel
+
+
+def sqlmodel_field_has_primary_key(field) -> bool:
+    """Determines if a field is a priamary.
+
+    Args:
+        field: a rx.model field
+
+    Returns:
+        If field is a primary key (Bool)
+    """
+    if getattr(field.field_info, "primary_key", None) is True:
+        return True
+    if getattr(field.field_info, "sa_column", None) is None:
+        return False
+    if getattr(field.field_info.sa_column, "primary_key", None) is True:
+        return True
+    return False