Allow preload grouping of through associations

Co-authored-by: Dinah Shi <dinahshi@github.com>

activerecord/lib/active_record/associations/preloader/association.rb

@@ -31,6 +31,10 @@ module ActiveRecord
           @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? }
         end
 
+        def runnable_loaders
+          [self]
+        end
+
         def run?
           @run
         end

activerecord/lib/active_record/associations/preloader/batch.rb

@@ -9,15 +9,19 @@ module ActiveRecord
         end
 
         def call
-          return if @preloaders.empty?
-
           branches = @preloaders.flat_map(&:branches)
           until branches.empty?
-            loaders = branches.flat_map(&:loaders)
+            loaders = branches.flat_map(&:runnable_loaders)
+
+            already_loaded, loaders = loaders.partition(&:already_loaded?)
+            already_loaded.each(&:run)
+
             group_and_load_similar(loaders)
             loaders.each(&:run)
 
-            branches = branches.flat_map(&:children)
+            finished, in_progress = branches.partition(&:done?)
+
+            branches = in_progress + finished.flat_map(&:children)
           end
         end
 
@@ -26,8 +30,6 @@ module ActiveRecord
 
           def group_and_load_similar(loaders)
             loaders.grep_v(ThroughAssociation).group_by(&:grouping_key).each do |(_, _, association_key_name), similar_loaders|
-              next if similar_loaders.all? { |l| l.already_loaded? }
-
               scope = similar_loaders.first.scope
               Association.load_records_in_batch(scope, association_key_name, similar_loaders)
             end

activerecord/lib/active_record/associations/preloader/branch.rb

@@ -33,6 +33,10 @@ module ActiveRecord
           loaders.all?(&:run?)
         end
 
+        def runnable_loaders
+          loaders.flat_map(&:runnable_loaders).reject(&:run?)
+        end
+
         def grouped_records
           h = {}
           source_records.each do |record|

activerecord/lib/active_record/associations/preloader/through_association.rb

@@ -35,9 +35,19 @@ module ActiveRecord
           end
         end
 
+        def runnable_loaders
+          if already_loaded?
+            [self]
+          elsif through_preloaders.all?(&:run?)
+            [self] + source_preloaders.flat_map(&:runnable_loaders)
+          else
+            through_preloaders.flat_map(&:runnable_loaders)
+          end
+        end
+
         private
           def source_preloaders
-            @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).call
+            @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).loaders
           end
 
           def middle_records
@@ -45,7 +55,7 @@ module ActiveRecord
           end
 
           def through_preloaders
-            @through_preloaders ||= ActiveRecord::Associations::Preloader.new(records: owners, associations: through_reflection.name, scope: through_scope, associate_by_default: false).call
+            @through_preloaders ||= ActiveRecord::Associations::Preloader.new(records: owners, associations: through_reflection.name, scope: through_scope, associate_by_default: false).loaders
           end
 
           def through_reflection

activerecord/test/cases/associations_test.rb

@@ -416,6 +416,25 @@ class PreloaderTest < ActiveRecord::TestCase
     end
   end
 
+  def test_preload_grouped_queries_of_middle_records
+    comments = [
+      comments(:eager_sti_on_associations_s_comment1),
+      comments(:eager_sti_on_associations_s_comment2),
+    ]
+
+    assert_queries(2) do
+      ActiveRecord::Associations::Preloader.new(records: comments, associations: [:author, :ordinary_post]).call
+    end
+  end
+
+  def test_preload_grouped_queries_of_through_records
+    author = authors(:david)
+
+    assert_queries(3) do
+      ActiveRecord::Associations::Preloader.new(records: [author], associations: [:hello_post_comments, :comments]).call
+    end
+  end
+
   def test_preload_through
     comments = [
       comments(:eager_sti_on_associations_s_comment1),

activerecord/test/models/comment.rb

@@ -64,6 +64,7 @@ class Comment < ActiveRecord::Base
 end
 
 class SpecialComment < Comment
+  belongs_to :ordinary_post, foreign_key: :post_id, class_name: "Post"
   has_one :author, through: :post
   default_scope { where(deleted_at: nil) }