Adds +as_relation+ flag to EnumeratorBuilder#active_record_on_batches

Currently, the resulting batch from #active_record_on_batches is an Array of records.
Many job authors end up converting this into a relation to perform batch operations
such as `#update_all`. By expanding the API of #active_record_on_batches to return
a relation if a flag is set, users who require their batch as an ActiveRecord::Relation
will no longer need to perform an additional query to turn records back into a relation.

lib/job-iteration/active_record_cursor.rb

@@ -55,7 +55,7 @@ def update_from_record(record)
       end
     end
 
-    def next_batch(batch_size)
+    def next_batch(batch_size, as_relation)
       return nil if @reached_end
 
       relation = @base_relation.limit(batch_size)
@@ -65,7 +65,7 @@ def next_batch(batch_size)
       end
 
       records = relation.uncached do
-        relation.to_a
+        as_relation ? relation : relation.to_a
       end
 
       update_from_record(records.last) unless records.empty?

lib/job-iteration/active_record_enumerator.rb

@@ -6,9 +6,10 @@ module JobIteration
   class ActiveRecordEnumerator
     SQL_DATETIME_WITH_NSEC = "%Y-%m-%d %H:%M:%S.%N"
 
-    def initialize(relation, columns: nil, batch_size: 100, cursor: nil)
+    def initialize(relation, columns: nil, batch_size: 100, as_relation: false, cursor: nil)
       @relation = relation
       @batch_size = batch_size
+      @as_relation = as_relation
       @columns = Array(columns || "#{relation.table_name}.#{relation.primary_key}")
       @cursor = cursor
     end
@@ -26,7 +27,7 @@ def records
     def batches
       cursor = finder_cursor
       Enumerator.new(method(:size)) do |yielder|
-        while (records = cursor.next_batch(@batch_size))
+        while (records = cursor.next_batch(@batch_size, @as_relation))
           yielder.yield(records, cursor_value(records.last)) if records.any?
         end
       end

lib/job-iteration/enumerator_builder.rb

@@ -99,6 +99,8 @@ def build_active_record_enumerator_on_records(scope, cursor:, **args)
     # Each Enumerator tick moves the cursor +batch_size+ rows forward.
     #
     # +batch_size:+ sets how many records will be fetched in one batch. Defaults to 100.
+    # +as_relation:+ when set to true, returns the batch as an ActiveRecord::Relation instead of a batch of records.
+    # Defaults to false.
     #
     # For the rest of arguments, see documentation for #build_active_record_enumerator_on_records
     def build_active_record_enumerator_on_batches(scope, cursor:, **args)

test/unit/active_job_iteration_test.rb

@@ -75,6 +75,21 @@ def each_iteration(record)
       end
     end
 
+    class BatchActiveRecordRelationIterationJob < SimpleIterationJob
+      def build_enumerator(cursor:)
+        enumerator_builder.active_record_on_batches(
+          Product.all,
+          cursor: cursor,
+          batch_size: 3,
+          as_relation: true
+        )
+      end
+
+      def each_iteration(relation)
+        self.class.records_performed << relation
+      end
+    end
+
     class AbortingActiveRecordIterationJob < ActiveRecordIterationJob
       def each_iteration(*)
         abort_strategy if self.class.records_performed.size == 2
@@ -412,6 +427,17 @@ def test_activerecord_batches_complete
       assert_equal(processed_records, BatchActiveRecordIterationJob.records_performed.flatten.map(&:id))
     end
 
+    def test_activerecord_batches_as_relation
+      push(BatchActiveRecordRelationIterationJob)
+
+      work_one_job
+      assert_jobs_in_queue(0)
+
+      records_performed = BatchActiveRecordIterationJob.records_performed
+      assert_equal([3, 3, 3, 1], records_performed.map(&:size))
+      assert(records_performed.all? { |record| record.is_a?(ActiveRecord::Relation) })
+    end
+
     def test_activerecord_batches
       iterate_exact_times(1.times)