sqlalchemy_polymorphic_union2.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))

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


class ModelA(Entity):
    __tablename__ = "model_a"
    id = Column(Integer, ForeignKey("entities.id"), primary_key=True)

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


class ModelB(Entity):
    __tablename__ = "model_b"
    id = Column(Integer, ForeignKey("entities.id"), primary_key=True)

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


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")


class Event(Base):
    __tablename__ = 'events'
    id = Column(Integer, primary_key=True)
    name = Column(String(100))
    entity_id = Column(Integer, ForeignKey("entities.id"))
    entity = relationship("Entity", back_populates="event")


Entity.__mapper__.add_property('task', relationship('Task', back_populates='entity'))
Entity.__mapper__.add_property('event', relationship('Event', back_populates='entity'))

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

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

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

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

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

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


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

events = session.query(Event).all()
for e in events:
    print(e.entity)


entities = session.query(Entity).all()
for e in entities:
    print(e.type, e.id)

session.close()