Add each_iteration callback

lib/job-iteration/iteration.rb

@@ -17,6 +17,7 @@ module Iteration
       define_callbacks :start
       define_callbacks :shutdown
       define_callbacks :complete
+      define_callbacks :each_iteration
     end
 
     module ClassMethods
@@ -36,6 +37,14 @@ def on_complete(*filters, &blk)
         set_callback(:complete, :after, *filters, &blk)
       end
 
+      def before_each_iteration(*filters, &blk)
+        set_callback(:each_iteration, :before, *filters, &blk)
+      end
+
+      def after_each_iteration(*filters, &blk)
+        set_callback(:each_iteration, :after, *filters, &blk)
+      end
+
       def supports_interruption?
         true
       end
@@ -123,7 +132,9 @@ def iterate_with_enumerator(enumerator, arguments)
       arguments = arguments.dup.freeze
       enumerator.each do |iteration, index|
         record_unit_of_work do
-          each_iteration(iteration, *arguments)
+          run_callbacks(:each_iteration) do
+            each_iteration(iteration, *arguments)
+          end
           self.cursor_position = index
         end
 

test/unit/active_job_iteration_test.rb

@@ -14,6 +14,10 @@ class SimpleIterationJob < ActiveJob::Base
       self.on_complete_called = 0
       cattr_accessor :on_shutdown_called, instance_accessor: false
       self.on_shutdown_called = 0
+      cattr_accessor :before_each_iteration_called, instance_accessor: false
+      self.before_each_iteration_called = 0
+      cattr_accessor :after_each_iteration_called, instance_accessor: false
+      self.after_each_iteration_called = 0
 
       on_start do
         self.class.on_start_called += 1
@@ -26,6 +30,14 @@ class SimpleIterationJob < ActiveJob::Base
       on_shutdown do
         self.class.on_shutdown_called += 1
       end
+
+      before_each_iteration do
+        self.class.before_each_iteration_called += 1
+      end
+
+      after_each_iteration do
+        self.class.after_each_iteration_called += 1
+      end
     end
 
     class MultiArgumentIterationJob < SimpleIterationJob
@@ -276,6 +288,9 @@ def setup
         klass.on_start_called = 0
         klass.on_complete_called = 0
         klass.on_shutdown_called = 0
+        klass.on_shutdown_called = 0
+        klass.before_each_iteration_called = 0
+        klass.after_each_iteration_called = 0
       end
       JobShouldExitJob.records_performed = []
       super
@@ -345,6 +360,16 @@ def test_failing_job
       assert_jobs_in_queue 0
     end
 
+    def test_failing_job_runs_each_iteration_callbacks
+      push(FailingIterationJob)
+
+      work_one_job
+      assert_jobs_in_queue 1
+
+      assert_equal 4, FailingIterationJob.before_each_iteration_called
+      assert_equal 3, FailingIterationJob.after_each_iteration_called
+    end
+
     def test_active_record_job
       iterate_exact_times(2.times)
 
@@ -440,6 +465,15 @@ def test_multiple_columns
       assert_equal first_products, MultipleColumnsActiveRecordIterationJob.records_performed
     end
 
+    def test_multiple_columns_callback_for_each_iteration
+      iterate_exact_times(3.times)
+      push(MultipleColumnsActiveRecordIterationJob)
+      work_one_job
+
+      assert_equal 3, MultipleColumnsActiveRecordIterationJob.before_each_iteration_called
+      assert_equal 3, MultipleColumnsActiveRecordIterationJob.after_each_iteration_called
+    end
+
     def test_single_iteration
       push(SingleIterationJob)
 
@@ -452,6 +486,15 @@ def test_single_iteration
       assert_equal 1, SingleIterationJob.on_complete_called
     end
 
+    def test_single_iteration_callback_for_each_iteration
+      push(SingleIterationJob)
+      work_one_job
+      assert_jobs_in_queue 0
+
+      assert_equal 1, SingleIterationJob.before_each_iteration_called
+      assert_equal 1, SingleIterationJob.after_each_iteration_called
+    end
+
     def test_relation_with_limit
       push(LimitActiveRecordIterationJob)
 
@@ -557,6 +600,14 @@ def test_log_completion_data
       end
     end
 
+    def test_aborting_in_each_iteration_job_will_execute_each_iteration_callback
+      push(AbortingActiveRecordIterationJob)
+      work_one_job
+      assert_equal 2, AbortingActiveRecordIterationJob.records_performed.size
+      assert_equal 3, AbortingActiveRecordIterationJob.before_each_iteration_called
+      assert_equal 2, AbortingActiveRecordIterationJob.after_each_iteration_called
+    end
+
     def test_aborting_in_each_iteration_job_will_execute_on_complete_callback
       push(AbortingActiveRecordIterationJob)
       work_one_job