From 9dfd0fc937b06a0cd82f79ae39557797be97375f Mon Sep 17 00:00:00 2001
From: Adrianna Chang <adrianna.chang@shopify.com>
Date: Thu, 18 Mar 2021 15:31:41 -0400
Subject: [PATCH] Handle no run when rescuing errors in TaskJob (#375)

* Handle no run when rescuing errors in TaskJob
---
 README.md                                    |  5 ++++-
 app/jobs/maintenance_tasks/task_job.rb       | 17 ++++++++++------
 test/jobs/maintenance_tasks/task_job_test.rb | 21 ++++++++++++++++++++
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index a84f7260..2164f864 100644
--- a/README.md
+++ b/README.md
@@ -339,12 +339,15 @@ end
 
 The error handler should be a lambda that accepts three arguments:
 
-* `error`: The object containing the exception that was raised.
+* `error`: The exception that was raised.
 * `task_context`: A hash with additional information about the Task and the
   error:
   * `task_name`: The name of the Task that errored
   * `started_at`: The time the Task started
   * `ended_at`: The time the Task errored
+  Note that `task_context` may be empty if the Task produced an error before any
+  context could be gathered (for example, if deserializing the job to process
+  your Task failed).
 * `errored_element`: The element, if any, that was being processed when the
   Task raised an exception. If you would like to pass this object to your
   exception monitoring service, make sure you **sanitize the object** to avoid
diff --git a/app/jobs/maintenance_tasks/task_job.rb b/app/jobs/maintenance_tasks/task_job.rb
index ea878168..373a48ee 100644
--- a/app/jobs/maintenance_tasks/task_job.rb
+++ b/app/jobs/maintenance_tasks/task_job.rb
@@ -122,13 +122,18 @@ def after_perform
 
     def on_error(error)
       @ticker.persist if defined?(@ticker)
-      @run.persist_error(error)
 
-      task_context = {
-        task_name: @run.task_name,
-        started_at: @run.started_at,
-        ended_at: @run.ended_at,
-      }
+      if defined?(@run)
+        @run.persist_error(error)
+
+        task_context = {
+          task_name: @run.task_name,
+          started_at: @run.started_at,
+          ended_at: @run.ended_at,
+        }
+      else
+        task_context = {}
+      end
       errored_element = @errored_element if defined?(@errored_element)
       MaintenanceTasks.error_handler.call(error, task_context, errored_element)
     end
diff --git a/test/jobs/maintenance_tasks/task_job_test.rb b/test/jobs/maintenance_tasks/task_job_test.rb
index 714d8493..c7e0d53b 100644
--- a/test/jobs/maintenance_tasks/task_job_test.rb
+++ b/test/jobs/maintenance_tasks/task_job_test.rb
@@ -316,5 +316,26 @@ class << self
     ensure
       MaintenanceTasks.error_handler = error_handler_before
     end
+
+    test '.perform_now handles case where run is not set and calls error handler' do
+      error_handler_before = MaintenanceTasks.error_handler
+      handled_error = nil
+      handled_task_context = nil
+      MaintenanceTasks.error_handler = ->(error, task_context, _errored_el) do
+        handled_error = error
+        handled_task_context = task_context
+      end
+
+      RaisingTaskJob = Class.new(TaskJob) do
+        before_perform(prepend: true) { raise 'Uh oh!' }
+      end
+
+      RaisingTaskJob.perform_now(@run)
+
+      assert_equal('Uh oh!', handled_error.message)
+      assert_empty(handled_task_context)
+    ensure
+      MaintenanceTasks.error_handler = error_handler_before
+    end
   end
 end