Skip to content

Commit e971caf

Browse files
authored
feat(celery): Support Celery abstract tasks (#1287)
Prior to this change, the Celery integration always instruments `task.run` and incorrectly instruments `task.__call__` (`task(...)` is equivalent to `type(task).__call__(...)`, not `task.__call__(...)`). After this change, we'll use the same logic as Celery to decide whether to instrument `task.__call__` or `task.run`. That change allows abstract tasks to catch/raise exceptions before the Sentry wrapper.
1 parent 5f2af2d commit e971caf

File tree

3 files changed

+31
-4
lines changed

3 files changed

+31
-4
lines changed

mypy.ini

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,3 +59,5 @@ ignore_missing_imports = True
5959
[mypy-sentry_sdk._queue]
6060
ignore_missing_imports = True
6161
disallow_untyped_defs = False
62+
[mypy-celery.app.trace]
63+
ignore_missing_imports = True

sentry_sdk/integrations/celery.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
Ignore,
3131
Reject,
3232
)
33+
from celery.app.trace import task_has_custom
3334
except ImportError:
3435
raise DidNotEnable("Celery not installed")
3536

@@ -57,10 +58,12 @@ def setup_once():
5758
def sentry_build_tracer(name, task, *args, **kwargs):
5859
# type: (Any, Any, *Any, **Any) -> Any
5960
if not getattr(task, "_sentry_is_patched", False):
60-
# Need to patch both methods because older celery sometimes
61-
# short-circuits to task.run if it thinks it's safe.
62-
task.__call__ = _wrap_task_call(task, task.__call__)
63-
task.run = _wrap_task_call(task, task.run)
61+
# determine whether Celery will use __call__ or run and patch
62+
# accordingly
63+
if task_has_custom(task, "__call__"):
64+
type(task).__call__ = _wrap_task_call(task, type(task).__call__)
65+
else:
66+
task.run = _wrap_task_call(task, task.run)
6467

6568
# `build_tracer` is apparently called for every task
6669
# invocation. Can't wrap every celery task for every invocation

tests/integrations/celery/test_celery.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,3 +407,25 @@ def walk_dogs(x, y):
407407
# passed as args or as kwargs, so make this generic
408408
DictionaryContaining({"celery_job": dict(task="dog_walk", **args_kwargs)})
409409
)
410+
411+
412+
def test_abstract_task(capture_events, celery, celery_invocation):
413+
events = capture_events()
414+
415+
class AbstractTask(celery.Task):
416+
abstract = True
417+
418+
def __call__(self, *args, **kwargs):
419+
try:
420+
return self.run(*args, **kwargs)
421+
except ZeroDivisionError:
422+
return None
423+
424+
@celery.task(name="dummy_task", base=AbstractTask)
425+
def dummy_task(x, y):
426+
return x / y
427+
428+
with start_transaction():
429+
celery_invocation(dummy_task, 1, 0)
430+
431+
assert not events

0 commit comments

Comments
 (0)