[celery] Patch via post-import hooks (#534)

* [core] revert argv patch

* [core] patch celery via post-import hooks

* [core] set patched to true for celery patching

* [core] remove patched module count side effect

* [celery] add tests verifying the import hook patching

Author: Kyle-Verhoog

ddtrace/bootstrap/sitecustomize.py

@@ -70,9 +70,6 @@ def update_patched_modules():
     if opts:
         tracer.configure(**opts)
 
-    if not hasattr(sys, 'argv'):
-        sys.argv = ['']
-
     if patch:
         update_patched_modules()
         from ddtrace import patch_all; patch_all(**EXTRA_PATCHED_MODULES) # noqa

ddtrace/monkey.py

@@ -77,16 +77,28 @@ def patch(raise_errors=True, **patch_modules):
         >>> patch(psycopg=True, elasticsearch=True)
     """
     modules = [m for (m, should_patch) in patch_modules.items() if should_patch]
-    count = 0
     for module in modules:
-        patched = patch_module(module, raise_errors=raise_errors)
-        if patched:
-            count += 1
-
+        # TODO: this is a temporary hack until we shift to using
+        # post-import hooks for everything.
+        if module == 'celery':
+            # if patch celery via post-import hooks
+            from wrapt.importer import when_imported
+
+            @when_imported('celery')
+            def patch_celery(hook):
+                from ddtrace.contrib.celery import patch
+                patch()
+
+            # manually add celery to patched modules
+            _PATCHED_MODULES.add(module)
+        else:
+            patch_module(module, raise_errors=raise_errors)
+
+    patched_modules = get_patched_modules()
     log.info("patched %s/%s modules (%s)",
-        count,
+        len(patched_modules),
         len(modules),
-        ",".join(get_patched_modules()))
+        ",".join(patched_modules))
 
 
 def patch_module(module, raise_errors=True):
@@ -115,7 +127,7 @@ def _patch_module(module):
     """
     path = 'ddtrace.contrib.%s' % module
     with _LOCK:
-        if module in _PATCHED_MODULES:
+        if module in _PATCHED_MODULES and module != 'celery':
             log.debug("already patched: %s", path)
             return False
 

tests/commands/ddtrace_run_patched_modules.py

@@ -6,5 +6,4 @@
 
 if __name__ == '__main__':
     ok_('redis' in monkey.get_patched_modules())
-    ok_('celery' in monkey.get_patched_modules())
     print("Test success")

tests/contrib/celery/autopatch.py

@@ -0,0 +1,13 @@
+from __future__ import print_function
+
+from nose.tools import ok_
+
+from ddtrace import Pin
+
+if __name__ == '__main__':
+    # have to import celery in order to have the post-import hooks run
+    import celery
+
+    # now celery.Celery should be patched and should have a pin
+    ok_(Pin.get_from(celery.Celery))
+    print("Test success")

tests/contrib/celery/test_autopatch.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+import subprocess
+import unittest
+
+
+class DdtraceRunTest(unittest.TestCase):
+    """Test that celery is patched successfully if run with ddtrace-run."""
+
+    def test_autopatch(self):
+        out = subprocess.check_output(
+            ['ddtrace-run', 'python', 'tests/contrib/celery/autopatch.py']
+        )
+        assert out.startswith(b"Test success")

tests/contrib/celery/test_patch.py

@@ -0,0 +1,22 @@
+import unittest
+from nose.tools import ok_
+from ddtrace import Pin
+
+
+class CeleryPatchTest(unittest.TestCase):
+    def test_patch_after_import(self):
+        import celery
+        from ddtrace import patch
+        patch(celery=True)
+
+        app = celery.Celery()
+        ok_(Pin.get_from(app) is not None)
+
+    def test_patch_before_import(self):
+        from ddtrace import patch
+        patch(celery=True)
+        import celery
+
+        app = celery.Celery()
+        ok_(Pin.get_from(app) is not None)
+

tox.ini

@@ -130,7 +130,6 @@ deps =
     celery41: celery>=4.1,<4.2
     celery42: celery>=4.2,<4.3
     ddtracerun: redis
-    ddtracerun: celery
     elasticsearch16: elasticsearch>=1.6,<1.7
     elasticsearch17: elasticsearch>=1.7,<1.8
     elasticsearch18: elasticsearch>=1.8,<1.9