diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb index 0897a85313..90166f1dfa 100644 --- a/activerecord/lib/active_record/associations/preloader.rb +++ b/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 diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb index 166cb01b85..fce2003764 100644 --- a/activerecord/lib/active_record/associations/preloader/association.rb +++ b/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) diff --git a/activerecord/lib/active_record/associations/preloader/batch.rb b/activerecord/lib/active_record/associations/preloader/batch.rb index 3b7ca74852..3af498f08c 100644 --- a/activerecord/lib/active_record/associations/preloader/batch.rb +++ b/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 diff --git a/activerecord/lib/active_record/associations/preloader/branch.rb b/activerecord/lib/active_record/associations/preloader/branch.rb new file mode 100644 index 0000000000..ebb9d37a7d --- /dev/null +++ b/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 diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb index 95a937c2cd..b7ac3b22b0 100644 --- a/activerecord/lib/active_record/associations/preloader/through_association.rb +++ b/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 diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb index 8a38aec2bc..d5aad8202b 100644 --- a/activerecord/test/cases/associations_test.rb +++ b/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), diff --git a/activerecord/test/models/comment.rb b/activerecord/test/models/comment.rb index 9834cd482a..68d4791fdf 100644 --- a/activerecord/test/models/comment.rb +++ b/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) } diff --git a/activerecord/test/models/discount.rb b/activerecord/test/models/discount.rb new file mode 100644 index 0000000000..06068df2fb --- /dev/null +++ b/activerecord/test/models/discount.rb @@ -0,0 +1,4 @@ +# frozen_string_literal: true + +class Discount < ActiveRecord::Base +end diff --git a/activerecord/test/models/invoice.rb b/activerecord/test/models/invoice.rb index 1851792ed5..1a62b7baa0 100644 --- a/activerecord/test/models/invoice.rb +++ b/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 diff --git a/activerecord/test/models/line_item.rb b/activerecord/test/models/line_item.rb index 3a51cf03b2..2c9876dac1 100644 --- a/activerecord/test/models/line_item.rb +++ b/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 diff --git a/activerecord/test/models/shipping_line.rb b/activerecord/test/models/shipping_line.rb new file mode 100644 index 0000000000..4044a5ee38 --- /dev/null +++ b/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 diff --git a/activerecord/test/schema/schema.rb b/activerecord/test/schema/schema.rb index f72cda1e3c..c9fd47e512 100644 --- a/activerecord/test/schema/schema.rb +++ b/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