sqlalchemy_polymorphic_union.py

from sqlalchemy import create_engine, ForeignKey, Column, Integer, String
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
from sqlalchemy.orm import Mapped, mapped_column

Base = declarative_base()


class Entity(Base):
    __tablename__ = "entities"
    id = Column(Integer, primary_key=True)
    type = Column(String(50))
    task = relationship("Task", back_populates="entity")

    __mapper_args__ = {"polymorphic_identity": "entity", "polymorphic_on": type}


class ModelA(Entity):
    __tablename__ = "model_a"
    id = Column(Integer, ForeignKey("entities.id"), primary_key=True)
    task = relationship("Task", back_populates="entity")
    __mapper_args__ = {
        "polymorphic_identity": "model_a",
    }


class ModelB(Entity):
    __tablename__ = "model_b"
    id = Column(Integer, ForeignKey("entities.id"), primary_key=True)
    task = relationship("Task", back_populates="entity")

    __mapper_args__ = {
        "polymorphic_identity": "model_b",
    }


class ModelC(Entity):
    __tablename__ = "model_c"
    id = Column(Integer, ForeignKey("entities.id"), primary_key=True)
    task = relationship("Task", back_populates="entity")
    __mapper_args__ = {
        "polymorphic_identity": "model_c",
    }


class Task(Base):
    __tablename__ = "tasks"
    id = Column(Integer, primary_key=True)
    name = Column(String(100))
    entity_id = Column(Integer, ForeignKey("entities.id"))
    entity = relationship("Entity", back_populates="task")


engine = create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)
session = Session()

task = Task(name="Example Task", entity=ModelA())
session.add(task)
session.commit()

task = Task(name="Another Task", entity=ModelB())
session.add(task)
session.commit()

task = Task(name="Yet Another Task", entity=ModelC())
session.add(task)
session.commit()

tasks = session.query(Task).all()
for t in tasks:
    print(t.entity)

session.close()