[celery] add patch_task and unpatch_task as a backward compatibility

This reverts commit 0440d39.

Author: Emanuele Palazzetti

ddtrace/contrib/celery/__init__.py

@@ -52,6 +52,7 @@ def run(self):
     if not missing_modules:
         from .app import patch_app, unpatch_app
         from .patch import patch, unpatch
+        from .task import patch_task, unpatch_task
 
         __all__ = [
             'patch',

ddtrace/contrib/celery/app.py

@@ -1,19 +1,44 @@
+from celery import signals
+
 from ddtrace import Pin
 from ddtrace.pin import _DD_PIN_NAME
 from ddtrace.ext import AppTypes
 
 from .util import APP, WORKER_SERVICE
+from .signals import (
+    trace_prerun,
+    trace_postrun,
+    trace_before_publish,
+    trace_after_publish,
+    trace_failure,
+)
 
 
 def patch_app(app, pin=None):
-    """Attach the Pin class to the application"""
+    """Attach the Pin class to the application and connect
+    our handlers to Celery signals.
+    """
     pin = pin or Pin(service=WORKER_SERVICE, app=APP, app_type=AppTypes.worker)
     pin.onto(app)
+
+    signals.task_prerun.connect(trace_prerun)
+    signals.task_postrun.connect(trace_postrun)
+    signals.before_task_publish.connect(trace_before_publish)
+    signals.after_task_publish.connect(trace_after_publish)
+    signals.task_failure.connect(trace_failure)
     return app
 
 
 def unpatch_app(app):
-    """ unpatch_app will remove tracing from a celery app """
+    """Remove the Pin instance from the application and disconnect
+    our handlers from Celery signal framework.
+    """
     pin = Pin.get_from(app)
     if pin is not None:
         delattr(app, _DD_PIN_NAME)
+
+    signals.task_prerun.disconnect(trace_prerun)
+    signals.task_postrun.disconnect(trace_postrun)
+    signals.before_task_publish.disconnect(trace_before_publish)
+    signals.after_task_publish.disconnect(trace_after_publish)
+    signals.task_failure.disconnect(trace_failure)

ddtrace/contrib/celery/patch.py

@@ -1,17 +1,7 @@
 import celery
 
-from celery import signals
-
 from .app import patch_app, unpatch_app
 
-from .signals import (
-    trace_prerun,
-    trace_postrun,
-    trace_before_publish,
-    trace_after_publish,
-    trace_failure,
-)
-
 
 def patch():
     """Instrument Celery base application and the `TaskRegistry` so
@@ -20,18 +10,8 @@ def patch():
     must be instrumented because Django doesn't use the Celery registry.
     """
     patch_app(celery.Celery)
-    signals.task_prerun.connect(trace_prerun)
-    signals.task_postrun.connect(trace_postrun)
-    signals.before_task_publish.connect(trace_before_publish)
-    signals.after_task_publish.connect(trace_after_publish)
-    signals.task_failure.connect(trace_failure)
 
 
 def unpatch():
     """Disconnect all signals and remove Tracing capabilities"""
     unpatch_app(celery.Celery)
-    signals.task_prerun.disconnect(trace_prerun)
-    signals.task_postrun.disconnect(trace_postrun)
-    signals.before_task_publish.disconnect(trace_before_publish)
-    signals.after_task_publish.disconnect(trace_after_publish)
-    signals.task_failure.disconnect(trace_failure)

ddtrace/contrib/celery/task.py

@@ -0,0 +1,31 @@
+from .app import patch_app
+
+from ...utils.deprecation import deprecation
+
+
+def patch_task(task, pin=None):
+    """Deprecated API. The new API uses signals that can be activated via
+    patch(celery=True) or through `ddtrace-run` script. Using this API
+    enables instrumentation on all tasks.
+    """
+    deprecation(
+        name='ddtrace.contrib.celery.patch_task',
+        message='Use `patch(celery=True)` or `ddtrace-run` script instead',
+        version='1.0.0',
+    )
+
+    # Enable instrumentation everywhere
+    patch_app(task.app)
+    return task
+
+def unpatch_task(task):
+    """Deprecated API. The new API uses signals that can be deactivated
+    via unpatch() API. This API is now a no-op implementation so it doesn't
+    affect instrumented tasks.
+    """
+    deprecation(
+        name='ddtrace.contrib.celery.patch_task',
+        message='Use `unpatch()` instead',
+        version='1.0.0',
+    )
+    return task

tests/contrib/celery/test_task_deprecation.py

@@ -0,0 +1,52 @@
+import warnings
+import unittest
+
+from celery import Celery
+
+from nose.tools import ok_
+
+from ddtrace.contrib.celery import patch_task, unpatch_task, unpatch
+
+
+class CeleryDeprecatedTaskPatch(unittest.TestCase):
+    """Ensures that the previous Task instrumentation is available
+    as Deprecated API.
+    """
+    def setUp(self):
+        # create a not instrumented Celery App
+        self.app = Celery('celery.test_app')
+
+    def tearDown(self):
+        # be sure the system is always unpatched
+        unpatch()
+        self.app = None
+
+    def test_patch_signals_connect(self):
+        # calling `patch_task` enables instrumentation globally
+        # while raising a Deprecation warning
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always')
+
+            @patch_task
+            @self.app.task
+            def fn_task():
+                return 42
+
+            ok_(len(w) == 1)
+            ok_(issubclass(w[-1].category, DeprecationWarning))
+            ok_('patch(celery=True)' in str(w[-1].message))
+
+    def test_unpatch_signals_diconnect(self):
+        # calling `unpatch_task` is a no-op that raises a Deprecation
+        # warning
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always')
+
+            @unpatch_task
+            @self.app.task
+            def fn_task():
+                return 42
+
+            ok_(len(w) == 1)
+            ok_(issubclass(w[-1].category, DeprecationWarning))
+            ok_('unpatch()' in str(w[-1].message))