Skip to content

Commit

Permalink
✨Properly support inheritance of Relationship attributes
Browse files Browse the repository at this point in the history
  • Loading branch information
earshinov committed Apr 8, 2024
1 parent c75743d commit 8f30325
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 2 deletions.
7 changes: 5 additions & 2 deletions sqlmodel/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,8 @@ def __new__(
**kwargs: Any,
) -> Any:
relationships: Dict[str, RelationshipInfo] = {}
for base in bases:
relationships.update(getattr(base, "__sqlmodel_relationships__", {}))
dict_for_pydantic = {}
original_annotations = get_annotations(class_dict)
pydantic_annotations = {}
Expand Down Expand Up @@ -471,8 +473,9 @@ def get_config(name: str) -> Any:
# If it was passed by kwargs, ensure it's also set in config
set_config_value(model=new_cls, parameter="table", value=config_table)
for k, v in get_model_fields(new_cls).items():
col = get_column_from_field(v)
setattr(new_cls, k, col)
if k not in relationships:
col = get_column_from_field(v)
setattr(new_cls, k, col)
# Set a config flag to tell FastAPI that this should be read with a field
# in orm_mode instead of preemptively converting it to a dict.
# This could be done by reading new_cls.model_config['table'] in FastAPI, but
Expand Down
62 changes: 62 additions & 0 deletions tests/test_relationship_inheritance.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
from typing import Optional

from sqlalchemy.orm import declared_attr, relationship
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select


def test_relationship_inheritance() -> None:
class User(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
name: str

class CreatedUpdatedMixin(SQLModel):
# With Pydantic V2, it is also possible to define `created_by` like this:
#
# ```python
# @declared_attr
# def _created_by(cls):
# return relationship(User, foreign_keys=cls.created_by_id)
#
# created_by: Optional[User] = Relationship(sa_relationship=_created_by))
# ```
#
# The difference from Pydantic V1 is that Pydantic V2 plucks attributes with names starting with '_' (but not '__')
# from class attributes and stores them separately as instances of `pydantic.ModelPrivateAttr` somewhere in depths of
# Pydantic internals. Under Pydantic V1 this doesn't happen, so SQLAlchemy ends up having two class attributes
# (`_created_by` and `created_by`) corresponding to one database attribute, causing a conflict and unreliable behavior.
# The approach with a lambda always works because it doesn't produce the second class attribute and thus eliminates
# the possibility of a conflict entirely.
#
created_by_id: Optional[int] = Field(default=None, foreign_key="user.id")
created_by: Optional[User] = Relationship(
sa_relationship=declared_attr(
lambda cls: relationship(User, foreign_keys=cls.created_by_id)
)
)

updated_by_id: Optional[int] = Field(default=None, foreign_key="user.id")
updated_by: Optional[User] = Relationship(
sa_relationship=declared_attr(
lambda cls: relationship(User, foreign_keys=cls.updated_by_id)
)
)

class Asset(CreatedUpdatedMixin, table=True):
id: Optional[int] = Field(default=None, primary_key=True)

engine = create_engine("sqlite://")

SQLModel.metadata.create_all(engine)

john = User(name="John")
jane = User(name="Jane")
asset = Asset(created_by=john, updated_by=jane)

with Session(engine) as session:
session.add(asset)
session.commit()

with Session(engine) as session:
asset = session.exec(select(Asset)).one()
assert asset.created_by.name == "John"
assert asset.updated_by.name == "Jane"

0 comments on commit 8f30325

Please sign in to comment.