Merge pull request #314 from mcanaves/django_flexible_configuration

[django] disable db or cache instrumentation

Author: Emanuele Palazzetti

ddtrace/contrib/django/__init__.py

@@ -13,14 +13,6 @@
         'ddtrace.contrib.django',
     ]
 
-    # or MIDDLEWARE_CLASSES for Django pre 1.10
-    MIDDLEWARE = [
-        # the tracer must be the first middleware
-        'ddtrace.contrib.django.TraceMiddleware',
-
-        # your middlewares...
-    ]
-
 The configuration for this integration is namespaced under the ``DATADOG_TRACE``
 Django setting. For example, your ``settings.py`` may contain::
 
@@ -54,10 +46,18 @@
   and a restart is required. By default the tracer is disabled when in ``DEBUG``
   mode, enabled otherwise.
 * ``AUTO_INSTRUMENT`` (default: ``True``): if set to false the code will not be
-  instrumented, while the tracer may be active for your internal usage. This could
-  be useful if you want to use the Django integration, but you want to trace only
-  particular functions or views. If set to False, the request middleware will be
-  disabled even if present.
+  instrumented (even if ``INSTRUMENT_DATABASE``, ``INSTRUMENT_CACHE`` or
+  ``INSTRUMENT_TEMPLATE`` are set to ``True``), while the tracer may be active
+  for your internal usage. This could be useful if you want to use the Django
+  integration, but you want to trace only particular functions or views. If set
+  to False, the request middleware will be disabled even if present.
+* ``INSTRUMENT_DATABASE`` (default: ``True``): if set to ``False`` database will not
+  be instrumented. Only configurable when ``AUTO_INSTRUMENT`` is set to ``True``.
+* ``INSTRUMENT_CACHE`` (default: ``True``): if set to ``False`` cache will not
+  be instrumented. Only configurable when ``AUTO_INSTRUMENT`` is set to ``True``.
+* ``INSTRUMENT_TEMPLATE`` (default: ``True``): if set to ``False`` template
+  rendering will not be instrumented. Only configurable when ``AUTO_INSTRUMENT``
+  is set to ``True``.
 * ``AGENT_HOSTNAME`` (default: ``localhost``): define the hostname of the trace agent.
 * ``AGENT_PORT`` (default: ``8126``): define the port of the trace agent.
 """

ddtrace/contrib/django/apps.py

@@ -8,7 +8,7 @@
 from .conf import settings
 from .cache import patch_cache
 from .templates import patch_template
-from .middleware import insert_exception_middleware
+from .middleware import insert_exception_middleware, insert_trace_middleware
 
 from ...ext import AppTypes
 
@@ -48,18 +48,23 @@ def ready(self):
 
         if settings.AUTO_INSTRUMENT:
             # trace Django internals
+            insert_trace_middleware()
             insert_exception_middleware()
-            try:
-                patch_db(tracer)
-            except Exception:
-                log.exception('error patching Django database connections')
 
-            try:
-                patch_template(tracer)
-            except Exception:
-                log.exception('error patching Django template rendering')
+            if settings.INSTRUMENT_TEMPLATE:
+                try:
+                    patch_template(tracer)
+                except Exception:
+                    log.exception('error patching Django template rendering')
 
-            try:
-                patch_cache(tracer)
-            except Exception:
-                log.exception('error patching Django cache')
+            if settings.INSTRUMENT_DATABASE:
+                try:
+                    patch_db(tracer)
+                except Exception:
+                    log.exception('error patching Django database connections')
+
+            if settings.INSTRUMENT_CACHE:
+                try:
+                    patch_cache(tracer)
+                except Exception:
+                    log.exception('error patching Django cache')

ddtrace/contrib/django/cache.py

@@ -89,3 +89,19 @@ def _wrap_method(cls, method_name):
 
         for method in TRACED_METHODS:
             _wrap_method(cache, method)
+
+def unpatch_method(cls, method_name):
+    method = getattr(cls, DATADOG_NAMESPACE.format(method=method_name), None)
+    if method is None:
+        log.debug('nothing to do, the class is not patched')
+        return
+    setattr(cls, method_name, method)
+    delattr(cls, DATADOG_NAMESPACE.format(method=method_name))
+
+def unpatch_cache():
+    cache_backends = [cache['BACKEND'] for cache in django_settings.CACHES.values()]
+    for cache_module in cache_backends:
+        cache = import_from_string(cache_module, cache_module)
+
+        for method in TRACED_METHODS:
+            unpatch_method(cache, method)

ddtrace/contrib/django/conf.py

@@ -25,6 +25,9 @@
     'AGENT_HOSTNAME': 'localhost',
     'AGENT_PORT': 8126,
     'AUTO_INSTRUMENT': True,
+    'INSTRUMENT_CACHE': True,
+    'INSTRUMENT_DATABASE': True,
+    'INSTRUMENT_TEMPLATE': True,
     'DEFAULT_DATABASE_PREFIX': '',
     'DEFAULT_SERVICE': 'django',
     'ENABLED': True,

ddtrace/contrib/django/db.py

@@ -12,25 +12,37 @@
 
 log = logging.getLogger(__name__)
 
+CURSOR_ATTR = '_datadog_original_cursor'
+
 
 def patch_db(tracer):
     for c in connections.all():
         patch_conn(tracer, c)
 
+def unpatch_db():
+    for c in connections.all():
+        unpatch_conn(c)
+
 
 def patch_conn(tracer, conn):
-    attr = '_datadog_original_cursor'
-    if hasattr(conn, attr):
+    if hasattr(conn, CURSOR_ATTR):
         log.debug("already patched")
         return
 
-    conn._datadog_original_cursor = conn.cursor
+    setattr(conn, CURSOR_ATTR, conn.cursor)
 
     def cursor():
         return TracedCursor(tracer, conn, conn._datadog_original_cursor())
 
     conn.cursor = cursor
 
+def unpatch_conn(conn):
+    cursor = getattr(conn, CURSOR_ATTR, None)
+    if cursor is None:
+        log.debug('nothing to do, the connection is not patched')
+        return
+    conn.cursor = cursor
+    delattr(conn, CURSOR_ATTR)
 
 class TracedCursor(object):
 

ddtrace/contrib/django/middleware.py

@@ -18,14 +18,35 @@
 
 log = logging.getLogger(__name__)
 
+EXCEPTION_MIDDLEWARE = 'ddtrace.contrib.django.TraceExceptionMiddleware'
+TRACE_MIDDLEWARE = 'ddtrace.contrib.django.TraceMiddleware'
+MIDDLEWARE_ATTRIBUTES = ['MIDDLEWARE', 'MIDDLEWARE_CLASSES']
+
+def insert_trace_middleware():
+    for middleware_attribute in MIDDLEWARE_ATTRIBUTES:
+        middleware = getattr(django_settings, middleware_attribute, None)
+        if middleware is not None and TRACE_MIDDLEWARE not in set(middleware):
+            setattr(django_settings, middleware_attribute, type(middleware)((TRACE_MIDDLEWARE,)) + middleware)
+            break
+
+def remove_trace_middleware():
+    for middleware_attribute in MIDDLEWARE_ATTRIBUTES:
+        middleware = getattr(django_settings, middleware_attribute, None)
+        if middleware and TRACE_MIDDLEWARE in set(middleware):
+            middleware.remove(TRACE_MIDDLEWARE)
+
 def insert_exception_middleware():
-    exception_middleware = 'ddtrace.contrib.django.TraceExceptionMiddleware'
-    middleware_attributes = ['MIDDLEWARE', 'MIDDLEWARE_CLASSES']
-    for middleware_attribute in middleware_attributes:
+    for middleware_attribute in MIDDLEWARE_ATTRIBUTES:
         middleware = getattr(django_settings, middleware_attribute, None)
-        if middleware and exception_middleware not in set(middleware):
-            setattr(django_settings, middleware_attribute, middleware + type(middleware)((exception_middleware,)))
+        if middleware is not None and EXCEPTION_MIDDLEWARE not in set(middleware):
+            setattr(django_settings, middleware_attribute, middleware + type(middleware)((EXCEPTION_MIDDLEWARE,)))
+            break
 
+def remove_exception_middleware():
+    for middleware_attribute in MIDDLEWARE_ATTRIBUTES:
+        middleware = getattr(django_settings, middleware_attribute, None)
+        if middleware and EXCEPTION_MIDDLEWARE in set(middleware):
+            middleware.remove(EXCEPTION_MIDDLEWARE)
 
 class InstrumentationMixin(MiddlewareClass):
     """

ddtrace/contrib/django/patch.py

@@ -22,16 +22,4 @@ def traced_setup(wrapped, instance, args, kwargs):
         else:
             settings.INSTALLED_APPS.append('ddtrace.contrib.django')
 
-    if hasattr(settings, 'MIDDLEWARE_CLASSES'):
-        if 'ddtrace.contrib.django.TraceMiddleware' not in settings.MIDDLEWARE_CLASSES:
-            if isinstance(settings.MIDDLEWARE_CLASSES, tuple):
-                # MIDDLEWARE_CLASSES is a tuple < 1.9
-                settings.MIDDLEWARE_CLASSES = ('ddtrace.contrib.django.TraceMiddleware', ) + settings.MIDDLEWARE_CLASSES
-            else:
-                settings.MIDDLEWARE_CLASSES.insert(0, 'ddtrace.contrib.django.TraceMiddleware')
-
-    if hasattr(settings, 'MIDDLEWARE'):
-        if 'ddtrace.contrib.django.TraceMiddleware' not in settings.MIDDLEWARE:
-            settings.MIDDLEWARE.insert(0, 'ddtrace.contrib.django.TraceMiddleware')
-
     wrapped(*args, **kwargs)

ddtrace/contrib/django/templates.py

@@ -15,6 +15,7 @@
 
 log = logging.getLogger(__name__)
 
+RENDER_ATTR = '_datadog_original_render'
 
 def patch_template(tracer):
     """ will patch django's template rendering function to include timing
@@ -24,12 +25,11 @@ def patch_template(tracer):
     # FIXME[matt] we're patching the template class here. ideally we'd only
     # patch so we can use multiple tracers at once, but i suspect this is fine
     # in practice.
-    attr = '_datadog_original_render'
-    if getattr(Template, attr, None):
+    if getattr(Template, RENDER_ATTR, None):
         log.debug("already patched")
         return
 
-    setattr(Template, attr, Template.render)
+    setattr(Template, RENDER_ATTR, Template.render)
 
     def traced_render(self, context):
         with tracer.trace('django.template', span_type=http.TEMPLATE) as span:
@@ -41,3 +41,11 @@ def traced_render(self, context):
                 span.set_tag('django.template_name', template_name)
 
     Template.render = traced_render
+
+def unpatch_template():
+    render = getattr(Template, RENDER_ATTR, None)
+    if render is None:
+        log.debug('nothing to do Template is already patched')
+        return
+    Template.render = render
+    delattr(Template, RENDER_ATTR)

tests/contrib/django/app/settings.py

@@ -4,6 +4,7 @@
 Django during tests
 """
 import os
+import django
 
 
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -68,28 +69,22 @@
     },
 ]
 
-# 1.10+ style
-MIDDLEWARE = [
-    # tracer middleware
-    'ddtrace.contrib.django.TraceMiddleware',
-
-    'django.contrib.sessions.middleware.SessionMiddleware',
-    'django.middleware.common.CommonMiddleware',
-    'django.middleware.csrf.CsrfViewMiddleware',
-    'django.contrib.auth.middleware.AuthenticationMiddleware',
-    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
-    'django.contrib.messages.middleware.MessageMiddleware',
-    'django.middleware.clickjacking.XFrameOptionsMiddleware',
-    'django.middleware.security.SecurityMiddleware',
-
-    'tests.contrib.django.app.middlewares.CatchExceptionMiddleware',
-]
+if django.VERSION >= (1, 10):
+    MIDDLEWARE = [
+        'django.contrib.sessions.middleware.SessionMiddleware',
+        'django.middleware.common.CommonMiddleware',
+        'django.middleware.csrf.CsrfViewMiddleware',
+        'django.contrib.auth.middleware.AuthenticationMiddleware',
+        'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
+        'django.contrib.messages.middleware.MessageMiddleware',
+        'django.middleware.clickjacking.XFrameOptionsMiddleware',
+        'django.middleware.security.SecurityMiddleware',
 
+        'tests.contrib.django.app.middlewares.CatchExceptionMiddleware',
+    ]
+# Always add the legacy conf to make sure we handle it properly
 # Pre 1.10 style
 MIDDLEWARE_CLASSES = [
-    # tracer middleware
-    'ddtrace.contrib.django.TraceMiddleware',
-
     'django.contrib.sessions.middleware.SessionMiddleware',
     'django.middleware.common.CommonMiddleware',
     'django.middleware.csrf.CsrfViewMiddleware',

tests/contrib/django/test_autopatching.py

@@ -1,41 +1,69 @@
+import django
+
 from ddtrace.monkey import patch
 from .utils import DjangoTraceTestCase
 from nose.tools import eq_, ok_
+from django.conf import settings
+from unittest import skipIf
+
 
 class DjangoAutopatchTest(DjangoTraceTestCase):
-    def test_autopatching(self):
+    def setUp(self):
+        super(DjangoAutopatchTest, self).setUp()
         patch(django=True)
-
-        import django
-        ok_(django._datadog_patch)
         django.setup()
 
-        from django.conf import settings
+    @skipIf(django.VERSION >= (1, 10), 'skip if version above 1.10')
+    def test_autopatching_middleware_classes(self):
+        ok_(django._datadog_patch)
         ok_('ddtrace.contrib.django' in settings.INSTALLED_APPS)
         eq_(settings.MIDDLEWARE_CLASSES[0], 'ddtrace.contrib.django.TraceMiddleware')
+        eq_(settings.MIDDLEWARE_CLASSES[-1], 'ddtrace.contrib.django.TraceExceptionMiddleware')
 
 
-    def test_autopatching_twice(self):
-        patch(django=True)
-
+    @skipIf(django.VERSION >= (1, 10), 'skip if version above 1.10')
+    def test_autopatching_twice_middleware_classes(self):
+        ok_(django._datadog_patch)
         # Call django.setup() twice and ensure we don't add a duplicate tracer
-        import django
-        django.setup()
         django.setup()
 
-        from django.conf import settings
-        found_app = 0
+        found_app = settings.INSTALLED_APPS.count('ddtrace.contrib.django')
+        eq_(found_app, 1)
 
-        for app in settings.INSTALLED_APPS:
-            if app == 'ddtrace.contrib.django':
-                found_app += 1
+        eq_(settings.MIDDLEWARE_CLASSES[0], 'ddtrace.contrib.django.TraceMiddleware')
+        eq_(settings.MIDDLEWARE_CLASSES[-1], 'ddtrace.contrib.django.TraceExceptionMiddleware')
 
+        found_mw = settings.MIDDLEWARE_CLASSES.count('ddtrace.contrib.django.TraceMiddleware')
+        eq_(found_mw, 1)
+        found_mw = settings.MIDDLEWARE_CLASSES.count('ddtrace.contrib.django.TraceExceptionMiddleware')
+        eq_(found_mw, 1)
+
+    @skipIf(django.VERSION < (1, 10), 'skip if version is below 1.10')
+    def test_autopatching_middleware(self):
+        ok_(django._datadog_patch)
+        ok_('ddtrace.contrib.django' in settings.INSTALLED_APPS)
+        eq_(settings.MIDDLEWARE[0], 'ddtrace.contrib.django.TraceMiddleware')
+        ok_('ddtrace.contrib.django.TraceMiddleware' not in settings.MIDDLEWARE_CLASSES)
+        eq_(settings.MIDDLEWARE[-1], 'ddtrace.contrib.django.TraceExceptionMiddleware')
+        ok_('ddtrace.contrib.django.TraceExceptionMiddleware' not in settings.MIDDLEWARE_CLASSES)
+
+
+    @skipIf(django.VERSION < (1, 10), 'skip if version is below 1.10')
+    def test_autopatching_twice_middleware(self):
+        ok_(django._datadog_patch)
+        # Call django.setup() twice and ensure we don't add a duplicate tracer
+        django.setup()
+
+        found_app = settings.INSTALLED_APPS.count('ddtrace.contrib.django')
         eq_(found_app, 1)
-        eq_(settings.MIDDLEWARE_CLASSES[0], 'ddtrace.contrib.django.TraceMiddleware')
 
-        found_mw = 0
-        for mw in settings.MIDDLEWARE_CLASSES:
-            if mw == 'ddtrace.contrib.django.TraceMiddleware':
-                found_mw += 1
+        eq_(settings.MIDDLEWARE[0], 'ddtrace.contrib.django.TraceMiddleware')
+        ok_('ddtrace.contrib.django.TraceMiddleware' not in settings.MIDDLEWARE_CLASSES)
+        eq_(settings.MIDDLEWARE[-1], 'ddtrace.contrib.django.TraceExceptionMiddleware')
+        ok_('ddtrace.contrib.django.TraceExceptionMiddleware' not in settings.MIDDLEWARE_CLASSES)
+
+        found_mw = settings.MIDDLEWARE.count('ddtrace.contrib.django.TraceMiddleware')
+        eq_(found_mw, 1)
 
+        found_mw = settings.MIDDLEWARE.count('ddtrace.contrib.django.TraceExceptionMiddleware')
         eq_(found_mw, 1)