Skip to content

Commit

Permalink
a friendly little helper (#4021)
Browse files Browse the repository at this point in the history
* a friendly little helper

* addressing comments

* update comment

---------

Co-authored-by: simon <simon@reflex.dev>
  • Loading branch information
Kastier1 and simon authored Oct 1, 2024
1 parent 9c3cc0c commit e96b4bf
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 3 deletions.
5 changes: 2 additions & 3 deletions reflex/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
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:
Expand Down Expand Up @@ -166,8 +166,7 @@ def __init_subclass__(cls):
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)
Expand Down
18 changes: 18 additions & 0 deletions reflex/utils/compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit e96b4bf

Please sign in to comment.