Merge pull request #41596 from jhawthorn/preloader_branch

Allow Preloader to group through association queries

activerecord/lib/active_record/associations/preloader.rb

@@ -47,11 +47,11 @@ module ActiveRecord
       eager_autoload do
         autoload :Association,        "active_record/associations/preloader/association"
         autoload :Batch,              "active_record/associations/preloader/batch"
+        autoload :Branch,             "active_record/associations/preloader/branch"
         autoload :ThroughAssociation, "active_record/associations/preloader/through_association"
       end
 
-      attr_reader :records, :associations, :scope, :associate_by_default, :polymorphic_parent
-      attr_reader :loaders
+      attr_reader :records, :associations, :scope, :associate_by_default
 
       # Eager loads the named associations for the given Active Record record(s).
       #
@@ -87,7 +87,7 @@ module ActiveRecord
       #   [ :books, :author ]
       #   { author: :avatar }
       #   [ :books, { author: :avatar } ]
-      def initialize(associate_by_default: true, polymorphic_parent: false, **kwargs)
+      def initialize(associate_by_default: true, **kwargs)
         if kwargs.empty?
           ActiveSupport::Deprecation.warn("Calling `Preloader#initialize` without arguments is deprecated and will be removed in Rails 7.0.")
         else
@@ -95,9 +95,15 @@ module ActiveRecord
           @associations = kwargs[:associations]
           @scope = kwargs[:scope]
           @associate_by_default = associate_by_default
-          @polymorphic_parent = polymorphic_parent
-          @child_preloader_args = []
-          @loaders = build_preloaders
+
+          @tree = Branch.new(
+            parent: nil,
+            association: nil,
+            children: associations,
+            associate_by_default: @associate_by_default,
+            scope: @scope
+          )
+          @tree.preloaded_records = records
         end
       end
 
@@ -108,7 +114,7 @@ module ActiveRecord
       def call
         Batch.new([self]).call
 
-        @loaders
+        loaders
       end
 
       def preload(records, associations, preload_scope = nil)
@@ -117,64 +123,13 @@ module ActiveRecord
         Preloader.new(records: records, associations: associations, scope: preload_scope).call
       end
 
-      def child_preloaders
-        @child_preloader_args.map do |reflection, child, parents|
-          build_child_preloader(reflection, child, parents)
-        end
+      def branches
+        @tree.children
       end
 
-      private
-        def build_preloaders
-          Array.wrap(associations).flat_map { |association|
-            Array(association).flat_map { |parent, child|
-              grouped_records(parent).flat_map do |reflection, reflection_records|
-                loaders = preloaders_for_reflection(reflection, reflection_records)
-
-                if child
-                  @child_preloader_args << [reflection, child, loaders]
-                end
-
-                loaders
-              end
-            }
-          }
-        end
-
-        def build_child_preloader(reflection, child, loaders)
-          child_polymorphic_parent = reflection && reflection.options[:polymorphic]
-          preloaded_records = loaders.flat_map(&:preloaded_records).uniq
-
-          Preloader.new(records: preloaded_records, associations: child, scope: scope, associate_by_default: associate_by_default, polymorphic_parent: child_polymorphic_parent)
-        end
-
-        def preloaders_for_reflection(reflection, reflection_records)
-          reflection_records.group_by { |record| record.association(reflection.name).klass }.map do |rhs_klass, rs|
-            preloader_for(reflection).new(rhs_klass, rs, reflection, scope, associate_by_default)
-          end
-        end
-
-        def grouped_records(association)
-          h = {}
-          records.each do |record|
-            reflection = record.class._reflect_on_association(association)
-            next if polymorphic_parent && !reflection || !record.association(association).klass
-            (h[reflection] ||= []) << record
-          end
-          h
-        end
-
-        # Returns a class containing the logic needed to load preload the data
-        # and attach it to a relation. The class returned implements a `run` method
-        # that accepts a preloader.
-        def preloader_for(reflection)
-          reflection.check_preloadable!
-
-          if reflection.options[:through]
-            ThroughAssociation
-          else
-            Association
-          end
-        end
+      def loaders
+        branches.flat_map(&:loaders)
+      end
     end
   end
 end

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

@@ -24,13 +24,25 @@ module ActiveRecord
           @preload_scope = preload_scope
           @associate     = associate_by_default || !preload_scope || preload_scope.empty_scope?
           @model         = owners.first && owners.first.class
+          @run = false
         end
 
         def already_loaded?
           @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? }
         end
 
+        def runnable_loaders
+          [self]
+        end
+
+        def run?
+          @run
+        end
+
         def run
+          return self if run?
+          @run = true
+
           if already_loaded?
             fetch_from_preloaded_records
             return self
@@ -46,17 +58,25 @@ module ActiveRecord
         end
 
         def records_by_owner
-          load_records unless defined?(@records_by_owner)
+          ensure_loaded unless defined?(@records_by_owner)
 
           @records_by_owner
         end
 
         def preloaded_records
-          load_records unless defined?(@preloaded_records)
+          ensure_loaded unless defined?(@preloaded_records)
 
           @preloaded_records
         end
 
+        def ensure_loaded
+          if already_loaded?
+            fetch_from_preloaded_records
+          else
+            load_records
+          end
+        end
+
         # The name of the key on the associated records
         def association_key_name
           reflection.join_primary_key(klass)

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

@@ -9,14 +9,20 @@ module ActiveRecord
         end
 
         def call
-          return if @preloaders.empty?
+          branches = @preloaders.flat_map(&:branches)
+          until branches.empty?
+            loaders = branches.flat_map(&:runnable_loaders)
 
-          loaders = @preloaders.flat_map(&:loaders)
-          group_and_load_similar(loaders)
-          loaders.each(&:run)
+            already_loaded, loaders = loaders.partition(&:already_loaded?)
+            already_loaded.each(&:run)
 
-          child_preloaders = @preloaders.flat_map(&:child_preloaders)
-          Batch.new(child_preloaders).call
+            group_and_load_similar(loaders)
+            loaders.each(&:run)
+
+            finished, in_progress = branches.partition(&:done?)
+
+            branches = in_progress + finished.flat_map(&:children)
+          end
         end
 
         private
@@ -24,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

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Associations
+    class Preloader
+      class Branch #:nodoc:
+        attr_reader :association, :children, :parent
+        attr_reader :scope, :associate_by_default
+        attr_writer :preloaded_records
+
+        def initialize(association:, children:, parent:, associate_by_default:, scope:)
+          @association = association
+          @parent = parent
+          @scope = scope
+          @associate_by_default = associate_by_default
+
+          @children = build_children(children)
+        end
+
+        def root?
+          parent.nil?
+        end
+
+        def source_records
+          @parent.preloaded_records
+        end
+
+        def preloaded_records
+          @preloaded_records ||= loaders.flat_map(&:preloaded_records)
+        end
+
+        def done?
+          loaders.all?(&:run?)
+        end
+
+        def runnable_loaders
+          loaders.flat_map(&:runnable_loaders).reject(&:run?)
+        end
+
+        def grouped_records
+          h = {}
+          source_records.each do |record|
+            reflection = record.class._reflect_on_association(association)
+            next if polymorphic_parent? && !reflection || !record.association(association).klass
+            (h[reflection] ||= []) << record
+          end
+          h
+        end
+
+        def preloaders_for_reflection(reflection, reflection_records)
+          reflection_records.group_by { |record| record.association(reflection.name).klass }.map do |rhs_klass, rs|
+            preloader_for(reflection).new(rhs_klass, rs, reflection, scope, associate_by_default)
+          end
+        end
+
+        def polymorphic_parent?
+          return false if root?
+
+          parent.polymorphic?
+        end
+
+        def polymorphic?
+          return false if root?
+
+          grouped_records.keys.any? do |reflection|
+            reflection.options[:polymorphic]
+          end
+        end
+
+        def loaders
+          @loaders ||=
+            grouped_records.flat_map do |reflection, reflection_records|
+              preloaders_for_reflection(reflection, reflection_records)
+            end
+        end
+
+        private
+          def build_children(children)
+            Array.wrap(children).flat_map { |association|
+              Array(association).flat_map { |parent, child|
+                Branch.new(
+                  parent: self,
+                  association: parent,
+                  children: child,
+                  associate_by_default: associate_by_default,
+                  scope: scope
+                )
+              }
+            }
+          end
+
+          # Returns a class containing the logic needed to load preload the data
+          # and attach it to a relation. The class returned implements a `run` method
+          # that accepts a preloader.
+          def preloader_for(reflection)
+            reflection.check_preloadable!
+
+            if reflection.options[:through]
+              ThroughAssociation
+            else
+              Association
+            end
+          end
+      end
+    end
+  end
+end

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

@@ -27,6 +27,10 @@ require "models/parrot"
 require "models/bird"
 require "models/treasure"
 require "models/price_estimate"
+require "models/invoice"
+require "models/discount"
+require "models/line_item"
+require "models/shipping_line"
 
 class AssociationsTest < ActiveRecord::TestCase
   fixtures :accounts, :companies, :developers, :projects, :developers_projects,
@@ -416,6 +420,79 @@ 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_some_already_loaded_associations
+    item_discount = Discount.create(amount: 5)
+    shipping_discount = Discount.create(amount: 20)
+
+    invoice = Invoice.new
+    line_item = LineItem.new(amount: 20)
+    line_item.discount_applications << LineItemDiscountApplication.new(discount: item_discount)
+    invoice.line_items << line_item
+
+    shipping_line = ShippingLine.new(amount: 50)
+    shipping_line.discount_applications << ShippingLineDiscountApplication.new(discount: shipping_discount)
+    invoice.shipping_lines << shipping_line
+
+    invoice.save!
+    invoice.reload
+
+    # SELECT "line_items".* FROM "line_items" WHERE "line_items"."invoice_id" = ?
+    # SELECT "shipping_lines".* FROM shipping_lines WHERE "shipping_lines"."invoice_id" = ?
+    # SELECT "line_item_discount_applications".* FROM "line_item_discount_applications" WHERE "line_item_discount_applications"."line_item_id" = ?
+    # SELECT "shipping_line_discount_applications".* FROM "shipping_line_discount_applications" WHERE "shipping_line_discount_applications"."shipping_line_id" = ?
+    # SELECT "discounts".* FROM "discounts" WHERE "discounts"."id" IN (?, ?).
+    assert_queries(5) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [invoice], associations: [
+        line_items: { discount_applications: :discount },
+        shipping_lines: { discount_applications: :discount },
+      ])
+      preloader.call
+    end
+
+    assert_no_queries do
+      assert_not_nil invoice.line_items.first.discount_applications.first.discount
+      assert_not_nil invoice.shipping_lines.first.discount_applications.first.discount
+    end
+
+    invoice.reload
+    invoice.line_items.map { |i| i.discount_applications.to_a }
+    # `line_items` and `line_item_discount_applications` are already preloaded, so we expect:
+    # SELECT "shipping_lines".* FROM shipping_lines WHERE "shipping_lines"."invoice_id" = ?
+    # SELECT "shipping_line_discount_applications".* FROM "shipping_line_discount_applications" WHERE "shipping_line_discount_applications"."shipping_line_id" = ?
+    # SELECT "discounts".* FROM "discounts" WHERE "discounts"."id" = ?.
+    assert_queries(3) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [invoice], associations: [
+        line_items: { discount_applications: :discount },
+        shipping_lines: { discount_applications: :discount },
+      ])
+      preloader.call
+    end
+
+    assert_no_queries do
+      assert_not_nil invoice.line_items.first.discount_applications.first.discount
+      assert_not_nil invoice.shipping_lines.first.discount_applications.first.discount
+    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) }
 

activerecord/test/models/discount.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class Discount < ActiveRecord::Base
+end

activerecord/test/models/invoice.rb

@@ -2,5 +2,6 @@
 
 class Invoice < ActiveRecord::Base
   has_many :line_items, autosave: true
+  has_many :shipping_lines, -> { from("shipping_lines") }, autosave: true
   before_save { |record| record.balance = record.line_items.map(&:amount).sum }
 end

activerecord/test/models/line_item.rb

@@ -2,4 +2,10 @@
 
 class LineItem < ActiveRecord::Base
   belongs_to :invoice, touch: true
+  has_many :discount_applications, class_name: "LineItemDiscountApplication"
+end
+
+class LineItemDiscountApplication < ActiveRecord::Base
+  belongs_to :line_item
+  belongs_to :discount
 end

activerecord/test/models/shipping_line.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+class ShippingLine < ActiveRecord::Base
+  belongs_to :invoice, touch: true
+  has_many :discount_applications, class_name: "ShippingLineDiscountApplication"
+end
+
+class ShippingLineDiscountApplication < ActiveRecord::Base
+  belongs_to :shipping_line
+  belongs_to :discount
+end

activerecord/test/schema/schema.rb

@@ -313,6 +313,10 @@ ActiveRecord::Schema.define do
     t.boolean :deleted
   end
 
+  create_table :discounts, force: true do |t|
+    t.integer :amount
+  end
+
   create_table :dl_keyed_belongs_tos, force: true, id: false do |t|
     t.primary_key :belongs_key
     t.references :destroy_async_parent
@@ -547,6 +551,11 @@ ActiveRecord::Schema.define do
     t.integer :amount
   end
 
+  create_table :line_item_discount_applications, force: true do |t|
+    t.integer :line_item_id
+    t.integer :discount_id
+  end
+
   create_table :lions, force: true do |t|
     t.integer :gender
     t.boolean :is_vegetarian, default: false
@@ -922,6 +931,16 @@ ActiveRecord::Schema.define do
     t.integer :shape_id
   end
 
+  create_table :shipping_lines, force: true do |t|
+    t.integer :invoice_id
+    t.integer :amount
+  end
+
+  create_table :shipping_line_discount_applications, force: true do |t|
+    t.integer :shipping_line_id
+    t.integer :discount_id
+  end
+
   create_table :ships, force: true do |t|
     t.string :name
     t.integer :pirate_id