Use Django 2.0 DB instrumentation too

contrib/opencensus-ext-django/examples/app/settings.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Django settings for test app."""
+import django
 
 # Build paths inside the project like this: os.path.join(BASE_DIR, ...)
 import os
@@ -32,18 +33,22 @@
     'opencensus.trace.ext.django',
 )
 
-MIDDLEWARE_CLASSES = (
+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',
     'opencensus.trace.ext.django.middleware.OpencensusMiddleware',
 )
 
+if django.VERSION < (2,):
+    MIDDLEWARE += (
+        'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
+    )
+
 ROOT_URLCONF = 'app.urls'
 
 TEMPLATES = [

contrib/opencensus-ext-django/opencensus/ext/django/middleware.py

@@ -28,6 +28,8 @@
 from opencensus.trace import utils
 from opencensus.trace.propagation import trace_context_http_header_format
 
+import django
+from django.db import connection
 try:
     from django.utils.deprecation import MiddlewareMixin
 except ImportError:  # pragma: NO COVER
@@ -182,9 +184,39 @@ def process_request(self, request):
                 SPAN_THREAD_LOCAL_KEY,
                 span)
 
+            if django.VERSION >= (2,):
+                connection.execute_wrappers.append(self._trace_db_call)
+
         except Exception:  # pragma: NO COVER
             log.error('Failed to trace request', exc_info=True)
 
+    def _trace_db_call(self, execute, sql, params, many, context):
+        _tracer = execution_context.get_opencensus_tracer()
+        if _tracer is not None:
+            if many:
+                method_name = 'executemany'
+            else:
+                method_name = 'execute'
+            db_type = context['connection'].vendor
+            # Note that although get_opencensus_tracer() returns a NoopTracer
+            # if no thread local has been set, set_opencensus_tracer() does NOT
+            # protect against setting None to the thread local - be defensive
+            # here
+            _span = _tracer.start_span()
+            _span.name = '{}.query'.format(db_type)
+            _span.span_kind = span_module.SpanKind.CLIENT
+            _tracer.add_attribute_to_current_span(
+                '{}.query'.format(db_type), sql)
+            _tracer.add_attribute_to_current_span(
+                '{}.cursor.method.name'.format(db_type), method_name)
+
+        result = execute(sql, params, many, context)
+
+        if _tracer is not None:
+            _tracer.end_span()
+
+        return result
+
     def process_view(self, request, view_func, *args, **kwargs):
         """Process view is executed before the view function, here we get the
         function name add set it as the span name.

contrib/opencensus-ext-django/setup.py

@@ -39,7 +39,7 @@
     include_package_data=True,
     long_description=open('README.rst').read(),
     install_requires=[
-        'Django >= 1.11.0, <= 1.11.20',
+        'Django >= 1.11.0',
         'opencensus >= 0.6.dev0, < 1.0.0',
     ],
     extras_require={},

contrib/opencensus-ext-django/tests/test_django_db_middleware.py

@@ -0,0 +1,74 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from collections import namedtuple
+
+import django
+import mock
+import pytest
+from django.test.utils import teardown_test_environment
+
+from opencensus.trace import execution_context
+
+
+class TestOpencensusDatabaseMiddleware(unittest.TestCase):
+    def setUp(self):
+        from django.conf import settings as django_settings
+        from django.test.utils import setup_test_environment
+
+        if not django_settings.configured:
+            django_settings.configure()
+        setup_test_environment()
+
+    def tearDown(self):
+        execution_context.clear()
+        teardown_test_environment()
+
+    def test_process_request(self):
+        if django.VERSION < (2, 0):
+            pytest.skip("Wrong version of Django")
+
+        from opencensus.ext.django import middleware
+
+        sql = "SELECT * FROM users"
+
+        MockConnection = namedtuple('Connection', ('vendor'))
+        connection = MockConnection('mysql')
+
+        mock_execute = mock.Mock()
+        mock_execute.return_value = "Mock result"
+
+        middleware_obj = middleware.OpencensusMiddleware()
+
+        result = middleware_obj._trace_db_call(
+            mock_execute, sql, params=[], many=False,
+            context={'connection': connection})
+
+        (mock_sql, mock_params, mock_many,
+            mock_context) = mock_execute.call_args[0]
+
+        self.assertEqual(mock_sql, sql)
+        self.assertEqual(mock_params, [])
+        self.assertEqual(mock_many, False)
+        self.assertEqual(mock_context, {'connection': connection})
+        self.assertEqual(result, "Mock result")
+
+        result = middleware_obj._trace_db_call(
+            mock_execute, sql, params=[], many=True,
+            context={'connection': connection})
+
+        (mock_sql, mock_params, mock_many,
+            mock_context) = mock_execute.call_args[0]
+        self.assertEqual(mock_many, True)

tests/system/trace/django/app/settings.py

@@ -34,8 +34,9 @@
     'opencensus.ext.django',
 )
 
-if django.VERSION >= (1, 10):
-    MIDDLEWARE = (
+if django.VERSION <= (1, 10):
+    # Middleware interface for Django version before 1.10
+    MIDDLEWARE_CLASSES = (
         'django.contrib.sessions.middleware.SessionMiddleware',
         'django.middleware.common.CommonMiddleware',
         'django.middleware.csrf.CsrfViewMiddleware',
@@ -47,19 +48,24 @@
         'opencensus.ext.django.middleware.OpencensusMiddleware',
     )
 
-# Middleware interface for Django version before 1.10
-MIDDLEWARE_CLASSES = (
+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',
     'opencensus.ext.django.middleware.OpencensusMiddleware',
 )
 
+# SessionAuthentication is unconditionally enabled for Django versions greater
+# than 2
+if django.VERSION < (2,):
+    MIDDLEWARE += (
+        'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
+    )
+
 ROOT_URLCONF = 'app.urls'
 
 TEMPLATES = [

tests/system/trace/django/app/urls.py

@@ -27,14 +27,14 @@
     1. Add an import:  from blog import urls as blog_urls
     2. Add a URL to urlpatterns:  url(r'^blog/', include(blog_urls))
 """
-from django.conf.urls import include, url
+from django.conf.urls import url
 from django.contrib import admin
 
 import app.views
 
 
 urlpatterns = [
-    url(r'^admin/', include(admin.site.urls)),
+    url(r'^admin/', admin.site.urls),
     url(r'^$', app.views.home),
     url(r'^greetings$', app.views.greetings),
     url(r'^_ah/health$', app.views.health_check),