Merge pull request rails#42553 from jhawthorn/preload_instance_scope

Support preloads on instance dependent associations

Author: jhawthorn

activerecord/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Allow preloading of associations with instance dependent scopes
+
+    *John Hawthorn*, *John Crepezzi*, *Adam Hess*, *Eileen M. Uchitelle*, *Dinah Shi*
+
 *   Do not try to rollback transactions that failed due to a `ActiveRecord::TransactionRollbackError`.
 
     *Jamie McCarthy*

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

@@ -42,11 +42,12 @@ def load_records_in_batch(loaders)
 
         attr_reader :klass
 
-        def initialize(klass, owners, reflection, preload_scope, associate_by_default = true)
+        def initialize(klass, owners, reflection, preload_scope, reflection_scope, associate_by_default)
           @klass         = klass
           @owners        = owners.uniq(&:__id__)
           @reflection    = reflection
           @preload_scope = preload_scope
+          @reflection_scope = reflection_scope
           @associate     = associate_by_default || !preload_scope || preload_scope.empty_scope?
           @model         = owners.first && owners.first.class
           @run = false

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

@@ -83,8 +83,19 @@ def grouped_records
         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)
+          reflection_records.group_by do |record|
+            klass = record.association(association).klass
+
+            if reflection.scope && reflection.scope.arity != 0
+              # For instance dependent scopes, the scope is potentially
+              # different for each record. To allow this we'll group each
+              # object separately into its own preloader
+              reflection_scope = reflection.join_scopes(klass.arel_table, klass.predicate_builder, klass, record).inject(&:merge!)
+            end
+
+            [klass, reflection_scope]
+          end.map do |(rhs_klass, reflection_scope), rs|
+            preloader_for(reflection).new(rhs_klass, rs, reflection, scope, reflection_scope, associate_by_default)
           end
         end
 
@@ -124,8 +135,6 @@ def build_children(children)
           # 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

activerecord/lib/active_record/reflection.rb

@@ -194,9 +194,9 @@ def join_scope(table, foreign_table, foreign_klass)
         klass_scope
       end
 
-      def join_scopes(table, predicate_builder, klass = self.klass) # :nodoc:
+      def join_scopes(table, predicate_builder, klass = self.klass, record = nil) # :nodoc:
         if scope
-          [scope_for(build_scope(table, predicate_builder, klass))]
+          [scope_for(build_scope(table, predicate_builder, klass), record)]
         else
           []
         end
@@ -483,18 +483,17 @@ def check_validity!
         check_validity_of_inverse!
       end
 
-      def check_preloadable!
+      def check_eager_loadable!
         return unless scope
 
         unless scope.arity == 0
           raise ArgumentError, <<-MSG.squish
             The association scope '#{name}' is instance dependent (the scope
-            block takes an argument). Preloading instance dependent scopes is
-            not supported.
+            block takes an argument). Eager loading instance dependent scopes
+            is not supported.
           MSG
         end
       end
-      alias :check_eager_loadable! :check_preloadable!
 
       def join_id_for(owner) # :nodoc:
         owner[join_foreign_key]
@@ -842,8 +841,8 @@ def scopes
         source_reflection.scopes + super
       end
 
-      def join_scopes(table, predicate_builder, klass = self.klass) # :nodoc:
-        source_reflection.join_scopes(table, predicate_builder, klass) + super
+      def join_scopes(table, predicate_builder, klass = self.klass, record = nil) # :nodoc:
+        source_reflection.join_scopes(table, predicate_builder, klass, record) + super
       end
 
       def has_scope?
@@ -1015,9 +1014,9 @@ def initialize(reflection, previous_reflection)
         @previous_reflection = previous_reflection
       end
 
-      def join_scopes(table, predicate_builder, klass = self.klass) # :nodoc:
-        scopes = @previous_reflection.join_scopes(table, predicate_builder) + super
-        scopes << build_scope(table, predicate_builder, klass).instance_exec(nil, &source_type_scope)
+      def join_scopes(table, predicate_builder, klass = self.klass, record = nil) # :nodoc:
+        scopes = @previous_reflection.join_scopes(table, predicate_builder, record) + super
+        scopes << build_scope(table, predicate_builder, klass).instance_exec(record, &source_type_scope)
       end
 
       def constraints

activerecord/test/cases/associations/eager_test.rb

@@ -1493,36 +1493,32 @@ def test_preloading_has_many_through_with_custom_scope
     assert_equal pets(:parrot), Owner.including_last_pet.first.last_pet
   end
 
-  test "preloading and eager loading of instance dependent associations is not supported" do
-    message = "association scope 'posts_with_signature' is"
-    error = assert_raises(ArgumentError) do
-      Author.includes(:posts_with_signature).to_a
+  test "preloading of instance dependent associations is supported" do
+    authors = Author.preload(:posts_with_signature).to_a
+    assert_not authors.empty?
+    authors.each do |author|
+      assert_predicate author.posts_with_signature, :loaded?
     end
-    assert_match message, error.message
-
-    error = assert_raises(ArgumentError) do
-      Author.preload(:posts_with_signature).to_a
-    end
-    assert_match message, error.message
+  end
 
+  test "eager loading of instance dependent associations is not supported" do
+    message = "association scope 'posts_with_signature' is"
     error = assert_raises(ArgumentError) do
       Author.eager_load(:posts_with_signature).to_a
     end
     assert_match message, error.message
   end
 
-  test "preloading and eager loading of optional instance dependent associations is not supported" do
-    message = "association scope 'posts_mentioning_author' is"
-    error = assert_raises(ArgumentError) do
-      Author.includes(:posts_mentioning_author).to_a
+  test "preloading of optional instance dependent associations is supported" do
+    authors = Author.includes(:posts_mentioning_author).to_a
+    assert_not authors.empty?
+    authors.each do |author|
+      assert_predicate author.posts_mentioning_author, :loaded?
     end
-    assert_match message, error.message
-
-    error = assert_raises(ArgumentError) do
-      Author.preload(:posts_mentioning_author).to_a
-    end
-    assert_match message, error.message
+  end
 
+  test "eager loading of optional instance dependent associations is not supported" do
+    message = "association scope 'posts_mentioning_author' is"
     error = assert_raises(ArgumentError) do
       Author.eager_load(:posts_mentioning_author).to_a
     end
@@ -1542,11 +1538,12 @@ def test_preloading_has_many_through_with_custom_scope
     end
   end
 
-  test "including associations with extensions and an instance dependent scope is not supported" do
-    e = assert_raises(ArgumentError) do
-      Author.includes(:posts_with_extension_and_instance).to_a
+  test "including associations with extensions and an instance dependent scope is supported" do
+    authors = Author.includes(:posts_with_extension_and_instance).to_a
+    assert_not authors.empty?
+    authors.each do |author|
+      assert_predicate author.posts_with_extension_and_instance, :loaded?
     end
-    assert_match(/Preloading instance dependent scopes is not supported/, e.message)
   end
 
   test "preloading readonly association" do

activerecord/test/cases/associations_test.rb

@@ -461,6 +461,93 @@ def test_preload_grouped_queries_of_through_records
     end
   end
 
+  def test_preload_with_instance_dependent_scope
+    david = authors(:david)
+    david2 = Author.create!(name: "David")
+    bob = authors(:bob)
+    post = Post.create!(
+      author: david,
+      title: "test post",
+      body: "this post is about David"
+    )
+    post2 = Post.create!(
+      author: david,
+      title: "test post 2",
+      body: "this post is also about David"
+    )
+
+    assert_queries(2) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [david, david2, bob], associations: :posts_mentioning_author)
+      preloader.call
+    end
+
+    assert_predicate david.posts_mentioning_author, :loaded?
+    assert_predicate david2.posts_mentioning_author, :loaded?
+    assert_predicate bob.posts_mentioning_author, :loaded?
+
+    assert_equal [post, post2].sort, david.posts_mentioning_author.sort
+    assert_equal [], david2.posts_mentioning_author
+    assert_equal [], bob.posts_mentioning_author
+  end
+
+  def test_preload_with_instance_dependent_through_scope
+    david = authors(:david)
+    david2 = Author.create!(name: "David")
+    bob = authors(:bob)
+    comment1 = david.posts.first.comments.create!(body: "Hi David!")
+    comment2 = david.posts.first.comments.create!(body: "This comment mentions david")
+
+    assert_queries(2) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [david, david2, bob], associations: :comments_mentioning_author)
+      preloader.call
+    end
+
+    assert_predicate david.comments_mentioning_author, :loaded?
+    assert_predicate david2.comments_mentioning_author, :loaded?
+    assert_predicate bob.comments_mentioning_author, :loaded?
+
+    assert_equal [comment1, comment2].sort, david.comments_mentioning_author.sort
+    assert_equal [], david2.comments_mentioning_author
+    assert_equal [], bob.comments_mentioning_author
+  end
+
+  def test_preload_with_through_instance_dependent_scope
+    david = authors(:david)
+    david2 = Author.create!(name: "David")
+    bob = authors(:bob)
+    post = Post.create!(
+      author: david,
+      title: "test post",
+      body: "this post is about David"
+    )
+    Post.create!(
+      author: david,
+      title: "test post 2",
+      body: "this post is also about David"
+    )
+    post3 = Post.create!(
+      author: bob,
+      title: "test post 3",
+      body: "this post is about Bob"
+    )
+    comment1 = post.comments.create!(body: "hi!")
+    comment2 = post.comments.create!(body: "hello!")
+    comment3 = post3.comments.create!(body: "HI BOB!")
+
+    assert_queries(3) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [david, david2, bob], associations: :comments_on_posts_mentioning_author)
+      preloader.call
+    end
+
+    assert_predicate david.comments_on_posts_mentioning_author, :loaded?
+    assert_predicate david2.comments_on_posts_mentioning_author, :loaded?
+    assert_predicate bob.comments_on_posts_mentioning_author, :loaded?
+
+    assert_equal [comment1, comment2].sort, david.comments_on_posts_mentioning_author.sort
+    assert_equal [], david2.comments_on_posts_mentioning_author
+    assert_equal [comment3], bob.comments_on_posts_mentioning_author
+  end
+
   def test_some_already_loaded_associations
     item_discount = Discount.create(amount: 5)
     shipping_discount = Discount.create(amount: 20)

activerecord/test/models/author.rb

@@ -205,8 +205,10 @@ def ratings
   has_many :posts_with_default_include, class_name: "PostWithDefaultInclude"
   has_many :comments_on_posts_with_default_include, through: :posts_with_default_include, source: :comments
 
-  has_many :posts_with_signature, ->(record) { where("posts.title LIKE ?", "%by #{record.name.downcase}%") }, class_name: "Post"
-  has_many :posts_mentioning_author, ->(record = nil) { where("posts.body LIKE ?", "%#{record&.name&.downcase}%") }, class_name: "Post"
+  has_many :posts_with_signature, ->(record) { where(arel_table[:title].matches("%by #{record.name.downcase}%")) }, class_name: "Post"
+  has_many :posts_mentioning_author, ->(record = nil) { where(arel_table[:body].matches("%#{record&.name&.downcase}%")) }, class_name: "Post"
+  has_many :comments_on_posts_mentioning_author, through: :posts_mentioning_author, source: :comments
+  has_many :comments_mentioning_author, ->(record) { where(arel_table[:body].matches("%#{record.name.downcase}%")) }, through: :posts, source: :comments
 
   has_one :recent_post, -> { order(id: :desc) }, class_name: "Post"
   has_one :recent_response, through: :recent_post, source: :comments