Set status on django request span even if child span was not finished (#308)

Author: Liudmila Molkova

opencensus/trace/ext/django/middleware.py

@@ -33,6 +33,7 @@
 HTTP_STATUS_CODE = attributes_helper.COMMON_ATTRIBUTES['HTTP_STATUS_CODE']
 
 REQUEST_THREAD_LOCAL_KEY = 'django_request'
+SPAN_THREAD_LOCAL_KEY = 'django_span'
 
 BLACKLIST_PATHS = 'BLACKLIST_PATHS'
 GCP_EXPORTER_PROJECT = 'GCP_EXPORTER_PROJECT'
@@ -70,12 +71,21 @@ def _get_django_request():
     return execution_context.get_opencensus_attr(REQUEST_THREAD_LOCAL_KEY)
 
 
+def _get_django_span():
+    """Get Django span from thread local.
+
+    :rtype: str
+    :returns: Django request.
+    """
+    return execution_context.get_opencensus_attr(SPAN_THREAD_LOCAL_KEY)
+
+
 def _get_current_tracer():
     """Get the current request tracer."""
     return execution_context.get_opencensus_tracer()
 
 
-def _set_django_attributes(tracer, request):
+def _set_django_attributes(span, request):
     """Set the django related attributes."""
     django_user = getattr(request, 'user', None)
 
@@ -87,10 +97,10 @@ def _set_django_attributes(tracer, request):
 
     # User id is the django autofield for User model as the primary key
     if user_id is not None:
-        tracer.add_attribute_to_current_span('django.user.id', str(user_id))
+        span.add_attribute('django.user.id', str(user_id))
 
     if user_name is not None:
-        tracer.add_attribute_to_current_span('django.user.name', user_name)
+        span.add_attribute('django.user.name', user_name)
 
 
 class OpencensusMiddleware(MiddlewareMixin):
@@ -185,13 +195,24 @@ def process_request(self, request):
             tracer.add_attribute_to_current_span(
                 attribute_key=HTTP_URL,
                 attribute_value=request.path)
+
+            # Add the span to thread local
+            # in some cases (exceptions, timeouts) currentspan in
+            # response event will be one of a child spans.
+            # let's keep reference to 'django' span and
+            # use it in response event
+            execution_context.set_opencensus_attr(
+                SPAN_THREAD_LOCAL_KEY,
+                span)
+
         except Exception:  # pragma: NO COVER
             log.error('Failed to trace request', exc_info=True)
 
     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.
         """
+
         # Do not trace if the url is blacklisted
         if utils.disable_tracing_url(request.path, self._blacklist_paths):
             return
@@ -211,13 +232,14 @@ def process_response(self, request, response):
             return response
 
         try:
-            tracer = _get_current_tracer()
-            tracer.add_attribute_to_current_span(
+            span = _get_django_span()
+            span.add_attribute(
                 attribute_key=HTTP_STATUS_CODE,
                 attribute_value=str(response.status_code))
 
-            _set_django_attributes(tracer, request)
+            _set_django_attributes(span, request)
 
+            tracer = _get_current_tracer()
             tracer.end_span()
             tracer.finish()
         except Exception:  # pragma: NO COVER

tests/unit/trace/ext/django/test_middleware.py

@@ -392,50 +392,91 @@ def test_process_response(self):
 
         self.assertEqual(span.attributes, expected_attributes)
 
+    def test_process_response_unfinished_child_span(self):
+        from opencensus.trace.ext.django import middleware
+
+        trace_id = '2dd43a1d6b2549c6bc2a1a54c2fc0b05'
+        span_id = '6e0c63257de34c92'
+        django_trace_id = '{}/{}'.format(trace_id, span_id)
+
+        django_request = RequestFactory().get('/', **{
+            google_cloud_format._TRACE_CONTEXT_HEADER_NAME: django_trace_id})
+
+        middleware_obj = middleware.OpencensusMiddleware()
+
+        middleware_obj.process_request(django_request)
+        tracer = middleware._get_current_tracer()
+        span = tracer.current_span()
+
+        exporter_mock = mock.Mock()
+        tracer.exporter = exporter_mock
+
+        django_response = mock.Mock()
+        django_response.status_code = 500
+
+        expected_attributes = {
+            'http.url': u'/',
+            'http.method': 'GET',
+            'http.status_code': '500',
+            'django.user.id': '123',
+            'django.user.name': 'test_name'
+        }
+
+        mock_user = mock.Mock()
+        mock_user.pk = 123
+        mock_user.get_username.return_value = 'test_name'
+        django_request.user = mock_user
+
+        tracer.start_span()
+        self.assertNotEqual(span, tracer.current_span())
+        middleware_obj.process_response(django_request, django_response)
+
+        self.assertEqual(span.attributes, expected_attributes)
+
 
 class Test__set_django_attributes(unittest.TestCase):
-    class Tracer(object):
+    class Span(object):
         def __init__(self):
             self.attributes = {}
 
-        def add_attribute_to_current_span(self, key, value):
+        def add_attribute(self, key, value):
             self.attributes[key] = value
 
     def test__set_django_attributes_no_user(self):
         from opencensus.trace.ext.django.middleware import \
             _set_django_attributes
-        tracer = self.Tracer()
+        span = self.Span()
         request = mock.Mock()
 
         request.user = None
 
-        _set_django_attributes(tracer, request)
+        _set_django_attributes(span, request)
 
         expected_attributes = {}
 
-        self.assertEqual(tracer.attributes, expected_attributes)
+        self.assertEqual(span.attributes, expected_attributes)
 
     def test__set_django_attributes_no_user_info(self):
         from opencensus.trace.ext.django.middleware import \
             _set_django_attributes
-        tracer = self.Tracer()
+        span = self.Span()
         request = mock.Mock()
         django_user = mock.Mock()
 
         request.user = django_user
         django_user.pk = None
         django_user.get_username.return_value = None
 
-        _set_django_attributes(tracer, request)
+        _set_django_attributes(span, request)
 
         expected_attributes = {}
 
-        self.assertEqual(tracer.attributes, expected_attributes)
+        self.assertEqual(span.attributes, expected_attributes)
 
     def test__set_django_attributes_with_user_info(self):
         from opencensus.trace.ext.django.middleware import \
             _set_django_attributes
-        tracer = self.Tracer()
+        span = self.Span()
         request = mock.Mock()
         django_user = mock.Mock()
 
@@ -445,10 +486,10 @@ def test__set_django_attributes_with_user_info(self):
         django_user.pk = test_id
         django_user.get_username.return_value = test_name
 
-        _set_django_attributes(tracer, request)
+        _set_django_attributes(span, request)
 
         expected_attributes = {
             'django.user.id': '123',
             'django.user.name': test_name}
 
-        self.assertEqual(tracer.attributes, expected_attributes)
+        self.assertEqual(span.attributes, expected_attributes)