SQLAlchemy Setup For Subclassed Association Table With Polymorphic Entities

Viewed 235

I am new to both python and sqlalchemy but have a rather complex modeling situation I am having trouble setting up. It involves an association table in which the association table has a polymorphic association for one of the entities.

I am very close to getting this working. When the data already exists in the DB, I can read and interact with it and the models as expected. The problem comes with writing, and I will address that after showing the code:

First, there is a shared base class defining tablename and id as postgres uuid

@as_declarative()
class Base(object):
    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()

    id = Column(pg.UUID(as_uuid=True), primary_key=True, default=uuid4)

There is an enum of defined types allowed in the polymorphic column. While I will eventually support Aaa and Bbb, this example only has Bbb defined so far for clarity.

class EntityTypes(Enum):
    AAA = Aaa.__name__.lower()
    BBB = Bbb.__name__.lower()

This is the model that represents the polymorphic association table. It has entity_id, entity_type, and ccc_id columns for the joins. "Entity" can be an Aaa or a Bbb, but Ccc is always a Ccc (non-polymorphic).

class EntityCcc(Base):
    """Polymorphic mapping between an EntityType and Ccc."""
    
    # entity_id is defined in subclasses, with foreign keys
    entity_type = Column(
        Enum(EntityTypes, values_callable=lambda x: [e.value for e in x]), nullable=False
    )

    ccc_id = Column(
        pg.UUID(as_uuid=True),
        ForeignKey(f"{Ccc.__tablename__}.id"),
        nullable=False
    )

    __mapper_args__ = {
        "polymorphic_on": entity_type
    }

This is a subclass of the polymorphic model which knows about Bbbs, so it sets the foreign key accordingly, and sets up the relationships. A similar AaaCcc class would be defined

class BbbCcc(EntityCcc):
    """
        Mapping between a Bbb and Ccc.
        Subclasses polymorphic join model to get specific
        bbb accessor and set appropriate foreign key.
    """

    __tablename__ = EntityCcc.__tablename__

    entity_id = Column(
        pg.UUID(as_uuid=True), ForeignKey(f"{Bbb.__tablename__}.id"), nullable=False
    )

    bbb = relationship("Bbb", back_populates='bbb_ccc')
    ccc = relationship("Ccc", back_populates='bbb_ccc')

    __mapper_args__ = {
        "polymorphic_identity": EntityTypes(Bbb.__name__.lower())
    }

This is a Bbb, which has a relationship setup to the subclassed join model, and a secondary relationship setup to point to its Ccc through the join table (using secondary)

class Bbb(Base):
    """Represents a Bbb entity."""

    name = Column(TEXT)

    bbb_ccc = relationship("BbbCcc", back_populates="bbb", uselist=False)

    ccc = relationship(
        "Ccc",
        secondary="entity_ccc",
        back_populates="bbb",
        uselist=False
    )

This is a Ccc, which has a relationship setup to the subclassed join model, and a secondary relationship setup to point to its Bbb through the join table (using secondary)

class Ccc(Base):
    """Represents a Ccc entity."""

    name = Column(TEXT)

    bbb_ccc = relationship("BbbCcc", back_populates="ccc", uselist=False)

    bbb = relationship(
        "Bbb",
        secondary="entity_ccc",
        back_populates="ccc",
        uselist=False
    )

So what's the problem?

With the appropriate entries already seeded in the database, I can interact with them as expected:

(Pdb) found_bbb_ccc = db.session.query(BbbCcc).get(uuid)
(Pdb) found_bbb_ccc
<app.models.mappings.bbb_ccc.BbbCcc object at 0x7f488ce6ebe0>
(Pdb) found_bbb_ccc.bbb
<app.models.entities.bbb.Bbb object at 0x7f488dd73f10>
(Pdb) found_bbb_ccc.ccc
<app.models.entities.ccc.Ccc object at 0x7f488ce6ec40>
(Pdb) found_bbb_ccc.bbb.ccc
<app.models.entities.ccc.Ccc object at 0x7f488ce6ec40>
(Pdb) found_bbb_ccc.ccc.bbb
<app.models.entities.bbb.Bbb object at 0x7f488dd73f10>

This shows that a bbb can reference and find its ccc through the join model, and vice versa. Reading through the association is fine. But creating a new association by writing is problematic:

new_bbb = Bbb(name='Bbb instance')
new_ccc = Ccc(name='Ccc instance')

new_bbb.ccc = new_ccc
db.session.commit()

*** sqlalchemy.exc.IntegrityError: (psycopg2.errors.NotNullViolation) null value in column "entity_type" violates not-null constraint
DETAIL:  Failing row contains (4b1f7ac7-16b0-4972-9577-bda1b5efe2aa, 2021-08-05 17:50:05.233465, 2021-08-05 17:50:05.233482, 63463492-0a9d-492f-b42a-72ec276f2768, null, a75d06af-33bd-4345-abbd-c6098e9a797d).

[SQL: INSERT INTO entity_ccc (created, updated, id, ccc_id, entity_id) VALUES (%(created)s, %(updated)s, %(id)s, %(ccc_id)s, %(entity_id)s)]
[parameters: {'created': datetime.datetime(2021, 8, 5, 17, 50, 5, 233465), 'updated': datetime.datetime(2021, 8, 5, 17, 50, 5, 233482), 'id': UUID('4b1f7ac7-16b0-4972-9577-bda1b5efe2aa'), 'ccc_id': UUID('a75d06af-33bd-4345-abbd-c6098e9a797d'), 'entity_id': UUID('63463492-0a9d-492f-b42a-72ec276f2768')}]
(Background on this error at: http://sqlalche.me/e/14/gkpj)

The error seen is that the writing of this data is not setting the polymorphic entity_type, which should be bbb here. I believe the problem is that the secondary relationship definition requires a table name, and not an object (passing BbbCcc would likely pick up the entity_type right?) but perhaps it is something else.

How can I tweak this code to allow the setting of the polymorphic association as described? Thanks!

1 Answers

try

from sqlalchemy import inspect

    [...]
    eng_mapper = inspect(Engineer)
    query.filter(
        eng_mapper.polymorphic_on.in_(
            m.polymorphic_identity
            for m in eng_mapper.polymorphic_iterator()
        ),
    )

I'd prefer a slightly less verbose incantation, but this works and doesn't require knowledge of the specific configuration of the polymorphic hierarchy.

Details When inspect() is called on an ORM mapped class, it returns the Mapper for that class. This is identical to the Model.mapper class attribute.

The Mapper contains all information needed to introspect the polymorphic hierarchy. In particular:

.polymorphic_on is the field (column) in the model at the top of the hierarchy that contains the polymorphic identity value for a record (e.g. for Engineer that would be the Employee.type field). .polymorphic_identity is the value that each instance of the mapped model will have in the .polymorphic_on field (e.g. for Engineer that would be "engineer"). .polymorphic_iterator() iterates over a collection of model Mappers that includes Model.mapper and the .mapper of all subclasses of Model recursively (e.g. for Engineer that would be an iterator containing only the Engineer.mapper). To make it more readable, one could easily turn the above filter expression into a function:

from sqlalchemy import inspect

def filter_instances_of(cls):
    mapper = inspect(cls)
    return mapper.polymorphic_on.in_(
        m.polymorphic_identity for m in mapper.polymorphic_iterator()
    )

And use it like:

query = query.filter(
    filter_instances_of(Engineer),
    [... other filter criteria ...]
)
Related