Add an option `end_at` to `find_in_batches`

that complements the `start`parameter to specify where to stop batch processing

1 files changed, 22 insertions, 12 deletions

diff --git a/activerecord/lib/active_record/relation/batches.rb b/activerecord/lib/active_record/relation/batches.rb
index 9d690af11d..a543341149 100644
--- a/activerecord/lib/active_record/relation/batches.rb
+++ b/activerecord/lib/active_record/relation/batches.rb
@@ -27,11 +27,12 @@ module ActiveRecord
     #
     # ==== Options
     # * <tt>:batch_size</tt> - Specifies the size of the batch. Default to 1000.
-    # * <tt>:start</tt> - Specifies the primary key value to start from.
+    # * <tt>:start</tt> - Specifies the primary key value to start from, inclusive of the value.
+    # * <tt>:end_at</tt> - Specifies the primary key value to end at, inclusive of the value.
     # This is especially useful if you want multiple workers dealing with
     # the same processing queue. You can make worker 1 handle all the records
     # between id 0 and 10,000 and worker 2 handle from 10,000 and beyond
-    # (by setting the +:start+ option on that worker).
+    # (by setting the +:start+ and +:end_at+ option on each worker).
     #
     #   # Let's process for a batch of 2000 records, skipping the first 2000 rows
     #   Person.find_each(start: 2000, batch_size: 2000) do |person|
@@ -45,14 +46,15 @@ module ActiveRecord
     #
     # NOTE: You can't set the limit either, that's used to control
     # the batch sizes.
-    def find_each(start: nil, batch_size: 1000)
+    def find_each(start: nil, end_at: nil, batch_size: 1000)
       if block_given?
-        find_in_batches(start: start, batch_size: batch_size) do |records|
+        find_in_batches(start: start, end_at: end_at, batch_size: batch_size) do |records|
           records.each { |record| yield record }
         end
       else
-        enum_for(:find_each, start: start, batch_size: batch_size) do
-          start ? where(table[primary_key].gteq(start)).size : size
+        enum_for(:find_each, start: start, end_at: end_at, batch_size: batch_size) do
+          relation = self
+          apply_limits(relation, start, end_at).size
         end
       end
     end
@@ -77,11 +79,12 @@ module ActiveRecord
     #
     # ==== Options
     # * <tt>:batch_size</tt> - Specifies the size of the batch. Default to 1000.
-    # * <tt>:start</tt> - Specifies the primary key value to start from.
+    # * <tt>:start</tt> - Specifies the primary key value to start from, inclusive of the value.
+    # * <tt>:end_at</tt> - Specifies the primary key value to end at, inclusive of the value.
     # This is especially useful if you want multiple workers dealing with
     # the same processing queue. You can make worker 1 handle all the records
     # between id 0 and 10,000 and worker 2 handle from 10,000 and beyond
-    # (by setting the +:start+ option on that worker).
+    # (by setting the +:start+ and +:end_at+ option on each worker).
     #
     #   # Let's process the next 2000 records
     #   Person.find_in_batches(start: 2000, batch_size: 2000) do |group|
@@ -95,12 +98,12 @@ module ActiveRecord
     #
     # NOTE: You can't set the limit either, that's used to control
     # the batch sizes.
-    def find_in_batches(start: nil, batch_size: 1000)
+    def find_in_batches(start: nil, end_at: nil, batch_size: 1000)
       relation = self
 
       unless block_given?
-        return to_enum(:find_in_batches, start: start, batch_size: batch_size) do
-          total = start ? where(table[primary_key].gteq(start)).size : size
+        return to_enum(:find_in_batches, start: start, end_at: end_at, batch_size: batch_size) do
+          total = apply_limits(relation, start, end_at).size
           (total - 1).div(batch_size) + 1
         end
       end
@@ -110,7 +113,8 @@ module ActiveRecord
       end
 
       relation = relation.reorder(batch_order).limit(batch_size)
-      records = start ? relation.where(table[primary_key].gteq(start)).to_a : relation.to_a
+      relation = apply_limits(relation, start, end_at)
+      records = relation.to_a
 
       while records.any?
         records_size = records.size
@@ -127,6 +131,12 @@ module ActiveRecord
 
     private
 
+    def apply_limits(relation, start, end_at)
+      relation = relation.where(table[primary_key].gteq(start)) if start
+      relation = relation.where(table[primary_key].lteq(end_at)) if end_at
+      relation
+    end
+
     def batch_order
       "#{quoted_table_name}.#{quoted_primary_key} ASC"
     end