Merge pull request rails#42654 from dinahshi/preloader-from-available

Add `available_records` argument to Associations::Preloader

Author: jhawthorn

activerecord/lib/active_record/associations/preloader.rb

@@ -87,13 +87,20 @@ 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.")
         else
           @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

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
 

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

@@ -4,15 +4,18 @@ 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
           branches = @preloaders.flat_map(&:branches)
           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)

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

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