fix: Fix celery integration with custom tag class

untitaker · untitaker · commit 2305feabdea9 · 2020-10-20T20:32:02.000+02:00
diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py
@@ -61,6 +61,7 @@ def sentry_build_tracer(name, task, *args, **kwargs):
                 # short-circuits to task.run if it thinks it's safe.
                 task.__call__ = _wrap_task_call(task, task.__call__)
                 task.run = _wrap_task_call(task, task.run)
+                task.apply_async = _wrap_apply_async(task, task.apply_async)
 
                 # `build_tracer` is apparently called for every task
                 # invocation. Can't wrap every celery task for every invocation
@@ -71,10 +72,6 @@ def sentry_build_tracer(name, task, *args, **kwargs):
 
         trace.build_tracer = sentry_build_tracer
 
-        from celery.app.task import Task  # type: ignore
-
-        Task.apply_async = _wrap_apply_async(Task.apply_async)
-
         _patch_worker_exit()
 
         # This logger logs every status of every task that ran on the worker.
@@ -88,15 +85,15 @@ def sentry_build_tracer(name, task, *args, **kwargs):
         ignore_logger("celery.redirected")
 
 
-def _wrap_apply_async(f):
-    # type: (F) -> F
+def _wrap_apply_async(task, f):
+    # type: (Any, F) -> F
     @wraps(f)
     def apply_async(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         hub = Hub.current
         integration = hub.get_integration(CeleryIntegration)
         if integration is not None and integration.propagate_traces:
-            with hub.start_span(op="celery.submit", description=args[0].name):
+            with hub.start_span(op="celery.submit", description=task.name):
                 with capture_internal_exceptions():
                     headers = dict(hub.iter_trace_propagation_headers())
 
diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py
@@ -99,8 +99,16 @@ def celery_invocation(request):
     """
     return request.param
 
+@pytest.mark.parametrize("custom_celery_task_cls", (True, False))
+def test_simple(capture_events, celery, celery_invocation, custom_celery_task_cls):
+
+    if custom_celery_task_cls:
+        class CustomTask(celery.Task):
+            def __call__(self, *args, **kwargs):
+                return self.run(*args, **kwargs)
+
+        celery.Task = CustomTask
 
-def test_simple(capture_events, celery, celery_invocation):
     events = capture_events()
 
     @celery.task(name="dummy_task")
