include columns from superclasses that indicate "selectin"

zzzeek · zzzeek · commit 40b00498e62d · 2023-02-27T01:35:26.000-05:00
Added support for the :paramref:`_orm.Mapper.polymorphic_load` parameter to be applied to each mapper in an inheritance hierarchy more than one level deep, allowing columns to load for all classes in the hierarchy that indicate ``"selectin"`` using a single statement, rather than ignoring elements on those intermediary classes that nonetheless indicate they also would participate in ``"selectin"`` loading and were not part of the base-most SELECT statement. Fixes: #9373 Change-Id: If8dcba0f0191f6c2818ecd15870bccfdf5ce1112
diff --git a/doc/build/changelog/unreleased_20/9373.rst b/doc/build/changelog/unreleased_20/9373.rst
@@ -0,0 +1,11 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 9373
+
+    Added support for the :paramref:`_orm.Mapper.polymorphic_load` parameter to
+    be applied to each mapper in an inheritance hierarchy more than one level
+    deep, allowing columns to load for all classes in the hierarchy that
+    indicate ``"selectin"`` using a single statement, rather than ignoring
+    elements on those intermediary classes that nonetheless indicate they also
+    would participate in ``"selectin"`` loading and were not part of the
+    base-most SELECT statement.
diff --git a/lib/sqlalchemy/orm/loading.py b/lib/sqlalchemy/orm/loading.py
@@ -972,26 +972,31 @@ def _instance_processor(
 
     if not refresh_state and _polymorphic_from is not None:
         key = ("loader", path.path)
+
         if key in context.attributes and context.attributes[key].strategy == (
             ("selectinload_polymorphic", True),
         ):
-            selectin_load_via = mapper._should_selectin_load(
-                context.attributes[key].local_opts["entities"],
-                _polymorphic_from,
-            )
+            option_entities = context.attributes[key].local_opts["entities"]
         else:
-            selectin_load_via = mapper._should_selectin_load(
-                None, _polymorphic_from
-            )
+            option_entities = None
+        selectin_load_via = mapper._should_selectin_load(
+            option_entities,
+            _polymorphic_from,
+        )
 
         if selectin_load_via and selectin_load_via is not _polymorphic_from:
             # only_load_props goes w/ refresh_state only, and in a refresh
             # we are a single row query for the exact entity; polymorphic
             # loading does not apply
             assert only_load_props is None
 
-            callable_ = _load_subclass_via_in(context, path, selectin_load_via)
-
+            callable_ = _load_subclass_via_in(
+                context,
+                path,
+                selectin_load_via,
+                _polymorphic_from,
+                option_entities,
+            )
             PostLoad.callable_for_path(
                 context,
                 load_path,
@@ -1212,17 +1217,42 @@ def ensure_no_pk(row):
     return _instance
 
 
-def _load_subclass_via_in(context, path, entity):
+def _load_subclass_via_in(
+    context, path, entity, polymorphic_from, option_entities
+):
     mapper = entity.mapper
 
+    # TODO: polymorphic_from seems to be a Mapper in all cases.
+    # this is likely not needed, but as we dont have typing in loading.py
+    # yet, err on the safe side
+    polymorphic_from_mapper = polymorphic_from.mapper
+    not_against_basemost = polymorphic_from_mapper.inherits is not None
+
     zero_idx = len(mapper.base_mapper.primary_key) == 1
 
-    if entity.is_aliased_class:
-        q, enable_opt, disable_opt = mapper._subclass_load_via_in(entity)
+    if entity.is_aliased_class or not_against_basemost:
+        q, enable_opt, disable_opt = mapper._subclass_load_via_in(
+            entity, polymorphic_from
+        )
     else:
         q, enable_opt, disable_opt = mapper._subclass_load_via_in_mapper
 
     def do_load(context, path, states, load_only, effective_entity):
+        if not option_entities:
+            # filter out states for those that would have selectinloaded
+            # from another loader
+            # TODO: we are currently ignoring the case where the
+            # "selectin_polymorphic" option is used, as this is much more
+            # complex / specific / very uncommon API use
+            states = [
+                (s, v)
+                for s, v in states
+                if s.mapper._would_selectin_load_only_from_given_mapper(mapper)
+            ]
+
+            if not states:
+                return
+
         orig_query = context.query
 
         options = (enable_opt,) + orig_query._with_options + (disable_opt,)
diff --git a/lib/sqlalchemy/orm/mapper.py b/lib/sqlalchemy/orm/mapper.py
@@ -3698,6 +3698,65 @@ def _iterate_to_target_viawpoly(self, mapper):
                 if m is mapper:
                     break
 
+    @HasMemoized.memoized_attribute
+    def _would_selectinload_combinations_cache(self):
+        return {}
+
+    def _would_selectin_load_only_from_given_mapper(self, super_mapper):
+        """return True if this mapper would "selectin" polymorphic load based
+        on the given super mapper, and not from a setting from a subclass.
+
+        given::
+
+            class A:
+                ...
+
+            class B(A):
+                __mapper_args__ = {"polymorphic_load": "selectin"}
+
+            class C(B):
+                ...
+
+            class D(B):
+                __mapper_args__ = {"polymorphic_load": "selectin"}
+
+        ``inspect(C)._would_selectin_load_only_from_given_mapper(inspect(B))``
+        returns True, because C does selectin loading because of B's setting.
+
+        OTOH, ``inspect(D)
+        ._would_selectin_load_only_from_given_mapper(inspect(B))``
+        returns False, because D does selectin loading because of its own
+        setting; when we are doing a selectin poly load from B, we want to
+        filter out D because it would already have its own selectin poly load
+        set up separately.
+
+        Added as part of #9373.
+
+        """
+        cache = self._would_selectinload_combinations_cache
+
+        try:
+            return cache[super_mapper]
+        except KeyError:
+            pass
+
+        # assert that given object is a supermapper, meaning we already
+        # strong reference it directly or indirectly.  this allows us
+        # to not worry that we are creating new strongrefs to unrelated
+        # mappers or other objects.
+        assert self.isa(super_mapper)
+
+        mapper = super_mapper
+        for m in self._iterate_to_target_viawpoly(mapper):
+            if m.polymorphic_load == "selectin":
+                retval = m is super_mapper
+                break
+        else:
+            retval = False
+
+        cache[super_mapper] = retval
+        return retval
+
     def _should_selectin_load(self, enabled_via_opt, polymorphic_from):
         if not enabled_via_opt:
             # common case, takes place for all polymorphic loads
@@ -3721,7 +3780,7 @@ def _should_selectin_load(self, enabled_via_opt, polymorphic_from):
         return None
 
     @util.preload_module("sqlalchemy.orm.strategy_options")
-    def _subclass_load_via_in(self, entity):
+    def _subclass_load_via_in(self, entity, polymorphic_from):
         """Assemble a that can load the columns local to
         this subclass as a SELECT with IN.
 
@@ -3739,6 +3798,16 @@ def _subclass_load_via_in(self, entity):
         disable_opt = strategy_options.Load(entity)
         enable_opt = strategy_options.Load(entity)
 
+        classes_to_include = {self}
+        m: Optional[Mapper[Any]] = self.inherits
+        while (
+            m is not None
+            and m is not polymorphic_from
+            and m.polymorphic_load == "selectin"
+        ):
+            classes_to_include.add(m)
+            m = m.inherits
+
         for prop in self.attrs:
 
             # skip prop keys that are not instrumented on the mapped class.
@@ -3747,7 +3816,7 @@ def _subclass_load_via_in(self, entity):
             if prop.key not in self.class_manager:
                 continue
 
-            if prop.parent is self or prop in keep_props:
+            if prop.parent in classes_to_include or prop in keep_props:
                 # "enable" options, to turn on the properties that we want to
                 # load by default (subject to options from the query)
                 if not isinstance(prop, StrategizedProperty):
@@ -3811,7 +3880,8 @@ def _subclass_load_via_in(self, entity):
 
     @HasMemoized.memoized_attribute
     def _subclass_load_via_in_mapper(self):
-        return self._subclass_load_via_in(self)
+        # the default is loading this mapper against the basemost mapper
+        return self._subclass_load_via_in(self, self.base_mapper)
 
     def cascade_iterator(
         self,
diff --git a/test/orm/inheritance/test_poly_loading.py b/test/orm/inheritance/test_poly_loading.py
