Subquery loading off of an "already seen" row

Given the following relations:

Company has many employees and a parent company
Employee has an employer

And the following instances:

Apple w/ employees Tim Cook & Jony Ive
Beats w/ employees Dr. Dre & Jimmy Iovine

The following query does not load everything:

employees = session.query(Employee).options(
    joinedload('employer').joinedload('parent').subqueryload('employees')
).filter().all()

It appears that it finds all employees and parents properly, but since all the companies get loaded into the session via the joined load, it seems to stop processing at this point and doesn't actually perform the subqueryload of employees. This example may seem contrived since it's circling back around through the relationships, but it's a simplified version of something that w/ filters applied is a very real query in our system.

Below is some code to run to show the problem. Just change the connection string to be able to connect to your database, the code will populate a few records for you.

import logging
from contextlib import contextmanager
from sqlalchemy import (
    Column,
    ForeignKey,
    Integer,
    MetaData,
    String,
    Table,
    create_engine,
    inspect,
)
from sqlalchemy.orm import (
    mapper,
    relationship,
    sessionmaker,
    subqueryload,
    joinedload,
)

class Model(object):
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)
class Company(Model): pass
class Employee(Model): pass

engine = create_engine(
    'mysql+pymysql://root:password@localhost/test?charset=utf8mb4',
    convert_unicode=True,
)
metadata = MetaData(bind=engine)
Session = sessionmaker(bind=engine)

company_table = Table(
    'companies', metadata,
    Column('id', Integer, primary_key=True, autoincrement=True),
    Column('name', String(255)),
    Column('parent_id', Integer, ForeignKey('companies.id')),
)
employee_table = Table(
    'employees', metadata,
    Column('id', Integer, primary_key=True, autoincrement=True),
    Column('name', String(255)),
    Column('employer_id', Integer, ForeignKey('companies.id')),
)
company_mapper = mapper(Company, company_table, properties={
    'parent': relationship(
        Company,
        uselist=False,
        remote_side=[company_table.columns.id],
    ),
    'employees': relationship(
        Employee,
        uselist=True,
        back_populates='employer',
    ),
})

employee_mapper = mapper(Employee, employee_table, properties={
    'employer': relationship(
        Company,
        uselist=False,
        back_populates='employees',
    )
})

connection = engine.connect()
metadata.create_all(connection)
connection.close()

@contextmanager
def session_manager():
    session = Session()
    try: yield session; session.commit()
    except: session.rollback(); raise
    finally: session.close()

def _ensure_loaded(obj, key_path):
    components = key_path.split('.')
    for i, key in enumerate(components):
        if key in inspect(obj).unloaded:
            raise AttributeError('Attribute {} not loaded during query'.format(
                '.'.join(components[:i+1])
            ))
        obj = getattr(obj, key)

        if obj is None:
            break

models = [
    Company(id=1, name='Apple'),
    Company(id=2, name='Beats', parent_id=1),
    Employee(id=1, name='Tim Cook', employer_id=1),
    Employee(id=2, name='Jony Ive', employer_id=1),
    Employee(id=3, name='Dr. Dre', employer_id=2),
    Employee(id=4, name='Jimmy Iovine', employer_id=2),
]

for model in models:
    with session_manager() as session:
        try:
            session.add(model)
            session.commit()
        except Exception as e: session.close()

logging.basicConfig()
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

with session_manager() as session:
    employees = session.query(Employee).options(
        joinedload('employer').joinedload('parent').subqueryload('employees')
    ).filter(Employee.id == 3).all()

    for employee in employees:
        _ensure_loaded(employee, 'employer.parent.employees')

# this block of code fails presumably because the parent employer is loaded into the session during
# the first joinedload of the employer which doesn't trigger the subqueryload for parent.employees
# for the company that actually does have a parent company.
with session_manager() as session:
    employees = session.query(Employee).options(
        joinedload('employer').joinedload('parent').subqueryload('employees')
    ).filter().all()

    for employee in employees:
        _ensure_loaded(employee, 'employer.parent.employees')

Comments (10)