Ensure mapper.polymorphic_on is polymorphic_prop.columns[0]

Fixed bug where joined eager loading would fail for a polymorphically-
loaded mapper, where the polymorphic_on was set to an un-mapped
expression such as a CASE expression.

Change-Id: Iffe68196aaac592165c89684f09f4c06cd78ce54
Fixes: #3800

3 files changed, 93 insertions, 12 deletions

diff --git a/doc/build/changelog/changelog_10.rst b/doc/build/changelog/changelog_10.rst
index 8808f6511..b3e76b475 100644
--- a/doc/build/changelog/changelog_10.rst
+++ b/doc/build/changelog/changelog_10.rst
@@ -20,6 +20,15 @@
 
     .. change::
         :tags: bug, orm
+        :tickets: 3800
+        :versions: 1.1.0
+
+        Fixed bug where joined eager loading would fail for a polymorphically-
+        loaded mapper, where the polymorphic_on was set to an un-mapped
+        expression such as a CASE expression.
+
+    .. change::
+        :tags: bug, orm
         :tickets: 3798
         :versions: 1.1.0
 
diff --git a/lib/sqlalchemy/orm/mapper.py b/lib/sqlalchemy/orm/mapper.py
index b8dc5b8c3..60848097c 100644
--- a/lib/sqlalchemy/orm/mapper.py
+++ b/lib/sqlalchemy/orm/mapper.py
@@ -1401,9 +1401,6 @@ class Mapper(InspectionAttr):
                 # polymorphic_on is a column that is already mapped
                 # to a ColumnProperty
                 prop = self._columntoproperty[self.polymorphic_on]
-                polymorphic_key = prop.key
-                self.polymorphic_on = prop.columns[0]
-                polymorphic_key = prop.key
             elif isinstance(self.polymorphic_on, MapperProperty):
                 # polymorphic_on is directly a MapperProperty,
                 # ensure it's a ColumnProperty
@@ -1414,8 +1411,6 @@ class Mapper(InspectionAttr):
                         "property or SQL expression "
                         "can be passed for polymorphic_on")
                 prop = self.polymorphic_on
-                self.polymorphic_on = prop.columns[0]
-                polymorphic_key = prop.key
             elif not expression._is_column(self.polymorphic_on):
                 # polymorphic_on is not a Column and not a ColumnProperty;
                 # not supported right now.
@@ -1477,12 +1472,14 @@ class Mapper(InspectionAttr):
                         col.label("_sa_polymorphic_on")
                     key = col.key
 
-                self._configure_property(
-                    key,
-                    properties.ColumnProperty(col,
-                                              _instrument=instrument),
-                    init=init, setparent=True)
-                polymorphic_key = key
+                prop = properties.ColumnProperty(col, _instrument=instrument)
+                self._configure_property(key, prop, init=init, setparent=True)
+
+            # the actual polymorphic_on should be the first public-facing
+            # column in the property
+            self.polymorphic_on = prop.columns[0]
+            polymorphic_key = prop.key
+
         else:
             # no polymorphic_on was set.
             # check inheriting mappers for one.
diff --git a/test/orm/inheritance/test_basic.py b/test/orm/inheritance/test_basic.py
index 42959a6e3..7b5f85b9b 100644
--- a/test/orm/inheritance/test_basic.py
+++ b/test/orm/inheritance/test_basic.py
@@ -1,5 +1,5 @@
 import warnings
-from sqlalchemy.testing import eq_, assert_raises, assert_raises_message
+from sqlalchemy.testing import eq_, is_, assert_raises, assert_raises_message
 from sqlalchemy import *
 from sqlalchemy import exc as sa_exc, util, event
 from sqlalchemy.orm import *
@@ -77,6 +77,65 @@ class O2MTest(fixtures.MappedTest):
         eq_(l[0].parent_foo.data, 'foo #1')
         eq_(l[1].parent_foo.data, 'foo #1')
 
+
+class PolyExpressionEagerLoad(fixtures.DeclarativeMappedTest):
+    run_setup_mappers = 'once'
+    __dialect__ = 'default'
+
+    @classmethod
+    def setup_classes(cls):
+        Base = cls.DeclarativeBasic
+
+        class A(fixtures.ComparableEntity, Base):
+            __tablename__ = 'a'
+
+            id = Column(Integer, primary_key=True,
+                        test_needs_autoincrement=True)
+            discriminator = Column(String(50), nullable=False)
+            child_id = Column(Integer, ForeignKey('a.id'))
+            child = relationship('A')
+
+            p_a = case([
+                (discriminator == "a", "a"),
+            ], else_="b")
+
+            __mapper_args__ = {
+                'polymorphic_identity': 'a',
+                "polymorphic_on": p_a,
+            }
+
+        class B(A):
+            __mapper_args__ = {
+                'polymorphic_identity': 'b'
+            }
+
+    @classmethod
+    def insert_data(cls):
+        A = cls.classes.A
+
+        session = Session(testing.db)
+        session.add_all([
+            A(id=1, discriminator='a'),
+            A(id=2, discriminator='b', child_id=1),
+            A(id=3, discriminator='c', child_id=1),
+        ])
+        session.commit()
+
+    def test_joinedload(self):
+        A = self.classes.A
+        B = self.classes.B
+
+        session = Session(testing.db)
+        result = session.query(A).filter_by(child_id=None).\
+            options(joinedload('child')).one()
+
+
+        eq_(
+            result,
+            A(id=1, discriminator='a', child=[B(id=2), B(id=3)]),
+        )
+
+
 class PolymorphicResolutionMultiLevel(fixtures.DeclarativeMappedTest,
                                         testing.AssertsCompiledSQL):
     run_setup_mappers = 'once'
@@ -396,6 +455,22 @@ class PolymorphicOnNotLocalTest(fixtures.MappedTest):
     def _roundtrip(self, set_event=True, parent_ident='parent', child_ident='child'):
         Parent, Child = self.classes.Parent, self.classes.Child
 
+        # locate the "polymorphic_on" ColumnProperty.   This isn't
+        # "officially" stored at the moment so do some heuristics to find it.
+        parent_mapper = inspect(Parent)
+        for prop in parent_mapper.column_attrs:
+            if not prop.instrument:
+                break
+        else:
+            prop = parent_mapper._columntoproperty[
+                parent_mapper.polymorphic_on]
+
+        # then make sure the column we will query on matches.
+        is_(
+            parent_mapper.polymorphic_on,
+            prop.columns[0]
+        )
+
         if set_event:
             @event.listens_for(Parent, "init", propagate=True)
             def set_identity(instance, *arg, **kw):