[celery] patch TaskRegistry to support old-style task with ddtrace-run (#484)

Author: Emanuele Palazzetti

ddtrace/contrib/celery/patch.py

@@ -1,15 +1,21 @@
-# Third party
 import celery
 
-# Project
+from wrapt import wrap_function_wrapper as _w
+
 from .app import patch_app, unpatch_app
+from .registry import _wrap_register
+from ...utils.wrappers import unwrap as _u
 
 
 def patch():
-    """ patch will add all available tracing to the celery library """
+    """Instrument Celery base application and the `TaskRegistry` so
+    that any new registered task is automatically instrumented
+    """
     setattr(celery, 'Celery', patch_app(celery.Celery))
+    _w('celery.app.registry', 'TaskRegistry.register', _wrap_register)
 
 
 def unpatch():
-    """ unpatch will remove tracing from the celery library """
+    """Removes instrumentation from Celery"""
     setattr(celery, 'Celery', unpatch_app(celery.Celery))
+    _u(celery.app.registry.TaskRegistry, 'register')

ddtrace/contrib/celery/registry.py

@@ -0,0 +1,15 @@
+from .task import patch_task
+
+
+def _wrap_register(func, instance, args, kwargs):
+    """Wraps the `TaskRegistry.register` function so that everytime
+    a `Task` is registered it is properly instrumented. This wrapper
+    is required because in old-style tasks (Celery 1.0+) we cannot
+    instrument the base class, otherwise a `Strategy` `KeyError`
+    exception is raised.
+    """
+    # the original signature requires one positional argument so the
+    # first and only parameter is the `Task` that must be instrumented
+    task = args[0]
+    patch_task(task)
+    func(*args, **kwargs)

tox.ini

@@ -128,8 +128,7 @@ deps =
     celery31: celery>=3.1,<3.2
     celery40: celery>=4.0,<4.1
     celery41: celery>=4.1,<4.2
-    # TODO[manu] update to a stable version of Celery
-    celery42: celery==4.2.0rc3
+    celery42: celery>=4.2,<4.3
     ddtracerun: redis
     ddtracerun: celery
     elasticsearch16: elasticsearch>=1.6,<1.7