Support request and resposne hooks for Django instrumentation (#407)

.github/workflows/test.yml

@@ -6,7 +6,7 @@ on:
     - 'release/*'
   pull_request:
 env:
-  CORE_REPO_SHA: 94bf80f3870bceefa72d9a61353eaf5d7dd30993
+  CORE_REPO_SHA: cad261e5dae1fe986c87e6965664b45cc9ab73c3
 
 jobs:
   build:

CHANGELOG.md

@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - `opentelemetry-instrumentation-urllib3` Add urllib3 instrumentation
   ([#299](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/299))
+- `opentelemetry-instrumenation-django` now supports request and response hooks.
+  ([#407](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/407))
 
 ## [0.19b0](https://github.com/open-telemetry/opentelemetry-python-contrib/releases/tag/v0.19b0) - 2021-03-26
 

instrumentation/opentelemetry-instrumentation-django/README.rst

@@ -45,6 +45,22 @@ will extract path_info and content_type attributes from every traced request and
 
 Django Request object reference: https://docs.djangoproject.com/en/3.1/ref/request-response/#attributes
 
+Request and Response hooks
+***************************
+The instrumentation supports specifying request and response hooks. These are functions that get called back by the instrumentation right after a Span is created for a request
+and right before the span is finished while processing a response. The hooks can be configured as follows:
+
+::
+
+    def request_hook(span, request):
+        pass
+
+    def response_hook(span, request, response):
+        pass
+
+    DjangoInstrumentation().instrument(request_hook=request_hook, response_hook=response_hook)
+
+
 References
 ----------
 

instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/__init__.py

@@ -11,6 +11,67 @@
 # 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.
+"""
+
+Instrument `django`_ to trace Django applications.
+
+.. _django: https://pypi.org/project/django/
+
+Usage
+-----
+
+.. code:: python
+
+    from opentelemetry.instrumentation.django import DjangoInstrumentor
+
+    DjangoInstrumentor().instrument()
+
+
+Configuration
+-------------
+
+Exclude lists
+*************
+To exclude certain URLs from being tracked, set the environment variable ``OTEL_PYTHON_DJANGO_EXCLUDED_URLS`` with comma delimited regexes representing which URLs to exclude.
+
+For example,
+
+::
+
+    export OTEL_PYTHON_DJANGO_EXCLUDED_URLS="client/.*/info,healthcheck"
+
+will exclude requests such as ``https://site/client/123/info`` and ``https://site/xyz/healthcheck``.
+
+Request attributes
+********************
+To extract certain attributes from Django's request object and use them as span attributes, set the environment variable ``OTEL_PYTHON_DJANGO_TRACED_REQUEST_ATTRS`` to a comma
+delimited list of request attribute names.
+
+For example,
+
+::
+
+    export OTEL_PYTHON_DJANGO_TRACED_REQUEST_ATTRS='path_info,content_type'
+
+will extract path_info and content_type attributes from every traced request and add them as span attritbues.
+
+Django Request object reference: https://docs.djangoproject.com/en/3.1/ref/request-response/#attributes
+
+Request and Response hooks
+***************************
+The instrumentation supports specifying request and response hooks. These are functions that get called back by the instrumentation right after a Span is created for a request
+and right before the span is finished while processing a response. The hooks can be configured as follows:
+
+.. code:: python
+
+    def request_hook(span, request):
+        pass
+
+    def response_hook(span, request, response):
+        pass
+
+    DjangoInstrumentation().instrument(request_hook=request_hook, response_hook=response_hook)
+"""
 
 from logging import getLogger
 from os import environ
@@ -44,6 +105,11 @@ class DjangoInstrumentor(BaseInstrumentor):
         if environ.get(OTEL_PYTHON_DJANGO_INSTRUMENT) == "False":
             return
 
+        _DjangoMiddleware._otel_request_hook = kwargs.pop("request_hook", None)
+        _DjangoMiddleware._otel_response_hook = kwargs.pop(
+            "response_hook", None
+        )
+
         # This can not be solved, but is an inherent problem of this approach:
         # the order of middleware entries matters, and here you have no control
         # on that:

instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py

@@ -14,6 +14,9 @@
 
 from logging import getLogger
 from time import time
+from typing import Callable
+
+from django.http import HttpRequest, HttpResponse
 
 from opentelemetry.context import attach, detach
 from opentelemetry.instrumentation.django.version import __version__
@@ -24,7 +27,7 @@ from opentelemetry.instrumentation.wsgi import (
     wsgi_getter,
 )
 from opentelemetry.propagate import extract
-from opentelemetry.trace import SpanKind, get_tracer, use_span
+from opentelemetry.trace import Span, SpanKind, get_tracer, use_span
 from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
 
 try:
@@ -62,6 +65,11 @@ class _DjangoMiddleware(MiddlewareMixin):
     _traced_request_attrs = get_traced_request_attrs("DJANGO")
     _excluded_urls = get_excluded_urls("DJANGO")
 
+    _otel_request_hook: Callable[[Span, HttpRequest], None] = None
+    _otel_response_hook: Callable[
+        [Span, HttpRequest, HttpResponse], None
+    ] = None
+
     @staticmethod
     def _get_span_name(request):
         try:
@@ -125,6 +133,11 @@ class _DjangoMiddleware(MiddlewareMixin):
         request.META[self._environ_span_key] = span
         request.META[self._environ_token] = token
 
+        if _DjangoMiddleware._otel_request_hook:
+            _DjangoMiddleware._otel_request_hook(  # pylint: disable=not-callable
+                span, request
+            )
+
     # pylint: disable=unused-argument
     def process_view(self, request, view_func, *args, **kwargs):
         # Process view is executed before the view function, here we get the
@@ -156,30 +169,30 @@ class _DjangoMiddleware(MiddlewareMixin):
         if self._excluded_urls.url_disabled(request.build_absolute_uri("?")):
             return response
 
-        if (
+        activation = request.META.pop(self._environ_activation_key, None)
-            self._environ_activation_key in request.META.keys()
+        span = request.META.pop(self._environ_span_key, None)
-            and self._environ_span_key in request.META.keys()
+
-        ):
+        if activation and span:
             add_response_attributes(
-                request.META[self._environ_span_key],
+                span,
                 "{} {}".format(response.status_code, response.reason_phrase),
                 response,
             )
 
-            request.META.pop(self._environ_span_key)
-
             exception = request.META.pop(self._environ_exception_key, None)
+            if _DjangoMiddleware._otel_response_hook:
+                _DjangoMiddleware._otel_response_hook(  # pylint: disable=not-callable
+                    span, request, response
+                )
+
             if exception:
-                request.META[self._environ_activation_key].__exit__(
+                activation.__exit__(
                     type(exception),
                     exception,
                     getattr(exception, "__traceback__", None),
                 )
             else:
-                request.META[self._environ_activation_key].__exit__(
+                activation.__exit__(None, None, None)
-                    None, None, None
-                )
-            request.META.pop(self._environ_activation_key)
 
         if self._environ_token in request.META.keys():
             detach(request.environ.get(self._environ_token))

instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py

@@ -18,10 +18,15 @@ from unittest.mock import Mock, patch
 from django import VERSION
 from django.conf import settings
 from django.conf.urls import url
+from django.http import HttpRequest, HttpResponse
 from django.test import Client
 from django.test.utils import setup_test_environment, teardown_test_environment
 
-from opentelemetry.instrumentation.django import DjangoInstrumentor
+from opentelemetry.instrumentation.django import (
+    DjangoInstrumentor,
+    _DjangoMiddleware,
+)
+from opentelemetry.sdk.trace import Span
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import SpanKind, StatusCode
@@ -268,3 +273,42 @@ class TestMiddleware(TestBase, WsgiTestBase):
         self.assertEqual(span.attributes["path_info"], "/span_name/1234/")
         self.assertEqual(span.attributes["content_type"], "test/ct")
         self.assertNotIn("non_existing_variable", span.attributes)
+
+    def test_hooks(self):
+        request_hook_args = ()
+        response_hook_args = ()
+
+        def request_hook(span, request):
+            nonlocal request_hook_args
+            request_hook_args = (span, request)
+
+        def response_hook(span, request, response):
+            nonlocal response_hook_args
+            response_hook_args = (span, request, response)
+            response["hook-header"] = "set by hook"
+
+        _DjangoMiddleware._otel_request_hook = request_hook
+        _DjangoMiddleware._otel_response_hook = response_hook
+
+        response = Client().get("/span_name/1234/")
+        _DjangoMiddleware._otel_request_hook = (
+            _DjangoMiddleware._otel_response_hook
+        ) = None
+
+        self.assertEqual(response["hook-header"], "set by hook")
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+        span = span_list[0]
+        self.assertEqual(span.attributes["path_info"], "/span_name/1234/")
+
+        self.assertEqual(len(request_hook_args), 2)
+        self.assertEqual(request_hook_args[0].name, span.name)
+        self.assertIsInstance(request_hook_args[0], Span)
+        self.assertIsInstance(request_hook_args[1], HttpRequest)
+
+        self.assertEqual(len(response_hook_args), 3)
+        self.assertEqual(request_hook_args[0], response_hook_args[0])
+        self.assertIsInstance(response_hook_args[1], HttpRequest)
+        self.assertIsInstance(response_hook_args[2], HttpResponse)
+        self.assertEqual(response_hook_args[2], response)