New style middleware

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

@@ -17,7 +17,6 @@
 import six
 
 import django.conf
-from django.utils.deprecation import MiddlewareMixin
 
 from opencensus.common import configuration
 from opencensus.trace import attributes_helper
@@ -33,9 +32,6 @@
 HTTP_URL = attributes_helper.COMMON_ATTRIBUTES['HTTP_URL']
 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'
 BLACKLIST_HOSTNAMES = 'BLACKLIST_HOSTNAMES'
 
@@ -48,59 +44,17 @@ class _DjangoMetaWrapper(object):
     Django request.META
     """
 
-    def __init__(self, meta=None):
-        self.meta = meta or _get_django_request().META
+    def __init__(self, meta):
+        self.meta = meta
 
     def get(self, key):
         return self.meta.get('HTTP_' + key.upper().replace('-', '_'))
 
 
-def _get_django_request():
-    """Get Django request from thread local.
-
-    :rtype: str
-    :returns: 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(span, request):
-    """Set the django related attributes."""
-    django_user = getattr(request, 'user', None)
-
-    if django_user is None:
-        return
-
-    user_id = django_user.pk
-    user_name = django_user.get_username()
-
-    # User id is the django autofield for User model as the primary key
-    if user_id is not None:
-        span.add_attribute('django.user.id', str(user_id))
-
-    if user_name is not None:
-        span.add_attribute('django.user.name', str(user_name))
-
-
-class OpencensusMiddleware(MiddlewareMixin):
-    """Saves the request in thread local"""
-
-    def __init__(self, get_response=None):
+class OpencensusMiddleware(object):
+    def __init__(self, get_response):
         self.get_response = get_response
+
         settings = getattr(django.conf.settings, 'OPENCENSUS', {})
         settings = settings.get('TRACE', {})
 
@@ -123,94 +77,80 @@ def __init__(self, get_response=None):
 
         self.blacklist_hostnames = settings.get(BLACKLIST_HOSTNAMES, None)
 
-    def process_request(self, request):
-        """Called on each request, before Django decides which view to execute.
-
+    def __call__(self, request):
+        """
         :type request: :class:`~django.http.request.HttpRequest`
         :param request: Django http request.
         """
-        # Do not trace if the url is blacklisted
-        if utils.disable_tracing_url(request.path, self.blacklist_paths):
-            return
+        span = None
 
-        # Add the request to thread local
-        execution_context.set_opencensus_attr(
-            REQUEST_THREAD_LOCAL_KEY,
-            request)
-
-        execution_context.set_opencensus_attr(
-            'blacklist_hostnames',
-            self.blacklist_hostnames)
-
-        try:
-            # Start tracing this request
-            span_context = self.propagator.from_headers(
-                _DjangoMetaWrapper(_get_django_request().META))
-
-            # Reload the tracer with the new span context
-            tracer = tracer_module.Tracer(
-                span_context=span_context,
-                sampler=self.sampler,
-                exporter=self.exporter,
-                propagator=self.propagator)
-
-            # Span name is being set at process_view
-            span = tracer.start_span()
-            span.span_kind = span_module.SpanKind.SERVER
-            tracer.add_attribute_to_current_span(
-                attribute_key=HTTP_METHOD,
-                attribute_value=request.method)
-            tracer.add_attribute_to_current_span(
-                attribute_key=HTTP_URL,
-                attribute_value=str(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
+        # Do not trace if the url is blacklisted
+        if not utils.disable_tracing_url(request.path, self.blacklist_paths):
+            # Add the request to thread local
             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.
+                'blacklist_hostnames',
+                self.blacklist_hostnames)
+
+            try:
+                # Start tracing this request
+                span_context = self.propagator.from_headers(
+                    _DjangoMetaWrapper(request.META))
+
+                # Reload the tracer with the new span context
+                tracer = tracer_module.Tracer(
+                    span_context=span_context,
+                    sampler=self.sampler,
+                    exporter=self.exporter,
+                    propagator=self.propagator)
+
+                # Span name is being set at process_view
+                span = tracer.start_span()
+                span.span_kind = span_module.SpanKind.SERVER
+                tracer.add_attribute_to_current_span(
+                    attribute_key=HTTP_METHOD,
+                    attribute_value=request.method)
+                tracer.add_attribute_to_current_span(
+                    attribute_key=HTTP_URL,
+                    attribute_value=request.path)
+
+            except Exception:  # pragma: NO COVER
+                log.error('Failed to trace request', exc_info=True)
+
+        response = self.get_response(request)
+
+        if span:
+            try:
+                span.name = utils.get_func_name(request.resolver_match.func)
+
+                span.add_attribute(
+                    attribute_key=HTTP_STATUS_CODE,
+                    attribute_value=str(response.status_code))
+
+                self.set_django_attributes(span, request, response)
+
+                tracer.end_span()
+                tracer.finish()
+            except Exception:  # pragma: NO COVER
+                log.error('Failed to trace request', exc_info=True)
+
+        return response
+
+    def set_django_attributes(self, span, request, response):  # noqa
+        """ The last method before the span finishes, allowing to set
+            django-related attributes on the span. Override this in a
+            subclass of the middleware to set custom attributes.
         """
+        django_user = getattr(request, 'user', None)
 
-        # Do not trace if the url is blacklisted
-        if utils.disable_tracing_url(request.path, self.blacklist_paths):
+        if django_user is None:
             return
 
-        try:
-            # Get the current span and set the span name to the current
-            # function name of the request.
-            tracer = _get_current_tracer()
-            span = tracer.current_span()
-            span.name = utils.get_func_name(view_func)
-        except Exception:  # pragma: NO COVER
-            log.error('Failed to trace request', exc_info=True)
+        user_id = django_user.pk
+        user_name = django_user.get_username()
 
-    def process_response(self, request, response):
-        # Do not trace if the url is blacklisted
-        if utils.disable_tracing_url(request.path, self.blacklist_paths):
-            return response
-
-        try:
-            span = _get_django_span()
-            span.add_attribute(
-                attribute_key=HTTP_STATUS_CODE,
-                attribute_value=str(response.status_code))
-
-            _set_django_attributes(span, request)
-
-            tracer = _get_current_tracer()
-            tracer.end_span()
-            tracer.finish()
-        except Exception:  # pragma: NO COVER
-            log.error('Failed to trace request', exc_info=True)
-        finally:
-            return response
+        # User id is the django autofield for User model as the primary key
+        if user_id is not None:
+            span.add_attribute('django.user.id', str(user_id))
+
+        if user_name is not None:
+            span.add_attribute('django.user.name', str(user_name))