Skip to content

Commit

Permalink
Fix Sentry handler from LocalTaskJob causing error (#18119)
Browse files Browse the repository at this point in the history
The `enrich_errors` method assumes the first argument to the function
its patch is a TaskInstance when infact it can also be a LocalTaskJob.

This is now handled by extracting the task_instance from the
LocalTaskJob

Closes #18118
  • Loading branch information
robinedwards authored Sep 9, 2021
1 parent c9d2946 commit f97ddf1
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions airflow/sentry.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,11 +144,14 @@ def add_breadcrumbs(self, task_instance, session=None):
sentry_sdk.add_breadcrumb(category="completed_tasks", data=data, level="info")

def enrich_errors(self, func):
"""Wrap TaskInstance._run_raw_task to support task specific tags and breadcrumbs."""
"""
Wrap TaskInstance._run_raw_task and LocalTaskJob._run_mini_scheduler_on_child_tasks
to support task specific tags and breadcrumbs.
"""
session_args_idx = find_session_idx(func)

@wraps(func)
def wrapper(task_instance, *args, **kwargs):
def wrapper(_self, *args, **kwargs):
# Wrapping the _run_raw_task function with push_scope to contain
# tags and breadcrumbs to a specific Task Instance

Expand All @@ -159,8 +162,14 @@ def wrapper(task_instance, *args, **kwargs):

with sentry_sdk.push_scope():
try:
return func(task_instance, *args, **kwargs)
return func(_self, *args, **kwargs)
except Exception as e:
# Is a LocalTaskJob get the task instance
if hasattr(_self, 'task_instance'):
task_instance = _self.task_instance
else:
task_instance = _self

self.add_tagging(task_instance)
self.add_breadcrumbs(task_instance, session=session)
sentry_sdk.capture_exception(e)
Expand Down

0 comments on commit f97ddf1

Please sign in to comment.