[django/celery] add shared_task decorator wrapper (#486)

Author: Emanuele Palazzetti

ddtrace/contrib/celery/patch.py

@@ -3,19 +3,24 @@
 from wrapt import wrap_function_wrapper as _w
 
 from .app import patch_app, unpatch_app
+from .task import _wrap_shared_task
 from .registry import _wrap_register
 from ...utils.wrappers import unwrap as _u
 
 
 def patch():
     """Instrument Celery base application and the `TaskRegistry` so
-    that any new registered task is automatically instrumented
+    that any new registered task is automatically instrumented. In the
+    case of Django-Celery integration, also the `@shared_task` decorator
+    must be instrumented because Django doesn't use the Celery registry.
     """
     setattr(celery, 'Celery', patch_app(celery.Celery))
     _w('celery.app.registry', 'TaskRegistry.register', _wrap_register)
+    _w('celery', 'shared_task', _wrap_shared_task)
 
 
 def unpatch():
     """Removes instrumentation from Celery"""
     setattr(celery, 'Celery', unpatch_app(celery.Celery))
     _u(celery.app.registry.TaskRegistry, 'register')
+    _u(celery, 'shared_task')

ddtrace/contrib/celery/task.py

@@ -79,6 +79,14 @@ def unpatch_task(task):
     return task
 
 
+def _wrap_shared_task(decorator, instance, args, kwargs):
+    """Wrapper for Django-Celery shared tasks. `shared_task` is a decorator
+    that returns a `Task` from the given function.
+    """
+    task = decorator(*args, **kwargs)
+    return patch_task(task)
+
+
 def _task_init(func, task, args, kwargs):
     func(*args, **kwargs)
 

tests/contrib/celery/test_task.py

@@ -488,3 +488,19 @@ class CelerySubClass(CelerySuperClass):
         t.run()
         spans = self.tracer.writer.pop()
         self.assertEqual(len(spans), 4)
+
+    def test_celery_shared_task(self):
+        @celery.shared_task
+        def add(x ,y):
+            return x + y
+
+        res = add.run(2, 2)
+        self.assertEqual(res, 4)
+        spans = self.tracer.writer.pop()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.service, 'celery-worker')
+        self.assertEqual(span.resource, 'tests.contrib.celery.test_task.add')
+        self.assertEqual(span.name, 'celery.run')
+        self.assertIsNone(span.parent_id)
+        self.assertEqual(span.error, 0)