Merge pull request #28867 from eugeneius/skip_query_cache_in_batches

Skip query cache for in_batches and friends

11 files changed, 224 insertions, 20 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 53b9752867..5b88722993 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,11 @@
+*   Skip query caching when working with batches of records (`find_each`, `find_in_batches`,
+    `in_batches`).
+
+    Previously, records would be fetched in batches, but all records would be retained in memory
+    until the end of the request or job.
+
+    *Eugene Kenny*
+
 *   Prevent errors raised by sql.active_record notification subscribers from being converted into
     ActiveRecord::StatementInvalid exceptions.
 
diff --git a/activerecord/lib/active_record/relation.rb b/activerecord/lib/active_record/relation.rb
index 52f5d5f3e3..76cf47a3ed 100644
--- a/activerecord/lib/active_record/relation.rb
+++ b/activerecord/lib/active_record/relation.rb
@@ -6,7 +6,7 @@ module ActiveRecord
                             :extending, :unscope]
 
     SINGLE_VALUE_METHODS = [:limit, :offset, :lock, :readonly, :reordering,
-                            :reverse_order, :distinct, :create_with]
+                            :reverse_order, :distinct, :create_with, :skip_query_cache]
     CLAUSE_METHODS = [:where, :having, :from]
     INVALID_METHODS_FOR_DELETE_ALL = [:limit, :distinct, :offset, :group, :having]
 
@@ -657,20 +657,32 @@ module ActiveRecord
       end
 
       def exec_queries(&block)
-        @records = eager_loading? ? find_with_associations.freeze : @klass.find_by_sql(arel, bound_attributes, &block).freeze
-
-        preload = preload_values
-        preload += includes_values unless eager_loading?
-        preloader = nil
-        preload.each do |associations|
-          preloader ||= build_preloader
-          preloader.preload @records, associations
-        end
+        skip_query_cache_if_necessary do
+          @records = eager_loading? ? find_with_associations.freeze : @klass.find_by_sql(arel, bound_attributes, &block).freeze
+
+          preload = preload_values
+          preload += includes_values unless eager_loading?
+          preloader = nil
+          preload.each do |associations|
+            preloader ||= build_preloader
+            preloader.preload @records, associations
+          end
 
-        @records.each(&:readonly!) if readonly_value
+          @records.each(&:readonly!) if readonly_value
 
-        @loaded = true
-        @records
+          @loaded = true
+          @records
+        end
+      end
+
+      def skip_query_cache_if_necessary
+        if skip_query_cache_value
+          uncached do
+            yield
+          end
+        else
+          yield
+        end
       end
 
       def build_preloader
diff --git a/activerecord/lib/active_record/relation/batches.rb b/activerecord/lib/active_record/relation/batches.rb
index ee1f25ec84..c7e4f8a88a 100644
--- a/activerecord/lib/active_record/relation/batches.rb
+++ b/activerecord/lib/active_record/relation/batches.rb
@@ -209,6 +209,7 @@ module ActiveRecord
 
       relation = relation.reorder(batch_order).limit(batch_limit)
       relation = apply_limits(relation, start, finish)
+      relation.skip_query_cache! # Retaining the results in the query cache would undermine the point of batching
       batch_relation = relation
 
       loop do
diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 8a54f8f2c3..aaba6c71f2 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -184,7 +184,7 @@ module ActiveRecord
         relation.select_values = column_names.map { |cn|
           @klass.has_attribute?(cn) || @klass.attribute_alias?(cn) ? arel_attribute(cn) : cn
         }
-        result = klass.connection.select_all(relation.arel, nil, bound_attributes)
+        result = skip_query_cache_if_necessary { klass.connection.select_all(relation.arel, nil, bound_attributes) }
         result.cast_values(klass.attribute_types)
       end
     end
@@ -260,7 +260,7 @@ module ActiveRecord
           query_builder = relation.arel
         end
 
-        result = @klass.connection.select_all(query_builder, nil, bound_attributes)
+        result = skip_query_cache_if_necessary { @klass.connection.select_all(query_builder, nil, bound_attributes) }
         row    = result.first
         value  = row && row.values.first
         type   = result.column_types.fetch(column_alias) do
@@ -311,7 +311,7 @@ module ActiveRecord
         relation.group_values  = group_fields
         relation.select_values = select_values
 
-        calculated_data = @klass.connection.select_all(relation.arel, nil, relation.bound_attributes)
+        calculated_data = skip_query_cache_if_necessary { @klass.connection.select_all(relation.arel, nil, relation.bound_attributes) }
 
         if association
           key_ids     = calculated_data.collect { |row| row[group_aliases.first] }
diff --git a/activerecord/lib/active_record/relation/finder_methods.rb b/activerecord/lib/active_record/relation/finder_methods.rb
index eee0f36f63..ac0b4f597e 100644
--- a/activerecord/lib/active_record/relation/finder_methods.rb
+++ b/activerecord/lib/active_record/relation/finder_methods.rb
@@ -315,7 +315,7 @@ module ActiveRecord
 
       relation = construct_relation_for_exists(relation, conditions)
 
-      connection.select_value(relation.arel, "#{name} Exists", relation.bound_attributes) ? true : false
+      skip_query_cache_if_necessary { connection.select_value(relation.arel, "#{name} Exists", relation.bound_attributes) } ? true : false
     rescue ::RangeError
       false
     end
@@ -376,7 +376,7 @@ module ActiveRecord
           if ActiveRecord::NullRelation === relation
             []
           else
-            rows = connection.select_all(relation.arel, "SQL", relation.bound_attributes)
+            rows = skip_query_cache_if_necessary { connection.select_all(relation.arel, "SQL", relation.bound_attributes) }
             join_dependency.instantiate(rows, aliases)
           end
         end
@@ -424,7 +424,7 @@ module ActiveRecord
 
         relation = relation.except(:select).select(values).distinct!
 
-        id_rows = @klass.connection.select_all(relation.arel, "SQL", relation.bound_attributes)
+        id_rows = skip_query_cache_if_necessary { @klass.connection.select_all(relation.arel, "SQL", relation.bound_attributes) }
         id_rows.map { |row| row[primary_key] }
       end
 
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index 9da8f96337..79495ead91 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -913,6 +913,11 @@ module ActiveRecord
       self
     end
 
+    def skip_query_cache! # :nodoc:
+      self.skip_query_cache_value = true
+      self
+    end
+
     # Returns the Arel object associated with the relation.
     def arel # :nodoc:
       @arel ||= build_arel
diff --git a/activerecord/test/cases/batches_test.rb b/activerecord/test/cases/batches_test.rb
index 1a66b82b2e..dcd3af487b 100644
--- a/activerecord/test/cases/batches_test.rb
+++ b/activerecord/test/cases/batches_test.rb
@@ -1,4 +1,5 @@
 require "cases/helper"
+require "models/comment"
 require "models/post"
 require "models/subscriber"
 
@@ -610,4 +611,64 @@ class EachTest < ActiveRecord::TestCase
     end
     assert_equal expected, actual
   end
+
+  test ".find_each bypasses the query cache for its own queries" do
+    Post.cache do
+      assert_queries(2) do
+        Post.find_each {}
+        Post.find_each {}
+      end
+    end
+  end
+
+  test ".find_each does not disable the query cache inside the given block" do
+    Post.cache do
+      Post.find_each(start: 1, finish: 1) do |post|
+        assert_queries(1) do
+          post.comments.count
+          post.comments.count
+        end
+      end
+    end
+  end
+
+  test ".find_in_batches bypasses the query cache for its own queries" do
+    Post.cache do
+      assert_queries(2) do
+        Post.find_in_batches {}
+        Post.find_in_batches {}
+      end
+    end
+  end
+
+  test ".find_in_batches does not disable the query cache inside the given block" do
+    Post.cache do
+      Post.find_in_batches(start: 1, finish: 1) do |batch|
+        assert_queries(1) do
+          batch.first.comments.count
+          batch.first.comments.count
+        end
+      end
+    end
+  end
+
+  test ".in_batches bypasses the query cache for its own queries" do
+    Post.cache do
+      assert_queries(2) do
+        Post.in_batches {}
+        Post.in_batches {}
+      end
+    end
+  end
+
+  test ".in_batches does not disable the query cache inside the given block" do
+    Post.cache do
+      Post.in_batches(start: 1, finish: 1) do |relation|
+        assert_queries(1) do
+          relation.count
+          relation.count
+        end
+      end
+    end
+  end
 end
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index 80baaac30a..7d6dc21e34 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -817,4 +817,46 @@ class CalculationsTest < ActiveRecord::TestCase
       assert_equal 6, Account.sum(:firm_id) { 1 }
     end
   end
+
+  test "#skip_query_cache! for #pluck" do
+    Account.cache do
+      assert_queries(1) do
+        Account.pluck(:credit_limit)
+        Account.pluck(:credit_limit)
+      end
+
+      assert_queries(2) do
+        Account.all.skip_query_cache!.pluck(:credit_limit)
+        Account.all.skip_query_cache!.pluck(:credit_limit)
+      end
+    end
+  end
+
+  test "#skip_query_cache! for a simple calculation" do
+    Account.cache do
+      assert_queries(1) do
+        Account.calculate(:sum, :credit_limit)
+        Account.calculate(:sum, :credit_limit)
+      end
+
+      assert_queries(2) do
+        Account.all.skip_query_cache!.calculate(:sum, :credit_limit)
+        Account.all.skip_query_cache!.calculate(:sum, :credit_limit)
+      end
+    end
+  end
+
+  test "#skip_query_cache! for a grouped calculation" do
+    Account.cache do
+      assert_queries(1) do
+        Account.group(:firm_id).calculate(:sum, :credit_limit)
+        Account.group(:firm_id).calculate(:sum, :credit_limit)
+      end
+
+      assert_queries(2) do
+        Account.all.skip_query_cache!.group(:firm_id).calculate(:sum, :credit_limit)
+        Account.all.skip_query_cache!.group(:firm_id).calculate(:sum, :credit_limit)
+      end
+    end
+  end
 end
diff --git a/activerecord/test/cases/finder_test.rb b/activerecord/test/cases/finder_test.rb
index 420f552ef6..af21cd529f 100644
--- a/activerecord/test/cases/finder_test.rb
+++ b/activerecord/test/cases/finder_test.rb
@@ -1232,6 +1232,34 @@ class FinderTest < ActiveRecord::TestCase
     assert_equal tyre2, zyke.tyres.custom_find_by(id: tyre2.id)
   end
 
+  test "#skip_query_cache! for #exists?" do
+    Topic.cache do
+      assert_queries(1) do
+        Topic.exists?
+        Topic.exists?
+      end
+
+      assert_queries(2) do
+        Topic.all.skip_query_cache!.exists?
+        Topic.all.skip_query_cache!.exists?
+      end
+    end
+  end
+
+  test "#skip_query_cache! for #exists? with a limited eager load" do
+    Topic.cache do
+      assert_queries(2) do
+        Topic.eager_load(:replies).limit(1).exists?
+        Topic.eager_load(:replies).limit(1).exists?
+      end
+
+      assert_queries(4) do
+        Topic.eager_load(:replies).limit(1).skip_query_cache!.exists?
+        Topic.eager_load(:replies).limit(1).skip_query_cache!.exists?
+      end
+    end
+  end
+
   private
     def table_with_custom_primary_key
       yield(Class.new(Toy) do
diff --git a/activerecord/test/cases/relation/mutation_test.rb b/activerecord/test/cases/relation/mutation_test.rb
index dea787c07f..8e73baa70a 100644
--- a/activerecord/test/cases/relation/mutation_test.rb
+++ b/activerecord/test/cases/relation/mutation_test.rb
@@ -90,7 +90,7 @@ module ActiveRecord
       assert_equal [], relation.extending_values
     end
 
-    (Relation::SINGLE_VALUE_METHODS - [:lock, :reordering, :reverse_order, :create_with]).each do |method|
+    (Relation::SINGLE_VALUE_METHODS - [:lock, :reordering, :reverse_order, :create_with, :skip_query_cache]).each do |method|
       test "##{method}!" do
         assert relation.public_send("#{method}!", :foo).equal?(relation)
         assert_equal :foo, relation.public_send("#{method}_value")
@@ -162,5 +162,10 @@ module ActiveRecord
       relation.distinct! :foo
       assert_equal :foo, relation.distinct_value
     end
+
+    test "skip_query_cache!" do
+      relation.skip_query_cache!
+      assert relation.skip_query_cache_value
+    end
   end
 end
diff --git a/activerecord/test/cases/relations_test.rb b/activerecord/test/cases/relations_test.rb
index 5767dec315..eb3449b331 100644
--- a/activerecord/test/cases/relations_test.rb
+++ b/activerecord/test/cases/relations_test.rb
@@ -2034,4 +2034,46 @@ class RelationTest < ActiveRecord::TestCase
 
     assert_equal 2, posts.to_a.length
   end
+
+  test "#skip_query_cache!" do
+    Post.cache do
+      assert_queries(1) do
+        Post.all.load
+        Post.all.load
+      end
+
+      assert_queries(2) do
+        Post.all.skip_query_cache!.load
+        Post.all.skip_query_cache!.load
+      end
+    end
+  end
+
+  test "#skip_query_cache! with an eager load" do
+    Post.cache do
+      assert_queries(1) do
+        Post.eager_load(:comments).load
+        Post.eager_load(:comments).load
+      end
+
+      assert_queries(2) do
+        Post.eager_load(:comments).skip_query_cache!.load
+        Post.eager_load(:comments).skip_query_cache!.load
+      end
+    end
+  end
+
+  test "#skip_query_cache! with a preload" do
+    Post.cache do
+      assert_queries(2) do
+        Post.preload(:comments).load
+        Post.preload(:comments).load
+      end
+
+      assert_queries(4) do
+        Post.preload(:comments).skip_query_cache!.load
+        Post.preload(:comments).skip_query_cache!.load
+      end
+    end
+  end
 end