Skip to content

PostgreSQLInheritance

Jump to bottom
mike bayer edited this page May 15, 2023 · 4 revisions

PostgreSQLInheritance

Here we demonstrate one way of mapping to a partitioned PostgreSQL structure, using the schema techniques described in:

https://www.postgresql.org/docs/current/ddl-partitioning.html#DDL-PARTITIONING-USING-INHERITANCE

What's notable about this example is that the "child" tables are not referred towards explicitly; SQL access is done through the base table only, PostgreSQL handles the "discriminator" aspect of these tables transparently using a trigger for the DML side and normal PG partitioning logic for the DQL (read) side.

Based on this behavior, SQLAlchemy would best refer to this model via the base table alone using single table inheritance.

To map this, we can use single inheritance configuration directly, and DDL events to emit the special PostgreSQL-specific commands to create "partitions" as well as the trigger.

The only caveat within the mapping configuration is that using INSERT..RETURNING statement against "parent.id" doesn't seem to obey the INHERITS mechanics otherwise in place. So we turn off "implicit returning" for the table so that parent_id_seq is pre-executed beforehand, allowing DML INSERT statements to proceed.

from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import func
from sqlalchemy import Integer
from sqlalchemy import select
from sqlalchemy import String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session

Base = declarative_base()


class Parent(Base):
    __tablename__ = "parent"

    id = Column(Integer, primary_key=True)
    discriminator = Column(String(50))
    data = Column(String(50))
    __mapper_args__ = {"polymorphic_on": discriminator}
    __table_args__ = {"implicit_returning": False}


class Child1(Parent):
    __tablename__ = None
    __mapper_args__ = {"polymorphic_identity": "ops"}


class Child2(Parent):
    __tablename__ = None
    __mapper_args__ = {"polymorphic_identity": "eng"}


@event.listens_for(Parent.__table__, "after_create")
def create_child_tables(target, connection, **kw):
    connection.exec_driver_sql(
        """
        CREATE TABLE child1(
            CHECK (discriminator='ops')
        ) INHERITS (parent)

    """
    )

    connection.exec_driver_sql(
        """
        CREATE TABLE child2(
            CHECK (discriminator='eng')
        ) INHERITS (parent)
    """
    )

    connection.exec_driver_sql(
        """
        CREATE OR REPLACE FUNCTION parent_insert_trigger()
        RETURNS TRIGGER AS $$
        BEGIN
            IF (NEW.discriminator='ops') THEN
            INSERT INTO child1 VALUES (NEW.*);
        ELSIF (NEW.discriminator='eng') THEN
            INSERT INTO child2 VALUES (NEW.*);
        ELSE
            RAISE EXCEPTION 'Discriminator out of range.  Expected "ops", "eng"';
        END IF;
        RETURN NULL;
        END;
        $$
        LANGUAGE plpgsql;
    """  # noqa: E501
    )

    connection.exec_driver_sql(
        """
    CREATE TRIGGER insert_parent
    BEFORE INSERT ON parent
    FOR EACH ROW EXECUTE PROCEDURE parent_insert_trigger();
    """
    )


@event.listens_for(Parent.__table__, "before_drop")
def drop_child_tables(target, connection, **kw):
    connection.exec_driver_sql("drop table child1")
    connection.exec_driver_sql("drop table child2")


e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)
Base.metadata.drop_all(e)
Base.metadata.create_all(e)


s = Session(e)

s.add_all(
    [
        Child1(data="c1"),
        Child2(data="c2"),
        Child1(data="c3"),
        Child2(data="c4"),
        Child1(data="c5"),
    ]
)
s.commit()
s.close()

assert s.scalar(select(func.count()).select_from(Child1)) == 3
assert s.scalar(select(func.count()).select_from(Child2)) == 2
print(s.scalars(select(Parent).order_by(Parent.data)).all())
Clone this wiki locally