- polymorphic loading, working very rudimentallyrow_proc_integration

running orm2010, the whole approach adds callcounts.
code is more complex, less reliable, pretty much a total
bust.

will try to see if create_row_processor can remain but
be called against either "result" or a context namespace
ahead of time, and if this somehow can be used to simplify
the need to call upon strategy lookup twice.

5 files changed, 173 insertions, 79 deletions

diff --git a/lib/sqlalchemy/orm/interfaces.py b/lib/sqlalchemy/orm/interfaces.py
index d12c20a2c..45ab9b48f 100644
--- a/lib/sqlalchemy/orm/interfaces.py
+++ b/lib/sqlalchemy/orm/interfaces.py
@@ -107,7 +107,9 @@ class MapperProperty(_MappedAttribute, InspectionAttr, util.MemoizedSlots):
         """
         return {}
 
-    def setup(self, context, entity, path, adapter, **kwargs):
+    def setup(
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, **kw):
         """Called by Query for the purposes of constructing a SQL statement.
 
         Each MapperProperty associated with the target mapper processes the
@@ -116,6 +118,17 @@ class MapperProperty(_MappedAttribute, InspectionAttr, util.MemoizedSlots):
 
         """
 
+    def setup_for_missing_attribute(
+            self, context, query_entity, path, mapper, populators, **kw):
+        """Setup a strategy for a Query where this property was not yet
+        included.
+
+        This function can do everything that setup() does, *except* attempt
+        to modify the SQL query; the method may be called after the query
+        has already been emitted and results are being received.
+
+        """
+
     def cascade_iterator(self, type_, state, visited_instances=None,
                          halt_on=None):
         """Iterate through instances related to the given instance for
@@ -481,13 +494,31 @@ class StrategizedProperty(MapperProperty):
     def _get_strategy_by_cls(self, cls):
         return self._get_strategy(cls._strategy_keys[0])
 
-    def setup(self, context, entity, path, adapter, **kwargs):
-        loader = self._get_context_loader(context, path)
-        if loader and loader.strategy:
-            strat = self._get_strategy(loader.strategy)
+    def setup(
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, **kw):
+
+        loadopt = self._get_context_loader(context, path)
+        if loadopt and loadopt.strategy:
+            strat = self._get_strategy(loadopt.strategy)
         else:
             strat = self.strategy
-        strat.setup_query(context, entity, path, loader, adapter, **kwargs)
+
+        strat.setup_query(
+            context, query_entity, path, mapper,
+            adapter, column_collection, populators, loadopt, **kw)
+
+    def setup_for_missing_attribute(
+            self, context, query_entity, path, mapper,
+            populators, **kw):
+
+        loadopt = self._get_context_loader(context, path)
+        if loadopt and loadopt.strategy:
+            strat = self._get_strategy(loadopt.strategy)
+        else:
+            strat = self.strategy
+        strat.setup_for_missing_attribute(
+            context, query_entity, path, mapper, populators, loadopt, **kw)
 
     def do_init(self):
         self._strategies = {}
@@ -585,7 +616,9 @@ class LoaderStrategy(object):
     def init_class_attribute(self, mapper):
         pass
 
-    def setup_query(self, context, entity, path, loadopt, adapter, **kwargs):
+    def setup_query(
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, loadopt, **kw):
         """Establish column and other state for a given QueryContext.
 
         This method fulfills the contract specified by MapperProperty.setup().
@@ -595,5 +628,14 @@ class LoaderStrategy(object):
 
         """
 
+    def setup_for_missing_attribute(
+            self, context, query_entity, path, mapper,
+            populators, loadopt, **kw):
+        """Establish loader behavior for an attribute that's not accommodated
+        by the query.
+
+        This is used for polymorphic loading when a subclass load is detected.
+        """
+
     def __str__(self):
         return str(self.parent_property)
diff --git a/lib/sqlalchemy/orm/loading.py b/lib/sqlalchemy/orm/loading.py
index 3d5aa4fb4..add2ae43b 100644
--- a/lib/sqlalchemy/orm/loading.py
+++ b/lib/sqlalchemy/orm/loading.py
@@ -43,7 +43,10 @@ def instances(query, cursor, context):
             def filter_fn(row):
                 return tuple(fn(x) for x, fn in zip(row, filter_fns))
 
-    if context.statement is not None:
+    if context._predefined_statement:
+        # if the Query didn't actually build the statement,
+        # we use the result set to determine where the columns
+        # we're looking for are located.
         context._setup_column_processors(
             [
                 (cursor._index_of(col), col)
@@ -223,13 +226,14 @@ def load_on_ident(query, key,
         return None
 
 
-def instance_processor(mapper, props_toload, context, column_collection,
-                       query_entity, path, adapter,
-                       only_load_props=None, refresh_state=None,
-                       polymorphic_discriminator=None,
-                       _polymorphic_from=None,
-                       _polymorphic_pk_getters=None,
-                       _polymorphic_from_populators=None):
+def _instance_processor(
+    mapper, props_toload, context, column_collection,
+    query_entity, path, adapter,
+    only_load_props=None, refresh_state=None,
+    polymorphic_discriminator=None,
+        _polymorphic_from=None,
+        _polymorphic_pk_getters=None,
+        _polymorphic_from_populators=None):
     """Produce a mapper level row processor callable
        which processes rows into mapped instances."""
 
@@ -240,6 +244,13 @@ def instance_processor(mapper, props_toload, context, column_collection,
         # of populators that were already set up for us.
 
         populators = _polymorphic_from_populators
+
+        for prop in props_toload:
+            prop.setup_for_missing_attribute(
+                context, query_entity, path, mapper,
+                _polymorphic_from_populators
+            )
+
     else:
         populators = collections.defaultdict(list)
 
@@ -257,12 +268,13 @@ def instance_processor(mapper, props_toload, context, column_collection,
                 context,
                 query_entity,
                 path,
+                mapper,
                 adapter,
-                only_load_props=only_load_props,
-                column_collection=column_collection,
+                column_collection,
                 populators=per_mapper_populators[prop.parent]
-                if load_is_polymorphic
-                else populators
+                if load_is_polymorphic and not mapper.isa(prop.parent)
+                else populators,
+                only_load_props=only_load_props,
             )
 
     if _polymorphic_pk_getters:
@@ -428,7 +440,8 @@ def instance_processor(mapper, props_toload, context, column_collection,
         # if we are doing polymorphic, dispatch to a different _instance()
         # method specific to the subclass mapper
         _instance = _decorate_polymorphic_switch(
-            _instance, context, mapper, per_mapper_populators, path,
+            _instance, context, mapper, props_toload,
+            per_mapper_populators, path,
             polymorphic_discriminator, adapter, pk_getters)
 
     return _instance
@@ -516,7 +529,8 @@ def _validate_version_id(mapper, state, dict_, row, adapter):
 
 
 def _decorate_polymorphic_switch(
-        instance_fn, context, mapper, per_mapper_populators, path,
+        instance_fn, context, mapper, props_toload,
+        per_mapper_populators, path,
         polymorphic_discriminator, adapter, pk_getters):
 
     if polymorphic_discriminator is not None:
@@ -535,6 +549,8 @@ def _decorate_polymorphic_switch(
     context.column_processors.append(
         (polymorphic_on, polymorphic_getter.setup))
 
+    props_setup = set(props_toload)
+
     def configure_subclass_mapper(discriminator):
         try:
             sub_mapper = mapper.polymorphic_map[discriminator]
@@ -554,27 +570,18 @@ def _decorate_polymorphic_switch(
                 if super_mapper is mapper:
                     break
 
-            # TODO!
-            # big problems:
-            # 1. "quick" is being multiply populated with redundant
-            # populators
-            # 2. columns like "golf_swing", which are not rendered in
-            # setup(), therefore have no populator at all, we normally
-            # are expecting an "expire" populator to set up for a deferred
-            # load.   We need to either make it so these populators aren't
-            # needed or
-            # that we in here do actually add more non-column populators,
-            # which may mean that we need some version of
-            # row_processor() again for this case.   It would be
-            # along the lines of missing_attribute_populator() and would be
-            # specific to those cases where we have to produce a subclass
-            # against a query that did not specify this class in its
-            # entities.
-            # if discriminator == 'boss':
-            #    import pdb
-            #    pdb.set_trace()
-            return instance_processor(
-                sub_mapper, None, context, None, None,
+            keys_setup = set(p.key for p in props_setup)
+
+            props_needed = set(
+                prop for prop in sub_mapper._props.values()
+            ).difference(props_setup)
+
+            props_needed = props_needed.difference(
+                p for p in props_needed if p.key in keys_setup
+            )
+
+            return _instance_processor(
+                sub_mapper, props_needed, context, None, None,
                 path, adapter, _polymorphic_from=mapper,
                 _polymorphic_pk_getters=pk_getters,
                 _polymorphic_from_populators=populators)
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 53cf227ee..147b4cf0c 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -18,7 +18,9 @@ ORM session, whereas the ``Select`` construct interacts directly with the
 database to return iterable result sets.
 
 """
+from __future__ import absolute_import
 
+import collections
 from itertools import chain
 
 from . import (
@@ -2942,6 +2944,7 @@ class Query(object):
             entity.setup_context(self, context)
 
         if context.statement is not None:
+            context._predefined_statement = True
             return context
 
         for rec in context.create_eager_joins:
@@ -3044,6 +3047,13 @@ class Query(object):
         statement.append_order_by(*context.eager_order_by)
 
         context._setup_column_processors(
+            enumerate(
+                col for (label, col) in statement._columns_plus_names
+            ),
+            outer_adapter
+        )
+
+        context._setup_column_processors(
             enumerate(context.primary_columns, 0)
         )
         context._setup_column_processors(
@@ -3093,7 +3103,7 @@ class Query(object):
         # that have been established
         context._setup_column_processors(
             enumerate(
-                context.primary_columns + context.secondary_columns
+                col for (label, col) in statement._columns_plus_names
             )
         )
 
@@ -3317,7 +3327,7 @@ class _MapperEntity(_QueryEntity):
         else:
             only_load_props = refresh_state = None
 
-        _instance = loading.instance_processor(
+        _instance = loading._instance_processor(
             self.mapper,
             props_toload,
             context,
@@ -3331,18 +3341,6 @@ class _MapperEntity(_QueryEntity):
         )
 
         context.loaders.append((self._label_name, _instance))
-        # TODO: this needs to be in instance_processor()
-        # and needs a getter fn.   a special entry in
-        # populators should be used here
-        # if self._polymorphic_discriminator is not None and \
-        #    self._polymorphic_discriminator \
-        #        is not self.mapper.polymorphic_on:
-        #
-        #    if adapter:
-        #        pd = adapter.columns[self._polymorphic_discriminator]
-        #    else:
-        #        pd = self._polymorphic_discriminator
-        #    context.primary_columns.append(pd)
 
     def __str__(self):
         return str(self.mapper)
@@ -3690,6 +3688,7 @@ class QueryContext(object):
     froms = ()
     for_update = None
     order_by = False
+    _predefined_statement = False
 
     def __init__(self, query):
 
@@ -3723,11 +3722,14 @@ class QueryContext(object):
         self.attributes = query._attributes.copy()
         self.loaders = []
 
-    def _setup_column_processors(self, cols):
-        d = dict(
-            (col, idx) for idx, col in cols
+    def _setup_column_processors(self, cols, adapter=None):
+        d = collections.defaultdict(
+            lambda: -1,
+            [(col, idx) for idx, col in cols]
         )
         for col, fn in self.column_processors:
+            if adapter:
+                col = adapter.columns[col]
             fn(d[col])
 
 
diff --git a/lib/sqlalchemy/orm/strategies.py b/lib/sqlalchemy/orm/strategies.py
index 085faaa2d..aab455771 100644
--- a/lib/sqlalchemy/orm/strategies.py
+++ b/lib/sqlalchemy/orm/strategies.py
@@ -26,6 +26,7 @@ from .session import _state_session
 import itertools
 import operator
 
+
 def _register_attribute(
     strategy, mapper, useobject,
     compare_function=None,
@@ -112,8 +113,8 @@ class UninstrumentedColumnLoader(LoaderStrategy):
         self.columns = self.parent_property.columns
 
     def setup_query(
-            self, context, entity, path, loadopt, adapter,
-            column_collection=None, **kwargs):
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, loadopt, **kw):
         for c in self.columns:
             if adapter:
                 c = adapter.columns[c]
@@ -133,8 +134,8 @@ class ColumnLoader(LoaderStrategy):
         self.is_composite = hasattr(self.parent_property, 'composite_class')
 
     def setup_query(
-            self, context, entity, path, loadopt,
-            adapter, column_collection, populators, **kwargs):
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, loadopt, **kw):
         for c in self.columns:
             if adapter:
                 c = adapter.columns[c]
@@ -148,11 +149,19 @@ class ColumnLoader(LoaderStrategy):
             # querying out every column.
 
         def quick_populate(index):
-            populators["quick"].append(
-                (self.key, operator.itemgetter(index))
-            )
+            if index == -1:
+                populators["expire"].append((self.key, True))
+            else:
+                populators["quick"].append(
+                    (self.key, operator.itemgetter(index))
+                )
         context.column_processors.append((self.columns[0], quick_populate))
 
+    def setup_for_missing_attribute(
+            self, context, query_entity, path, mapper,
+            populators, loadopt, **kw):
+        populators["expire"].append((self.key, True))
+
     def init_class_attribute(self, mapper):
         self.is_class_level = True
         coltype = self.columns[0].type
@@ -194,8 +203,9 @@ class DeferredColumnLoader(LoaderStrategy):
         )
 
     def setup_query(
-            self, context, entity, path, loadopt, adapter,
-            populators, only_load_props=None, **kwargs):
+            self, context, query_entity, path, mapper,
+            adapter, column_collection, populators, loadopt,
+            only_load_props=None, **kw):
 
         if (
             (
@@ -216,10 +226,19 @@ class DeferredColumnLoader(LoaderStrategy):
             )
         ):
             self.parent_property._get_strategy_by_cls(ColumnLoader).\
-                setup_query(context, entity,
-                            path, loadopt, adapter,
-                            populators=populators, **kwargs)
-        elif not self.is_class_level:
+                setup_query(
+                    context, query_entity,
+                    path, mapper, adapter,
+                    populators, loadopt, **kw)
+        else:
+            self.setup_for_missing_attribute(
+                context, query_entity, path, mapper, populators, loadopt, **kw
+            )
+
+    def setup_for_missing_attribute(
+            self, context, query_entity, path, mapper, populators,
+            loadopt, **kw):
+        if not self.is_class_level:
             set_deferred_for_local_state = \
                 InstanceState._instance_level_callable_processor(
                     self.parent.class_manager,
@@ -1432,7 +1451,7 @@ class JoinedLoader(AbstractRelationshipLoader):
         if eager_adapter is not False:
             key = self.key
 
-            _instance = loading.instance_processor(
+            _instance = loading._instance_processor(
                 self.mapper,
                 context,
                 result,
diff --git a/test/orm/test_loading.py b/test/orm/test_loading.py
index 2ba741b53..bfc1ad0c8 100644
--- a/test/orm/test_loading.py
+++ b/test/orm/test_loading.py
@@ -124,7 +124,7 @@ class InstancesTest(_fixtures.FixtureTest):
         )
 
 
-class PolymorphicInstancesTest(_poly_fixtures._Polymorphic):
+class _PolymorphicInstancesTest(_poly_fixtures._PolymorphicFixtureBase):
     def test_query_load_entity(self):
         Person, Engineer, Manager, Boss = (
             _poly_fixtures.Person, _poly_fixtures.Engineer,
@@ -138,15 +138,39 @@ class PolymorphicInstancesTest(_poly_fixtures._Polymorphic):
         eq_(
             rows,
             [
-                Engineer(name='dilbert'),
-                Engineer(name='wally'),
-                Boss(name='pointy haired boss', golf_swing='fore!'),
-                Manager(manager_name='dogbert'),
-                Engineer(name='vlad')
+                Engineer(
+                    name='dilbert', primary_language='java',
+                    status='regular engineer'),
+                Engineer(
+                    name='wally', primary_language='c++',
+                    status='regular engineer'),
+                Boss(
+                    name='pointy haired boss', golf_swing='fore',
+                    manager_name='pointy', status='da boss'),
+                Manager(
+                    name='dogbert', manager_name='dogbert',
+                    status='regular manager'),
+                Engineer(
+                    name='vlad', primary_language='cobol',
+                    status='elbonian engineer')
             ]
         )
 
 
+class PolymorphicInstancesDeferredTest(
+        _poly_fixtures._Polymorphic, _PolymorphicInstancesTest):
+    """test polymorphic loading with missing attributes on subclasses.
+
+    """
+
+
+class PolymorphicInstancesPolymorphicTest(
+        _poly_fixtures._PolymorphicPolymorphic, _PolymorphicInstancesTest):
+    """test polymorphic loading with all attributes in the query.
+
+    """
+
+
 class MergeResultTest(_fixtures.FixtureTest):
     run_setup_mappers = 'once'
     run_inserts = 'once'