Merge pull request #37747 from bradleyprice/check-association-loaded-across-collection-on-preload

Check that entire collection has been loaded before short circuiting

activerecord/lib/active_record/associations/preloader.rb

@@ -177,7 +177,7 @@ def preloaded_records
         # and attach it to a relation. The class returned implements a `run` method
         # that accepts a preloader.
         def preloader_for(reflection, owners)
-          if owners.first.association(reflection.name).loaded?
+          if owners.all? { |o| o.association(reflection.name).loaded? }
             return AlreadyLoaded
           end
           reflection.check_preloadable!

activerecord/test/cases/associations/eager_test.rb

@@ -626,6 +626,21 @@ def test_eager_with_has_many_through_an_sti_join_model
     assert_equal [comments(:does_it_hurt)], assert_no_queries { author.special_post_comments }
   end
 
+  def test_preloading_with_has_one_through_an_sti_with_after_initialize
+    author_a = Author.create!(name: "A")
+    author_b = Author.create!(name: "B")
+    post_a = StiPost.create!(author: author_a, title: "TITLE", body: "BODY")
+    post_b = SpecialPost.create!(author: author_b, title: "TITLE", body: "BODY")
+    comment_a = SpecialComment.create!(post: post_a, body: "TEST")
+    comment_b = SpecialComment.create!(post: post_b, body: "TEST")
+    reset_callbacks(StiPost, :initialize) do
+      StiPost.after_initialize { author }
+      SpecialComment.where(id: [comment_a.id, comment_b.id]).includes(:author).each do |comment|
+        assert comment.author
+      end
+    end
+  end
+
   def test_preloading_has_many_through_with_implicit_source
     authors = Author.includes(:very_special_comments).to_a
     assert_no_queries do

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -2636,7 +2636,7 @@ def self.name
   end
 
   test "prevent double insertion of new object when the parent association loaded in the after save callback" do
-    reset_callbacks(:save, Bulb) do
+    reset_callbacks(Bulb, :save) do
       Bulb.after_save { |record| record.car.bulbs.load }
 
       car = Car.create!
@@ -2647,7 +2647,7 @@ def self.name
   end
 
   test "prevent double firing the before save callback of new object when the parent association saved in the callback" do
-    reset_callbacks(:save, Bulb) do
+    reset_callbacks(Bulb, :save) do
       count = 0
       Bulb.before_save { |record| record.car.save && count += 1 }
 
@@ -2702,7 +2702,7 @@ def test_ids_reader_memoization
   end
 
   def test_loading_association_in_validate_callback_doesnt_affect_persistence
-    reset_callbacks(:validation, Bulb) do
+    reset_callbacks(Bulb, :validation) do
       Bulb.after_validation { |record| record.car.bulbs.load }
 
       car = Car.create!(name: "Car")
@@ -2728,18 +2728,4 @@ def test_create_children_could_be_rolled_back_by_after_save
     def force_signal37_to_load_all_clients_of_firm
       companies(:first_firm).clients_of_firm.load_target
     end
-
-    def reset_callbacks(kind, klass)
-      old_callbacks = {}
-      old_callbacks[klass] = klass.send("_#{kind}_callbacks").dup
-      klass.subclasses.each do |subclass|
-        old_callbacks[subclass] = subclass.send("_#{kind}_callbacks").dup
-      end
-      yield
-    ensure
-      klass.send("_#{kind}_callbacks=", old_callbacks[klass])
-      klass.subclasses.each do |subclass|
-        subclass.send("_#{kind}_callbacks=", old_callbacks[subclass])
-      end
-    end
 end

activerecord/test/cases/associations/inverse_associations_test.rb

@@ -536,13 +536,6 @@ def test_inverse_instance_should_be_set_before_initialize_callbacks_are_run
       assert_predicate Man.joins(:interests).includes(:interests).first.interests, :any?
     end
   end
-
-  def reset_callbacks(target, type)
-    old_callbacks = target.send(:get_callbacks, type).deep_dup
-    yield
-  ensure
-    target.send(:set_callbacks, type, old_callbacks) if old_callbacks
-  end
 end
 
 class InverseBelongsToTests < ActiveRecord::TestCase

activerecord/test/cases/test_case.rb

@@ -81,6 +81,20 @@ def has_column?(model, column_name)
     def frozen_error_class
       Object.const_defined?(:FrozenError) ? FrozenError : RuntimeError
     end
+
+    def reset_callbacks(klass, kind)
+      old_callbacks = {}
+      old_callbacks[klass] = klass.send("_#{kind}_callbacks").dup
+      klass.subclasses.each do |subclass|
+        old_callbacks[subclass] = subclass.send("_#{kind}_callbacks").dup
+      end
+      yield
+    ensure
+      klass.send("_#{kind}_callbacks=", old_callbacks[klass])
+      klass.subclasses.each do |subclass|
+        subclass.send("_#{kind}_callbacks=", old_callbacks[subclass])
+      end
+    end
   end
 
   class PostgreSQLTestCase < TestCase

activerecord/test/models/comment.rb

@@ -59,6 +59,7 @@ def to_s
 end
 
 class SpecialComment < Comment
+  has_one :author, through: :post
   default_scope { where(deleted_at: nil) }
 
   def self.what_are_you