Enable multi-level selectin polymorphic loading

Change-Id: Icc742bbeecdb7448ce84caccd63e086af16e81c1
Fixes: #4026

2 files changed, 342 insertions, 13 deletions

diff --git a/test/orm/inheritance/_poly_fixtures.py b/test/orm/inheritance/_poly_fixtures.py
index 79ff456e4..f1f9cd6f3 100644
--- a/test/orm/inheritance/_poly_fixtures.py
+++ b/test/orm/inheritance/_poly_fixtures.py
@@ -1,15 +1,10 @@
-from sqlalchemy import Integer, String, ForeignKey, func, desc, and_, or_
-from sqlalchemy.orm import interfaces, relationship, mapper, \
-    clear_mappers, create_session, joinedload, joinedload_all, \
-    subqueryload, subqueryload_all, polymorphic_union, aliased,\
-    class_mapper
-from sqlalchemy import exc as sa_exc
-from sqlalchemy.engine import default
+from sqlalchemy import Integer, String, ForeignKey
+from sqlalchemy.orm import relationship, mapper, \
+    create_session, polymorphic_union
 
 from sqlalchemy.testing import AssertsCompiledSQL, fixtures
-from sqlalchemy import testing
 from sqlalchemy.testing.schema import Table, Column
-from sqlalchemy.testing import assert_raises, eq_
+from sqlalchemy.testing import config
 
 
 class Company(fixtures.ComparableEntity):
@@ -370,3 +365,127 @@ class _PolymorphicJoins(_PolymorphicFixtureBase):
         manager_with_polymorphic = ('*', manager_join)
         return person_with_polymorphic,\
             manager_with_polymorphic
+
+
+class GeometryFixtureBase(fixtures.DeclarativeMappedTest):
+    """Provides arbitrary inheritance hierarchies based on a dictionary
+    structure.
+
+    e.g.::
+
+        self._fixture_from_geometry(
+            "a": {
+                "subclasses": {
+                    "b": {"polymorphic_load": "selectin"},
+                    "c": {
+                        "subclasses": {
+                            "d": {
+                                "polymorphic_load": "inlne", "single": True
+                            },
+                            "e": {
+                                "polymorphic_load": "inline", "single": True
+                            },
+                        },
+                        "polymorphic_load": "selectin",
+                    }
+                }
+            }
+        )
+
+    would provide the equivalent of::
+
+        class a(Base):
+            __tablename__ = 'a'
+
+            id = Column(Integer, primary_key=True)
+            a_data = Column(String(50))
+            type = Column(String(50))
+            __mapper_args__ = {
+                "polymorphic_on": type,
+                "polymorphic_identity": "a"
+            }
+
+        class b(a):
+            __tablename__ = 'b'
+
+            id = Column(ForeignKey('a.id'), primary_key=True)
+            b_data = Column(String(50))
+
+            __mapper_args__ = {
+                "polymorphic_identity": "b",
+                "polymorphic_load": "selectin"
+            }
+
+            # ...
+
+        class c(a):
+            __tablename__ = 'c'
+
+        class d(c):
+            # ...
+
+        class e(c):
+            # ...
+
+    Declarative is used so that we get extra behaviors of declarative,
+    such as single-inheritance column masking.
+
+    """
+
+    run_create_tables = 'each'
+    run_define_tables = 'each'
+    run_setup_classes = 'each'
+    run_setup_mappers = 'each'
+
+    def _fixture_from_geometry(self, geometry, base=None):
+        if not base:
+            is_base = True
+            base = self.DeclarativeBasic
+        else:
+            is_base = False
+
+        for key, value in geometry.items():
+            if is_base:
+                type_ = Column(String(50))
+                items = {
+                    "__tablename__": key,
+                    "id": Column(Integer, primary_key=True),
+                    "type": type_,
+                    "__mapper_args__": {
+                        "polymorphic_on": type_,
+                        "polymorphic_identity": key
+                    }
+
+                }
+            else:
+                items = {
+                    "__mapper_args__": {
+                        "polymorphic_identity": key
+                    }
+                }
+
+                if not value.get("single", False):
+                    items["__tablename__"] = key
+                    items["id"] = Column(
+                        ForeignKey("%s.id" % base.__tablename__),
+                        primary_key=True)
+
+            items["%s_data" % key] = Column(String(50))
+
+            # add other mapper options to be transferred here as needed.
+            for mapper_opt in ("polymorphic_load", ):
+                if mapper_opt in value:
+                    items["__mapper_args__"][mapper_opt] = value[mapper_opt]
+
+            if is_base:
+                klass = type(key, (fixtures.ComparableEntity, base, ), items)
+            else:
+                klass = type(key, (base, ), items)
+
+            if "subclasses" in value:
+                self._fixture_from_geometry(value["subclasses"], klass)
+
+        if is_base and self.metadata.tables and self.run_create_tables:
+            self.tables.update(self.metadata.tables)
+            self.metadata.create_all(config.db)
+
diff --git a/test/orm/inheritance/test_poly_loading.py b/test/orm/inheritance/test_poly_loading.py
index ab807b45c..f6046b3b2 100644
--- a/test/orm/inheritance/test_poly_loading.py
+++ b/test/orm/inheritance/test_poly_loading.py
@@ -1,12 +1,12 @@
 from sqlalchemy import String, Integer, Column, ForeignKey
 from sqlalchemy.orm import relationship, Session, \
-    selectin_polymorphic, selectinload
+    selectin_polymorphic, selectinload, with_polymorphic
 from sqlalchemy.testing import fixtures
 from sqlalchemy import testing
 from sqlalchemy.testing import eq_
-from sqlalchemy.testing.assertsql import AllOf, CompiledSQL, EachOf
-from ._poly_fixtures import Company, Person, Engineer, Manager, Boss, \
-    Machine, Paperwork, _Polymorphic
+from sqlalchemy.testing.assertsql import AllOf, CompiledSQL, EachOf, Or
+from ._poly_fixtures import Company, Person, Engineer, Manager, \
+    _Polymorphic, GeometryFixtureBase
 
 
 class BaseAndSubFixture(object):
@@ -258,3 +258,213 @@ class FixtureLoadTest(_Polymorphic, testing.AssertsExecutionResults):
         )
         eq_(result, [self.c1, self.c2])
 
+
+class TestGeometries(GeometryFixtureBase):
+
+    def test_threelevel_selectin_to_inline_mapped(self):
+        self._fixture_from_geometry({
+            "a": {
+                "subclasses": {
+                    "b": {"polymorphic_load": "selectin"},
+                    "c": {
+                        "subclasses": {
+                            "d": {
+                                "polymorphic_load": "inline", "single": True
+                            },
+                            "e": {
+                                "polymorphic_load": "inline", "single": True
+                            },
+                        },
+                        "polymorphic_load": "selectin",
+                    }
+                }
+            }
+        })
+
+        a, b, c, d, e = self.classes("a", "b", "c", "d", "e")
+        sess = Session()
+        sess.add_all([d(d_data="d1"), e(e_data="e1")])
+        sess.commit()
+
+        q = sess.query(a)
+
+        result = self.assert_sql_execution(
+            testing.db,
+            q.all,
+            CompiledSQL(
+                "SELECT a.type AS a_type, a.id AS a_id, "
+                "a.a_data AS a_a_data FROM a",
+                {}
+            ),
+            Or(
+                CompiledSQL(
+                    "SELECT a.type AS a_type, c.id AS c_id, a.id AS a_id, "
+                    "c.c_data AS c_c_data, c.e_data AS c_e_data, "
+                    "c.d_data AS c_d_data "
+                    "FROM a JOIN c ON a.id = c.id "
+                    "WHERE a.id IN ([EXPANDING_primary_keys]) ORDER BY a.id",
+                    [{'primary_keys': [1, 2]}]
+                ),
+                CompiledSQL(
+                    "SELECT a.type AS a_type, c.id AS c_id, a.id AS a_id, "
+                    "c.c_data AS c_c_data, "
+                    "c.d_data AS c_d_data, c.e_data AS c_e_data "
+                    "FROM a JOIN c ON a.id = c.id "
+                    "WHERE a.id IN ([EXPANDING_primary_keys]) ORDER BY a.id",
+                    [{'primary_keys': [1, 2]}]
+                )
+            )
+        )
+        with self.assert_statement_count(testing.db, 0):
+            eq_(
+                result,
+                [d(d_data="d1"), e(e_data="e1")]
+            )
+
+    def test_threelevel_selectin_to_inline_options(self):
+        self._fixture_from_geometry({
+            "a": {
+                "subclasses": {
+                    "b": {},
+                    "c": {
+                        "subclasses": {
+                            "d": {
+                                "single": True
+                            },
+                            "e": {
+                                "single": True
+                            },
+                        },
+                    }
+                }
+            }
+        })
+
+        a, b, c, d, e = self.classes("a", "b", "c", "d", "e")
+        sess = Session()
+        sess.add_all([d(d_data="d1"), e(e_data="e1")])
+        sess.commit()
+
+        c_alias = with_polymorphic(c, (d, e))
+        q = sess.query(a).options(
+            selectin_polymorphic(a, [b, c_alias])
+        )
+
+        result = self.assert_sql_execution(
+            testing.db,
+            q.all,
+            CompiledSQL(
+                "SELECT a.type AS a_type, a.id AS a_id, "
+                "a.a_data AS a_a_data FROM a",
+                {}
+            ),
+            Or(
+                CompiledSQL(
+                    "SELECT a.type AS a_type, c.id AS c_id, a.id AS a_id, "
+                    "c.c_data AS c_c_data, c.e_data AS c_e_data, "
+                    "c.d_data AS c_d_data "
+                    "FROM a JOIN c ON a.id = c.id "
+                    "WHERE a.id IN ([EXPANDING_primary_keys]) ORDER BY a.id",
+                    [{'primary_keys': [1, 2]}]
+                ),
+                CompiledSQL(
+                    "SELECT a.type AS a_type, c.id AS c_id, a.id AS a_id, "
+                    "c.c_data AS c_c_data, c.d_data AS c_d_data, "
+                    "c.e_data AS c_e_data "
+                    "FROM a JOIN c ON a.id = c.id "
+                    "WHERE a.id IN ([EXPANDING_primary_keys]) ORDER BY a.id",
+                    [{'primary_keys': [1, 2]}]
+                ),
+            )
+        )
+        with self.assert_statement_count(testing.db, 0):
+            eq_(
+                result,
+                [d(d_data="d1"), e(e_data="e1")]
+            )
+
+    def test_threelevel_selectin_to_inline_awkward_alias_options(self):
+        self._fixture_from_geometry({
+            "a": {
+                "subclasses": {
+                    "b": {},
+                    "c": {
+                        "subclasses": {
+                            "d": {},
+                            "e": {},
+                        },
+                    }
+                }
+            }
+        })
+
+        a, b, c, d, e = self.classes("a", "b", "c", "d", "e")
+        sess = Session()
+        sess.add_all([d(d_data="d1"), e(e_data="e1")])
+        sess.commit()
+
+        from sqlalchemy import select
+
+        a_table, c_table, d_table, e_table = self.tables("a", "c", "d", "e")
+
+        poly = select([
+            a_table.c.id, a_table.c.type, c_table, d_table, e_table
+        ]).select_from(
+            a_table.join(c_table).outerjoin(d_table).outerjoin(e_table)
+        ).apply_labels().alias('poly')
+
+        c_alias = with_polymorphic(c, (d, e), poly)
+        q = sess.query(a).options(
+            selectin_polymorphic(a, [b, c_alias])
+        ).order_by(a.id)
+
+        result = self.assert_sql_execution(
+            testing.db,
+            q.all,
+            CompiledSQL(
+                "SELECT a.type AS a_type, a.id AS a_id, "
+                "a.a_data AS a_a_data FROM a ORDER BY a.id",
+                {}
+            ),
+            Or(
+                # here, the test is that the adaptation of "a" takes place
+                CompiledSQL(
+                    "SELECT poly.a_type AS poly_a_type, "
+                    "poly.c_id AS poly_c_id, "
+                    "poly.a_id AS poly_a_id, poly.c_c_data AS poly_c_c_data, "
+                    "poly.e_id AS poly_e_id, poly.e_e_data AS poly_e_e_data, "
+                    "poly.d_id AS poly_d_id, poly.d_d_data AS poly_d_d_data "
+                    "FROM (SELECT a.id AS a_id, a.type AS a_type, "
+                    "c.id AS c_id, "
+                    "c.c_data AS c_c_data, d.id AS d_id, "
+                    "d.d_data AS d_d_data, "
+                    "e.id AS e_id, e.e_data AS e_e_data FROM a JOIN c "
+                    "ON a.id = c.id LEFT OUTER JOIN d ON c.id = d.id "
+                    "LEFT OUTER JOIN e ON c.id = e.id) AS poly "
+                    "WHERE poly.a_id IN ([EXPANDING_primary_keys]) "
+                    "ORDER BY poly.a_id",
+                    [{'primary_keys': [1, 2]}]
+                ),
+                CompiledSQL(
+                    "SELECT poly.a_type AS poly_a_type, "
+                    "poly.c_id AS poly_c_id, "
+                    "poly.a_id AS poly_a_id, poly.c_c_data AS poly_c_c_data, "
+                    "poly.d_id AS poly_d_id, poly.d_d_data AS poly_d_d_data, "
+                    "poly.e_id AS poly_e_id, poly.e_e_data AS poly_e_e_data "
+                    "FROM (SELECT a.id AS a_id, a.type AS a_type, "
+                    "c.id AS c_id, c.c_data AS c_c_data, d.id AS d_id, "
+                    "d.d_data AS d_d_data, e.id AS e_id, "
+                    "e.e_data AS e_e_data FROM a JOIN c ON a.id = c.id "
+                    "LEFT OUTER JOIN d ON c.id = d.id "
+                    "LEFT OUTER JOIN e ON c.id = e.id) AS poly "
+                    "WHERE poly.a_id IN ([EXPANDING_primary_keys]) "
+                    "ORDER BY poly.a_id",
+                    [{'primary_keys': [1, 2]}]
+                )
+            )
+        )
+        with self.assert_statement_count(testing.db, 0):
+            eq_(
+                result,
+                [d(d_data="d1"), e(e_data="e1")]
+            )