From 6074c597576ceef3f9368ecb7fd48ca2d4dfb759 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Fri, 19 Feb 2021 10:57:45 +0100
Subject: [PATCH] Call `run` after preloading records

Ref: https://github.com/rails/rails/pull/41385

Otherwise the association isn't marked as loaded, and a ThroughAssociation
might perform the same query uselessly.

`already_loaded?` also has to be made lazy, because #41285
made it so that alls the preloaders are instiated together before being ran.
---
 .../active_record/associations/preloader.rb   |  2 +-
 .../associations/preloader/association.rb     | 11 +++----
 activerecord/test/cases/associations_test.rb  | 29 ++++++++++++++++++-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb
index 2275de85a6..336f5bd5a4 100644
--- a/activerecord/lib/active_record/associations/preloader.rb
+++ b/activerecord/lib/active_record/associations/preloader.rb
@@ -103,7 +103,7 @@ module ActiveRecord
 
         loaders = build_preloaders
         group_and_load_similar(loaders)
-        loaders.map(&:run)
+        loaders.each(&:run)
 
         child_preloaders.each { |reflection, child, parents| build_child_preloader(reflection, child, parents) }
 
diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index b30ec816a6..166cb01b85 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -11,7 +11,10 @@ module ActiveRecord
             loaders.each { |l| l.set_inverse(record) }
           end
 
-          loaders.each { |l| l.load_records(raw_records) }
+          loaders.each do |loader|
+            loader.load_records(raw_records)
+            loader.run
+          end
         end
 
         def initialize(klass, owners, reflection, preload_scope, associate_by_default = true)
@@ -21,16 +24,14 @@ module ActiveRecord
           @preload_scope = preload_scope
           @associate     = associate_by_default || !preload_scope || preload_scope.empty_scope?
           @model         = owners.first && owners.first.class
-
-          @already_loaded = owners.all? { |o| o.association(reflection.name).loaded? }
         end
 
         def already_loaded?
-          @already_loaded
+          @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? }
         end
 
         def run
-          if @already_loaded
+          if already_loaded?
             fetch_from_preloaded_records
             return self
           end
diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb
index 84d1d61304..9747a24aee 100644
--- a/activerecord/test/cases/associations_test.rb
+++ b/activerecord/test/cases/associations_test.rb
@@ -30,7 +30,8 @@ require "models/price_estimate"
 
 class AssociationsTest < ActiveRecord::TestCase
   fixtures :accounts, :companies, :developers, :projects, :developers_projects,
-           :computers, :people, :readers, :authors, :author_addresses, :author_favorites
+           :computers, :people, :readers, :authors, :author_addresses, :author_favorites,
+           :comments, :posts
 
   def test_eager_loading_should_not_change_count_of_children
     liquid = Liquid.create(name: "salty")
@@ -407,12 +408,30 @@ class PreloaderTest < ActiveRecord::TestCase
     assert_queries(1) do
       preloader = ActiveRecord::Associations::Preloader.new(records: [book, post], associations: :author)
       preloader.call
+    end
 
+    assert_no_queries do
       book.author
       post.author
     end
   end
 
+  def test_preload_through
+    comments = [
+      comments(:eager_sti_on_associations_s_comment1),
+      comments(:eager_sti_on_associations_s_comment2),
+    ]
+
+    assert_queries(2) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: comments, associations: [:author, :post])
+      preloader.call
+    end
+
+    assert_no_queries do
+      comments.each(&:author)
+    end
+  end
+
   def test_preload_with_grouping_sets_inverse_association
     mary = authors(:mary)
     bob = authors(:bob)
@@ -423,7 +442,9 @@ class PreloaderTest < ActiveRecord::TestCase
     assert_queries(1) do
       preloader = ActiveRecord::Associations::Preloader.new(records: favorites, associations: [:author, :favorite_author])
       preloader.call
+    end
 
+    assert_no_queries do
       favorites.first.author
       favorites.first.favorite_author
     end
@@ -440,7 +461,9 @@ class PreloaderTest < ActiveRecord::TestCase
     assert_queries(2) do
       preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author_with_the_letter_a)
       preloader.call
+    end
 
+    assert_no_queries do
       post.author_with_the_letter_a
       postesque.author_with_the_letter_a
     end
@@ -452,7 +475,9 @@ class PreloaderTest < ActiveRecord::TestCase
     assert_queries(3) do
       preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author_with_address)
       preloader.call
+    end
 
+    assert_no_queries do
       post.author_with_address
       postesque.author_with_address
     end
@@ -466,7 +491,9 @@ class PreloaderTest < ActiveRecord::TestCase
     assert_queries(2) do
       preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author)
       preloader.call
+    end
 
+    assert_no_queries do
       post.author
       postesque.author
     end