diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb index 4785c838ef..e741b04e42 100644 --- a/activerecord/lib/active_record/associations/preloader/association.rb +++ b/activerecord/lib/active_record/associations/preloader/association.rb @@ -23,11 +23,7 @@ module ActiveRecord end def records_for(loaders) - ids = loaders.flat_map(&:owner_keys).uniq - - scope.where(association_key_name => ids).load do |record| - loaders.each { |l| l.set_inverse(record) } - end + LoaderRecords.new(loaders, self).records end def load_records_in_batch(loaders) @@ -38,6 +34,52 @@ module ActiveRecord loader.run end end + + def load_records_for_keys(keys, &block) + scope.where(association_key_name => keys).load(&block) + end + end + + class LoaderRecords + def initialize(loaders, loader_query) + @loader_query = loader_query + @loaders = loaders + @keys_to_load = Set.new + @already_loaded_records_by_key = {} + + populate_keys_to_load_and_already_loaded_records + end + + def records + load_records + already_loaded_records + end + + private + attr_reader :loader_query, :loaders, :keys_to_load, :already_loaded_records_by_key + + def populate_keys_to_load_and_already_loaded_records + loaders.each do |loader| + loader.owners_by_key.each do |key, owners| + if loaded_owner = owners.find { |owner| loader.loaded?(owner) } + already_loaded_records_by_key[key] = loader.target_for(loaded_owner) + else + keys_to_load << key + end + end + end + + @keys_to_load.subtract(already_loaded_records_by_key.keys) + end + + def load_records + loader_query.load_records_for_keys(keys_to_load) do |record| + loaders.each { |l| l.set_inverse(record) } + end + end + + def already_loaded_records + already_loaded_records_by_key.values.flatten + end end attr_reader :klass @@ -57,12 +99,8 @@ module ActiveRecord @klass.table_name end - def data_available? - already_loaded? - end - def future_classes - if run? || already_loaded? + if run? [] else [@klass] @@ -81,11 +119,6 @@ module ActiveRecord return self if run? @run = true - if already_loaded? - fetch_from_preloaded_records - return self - end - records = records_by_owner owners.each do |owner| @@ -96,25 +129,17 @@ module ActiveRecord end def records_by_owner - ensure_loaded unless defined?(@records_by_owner) + load_records unless defined?(@records_by_owner) @records_by_owner end def preloaded_records - ensure_loaded unless defined?(@preloaded_records) + load_records 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) @@ -124,8 +149,19 @@ module ActiveRecord LoaderQuery.new(scope, association_key_name) end - def owner_keys - @owner_keys ||= owners_by_key.keys + def owners_by_key + @owners_by_key ||= owners.each_with_object({}) do |owner, result| + key = convert_key(owner[owner_key_name]) + (result[key] ||= []) << owner if key + end + end + + def loaded?(owner) + owner.association(reflection.name).loaded? + end + + def target_for(owner) + Array.wrap(owner.association(reflection.name).target) end def scope @@ -185,25 +221,16 @@ module ActiveRecord private 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| - Array(owner.association(reflection.name).target) - end - - @preloaded_records = records_by_owner.flat_map(&:last) - end - # The name of the key on the model which declares the association def owner_key_name reflection.join_foreign_key end def associate_records_to_owner(owner, records) + return if loaded?(owner) + association = owner.association(reflection.name) + if reflection.collection? association.target = records else @@ -211,13 +238,6 @@ module ActiveRecord end end - def owners_by_key - @owners_by_key ||= owners.each_with_object({}) do |owner, result| - key = convert_key(owner[owner_key_name]) - (result[key] ||= []) << owner if key - end - end - def key_conversion_required? unless defined?(@key_conversion_required) @key_conversion_required = (association_key_type != owner_key_type) diff --git a/activerecord/lib/active_record/associations/preloader/batch.rb b/activerecord/lib/active_record/associations/preloader/batch.rb index 97304625d4..17cf01f828 100644 --- a/activerecord/lib/active_record/associations/preloader/batch.rb +++ b/activerecord/lib/active_record/associations/preloader/batch.rb @@ -16,10 +16,7 @@ module ActiveRecord loaders.each { |loader| loader.associate_records_from_unscoped(@available_records[loader.klass]) } - already_loaded = loaders.select(&:data_available?) - if already_loaded.any? - already_loaded.each(&:run) - elsif loaders.any? + if loaders.any? future_tables = branches.flat_map do |branch| branch.future_classes - branch.runnable_loaders.map(&:klass) end.map(&:table_name).uniq diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb index 212435a6a0..67f6bc4293 100644 --- a/activerecord/lib/active_record/associations/preloader/through_association.rb +++ b/activerecord/lib/active_record/associations/preloader/through_association.rb @@ -10,10 +10,13 @@ module ActiveRecord def records_by_owner return @records_by_owner if defined?(@records_by_owner) - source_records_by_owner = source_preloaders.map(&:records_by_owner).reduce(:merge) - through_records_by_owner = through_preloaders.map(&:records_by_owner).reduce(:merge) @records_by_owner = owners.each_with_object({}) do |owner, result| + if loaded?(owner) + result[owner] = target_for(owner) + next + end + through_records = through_records_by_owner[owner] || [] if owners.first.association(through_reflection.name).loaded? @@ -35,12 +38,6 @@ module ActiveRecord end end - def data_available? - return true if super() - through_preloaders.all?(&:run?) && - source_preloaders.all?(&:run?) - end - def runnable_loaders if data_available? [self] @@ -52,7 +49,7 @@ module ActiveRecord end def future_classes - if run? || data_available? + if run? [] elsif through_preloaders.all?(&:run?) source_preloaders.flat_map(&:future_classes).uniq @@ -67,6 +64,11 @@ module ActiveRecord end private + def data_available? + owners.all? { |owner| loaded?(owner) } || + through_preloaders.all?(&:run?) && source_preloaders.all?(&:run?) + end + def source_preloaders @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).loaders end @@ -87,6 +89,14 @@ module ActiveRecord reflection.source_reflection end + def source_records_by_owner + @source_records_by_owner ||= source_preloaders.map(&:records_by_owner).reduce(:merge) + end + + def through_records_by_owner + @through_records_by_owner ||= through_preloaders.map(&:records_by_owner).reduce(:merge) + end + def preload_index @preload_index ||= preloaded_records.each_with_object({}).with_index do |(record, result), index| result[record] = index diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb index ba236fe979..189b634c00 100644 --- a/activerecord/test/cases/associations_test.rb +++ b/activerecord/test/cases/associations_test.rb @@ -443,6 +443,18 @@ class PreloaderTest < ActiveRecord::TestCase end end + def test_preload_grouped_queries_with_already_loaded_records + book = books(:awdr) + post = posts(:welcome) + book.author + + assert_no_queries do + ActiveRecord::Associations::Preloader.new(records: [book, post], associations: :author).call + book.author + post.author + end + end + def test_preload_grouped_queries_of_middle_records comments = [ comments(:eager_sti_on_associations_s_comment1), @@ -792,6 +804,41 @@ class PreloaderTest < ActiveRecord::TestCase end end + def test_preload_with_only_some_records_available + bob_post = posts(:misc_by_bob) + mary_post = posts(:misc_by_mary) + bob = authors(:bob) + mary = authors(:mary) + + assert_queries(1) do + ActiveRecord::Associations::Preloader.new(records: [bob_post, mary_post], associations: :author, available_records: [bob]).call + end + + assert_no_queries do + assert_same bob, bob_post.author + assert_equal mary, mary_post.author + end + end + + def test_preload_with_some_records_already_loaded + bob_post = posts(:misc_by_bob) + mary_post = posts(:misc_by_mary) + bob = bob_post.author + mary = authors(:mary) + + assert bob_post.association(:author).loaded? + assert_not mary_post.association(:author).loaded? + + assert_queries(1) do + ActiveRecord::Associations::Preloader.new(records: [bob_post, mary_post], associations: :author).call + end + + assert_no_queries do + assert_same bob, bob_post.author + assert_equal mary, mary_post.author + end + end + def test_preload_with_available_records_with_through_association author = authors(:david) categories = Category.all.to_a @@ -805,6 +852,25 @@ class PreloaderTest < ActiveRecord::TestCase assert categories.map(&:object_id).include?(author.essay_category.object_id) end + def test_preload_with_only_some_records_available_with_through_associations + mary = authors(:mary) + mary_essay = essays(:mary_stay_home) + mary_category = categories(:technology) + mary_essay.update!(category: mary_category) + + dave = authors(:david) + dave_category = categories(:general) + + assert_queries(2) do + ActiveRecord::Associations::Preloader.new(records: [mary, dave], associations: :essay_category, available_records: [mary_category]).call + end + + assert_no_queries do + assert_same mary_category, mary.essay_category + assert_equal dave_category, dave.essay_category + end + end + def test_preload_with_available_records_with_multiple_classes essay = essays(:david_modest_proposal) general = categories(:general) @@ -858,6 +924,20 @@ class PreloaderTest < ActiveRecord::TestCase assert_equal david, post.author end end + + def test_preload_with_unpersisted_records_no_ops + author = Author.new + new_post_with_author = Post.new(author: author) + new_post_without_author = Post.new + posts = [new_post_with_author, new_post_without_author] + + assert_no_queries do + ActiveRecord::Associations::Preloader.new(records: posts, associations: :author).call + + assert_same author, new_post_with_author.author + assert_nil new_post_without_author.author + end + end end class GeneratedMethodsTest < ActiveRecord::TestCase