Merge pull request #32655 from robotdana/preload-through-polymorphic-associations

Preloading through polymorphic associations (still raises for typos on regular associations)

4 files changed, 51 insertions, 10 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index c4a4e767cd..0d8fa48235 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,7 @@
+*   Add support to preload associations of polymorphic associations when not all the records have the requested associations.
+
+    *Dana Sherson*
+
 *   Add `touch_all` method to `ActiveRecord::Relation`.
 
     Example:
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] ||= {}
diff --git a/activerecord/test/cases/associations/eager_test.rb b/activerecord/test/cases/associations/eager_test.rb
index 6e0cf30092..f46be8734b 100644
--- a/activerecord/test/cases/associations/eager_test.rb
+++ b/activerecord/test/cases/associations/eager_test.rb
@@ -1515,6 +1515,35 @@ class EagerAssociationTest < ActiveRecord::TestCase
     Author.preload(:readonly_comments).first!
   end
 
+  test "preloading through a polymorphic association doesn't require the association to exist" do
+    sponsors = []
+    assert_queries 5 do
+      sponsors = Sponsor.where(sponsorable_id: 1).preload(sponsorable: [:post, :membership]).to_a
+    end
+    # check the preload worked
+    assert_queries 0 do
+      sponsors.map(&:sponsorable).map { |s| s.respond_to?(:posts) ? s.post.author : s.membership }
+    end
+  end
+
+  test "preloading a regular association through a polymorphic association doesn't require the association to exist on all types" do
+    sponsors = []
+    assert_queries 6 do
+      sponsors = Sponsor.where(sponsorable_id: 1).preload(sponsorable: [{ post: :first_comment }, :membership]).to_a
+    end
+    # check the preload worked
+    assert_queries 0 do
+      sponsors.map(&:sponsorable).map { |s| s.respond_to?(:posts) ? s.post.author : s.membership }
+    end
+  end
+
+  test "preloading a regular association with a typo through a polymorphic association still raises" do
+    # this test contains an intentional typo of first -> fist
+    assert_raises(ActiveRecord::AssociationNotFoundError) do
+      Sponsor.where(sponsorable_id: 1).preload(sponsorable: [{ post: :fist_comment }, :membership]).to_a
+    end
+  end
+
   private
     def find_all_ordered(klass, include = nil)
       klass.order("#{klass.table_name}.#{klass.primary_key}").includes(include).to_a
diff --git a/activerecord/test/fixtures/sponsors.yml b/activerecord/test/fixtures/sponsors.yml
index 2da541c539..02ddb8dd38 100644
--- a/activerecord/test/fixtures/sponsors.yml
+++ b/activerecord/test/fixtures/sponsors.yml
@@ -10,3 +10,6 @@ crazy_club_sponsor_for_groucho:
   sponsor_club: crazy_club
   sponsorable_id: 3
   sponsorable_type: Member
+sponsor_for_author_david:
+  sponsorable_id: 1
+  sponsorable_type: Author