Merge pull request #28867 from eugeneius/skip_query_cache_in_batches

Skip query cache for in_batches and friends

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.
 

activerecord/lib/active_record/relation.rb

@@ -6,7 +6,7 @@ class Relation
                             :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 @@ def has_join_values?
       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

activerecord/lib/active_record/relation/batches.rb

@@ -209,6 +209,7 @@ def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore:
 
       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

activerecord/lib/active_record/relation/calculations.rb

@@ -184,7 +184,7 @@ def pluck(*column_names)
         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 @@ def execute_simple_calculation(operation, column_name, distinct) #:nodoc:
           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 @@ def execute_grouped_calculation(operation, column_name, distinct) #:nodoc:
         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] }

activerecord/lib/active_record/relation/finder_methods.rb

@@ -315,7 +315,7 @@ def exists?(conditions = :none)
 
       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 @@ def find_with_associations
           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 @@ def limited_ids_for(relation)
 
         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
 

activerecord/lib/active_record/relation/query_methods.rb

@@ -913,6 +913,11 @@ def reverse_order! # :nodoc:
       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

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 @@ def test_find_in_batches_should_return_a_sized_enumerator
     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

activerecord/test/cases/calculations_test.rb

@@ -817,4 +817,46 @@ def test_deprecate_sum_with_block_and_column_name
       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

activerecord/test/cases/finder_test.rb

@@ -1232,6 +1232,34 @@ def test_finder_with_offset_string
     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

activerecord/test/cases/relation/mutation_test.rb

@@ -90,7 +90,7 @@ def relation
       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 @@ def relation
       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

activerecord/test/cases/relations_test.rb

@@ -2034,4 +2034,46 @@ def stubbed_connection.combine_bind_parameters(**kwargs)
 
     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