Merge pull request #458 from palazzem/sitecustomize

[core] ensure users sitecustomize.py is called

Author: Emanuele Palazzetti

ddtrace/bootstrap/sitecustomize.py

@@ -4,6 +4,8 @@
 """
 
 import os
+import imp
+import sys
 import logging
 
 from ddtrace.utils.formats import asbool
@@ -46,6 +48,8 @@ def update_patched_modules():
     patch = True
 
     # Respect DATADOG_* environment variables in global tracer configuration
+    # TODO: these variables are deprecated; use utils method and update our documentation
+    # correct prefix should be DD_*
     enabled = os.environ.get("DATADOG_TRACE_ENABLED")
     hostname = os.environ.get("DATADOG_TRACE_AGENT_HOSTNAME")
     port = os.environ.get("DATADOG_TRACE_AGENT_PORT")
@@ -76,5 +80,24 @@ def update_patched_modules():
 
     if 'DATADOG_ENV' in os.environ:
         tracer.set_tags({"env": os.environ["DATADOG_ENV"]})
+
+    # Ensure sitecustomize.py is properly called if available in application directories:
+    # * exclude `bootstrap_dir` from the search
+    # * find a user `sitecustomize.py` module
+    # * import that module via `imp`
+    bootstrap_dir = os.path.dirname(__file__)
+    path = list(sys.path)
+    path.remove(bootstrap_dir)
+
+    try:
+        (f, path, description) = imp.find_module('sitecustomize', path)
+    except ImportError:
+        pass
+    else:
+        # `sitecustomize.py` found, load it
+        log.debug('sitecustomize from user found in: %s', path)
+        imp.load_module('sitecustomize', f, path, description)
+
+
 except Exception as e:
     log.warn("error configuring Datadog tracing", exc_info=True)

tests/commands/bootstrap/__init__.py

@@ -0,0 +1 @@
+CORRECT_IMPORT = True

tests/commands/bootstrap/sitecustomize.py

@@ -0,0 +1,19 @@
+from __future__ import print_function
+
+import sys
+from ddtrace import tracer
+from nose.tools import ok_
+
+
+if __name__ == '__main__':
+    # detect if `-S` is used
+    suppress = len(sys.argv) == 2 and sys.argv[1] is '-S'
+    if suppress:
+        ok_('sitecustomize' not in sys.modules)
+    else:
+        ok_('sitecustomize' in sys.modules)
+
+    # ensure the right `sitecustomize` will be imported
+    import sitecustomize
+    ok_(sitecustomize.CORRECT_IMPORT)
+    print('Test success')

tests/commands/ddtrace_run_sitecustomize.py

@@ -5,6 +5,8 @@
 import subprocess
 import unittest
 
+from ..util import inject_sitecustomize
+
 
 class DdtraceRunTest(unittest.TestCase):
     def tearDown(self):
@@ -147,3 +149,23 @@ def test_patch_modules_from_env(self):
         update_patched_modules()
         assert EXTRA_PATCHED_MODULES["boto"] == True
         assert EXTRA_PATCHED_MODULES["django"] == False
+
+    def test_sitecustomize_run(self):
+        # [Regression test]: ensure users `sitecustomize.py` is properly loaded,
+        # so that our `bootstrap/sitecustomize.py` doesn't override the one
+        # defined in users' PYTHONPATH.
+        env = inject_sitecustomize('tests/commands/bootstrap')
+        out = subprocess.check_output(
+            ['ddtrace-run', 'python', 'tests/commands/ddtrace_run_sitecustomize.py'],
+            env=env,
+        )
+        assert out.startswith(b"Test success")
+
+    def test_sitecustomize_run_suppressed(self):
+        # ensure `sitecustomize.py` is not loaded if `-S` is used
+        env = inject_sitecustomize('tests/commands/bootstrap')
+        out = subprocess.check_output(
+            ['ddtrace-run', 'python', 'tests/commands/ddtrace_run_sitecustomize.py', '-S'],
+            env=env,
+        )
+        assert out.startswith(b"Test success")

tests/commands/test_runner.py

@@ -1,7 +1,9 @@
 import os
+import sys
 import mock
 import ddtrace
 
+from ddtrace import __file__ as root_file
 from nose.tools import ok_
 from contextlib import contextmanager
 
@@ -75,3 +77,26 @@ def set_env(**environ):
     finally:
         os.environ.clear()
         os.environ.update(old_environ)
+
+
+def inject_sitecustomize(path):
+    """Creates a new environment, injecting a ``sitecustomize.py`` module in
+    the current PYTHONPATH.
+
+    :param path: package path containing ``sitecustomize.py`` module, starting
+                 from the ddtrace root folder
+    :returns: a cloned environment that includes an altered PYTHONPATH with
+              the given `sitecustomize.py`
+    """
+    root_folder = os.path.dirname(root_file)
+    # Copy the current environment and replace the PYTHONPATH. This is
+    # required otherwise `ddtrace` scripts are not found when `env` kwarg is
+    # passed
+    env = os.environ.copy()
+    sitecustomize = os.path.join(root_folder, '..', path)
+
+    # Add `boostrap` module so that `sitecustomize.py` is at the bottom
+    # of the PYTHONPATH
+    python_path = list(sys.path) + [sitecustomize]
+    env['PYTHONPATH'] = ':'.join(python_path)[1:]
+    return env