correlation fails with with_polymorphic selectables

we probably need to get correlate to find the actual tables to correlate here. two workarounds presented. Also why doesn't query have correlate_except() ?!
from sqlalchemy import Column, Integer, ForeignKey, String, DateTime
from sqlalchemy import and_, func
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import arrow

engine = create_engine('sqlite:///:memory:', echo=True)
Session = sessionmaker(bind=engine)
Base = declarative_base()


class Parent(Base):
    __tablename__ = 'parent_objs'

    id = Column(Integer, primary_key=True)
    obj_type = Column(String, nullable=False)
    user_id = Column(Integer, ForeignKey('user_objs.id'))
    time = Column(DateTime)

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


class Child(Parent):
    __tablename__ = 'child_objs'

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

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


class User(Base):
    __tablename__ = 'user_objs'

    id = Column(Integer, primary_key=True)


class UserEvent(Base):
    __tablename__ = 'user_events'

    id = Column(Integer, primary_key=True)
    time = Column(DateTime)
    user_id = Column(Integer, ForeignKey('user_objs.id'))


Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)

session = Session()
twenty_hours_ago = arrow.utcnow().replace(hours=-20).naive
ten_hours_ago = arrow.utcnow().replace(hours=-10).naive
just_after_ten_hours_ago = arrow.utcnow().replace(hours=-10, seconds=+1).naive
five_hours_ago = arrow.utcnow().replace(hours=-5).naive
just_after_five_hours_ago = arrow.utcnow().replace(hours=-5, seconds=+1).naive
now = arrow.utcnow().naive


u = User(id=1)
session.add(u)
session.commit()

ue_1 = UserEvent(id=1, user_id=u.id, time=twenty_hours_ago)
ue_2 = UserEvent(id=2, user_id=u.id, time=ten_hours_ago)
ue_3 = UserEvent(id=3, user_id=u.id, time=five_hours_ago)
ue_4 = UserEvent(id=4, user_id=u.id, time=now)

child_1 = Child(id=1, time=just_after_ten_hours_ago, user_id=u.id)
child_2 = Child(id=2, time=just_after_five_hours_ago, user_id=u.id)


map(session.add, [
    ue_1, ue_2, ue_3, ue_4, child_1, child_2])
session.commit()


# failure one
q_with_declarative_one = (
    session.query(Parent.id, UserEvent.id)
    .join(User, Parent.user_id == User.id)
    .join(UserEvent, and_(
        UserEvent.user_id == User.id,
        UserEvent.time == session.query(
            func.max(UserEvent.time)
        ).filter(UserEvent.user_id == User.id)
         .filter(UserEvent.time <= Parent.time)
         .correlate(Parent).correlate(User)
         .as_scalar()
    ))
)

# workaround one
q_with_declarative_two = (
    session.query(Parent.id, UserEvent.id)
    .join(User, Parent.user_id == User.id)
    .join(UserEvent, and_(
        UserEvent.user_id == User.id,
        UserEvent.time == session.query(
            func.max(UserEvent.time)
        ).filter(UserEvent.user_id == User.id)
         .filter(UserEvent.time <= Parent.time)
         .as_scalar().correlate_except(UserEvent)
    ))
)

# workaround two
from sqlalchemy.orm import with_polymorphic
parent_poly = with_polymorphic(Parent, [], Parent.__table__)

q_with_declarative_three = (
    session.query(parent_poly.id, UserEvent.id)
    .join(User, parent_poly.user_id == User.id)
    .join(UserEvent, and_(
        UserEvent.user_id == User.id,
        UserEvent.time == session.query(
            func.max(UserEvent.time)
        ).filter(UserEvent.user_id == User.id)
         .filter(UserEvent.time <= parent_poly.time)
         .correlate(parent_poly).correlate(User)
         .as_scalar()
    ))
)


print q_with_declarative_one.all()
print q_with_declarative_two.all()
print q_with_declarative_three.all()
Comments (1)
