Allow Active Record collection tasks to specify a batch size.

README.md

@@ -156,6 +156,31 @@ module Maintenance
 end
 ```
 
+#### Customizing the Batch Size
+
+When processing records from an Active Record Relation, records are fetched in
+batches internally, and then each record is passed to the `#process` method.
+Maintenance Tasks will query the database to fetch records in batches of 100 by
+default, but the batch size can be modified using the `collection_batch_size` macro:
+
+```ruby
+# app/tasks/maintenance/update_posts_task.rb
+
+module Maintenance
+  class UpdatePostsTask < MaintenanceTasks::Task
+    # Fetch records in batches of 1000
+    collection_batch_size(1000)
+
+    def collection
+      Post.all
+    end
+
+    def process(post)
+      post.update!(content: "New content!")
+    end
+  end
+end
+
 ### Creating a CSV Task
 
 You can also write a Task that iterates on a CSV file. Note that writing CSV

app/jobs/concerns/maintenance_tasks/task_job_concern.rb

@@ -38,7 +38,9 @@ def build_enumerator(_run, cursor:)
       when :no_collection
         enumerator_builder.build_once_enumerator(cursor: nil)
       when ActiveRecord::Relation
-        enumerator_builder.active_record_on_records(collection, cursor: cursor, columns: @task.cursor_columns)
+        options = { cursor: cursor, columns: @task.cursor_columns }
+        options[:batch_size] = @task.active_record_enumerator_batch_size if @task.active_record_enumerator_batch_size
+        enumerator_builder.active_record_on_records(collection, **options)
       when ActiveRecord::Batches::BatchEnumerator
         if collection.start || collection.finish
           raise ArgumentError, <<~MSG.squish

app/models/maintenance_tasks/task.rb

@@ -18,6 +18,12 @@ class NotFoundError < NameError; end
     # @api private
     class_attribute :throttle_conditions, default: []
 
+    # The number of active records to fetch in a single query when iterating
+    # over an Active Record collection task.
+    #
+    # @api private
+    class_attribute :active_record_enumerator_batch_size
+
     # @api private
     class_attribute :collection_builder_strategy, default: NullCollectionBuilder.new
 
@@ -137,6 +143,14 @@ def throttle_on(backoff: 30.seconds, &condition)
         self.throttle_conditions += [{ throttle_on: condition, backoff: backoff_as_proc }]
       end
 
+      # Limit the number of records that will be fetched in a single query when
+      # iterating over an Active Record collection task.
+      #
+      # @param size [Integer] the number of records to fetch in a single query.
+      def collection_batch_size(size)
+        self.active_record_enumerator_batch_size = size
+      end
+
       # Initialize a callback to run after the task starts.
       #
       # @param filter_list apply filters to the callback

test/dummy/app/tasks/maintenance/import_posts_with_options_task.rb

@@ -2,7 +2,7 @@
 
 module Maintenance
   class ImportPostsWithOptionsTask < MaintenanceTasks::Task
-    csv_collection(skip_lines: /^#/, converters: ->(field) { field.upcase })
+    csv_collection(skip_lines: /^#/, converters: ->(field) { field.to_s.upcase })
 
     def process(row)
       Post.create!(title: row["title"], content: row["content"])

test/dummy/app/tasks/maintenance/update_posts_task.rb

@@ -6,6 +6,8 @@ class << self
       attr_accessor :fast_task
     end
 
+    collection_batch_size 1000
+
     def collection
       Post.all
     end

test/jobs/maintenance_tasks/task_job_test.rb

@@ -613,6 +613,17 @@ class << self
       assert_equal 2, run.reload.tick_total
     end
 
+    test "Active Record Relation tasks can specify a relation batch size" do
+      run = Run.create!(task_name: "Maintenance::UpdatePostsTask")
+
+      JobIteration::EnumeratorBuilder
+        .any_instance
+        .expects(:active_record_on_records)
+        .with(anything, has_entry(batch_size: 1000))
+
+      TaskJob.perform_now(run)
+    end
+
     test "MaintenanceTasks::TaskJobConcern#build_enumerator provides cursor_columns as the column argument to active_record_on_records" do
       cursor_columns = [:created_at, :id]