Error when selectin eager loading on polymorphic subtypes

Issue #4156 resolved
Theron Luhn
created an issue

I'm attempting to use selectin eager loading on polymorphic subtypes using the strategy laid out here: http://docs.sqlalchemy.org/en/latest/orm/inheritance_loading.html#eager-loading-of-specific-or-polymorphic-subtypes

Here's some example code:

from sqlalchemy import String, Integer, Column, create_engine, ForeignKey, inspect
from sqlalchemy.orm import relationship, Session, subqueryload, selectinload
from sqlalchemy.ext.declarative import declarative_base


Base = declarative_base()


class Company(Base):
    __tablename__ = 'company'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)
    employees = relationship('Employee')


class Employee(Base):
    __tablename__ = 'employee'
    id = Column(Integer, primary_key=True)
    type = Column(String, nullable=False)
    name = Column(String, nullable=False)
    company_id = Column(Integer, ForeignKey('company.id'), nullable=False)

    __mapper_args__ = {
        'polymorphic_on': 'type',
        'with_polymorphic': '*',
    }


class Programmer(Employee):
    __tablename__ = 'programmer'
    id = Column(Integer, ForeignKey('employee.id'), primary_key=True)
    languages = relationship('Language')

    __mapper_args__ = {
        'polymorphic_identity': 'programmer',
    }


class Manager(Employee):
    __tablename__ = 'manager'
    id = Column(Integer, ForeignKey('employee.id'), primary_key=True)
    something = Column(String)

    __mapper_args__ = {
        'polymorphic_identity': 'manager',
    }


class Language(Base):
    __tablename__ = 'language'
    id = Column(Integer, primary_key=True)
    programmer_id = Column(
        Integer,
        ForeignKey('programmer.id'),
        nullable=False,
    )
    name = Column(String, nullable=False)


engine = create_engine('postgresql://localhost:5432/sa', )
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
db = Session(engine, enable_baked_queries=False)


company = Company(
    id=1,
    name='Foobar Corp',
    employees=[Programmer(
        id=1,
        name='John Smith',
        languages=[Language(id=1, name='Python')],
    ), Manager(
        id=2,
        name='Foo',
        something='foo',
    )],
)
db.add(company)
db.flush()
db.expunge_all()

company = db.query(Company).filter(
    Company.id == 1,
).options(
    selectinload(Company.employees.of_type(Programmer)).selectinload(Programmer.languages),
).one()
print(company.employees)

This results in:

Traceback (most recent call last):
  File "scratchpad/sqlalchemy_polymorphic_loading_bug.py", line 86, in <module>
    selectinload(Company.employees.of_type(Programmer)).selectinload(Programmer.languages),
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2837, in one
    ret = self.one_or_none()
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/query.py", line 2807, in one_or_none
    ret = list(self)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 97, in instances
    util.raise_from_cause(err)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 203, in raise_from_cause
    reraise(type(exception), exception, tb=exc_tb, cause=cause)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 187, in reraise
    raise value
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 85, in instances
    post_load.invoke(context, path)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 750, in invoke
    self.load_keys, *arg, **kw)
  File "<string>", line 1, in <lambda>
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/strategies.py", line 1966, in _load_for_path
    lambda x: x[0]
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/strategies.py", line 1958, in <dictcomp>
    data = {
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 97, in instances
    util.raise_from_cause(err)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 203, in raise_from_cause
    reraise(type(exception), exception, tb=exc_tb, cause=cause)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 187, in reraise
    raise value
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 85, in instances
    post_load.invoke(context, path)
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/loading.py", line 750, in invoke
    self.load_keys, *arg, **kw)
  File "<string>", line 1, in <lambda>
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/strategies.py", line 1987, in _load_for_path
    state.get_impl(self.key).set_committed_value(
  File "/Users/luhn/Code/revenue/.env/lib/python3.6/site-packages/sqlalchemy/orm/state.py", line 423, in get_impl
    return self.manager[key].impl
KeyError: 'languages'

This only happens when both a Manager and Programmer object are in the employees relation. If only a Manager or only a Programmer are attached to a company, it works fine.

This only happens when using selectinload, using subqueryload works fine.

I'm running the current master branch (1.2.0 doesn't work due to #4153)

Comments (7)

  1. Theron Luhn reporter

    It seems that the following query works fine, and is actually closer to what I want to achieve.

    all_employee_types = with_polymorphic(Employee, '*')
    company = db.query(Company).filter(
        Company.id == 1,
    ).options(
        selectinload(Company.employees.of_type(all_employee_types))
        .selectinload(all_employee_types.Programmer.languages),
    ).one()
    

    The polymorphic loading API is a bit confusing... was I using it wrong in my original code sample?

  2. Theron Luhn reporter

    Okay, back again: When using the "working" query above, but trying to eager load other relationships as well, it breaks again. Here's a couple queries I tried, the first executes successfully but doesn't eager load Employee.benefits, the second one throws the same error as above.

    from sqlalchemy import String, Integer, Column, create_engine, ForeignKey, inspect
    from sqlalchemy.orm import relationship, Session, subqueryload, selectinload, with_polymorphic
    from sqlalchemy.ext.declarative import declarative_base
    
    
    Base = declarative_base()
    
    
    class Company(Base):
        __tablename__ = 'company'
        id = Column(Integer, primary_key=True)
        name = Column(String, nullable=False)
        employees = relationship('Employee')
    
    
    class Employee(Base):
        __tablename__ = 'employee'
        id = Column(Integer, primary_key=True)
        type = Column(String, nullable=False)
        name = Column(String, nullable=False)
        benefits = relationship('Benefit')
        company_id = Column(Integer, ForeignKey('company.id'), nullable=False)
    
        __mapper_args__ = {
            'polymorphic_on': 'type',
            'with_polymorphic': '*',
        }
    
    
    class Programmer(Employee):
        __tablename__ = 'programmer'
        id = Column(Integer, ForeignKey('employee.id'), primary_key=True)
        languages = relationship('Language')
    
        __mapper_args__ = {
            'polymorphic_identity': 'programmer',
        }
    
    
    class Manager(Employee):
        __tablename__ = 'manager'
        id = Column(Integer, ForeignKey('employee.id'), primary_key=True)
        something = Column(String)
    
        __mapper_args__ = {
            'polymorphic_identity': 'manager',
        }
    
    
    class Language(Base):
        __tablename__ = 'language'
        id = Column(Integer, primary_key=True)
        programmer_id = Column(
            Integer,
            ForeignKey('programmer.id'),
            nullable=False,
        )
        name = Column(String, nullable=False)
    
    
    class Benefit(Base):
        __tablename__ = 'benefit'
        id = Column(Integer, primary_key=True)
        name = Column(String, nullable=False)
        employee_id = Column(Integer, ForeignKey('employee.id'), nullable=False)
    
    
    engine = create_engine('postgresql://localhost:5432/sa', echo=True)
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)
    db = Session(engine, enable_baked_queries=False)
    
    
    company = Company(
        id=1,
        name='Foobar Corp',
        employees=[Programmer(
            id=1,
            name='John Smith',
            languages=[Language(id=1, name='Python')],
        ), Manager(
            id=2,
            name='Foo',
            something='foo',
            benefits=[Benefit(id=1, name='401k')],
        )],
    )
    db.add(company)
    db.flush()
    db.expunge_all()
    
    all_employee_types = with_polymorphic(Employee, '*')
    company = db.query(Company).filter(
        Company.id == 1,
    ).options(
        # This doesn't give an error, but neither does it eagerly load Employee.benefits
        selectinload(Company.employees).selectinload(Employee.benefits),
        selectinload(Company.employees.of_type(all_employee_types))
        .selectinload(all_employee_types.Programmer.languages),
    ).one()
    company = db.query(Company).filter(
        Company.id == 1,
    ).options(
        # This does give an error
        # However, if I remove the Manager object, the error goes away and
        # Employee.benefits is properly eager loaded.
        selectinload('employees.benefits'),
        selectinload(Company.employees.of_type(all_employee_types))
        .selectinload(all_employee_types.Programmer.languages),
    ).one()
    
  3. Michael Bayer repo owner

    if we fix the first bug with the one liner (even though im not sure that's the place this should be fixed yet), the second example w/ the benefits can work if all the of_type stuff is removed, I added a Benefit to both and made sure I can get to all benefits / languages without extra SQL emitted:

    company = db.query(Company).filter(
        Company.id == 1,
    ).options(
        selectinload(Company.employees).selectinload(Employee.benefits),
        selectinload(Company.employees).selectinload(Programmer.languages),
    ).one()
    

    it works this way also:

        selectinload(Company.employees).selectinload(Employee.benefits),
        selectinload(Company.employees.of_type(Programmer))
        .selectinload(Programmer.languages),
    

    as far as throwing in the extra with_polymorphic(), this relationship is already polymorphically loading, I'm not sure why adding it in seems to interfere with everything, that is a bug but the overlapping of options in that case is so confusing and unintuitive that I'm not inclined to get into that case right now.

  4. Theron Luhn reporter

    I think your updated query makes more sense and is easier to read, so I'm quite happy forgoing my awkward second example. I don't fully understand polymorphic loading yet, so I'm just throwing things against the wall.

  5. Michael Bayer repo owner

    Limit select in loading for correct types

    Fixed bug in new "selectin" relationship loader where the loader could try to load a non-existent relationship when loading a collection of polymorphic objects, where only some of the mappers include that relationship, typically when :meth:.PropComparator.of_type is being used.

    This generalizes the mapper limiting that was present in _load_subclass_via_in() to be part of the PostLoad object itself, and is used by both polymorphic selectin loading and relationship selectin loading.

    Change-Id: I31416550e27bc8374b673860f57d9dcf96abe87d Fixes: #4156

    → <<cset a216625bd033>>

  6. Log in to comment