Introduce MaintenanceTasks::Task.rescue_from

By default any errors encountered while processing an iteration will be raised and
the task with error. Errors can be rescued using the `rescue_from` handler. Errors
rescued this way will be reported by the handler and iteration will continue.

The default handler just calls `Rails.error.report` with the raised error.

```ruby
module Maintenance
  class UpdatePostRescueErrorTask < MaintenanceTasks::Task
    rescue_from ActiveRecord::RecordNotUnique
  end
end
```

Or you can pass a custom handler with a block.

```ruby
module Maintenance
  class UpdatePostRescueErrorTask < MaintenanceTasks::Task
    rescue_from ActiveRecord::RecordNotUnique do |error|
      CustomErrorHandler.report(error)
    end
  end
end
```

Co-authored-by: Adrianna Chang <[email protected]>

README.md

@@ -407,6 +407,34 @@ module Maintenance
 end
 ```
 
+### Handling errors during iteration
+
+By default any errors encountered while processing an iteration will be raised and
+the task with error. Errors can be rescued using the `rescue_from` handler. Errors
+rescued this way will be reported by the handler and iteration will continue.
+
+The default handler just calls `Rails.error.report` with the raised error.
+
+```ruby
+module Maintenance
+  class UpdatePostRescueErrorTask < MaintenanceTasks::Task
+    rescue_from ActiveRecord::RecordNotUnique
+  end
+end
+```
+
+Or you can pass a custom handler with a block.
+
+```ruby
+module Maintenance
+  class UpdatePostRescueErrorTask < MaintenanceTasks::Task
+    rescue_from ActiveRecord::RecordNotUnique do |error|
+      CustomErrorHandler.report(error)
+    end
+  end
+end
+```
+
 ### Throttling
 
 Maintenance tasks often modify a lot of data and can be taxing on your database.

app/jobs/concerns/maintenance_tasks/task_job_concern.rb

@@ -112,7 +112,11 @@ def task_iteration(input)
       end
     rescue => error
       @errored_element = input
-      raise error
+      if @task.exception_handlers.keys.include?(error.class)
+        @task.exception_handlers[error.class].call(error)
+      else
+        raise error
+      end
     end
 
     def before_perform

app/models/maintenance_tasks/task.rb

@@ -27,6 +27,8 @@ class NotFoundError < NameError; end
     # @api private
     class_attribute :collection_builder_strategy, default: NullCollectionBuilder.new
 
+    class_attribute :exception_handlers, default: {}
+
     define_callbacks :start, :complete, :error, :cancel, :pause, :interrupt
 
     class << self
@@ -200,6 +202,22 @@ def after_error(*filter_list, &block)
         set_callback(:error, :after, *filter_list, &block)
       end
 
+      # Rescue listed exceptions during an iteration. By default the rescued exception will be logged with
+      # Rails.error.report. A block can be supplied for custom handling behaviour.
+      #
+      # @param exceptions [Class] list of exceptions to rescue
+      # @yield [error] Block to be invoked with the exception instance
+      # @yieldparam error [StandardError] Exception instance
+      def rescue_from(*exceptions, &block)
+        handler = block_given? ? block : ->(e) { default_handler(e) }
+
+        handlers = exceptions.to_h do |exception|
+          [exception, handler]
+        end
+
+        self.exception_handlers = exception_handlers.merge(handlers)
+      end
+
       private
 
       def load_constants
@@ -208,6 +226,14 @@ def load_constants
 
         Rails.autoloaders.main.eager_load_namespace(namespace)
       end
+
+      def default_handler(error)
+        if Gem::Version.new(Rails::VERSION::STRING) > "7.2"
+          Rails.error.report(error)
+        else
+          Rails.logger.error("[#{self}] iteration failed with: #{error.class}: #{error.message}")
+        end
+      end
     end
 
     # The contents of a CSV file to be processed by a Task.

test/jobs/maintenance_tasks/task_job_test.rb

@@ -707,5 +707,45 @@ class << self
 
       TaskJob.perform_now(run)
     end
+
+    test ".rescue_on" do
+      run = Run.create!(task_name: "Maintenance::UpdatePostsTask")
+      block_called = false
+
+      Maintenance::UpdatePostsTask.rescue_from(ActiveRecord::RecordInvalid) do
+        block_called = true
+      end
+
+      Maintenance::UpdatePostsTask.any_instance.expects(:process).raises(ActiveRecord::RecordInvalid).twice
+
+      assert_nothing_raised do
+        TaskJob.perform_now(run)
+      end
+
+      assert_predicate(run.reload, :succeeded?)
+      assert_equal(run.tick_count, 2)
+      assert(block_called)
+    ensure
+      Maintenance::UpdatePostsTask.exception_handlers = {}
+    end
+
+    test ".rescue_on with handler that raises" do
+      run = Run.create!(task_name: "Maintenance::UpdatePostsTask")
+
+      Maintenance::UpdatePostsTask.rescue_from(ActiveRecord::RecordInvalid) do
+        raise
+      end
+
+      Maintenance::UpdatePostsTask.any_instance.expects(:process).raises(ActiveRecord::RecordInvalid)
+
+      assert_nothing_raised do
+        TaskJob.perform_now(run)
+      end
+
+      assert_predicate(run.reload, :errored?)
+      assert_equal(run.tick_count, 0)
+    ensure
+      Maintenance::UpdatePostsTask.exception_handlers = {}
+    end
   end
 end

test/models/maintenance_tasks/task_test.rb

@@ -116,5 +116,35 @@ class TaskTest < ActiveSupport::TestCase
       task = Task.new
       assert_nil task.cursor_columns
     end
+
+    test ".rescue_from reports errors by default" do
+      Maintenance::TestTask.rescue_from(StandardError)
+      task = Maintenance::TestTask.new
+
+      assert_equal([StandardError], task.exception_handlers.keys)
+      if Gem::Version.new(Rails::VERSION::STRING) > "7.2"
+        assert_error_reported(StandardError) do
+          task.exception_handlers[StandardError].call(StandardError.new)
+        end
+      else
+        Rails.logger.expects(:error).with("[Maintenance::TestTask] iteration failed with: StandardError: oh no!!")
+        task.exception_handlers[StandardError].call(StandardError.new("oh no!!"))
+      end
+    ensure
+      Maintenance::TestTask.exception_handlers = {}
+    end
+
+    test ".rescue_from with custom handler" do
+      handler_called = false
+      Maintenance::TestTask.rescue_from(StandardError) do
+        handler_called = true
+      end
+      task = Maintenance::TestTask.new
+      task.exception_handlers[StandardError].call(StandardError.new)
+
+      assert(handler_called)
+    ensure
+      Maintenance::TestTask.exception_handlers = {}
+    end
   end
 end