✏️Fix model_validate in presence of inherited Relationship fields, add unit test

Author: earshinov

sqlmodel/main.py

@@ -503,8 +503,28 @@ def __new__(
         **kwargs: Any,
     ) -> Any:
         relationships: Dict[str, RelationshipInfo] = {}
+        backup_base_annotations: Dict[Type[Any], Dict[str, Any]] = {}
         for base in bases:
-            relationships.update(getattr(base, "__sqlmodel_relationships__", {}))
+            base_relationships = getattr(base, "__sqlmodel_relationships__", None)
+            if base_relationships:
+                relationships.update(base_relationships)
+                #
+                # Temporarily pluck out `__annotations__` corresponding to relationships from base classes, otherwise these annotations
+                # make their way into `cls.model_fields` as `FieldInfo(..., required=True)`, even when the relationships are declared
+                # optional.  When a model instance is then constructed using `model_validate` and an optional relationship field is not
+                # passed, this leads to an incorrect `pydantic.ValidationError`.
+                #
+                # We can't just clean up `new_cls.model_fields` after `new_cls` is constructed because by this time
+                # Pydantic has created model schema and validation rules, so this won't fix the problem.
+                #
+                base_annotations = getattr(base, "__annotations__", None)
+                if base_annotations:
+                    backup_base_annotations[base] = base_annotations
+                    base.__annotations__ = {
+                        name: typ
+                        for name, typ in base_annotations.items()
+                        if name not in base_relationships
+                    }
         dict_for_pydantic = {}
         original_annotations = get_annotations(class_dict)
         pydantic_annotations = {}
@@ -539,6 +559,9 @@ def __new__(
             key: kwargs[key] for key in kwargs.keys() & allowed_config_kwargs
         }
         new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs)
+        # Restore base annotations
+        for base, annotations in backup_base_annotations.items():
+            base.__annotations__ = annotations
         new_cls.__annotations__ = {
             **relationship_annotations,
             **pydantic_annotations,

tests/test_relationship_inheritance.py renamed to tests/test_inherit_relationship.py

@@ -1,12 +1,14 @@
 import datetime
 from typing import Optional
 
+import pydantic
 from sqlalchemy import DateTime, func
 from sqlalchemy.orm import declared_attr, relationship
 from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select
+from sqlmodel._compat import IS_PYDANTIC_V2
 
 
-def test_relationship_inheritance() -> None:
+def test_inherit_relationship(clear_sqlmodel) -> None:
     def now():
         return datetime.datetime.now(tz=datetime.timezone.utc)
 
@@ -90,3 +92,57 @@ class Document(CreatedUpdatedMixin, table=True):
         doc = session.exec(select(Document)).one()
         assert doc.created_by.name == "Jane"
         assert doc.updated_by.name == "John"
+
+
+def test_inherit_relationship_model_validate(clear_sqlmodel) -> None:
+    class User(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    class Mixin(SQLModel):
+        owner_id: Optional[int] = Field(default=None, foreign_key="user.id")
+        owner: Optional[User] = Relationship(
+            sa_relationship=declared_attr(
+                lambda cls: relationship(User, foreign_keys=cls.owner_id)
+            )
+        )
+
+    class Asset(Mixin, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+
+    class AssetCreate(pydantic.BaseModel):
+        pass
+
+    asset_create = AssetCreate()
+
+    engine = create_engine("sqlite://")
+
+    SQLModel.metadata.create_all(engine)
+
+    user = User()
+
+    # Owner must be optional
+    asset = Asset.model_validate(asset_create)
+    with Session(engine) as session:
+        session.add(asset)
+        session.commit()
+        session.refresh(asset)
+        assert asset.id is not None
+        assert asset.owner_id is None
+        assert asset.owner is None
+
+    # When set, owner must be saved
+    #
+    # Under Pydantic V2, relationship fields set it `model_validate` are not saved,
+    # with or without inheritance.  Consider it a known issue.
+    #
+    if IS_PYDANTIC_V2:
+        asset = Asset.model_validate(asset_create, update={"owner": user})
+        with Session(engine) as session:
+            session.add(asset)
+            session.commit()
+            session.refresh(asset)
+            session.refresh(user)
+            assert asset.id is not None
+            assert user.id is not None
+            assert asset.owner_id == user.id
+            assert asset.owner.id == user.id