Error when selectin eager loading on polymorphic subtypes

[sqlalchemy-bot]

Migrated issue, originally created by Theron Luhn

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)

[sqlalchemy-bot]

Theron Luhn wrote:

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?

[sqlalchemy-bot]

Theron Luhn wrote:

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

[sqlalchemy-bot]

Changes by Theron Luhn:

• edited description

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

first guess for first problem at https://gerrit.sqlalchemy.org/#/c/zzzeek/sqlalchemy/+/630. second example looks like a separate issue.

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

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.

[sqlalchemy-bot]

Theron Luhn wrote:

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.

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

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

→ a216625

[sqlalchemy-bot]

Changes by Michael Bayer (@zzzeek):

• changed status to closed