Can preload associations through polymorphic associations

1 files changed, 15 insertions, 10 deletions

diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb
index 1ea0aeac3a..5c2ac5b374 100644
--- a/activerecord/lib/active_record/associations/preloader.rb
+++ b/activerecord/lib/active_record/associations/preloader.rb
@@ -98,26 +98,30 @@ module ActiveRecord
       private
 
         # Loads all the given data into +records+ for the +association+.
-        def preloaders_on(association, records, scope)
+        def preloaders_on(association, records, scope, polymorphic_parent = false)
           case association
           when Hash
-            preloaders_for_hash(association, records, scope)
+            preloaders_for_hash(association, records, scope, polymorphic_parent)
           when Symbol
-            preloaders_for_one(association, records, scope)
+            preloaders_for_one(association, records, scope, polymorphic_parent)
           when String
-            preloaders_for_one(association.to_sym, records, scope)
+            preloaders_for_one(association.to_sym, records, scope, polymorphic_parent)
           else
             raise ArgumentError, "#{association.inspect} was not recognized for preload"
           end
         end
 
-        def preloaders_for_hash(association, records, scope)
+        def preloaders_for_hash(association, records, scope, polymorphic_parent)
           association.flat_map { |parent, child|
-            loaders = preloaders_for_one parent, records, scope
+            loaders = preloaders_for_one parent, records, scope, polymorphic_parent
 
             recs = loaders.flat_map(&:preloaded_records).uniq
+
+            reflection = records.first.class._reflect_on_association(parent)
+            polymorphic_parent = reflection && reflection.options[:polymorphic]
+
             loaders.concat Array.wrap(child).flat_map { |assoc|
-              preloaders_on assoc, recs, scope
+              preloaders_on assoc, recs, scope, polymorphic_parent
             }
             loaders
           }
@@ -135,8 +139,8 @@ module ActiveRecord
         # Additionally, polymorphic belongs_to associations can have multiple associated
         # classes, depending on the polymorphic_type field. So we group by the classes as
         # well.
-        def preloaders_for_one(association, records, scope)
-          grouped_records(association, records).flat_map do |reflection, klasses|
+        def preloaders_for_one(association, records, scope, polymorphic_parent)
+          grouped_records(association, records, polymorphic_parent).flat_map do |reflection, klasses|
             klasses.map do |rhs_klass, rs|
               loader = preloader_for(reflection, rs).new(rhs_klass, rs, reflection, scope)
               loader.run self
@@ -145,10 +149,11 @@ module ActiveRecord
           end
         end
 
-        def grouped_records(association, records)
+        def grouped_records(association, records, polymorphic_parent)
           h = {}
           records.each do |record|
             next unless record
+            next if polymorphic_parent && !record.class._reflect_on_association(association)
             assoc = record.association(association)
             next unless assoc.klass
             klasses = h[assoc.reflection] ||= {}