query._adapt_polymorphic_element should try to adapt based on entities before tables

from sqlalchemy import *
from sqlalchemy.ext.declarative.api import declarative_base
from sqlalchemy.orm import *

Base = declarative_base()

class A(Base):
    __tablename__ = 'a'
    id = Column(Integer, primary_key=True)
    name = Column(String(20), nullable=False)

class B(A):
    __tablename__ = 'b'
    id = Column(Integer, ForeignKey('a.id'), primary_key=True)

class C(A):
    __tablename__ = 'c'
    id = Column(Integer, ForeignKey('a.id'), primary_key=True)
    bid = Column(Integer, ForeignKey('b.id'))

class D(A):
    __tablename__ = 'd'
    id = Column(Integer, ForeignKey('a.id'), primary_key=True)
    cid = Column(Integer, ForeignKey('b.id'))

s = Session()

q = s.query(B.name, C.name, D.name).select_from(B).\
                join(C, C.bid == B.id).\
                join(D, D.cid == C.id)

btoc = q._from_obj[0](0).left

ac_adapted = btoc.right.element.left
c_adapted = btoc.right.element.right

assert ac_adapted.element is A.__table__
assert c_adapted.element is C.__table__


ctod = q._from_obj[0](0).right
ad_adapted = ctod.left
d_adapted = ctod.right
assert ad_adapted.element is A.__table__
assert d_adapted.element is D.__table__

bname, cname, dname = q._entities

b_name_adapted = bname._resolve_expr_against_query_aliases(q, bname.column, None)
c_name_adapted = cname._resolve_expr_against_query_aliases(q, cname.column, None)
d_name_adapted = dname._resolve_expr_against_query_aliases(q, dname.column, None)

assert bool(b_name_adapted == A.__table__.c.name)
assert bool(c_name_adapted == ac_adapted.c.name)
assert bool(d_name_adapted == ad_adapted.c.name)

assert str(q.with_labels().statement.compile()).startswith(
            "SELECT a.name AS a_name, a_1.name AS a_1_name, a_2.name AS a_2_name")

patch:

--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -203,12 +203,17 @@ class Query(object):
                 self._polymorphic_adapters.pop(m.local_table, None)

     def _adapt_polymorphic_element(self, element):
-        if isinstance(element, expression.FromClause):
-            search = element
-        elif hasattr(element, 'table'):
-            search = element.table
-        else:
-            search = None
+        search = None
+        if "parententity" in element._annotations:
+            search = element._annotations['parententity']('parententity')
+            if search not in self._polymorphic_adapters:
+                search = None
+
+        if search is None:
+            if isinstance(element, expression.FromClause):
+                search = element
+            elif hasattr(element, 'table'):
+                search = element.table

         if search is not None:
             alias = self._polymorphic_adapters.get(search, None)

Comments (2)