SQLAlchemy concrete table inheritance with relationship and imperative mapping

example.py

import abc
from dataclasses import dataclass

from sqlalchemy import Column, ForeignKey, Integer, String, Table
from sqlalchemy.engine.create import create_engine
from sqlalchemy.orm import polymorphic_union, registry, relationship
from sqlalchemy.orm.session import sessionmaker


# Create abstract so the parent can't be instantiated
@dataclass
class AbstractDataclass(abc.ABC):
    def __new__(cls, *args, **kwargs):
        if cls == AbstractDataclass or cls.__bases__[0] == AbstractDataclass:
            raise TypeError("Cannot instantiate abstract class.")
        return super().__new__(cls)


# Models
class FinancialInfo(AbstractDataclass):
    pass


@dataclass
class PixInfo(FinancialInfo):
    pix_number: str


@dataclass
class AccountInfo(FinancialInfo):
    bank_number: int


@dataclass
class User:
    name: str
    financial_info: FinancialInfo


# Databse definitions
# engine = create_engine("sqlite://", echo=True)
engine = create_engine("sqlite://")

Session = sessionmaker(
    bind=engine,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False,
)

mapper_registry = registry()

metadata_obj = mapper_registry.metadata

# Tables

user_table = Table(
    "user",
    metadata_obj,
    Column("id", Integer, primary_key=True),
    Column("name", String(50)),
)

pix_info_table = Table(
    "pix_info",
    metadata_obj,
    Column("id", Integer, primary_key=True),
    Column("pix_number", String(50), nullable=False),
    Column(
        "user_id",
        Integer,
        ForeignKey("user.id"),
        nullable=False,
    ),
)

account_info_table = Table(
    "account_info",
    metadata_obj,
    Column("id", Integer, primary_key=True),
    Column("bank_number", String(50), nullable=False),
    Column(
        "user_id",
        Integer,
        ForeignKey("user.id"),
        nullable=False,
    ),
)

# Mapping
pjoin = polymorphic_union(
    {
        "pix_info": pix_info_table,
        "account_info": account_info_table,
    },
    "type",
    "pjoin",
)
employee_mapper = mapper_registry.map_imperatively(
    FinancialInfo,
    pjoin,
    with_polymorphic=("*", pjoin),
    polymorphic_on=pjoin.c.type,
)
mapper_registry.map_imperatively(
    PixInfo,
    pix_info_table,
    polymorphic_identity="pix_info",
    inherits=FinancialInfo,
    concrete=True,
)
mapper_registry.map_imperatively(
    AccountInfo,
    account_info_table,
    polymorphic_identity="account_info",
    inherits=FinancialInfo,
    concrete=True,
)


mapper_registry.map_imperatively(
    User,
    user_table,
    properties={"financial_info": relationship(FinancialInfo, uselist=False)},
)

# Create everything on the database
metadata_obj.create_all(engine)


# add data to the DB
with Session() as session:
    session.add(
        User(name="Ricardo", financial_info=PixInfo(pix_number="[email protected]"))
    )
    session.add(User(name="Marlon", financial_info=AccountInfo(bank_number=123)))
    # session.add(User(name="Felipe", financial_info=FinancialInfo()))
    session.commit()

# query data
with Session() as session:
    print(session.query(FinancialInfo).all())
    print("\n")
    print(session.query(PixInfo).all())
    print("\n")
    print(session.query(AccountInfo).all())
    print("\n")
    print(session.query(User).all())