SQLAlchemy Subqueryload Poly

sqlalchemy_subqueryload_poly.py

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import (relationship, sessionmaker, subqueryload,
                            with_polymorphic)
from sqlalchemy import create_engine, Column, String, Integer
from sqlalchemy.schema import ForeignKey

Base = declarative_base()


class A(Base):
    __tablename__ = 't_a'

    id = Column(Integer, primary_key=True)


class B(Base):
    __tablename__ = 't_b'

    type = Column(String(2))
    __mapper_args__ = {
        'polymorphic_identity': 'b',
        'polymorphic_on': type
    }

    id = Column(Integer, primary_key=True)

    # Relationship to A
    a_id = Column(Integer, ForeignKey('t_a.id'))
    a = relationship('A', backref='bs')


class B2(B):
    __tablename__ = 't_b2'

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

    id = Column(Integer, ForeignKey('t_b.id'), primary_key=True)



class C(Base):
    __tablename__ = 't_c'

    type = Column(String(2))
    __mapper_args__ = {
        'polymorphic_identity': 'c',
        'polymorphic_on': type
    }

    id = Column(Integer, primary_key=True)

    # Relationship to B
    b_id = Column(Integer, ForeignKey('t_b.id'))
    b = relationship('B', backref='cs')


class C2(C):
    __tablename__ = 't_c2'

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

    id = Column(Integer, ForeignKey('t_c.id'), primary_key=True)


class D(Base):
    __tablename__ = 't_d'

    id = Column(Integer, primary_key=True)

    # Relationship to B
    c_id = Column(Integer, ForeignKey('t_c.id'))
    c = relationship('C', backref='ds')

engine = create_engine('sqlite://', echo=True)

Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
session = Session()

for i in xrange(2):
    a = A()
    session.add(a)

    b = B2(a=a)
    session.add(b)

    c = C2(b=b)
    session.add(c)

    d = D(c=c)
    session.add(d)

session.commit()

b_b2 = with_polymorphic(B, [B2], flat=True)
c_c2 = with_polymorphic(C, [C2], flat=True)

# Broken -- the query on D has a cross join between
# (A join B) and (B join C join D).
r = session.query(
    A
).options(
    subqueryload(
        A.bs.of_type(b_b2)
    ).subqueryload(
        b_b2.cs.of_type(c_c2)
    ).subqueryload(
        c_c2.ds
    )
).all()

# Note that the output is correct even though the generated query is wrong.
for a in r:
    print a.id
    for b in a.bs:
        print '\t', b.id
        for c in b.cs:
            print '\t\t', c.id
            for d in c.ds:
                print '\t\t\t', d.id