fix relationship problem of parent class during polymorphic inherit

Author: John Lyu

sqlmodel/main.py

@@ -548,6 +548,14 @@ def __new__(
                 if issubclass(base, BaseModel):
                     base_fields.update(get_model_fields(base))
                     base_annotations.update(base.__annotations__)
+                    if hasattr(base, "__sqlmodel_relationships__"):
+                        for k in base.__sqlmodel_relationships__:
+                            # create a dummy attribute to avoid inherit
+                            # pydantic will treat it as class variables, and will not become fields on model instances
+                            anno = base_annotations.get(k, Any)
+                            dummy_anno = ClassVar[anno]
+                            dict_used["__annotations__"][k] = dummy_anno
+
                     if hasattr(base, "__tablename__"):
                         is_polymorphic = True
             # use base_fields overwriting the ones from the class for inherit

tests/test_polymorphic_model.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy import ForeignKey
 from sqlalchemy.orm import mapped_column
-from sqlmodel import Field, Session, SQLModel, create_engine, select
+from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select
 
 from tests.conftest import needs_pydanticv2
 
@@ -16,7 +16,7 @@ class Hero(SQLModel, table=True):
 
         __mapper_args__ = {
             "polymorphic_on": "hero_type",
-            "polymorphic_identity": "hero",
+            "polymorphic_identity": "normal_hero",
         }
 
     class DarkHero(Hero):
@@ -59,7 +59,7 @@ class Hero(SQLModel, table=True):
 
         __mapper_args__ = {
             "polymorphic_on": "hero_type",
-            "polymorphic_identity": "hero",
+            "polymorphic_identity": "normal_hero",
         }
 
     class DarkHero(Hero):
@@ -103,7 +103,7 @@ class Hero(SQLModel, table=True):
 
         __mapper_args__ = {
             "polymorphic_on": "hero_type",
-            "polymorphic_identity": "hero",
+            "polymorphic_identity": "normal_hero",
         }
 
     class DarkHero(Hero):
@@ -130,3 +130,48 @@ class DarkHero(Hero):
         result = db.exec(statement).all()
     assert len(result) == 1
     assert isinstance(result[0].dark_power, str)
+
+
+@needs_pydanticv2
+def test_polymorphic_relationship(clear_sqlmodel) -> None:
+    class Tool(SQLModel, table=True):
+        __tablename__ = "tool_table"
+
+        id: int = Field(primary_key=True)
+
+        name: str
+
+    class Person(SQLModel, table=True):
+        __tablename__ = "person_table"
+
+        id: int = Field(primary_key=True)
+
+        discriminator: str
+        name: str
+
+        tool_id: int = Field(foreign_key="tool_table.id")
+        tool: Tool = Relationship()
+
+        __mapper_args__ = {
+            "polymorphic_on": "discriminator",
+            "polymorphic_identity": "simple_person",
+        }
+
+    class Worker(Person):
+        __mapper_args__ = {
+            "polymorphic_identity": "worker",
+        }
+
+    engine = create_engine("sqlite:///:memory:", echo=True)
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as db:
+        tool = Tool(id=1, name="Hammer")
+        db.add(tool)
+        worker = Worker(id=2, name="Bob", tool_id=1)
+        db.add(worker)
+        db.commit()
+
+        statement = select(Worker).where(Worker.tool_id == 1)
+        result = db.exec(statement).all()
+        assert len(result) == 1
+        assert isinstance(result[0].tool, Tool)