diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb index fc92e92a7e9fc..ad657e5d1fcd4 100644 --- a/activerecord/lib/active_record/associations/preloader.rb +++ b/activerecord/lib/active_record/associations/preloader.rb @@ -87,6 +87,12 @@ class Preloader #:nodoc: # [ :books, :author ] # { author: :avatar } # [ :books, { author: :avatar } ] + # + # +available_records+ is an array of ActiveRecord::Base. The Preloader + # will try to use the objects in this array to preload the requested + # associations before querying the database. This can save database + # queries by reusing in-memory objects. The optimization is only applied + # to single associations (i.e. :belongs_to, :has_one) with no scopes. 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.") @@ -94,6 +100,7 @@ def initialize(associate_by_default: true, **kwargs) @records = kwargs[:records] @associations = kwargs[:associations] @scope = kwargs[:scope] + @available_records = kwargs[:available_records] || [] @associate_by_default = associate_by_default @tree = Branch.new( @@ -112,7 +119,7 @@ def empty? end def call - Batch.new([self]).call + Batch.new([self], available_records: @available_records).call loaders end diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb index a9f8a7cce3466..5fc036e22f0d4 100644 --- a/activerecord/lib/active_record/associations/preloader/association.rb +++ b/activerecord/lib/active_record/associations/preloader/association.rb @@ -163,6 +163,25 @@ def load_records(raw_records = nil) end end + def associate_records_from_unscoped(unscoped_records) + return if unscoped_records.nil? || unscoped_records.empty? + return if !reflection_scope.empty_scope? + return if preload_scope && !preload_scope.empty_scope? + return if reflection.collection? + + unscoped_records.each do |record| + owners = owners_by_key[convert_key(record[association_key_name])] + owners&.each_with_index do |owner, i| + association = owner.association(reflection.name) + association.target = record + + if i == 0 # Set inverse on first owner + association.set_inverse_instance(record) + end + end + end + end + private attr_reader :owners, :reflection, :preload_scope, :model diff --git a/activerecord/lib/active_record/associations/preloader/batch.rb b/activerecord/lib/active_record/associations/preloader/batch.rb index cfcd9344ce80b..556650f184a16 100644 --- a/activerecord/lib/active_record/associations/preloader/batch.rb +++ b/activerecord/lib/active_record/associations/preloader/batch.rb @@ -4,8 +4,9 @@ module ActiveRecord module Associations class Preloader class Batch #:nodoc: - def initialize(preloaders) + def initialize(preloaders, available_records:) @preloaders = preloaders.reject(&:empty?) + @available_records = available_records.flatten.group_by(&:class) end def call @@ -13,6 +14,8 @@ def call until branches.empty? loaders = branches.flat_map(&:runnable_loaders) + 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) diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb index 8cdc4db679bbe..06050fd129de2 100644 --- a/activerecord/test/cases/associations_test.rb +++ b/activerecord/test/cases/associations_test.rb @@ -31,6 +31,7 @@ require "models/discount" require "models/line_item" require "models/shipping_line" +require "models/essay" class AssociationsTest < ActiveRecord::TestCase fixtures :accounts, :companies, :developers, :projects, :developers_projects, @@ -384,7 +385,7 @@ def test_associations_raise_with_name_error_if_associated_to_classes_that_do_not end class PreloaderTest < ActiveRecord::TestCase - fixtures :posts, :comments, :books, :authors, :tags, :taggings + fixtures :posts, :comments, :books, :authors, :tags, :taggings, :essays, :categories def test_preload_with_scope post = posts(:welcome) @@ -760,6 +761,85 @@ def test_preload_does_not_group_same_scope_different_key_name postesque.author end end + + def test_preload_with_available_records + post = posts(:welcome) + david = authors(:david) + + assert_no_queries do + ActiveRecord::Associations::Preloader.new(records: [post], associations: :author, available_records: [[david]]).call + + assert_predicate post.association(:author), :loaded? + assert_same david, post.author + end + end + + def test_preload_with_available_records_with_through_association + author = authors(:david) + categories = Category.all.to_a + + assert_queries(1) do + # One query to get the middle records (i.e. essays) + ActiveRecord::Associations::Preloader.new(records: [author], associations: :essay_category, available_records: categories).call + end + + assert_predicate author.association(:essay_category), :loaded? + assert categories.map(&:object_id).include?(author.essay_category.object_id) + end + + def test_preload_with_available_records_with_multiple_classes + essay = essays(:david_modest_proposal) + general = categories(:general) + david = authors(:david) + + assert_no_queries do + ActiveRecord::Associations::Preloader.new(records: [essay], associations: [:category, :author], available_records: [general, david]).call + + assert_predicate essay.association(:category), :loaded? + assert_predicate essay.association(:author), :loaded? + assert_same general, essay.category + assert_same david, essay.author + end + end + + def test_preload_with_available_records_queries_when_scoped + post = posts(:welcome) + david = authors(:david) + + assert_queries(1) do + ActiveRecord::Associations::Preloader.new(records: [post], associations: :author, scope: Author.where(name: "David"), available_records: [david]).call + end + + assert_predicate post.association(:author), :loaded? + assert_not_equal david.object_id, post.author.object_id + end + + def test_preload_with_available_records_queries_when_collection + post = posts(:welcome) + comments = Comment.all.to_a + + assert_queries(1) do + ActiveRecord::Associations::Preloader.new(records: [post], associations: :comments, available_records: comments).call + end + + assert_predicate post.association(:comments), :loaded? + assert_empty post.comments.map(&:object_id) & comments.map(&:object_id) + end + + def test_preload_with_available_records_queries_when_incomplete + post = posts(:welcome) + bob = authors(:bob) + david = authors(:david) + + assert_queries(1) do + ActiveRecord::Associations::Preloader.new(records: [post], associations: :author, available_records: [bob]).call + end + + assert_no_queries do + assert_predicate post.association(:author), :loaded? + assert_equal david, post.author + end + end end class GeneratedMethodsTest < ActiveRecord::TestCase diff --git a/activerecord/test/models/essay.rb b/activerecord/test/models/essay.rb index e59db4d877880..06764250de607 100644 --- a/activerecord/test/models/essay.rb +++ b/activerecord/test/models/essay.rb @@ -1,7 +1,7 @@ # frozen_string_literal: true class Essay < ActiveRecord::Base - belongs_to :author + belongs_to :author, primary_key: :name belongs_to :writer, primary_key: :name, polymorphic: true belongs_to :category, primary_key: :name has_one :owner, primary_key: :name