diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb index c907a3017a..2d80f912c8 100644 --- a/activerecord/lib/active_record/associations/preloader/association.rb +++ b/activerecord/lib/active_record/associations/preloader/association.rb @@ -40,6 +40,8 @@ module ActiveRecord end end + attr_reader :klass + def initialize(klass, owners, reflection, preload_scope, associate_by_default = true) @klass = klass @owners = owners.uniq(&:__id__) @@ -50,8 +52,20 @@ module ActiveRecord @run = false end - def already_loaded? - @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? } + def table_name + @klass.table_name + end + + def data_available? + already_loaded? + end + + def future_classes + if run? || already_loaded? + [] + else + [@klass] + end end def runnable_loaders @@ -149,7 +163,11 @@ module ActiveRecord end private - attr_reader :owners, :reflection, :preload_scope, :model, :klass + attr_reader :owners, :reflection, :preload_scope, :model + + def already_loaded? + @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? } + end def fetch_from_preloaded_records @records_by_owner = owners.index_with do |owner| diff --git a/activerecord/lib/active_record/associations/preloader/batch.rb b/activerecord/lib/active_record/associations/preloader/batch.rb index a7c4bc63b7..cfcd9344ce 100644 --- a/activerecord/lib/active_record/associations/preloader/batch.rb +++ b/activerecord/lib/active_record/associations/preloader/batch.rb @@ -13,11 +13,20 @@ module ActiveRecord until branches.empty? loaders = branches.flat_map(&:runnable_loaders) - already_loaded, loaders = loaders.partition(&:already_loaded?) - already_loaded.each(&:run) + already_loaded = loaders.select(&:data_available?) + if already_loaded.any? + already_loaded.each(&:run) + elsif loaders.any? + future_tables = branches.flat_map do |branch| + branch.future_classes - branch.runnable_loaders.map(&:klass) + end.map(&:table_name).uniq - group_and_load_similar(loaders) - loaders.each(&:run) + target_loaders = loaders.reject { |l| future_tables.include?(l.table_name) } + target_loaders = loaders if target_loaders.empty? + + group_and_load_similar(target_loaders) + target_loaders.each(&:run) + end finished, in_progress = branches.partition(&:done?) diff --git a/activerecord/lib/active_record/associations/preloader/branch.rb b/activerecord/lib/active_record/associations/preloader/branch.rb index bf1d077cfc..c7d43ff295 100644 --- a/activerecord/lib/active_record/associations/preloader/branch.rb +++ b/activerecord/lib/active_record/associations/preloader/branch.rb @@ -15,6 +15,40 @@ module ActiveRecord @associate_by_default = associate_by_default @children = build_children(children) + @loaders = nil + end + + def future_classes + (immediate_future_classes + children.flat_map(&:future_classes)).uniq + end + + def immediate_future_classes + if parent.done? + loaders.flat_map(&:future_classes).uniq + else + likely_reflections.reject(&:polymorphic?).flat_map do |reflection| + reflection. + chain. + map(&:klass) + end.uniq + end + end + + def target_classes + if done? + preloaded_records.map(&:klass).uniq + elsif parent.done? + loaders.map(&:klass).uniq + else + likely_reflections.reject(&:polymorphic?).map(&:klass).uniq + end + end + + def likely_reflections + parent_classes = parent.target_classes + parent_classes.map do |parent_klass| + parent_klass._reflect_on_association(@association) + end.compact end def root? @@ -30,7 +64,7 @@ module ActiveRecord end def done? - loaders.all?(&:run?) + root? || (@loaders && @loaders.all?(&:run?)) end def runnable_loaders diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb index b7ac3b22b0..9826253d4a 100644 --- a/activerecord/lib/active_record/associations/preloader/through_association.rb +++ b/activerecord/lib/active_record/associations/preloader/through_association.rb @@ -35,16 +35,37 @@ module ActiveRecord end end + def data_available? + return true if super() + through_preloaders.all?(&:run?) && + source_preloaders.all?(&:run?) + end + def runnable_loaders - if already_loaded? + if data_available? [self] elsif through_preloaders.all?(&:run?) - [self] + source_preloaders.flat_map(&:runnable_loaders) + source_preloaders.flat_map(&:runnable_loaders) else through_preloaders.flat_map(&:runnable_loaders) end end + def future_classes + if run? || data_available? + [] + elsif through_preloaders.all?(&:run?) + source_preloaders.flat_map(&:future_classes).uniq + else + through_classes = through_preloaders.flat_map(&:future_classes) + source_classes = source_reflection. + chain. + reject { |reflection| reflection.respond_to?(:polymorphic?) && reflection.polymorphic? }. + map(&:klass) + (through_classes + source_classes).uniq + 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).loaders diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb index 3fb3231798..5992f3f181 100644 --- a/activerecord/test/cases/associations_test.rb +++ b/activerecord/test/cases/associations_test.rb @@ -362,7 +362,7 @@ class OverridingAssociationsTest < ActiveRecord::TestCase end class PreloaderTest < ActiveRecord::TestCase - fixtures :posts, :comments, :books, :authors + fixtures :posts, :comments, :books, :authors, :tags, :taggings def test_preload_with_scope post = posts(:welcome) @@ -569,6 +569,40 @@ class PreloaderTest < ActiveRecord::TestCase end end + def test_preload_can_group_separate_levels + mary = authors(:mary) + bob = authors(:bob) + + AuthorFavorite.create!(author: mary, favorite_author: bob) + + assert_queries(3) do + preloader = ActiveRecord::Associations::Preloader.new(records: [mary], associations: [:posts, favorite_authors: :posts]) + preloader.call + end + + assert_no_queries do + mary.posts + mary.favorite_authors.map(&:posts) + end + end + + def test_preload_can_group_multi_level_ping_pong_through + mary = authors(:mary) + bob = authors(:bob) + + AuthorFavorite.create!(author: mary, favorite_author: bob) + + assert_queries(9) do + preloader = ActiveRecord::Associations::Preloader.new(records: [mary], associations: { similar_posts: :comments, favorite_authors: { similar_posts: :comments } }) + preloader.call + end + + assert_no_queries do + mary.similar_posts.map(&:comments).each(&:to_a) + mary.favorite_authors.flat_map(&:similar_posts).map(&:comments).each(&:to_a) + end + end + def test_preload_does_not_group_same_class_different_scope post = posts(:welcome) postesque = Postesque.create(author: Author.last)