Patch Rake::Task#execute instead of Rake::Application#display_error_message

lib/exception_notifier/rake.rb

@@ -1,3 +1,4 @@
+require 'exception_notifier/rake/rails' if defined?(Rails)
 require 'exception_notifier/rake/rake'
 require 'exception_notifier/rake/rake_patch'
 require 'exception_notifier/rake/version'

lib/exception_notifier/rake/rails.rb

@@ -0,0 +1,17 @@
+# Based on/adapted from https://github.com/airbrake/airbrake/blob/master/lib/airbrake/rails.rb
+
+module ExceptionNotifier
+  class Rake
+    class Railtie < ::Rails::Railtie
+      rake_tasks do
+        # Report exceptions occurring in Rake tasks.
+        require 'exception_notifier/rake/rake_patch'
+        # Work around https://github.com/nikhaldi/exception_notification-rake/issues/26
+        # Rake::TaskManager won't have been defined when rake_patch.rb was first loaded.
+        if Rails.env.development?
+          load 'exception_notifier/rake/rake_patch.rb'
+        end
+      end
+    end
+  end
+end

lib/exception_notifier/rake/rake_patch.rb

@@ -1,25 +1,61 @@
-# Monkey patching patterns lifted from
-# https://github.com/thoughtbot/airbrake/blob/master/lib/airbrake/rake_handler.rb
+# Copied/adapted from https://github.com/airbrake/airbrake/blob/master/lib/airbrake/rake.rb
+
+if Rake.const_defined?(:TaskManager)
+  Rake::TaskManager.record_task_metadata = true
+end
+
 module ExceptionNotifier
-  module RakePatch
-    def display_error_message(ex)
-      super(ex)
-      ExceptionNotifier::Rake.maybe_deliver_notification(ex,
-        :rake_command_line => reconstruct_command_line)
+  module RakeTaskPatch
+    # A wrapper around the original +#execute+, that catches all errors and
+    # passes them on to ExceptionNotifier.
+    #
+    # rubocop:disable Lint/RescueException
+    def execute(args = nil)
+      super(args)
+    rescue Exception => ex
+      ExceptionNotifier::Rake.maybe_deliver_notification(
+        ex,
+        task_info,
+      )
+      raise ex
     end
+    # rubocop:enable Lint/RescueException
+
+  private
+
+    # rubocop:disable Metrics/CyclomaticComplexity, Metrics/AbcSize
+    def task_info
+      info = {}
+
+      info[:rake_command_line] = reconstruct_command_line
+      info[:name] = name
+      info[:timestamp] = timestamp.to_s
+      info[:investigation] = investigation
+
+      info[:full_comment] = full_comment if full_comment
+      info[:arg_names] = arg_names if arg_names.any?
+      info[:arg_description] = arg_description if arg_description
+      info[:locations] = locations if locations.any?
+      info[:sources] = sources if sources.any?
+
+      if prerequisite_tasks.any?
+        info[:prerequisite_tasks] = prerequisite_tasks.map do |p|
+          p.__send__(:task_info)
+        end
+      end
+
+      info
+    end
+    # rubocop:enable Metrics/CyclomaticComplexity, Metrics/AbcSize
 
     def reconstruct_command_line
       "rake #{ARGV.join(' ')}"
     end
   end
 end
 
-# Only do this if we're actually in a Rake context. In some contexts (e.g.,
-# in the Rails console) Rake might not be defined.
-if Object.const_defined?(:Rake) && Rake.respond_to?(:application)
-  Rake.application.instance_eval do
-    class << self
-      prepend ExceptionNotifier::RakePatch
-    end
+module Rake
+  class Task
+    prepend ExceptionNotifier::RakeTaskPatch
   end
 end