Django: Capture custom request/response headers (open-telemetry#1024)

Author: sanketmehta28

CHANGELOG.md

@@ -8,7 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.10.0-0.29b0...HEAD)
 
 ### Added
-
+- `opentelemetry-instrumentation-django` Capture custom request/response headers in span attributes
+  ([#1024])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1024)
+- `opentelemetry-instrumentation-asgi` Capture custom request/response headers in span attributes
+  ([#1004])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1004)
 - `opentelemetry-instrumentation-psycopg2` extended the sql commenter support of dbapi into psycopg2
   ([#940](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/940))
 - `opentelemetry-instrumentation-flask` Fix non-recording span bug
@@ -24,9 +27,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [1.10.0-0.29b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.10.0-0.29b0) - 2022-03-10
 
-- `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes
-  ([#1004])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1004)
-
 - `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes
   ([#925])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/925)
 - `opentelemetry-instrumentation-flask` Flask: Capture custom request/response headers in span attributes

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

@@ -76,6 +76,52 @@ def response_hook(span, request, response):
 Django Request object: https://docs.djangoproject.com/en/3.1/ref/request-response/#httprequest-objects
 Django Response object: https://docs.djangoproject.com/en/3.1/ref/request-response/#httpresponse-objects
 
+Capture HTTP request and response headers
+*****************************************
+You can configure the agent to capture predefined HTTP headers as span attributes, according to the `semantic convention <https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers>`_.
+
+Request headers
+***************
+To capture predefined HTTP request headers as span attributes, set the environment variable ``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST``
+to a comma-separated list of HTTP header names.
+
+For example,
+::
+
+    export OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST="content_type,custom_request_header"
+
+will extract content_type and custom_request_header from request headers and add them as span attributes.
+
+It is recommended that you should give the correct names of the headers to be captured in the environment variable.
+Request header names in django are case insensitive. So, giving header name as ``CUStom_Header`` in environment variable will be able capture header with name ``custom-header``.
+
+The name of the added span attribute will follow the format ``http.request.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
+The value of the attribute will be single item list containing all the header values.
+
+Example of the added span attribute,
+``http.request.header.custom_request_header = ["<value1>,<value2>"]``
+
+Response headers
+****************
+To capture predefined HTTP response headers as span attributes, set the environment variable ``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE``
+to a comma-separated list of HTTP header names.
+
+For example,
+::
+
+    export OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE="content_type,custom_response_header"
+
+will extract content_type and custom_response_header from response headers and add them as span attributes.
+
+It is recommended that you should give the correct names of the headers to be captured in the environment variable.
+Response header names captured in django are case insensitive. So, giving header name as ``CUStomHeader`` in environment variable will be able capture header with name ``customheader``.
+
+The name of the added span attribute will follow the format ``http.response.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
+The value of the attribute will be single item list containing all the header values.
+
+Example of the added span attribute,
+``http.response.header.custom_response_header = ["<value1>,<value2>"]``
+
 API
 ---
 """

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

@@ -28,13 +28,19 @@
     _start_internal_or_server_span,
     extract_attributes_from_object,
 )
+from opentelemetry.instrumentation.wsgi import (
+    add_custom_request_headers as wsgi_add_custom_request_headers,
+)
+from opentelemetry.instrumentation.wsgi import (
+    add_custom_response_headers as wsgi_add_custom_response_headers,
+)
 from opentelemetry.instrumentation.wsgi import add_response_attributes
 from opentelemetry.instrumentation.wsgi import (
     collect_request_attributes as wsgi_collect_request_attributes,
 )
 from opentelemetry.instrumentation.wsgi import wsgi_getter
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import Span, use_span
+from opentelemetry.trace import Span, SpanKind, use_span
 from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
 
 try:
@@ -77,7 +83,13 @@ def __call__(self, request):
 
 # try/except block exclusive for optional ASGI imports.
 try:
-    from opentelemetry.instrumentation.asgi import asgi_getter
+    from opentelemetry.instrumentation.asgi import asgi_getter, asgi_setter
+    from opentelemetry.instrumentation.asgi import (
+        collect_custom_request_headers_attributes as asgi_collect_custom_request_attributes,
+    )
+    from opentelemetry.instrumentation.asgi import (
+        collect_custom_response_headers_attributes as asgi_collect_custom_response_attributes,
+    )
     from opentelemetry.instrumentation.asgi import (
         collect_request_attributes as asgi_collect_request_attributes,
     )
@@ -213,6 +225,13 @@ def process_request(self, request):
                     self._traced_request_attrs,
                     attributes,
                 )
+                if span.is_recording() and span.kind == SpanKind.SERVER:
+                    attributes.update(
+                        asgi_collect_custom_request_attributes(carrier)
+                    )
+            else:
+                if span.is_recording() and span.kind == SpanKind.SERVER:
+                    wsgi_add_custom_request_headers(span, carrier)
 
             for key, value in attributes.items():
                 span.set_attribute(key, value)
@@ -257,6 +276,7 @@ def process_exception(self, request, exception):
         if self._environ_activation_key in request.META.keys():
             request.META[self._environ_exception_key] = exception
 
+    # pylint: disable=too-many-branches
     def process_response(self, request, response):
         if self._excluded_urls.url_disabled(request.build_absolute_uri("?")):
             return response
@@ -271,12 +291,25 @@ def process_response(self, request, response):
         if activation and span:
             if is_asgi_request:
                 set_status_code(span, response.status_code)
+
+                if span.is_recording() and span.kind == SpanKind.SERVER:
+                    custom_headers = {}
+                    for key, value in response.items():
+                        asgi_setter.set(custom_headers, key, value)
+
+                    custom_res_attributes = (
+                        asgi_collect_custom_response_attributes(custom_headers)
+                    )
+                    for key, value in custom_res_attributes.items():
+                        span.set_attribute(key, value)
             else:
                 add_response_attributes(
                     span,
                     f"{response.status_code} {response.reason_phrase}",
                     response.items(),
                 )
+                if span.is_recording() and span.kind == SpanKind.SERVER:
+                    wsgi_add_custom_response_headers(span, response.items())
 
             propagator = get_global_response_propagator()
             if propagator:

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

@@ -43,14 +43,20 @@
     format_span_id,
     format_trace_id,
 )
-from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_excluded_urls,
+    get_traced_request_attrs,
+)
 
 # pylint: disable=import-error
 from .views import (
     error,
     excluded,
     excluded_noarg,
     excluded_noarg2,
+    response_with_custom_header,
     route_span_name,
     traced,
     traced_template,
@@ -67,6 +73,7 @@
 
 urlpatterns = [
     re_path(r"^traced/", traced),
+    re_path(r"^traced_custom_header/", response_with_custom_header),
     re_path(r"^route/(?P<year>[0-9]{4})/template/$", traced_template),
     re_path(r"^error/", error),
     re_path(r"^excluded_arg/", excluded),
@@ -451,3 +458,107 @@ def test_django_with_wsgi_instrumented(self):
                 parent_span.get_span_context().span_id,
                 span_list[0].parent.span_id,
             )
+
+
+class TestMiddlewareWsgiWithCustomHeaders(TestBase, WsgiTestBase):
+    @classmethod
+    def setUpClass(cls):
+        conf.settings.configure(ROOT_URLCONF=modules[__name__])
+        super().setUpClass()
+
+    def setUp(self):
+        super().setUp()
+        setup_test_environment()
+        tracer_provider, exporter = self.create_tracer_provider()
+        self.exporter = exporter
+        _django_instrumentor.instrument(tracer_provider=tracer_provider)
+        self.env_patch = patch.dict(
+            "os.environ",
+            {
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+            },
+        )
+        self.env_patch.start()
+
+    def tearDown(self):
+        super().tearDown()
+        self.env_patch.stop()
+        teardown_test_environment()
+        _django_instrumentor.uninstrument()
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        conf.settings = conf.LazySettings()
+
+    def test_http_custom_request_headers_in_span_attributes(self):
+        expected = {
+            "http.request.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.request.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        Client(
+            HTTP_CUSTOM_TEST_HEADER_1="test-header-value-1",
+            HTTP_CUSTOM_TEST_HEADER_2="test-header-value-2",
+        ).get("/traced/")
+        spans = self.exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+        self.memory_exporter.clear()
+
+    def test_http_custom_request_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.request.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        Client(HTTP_CUSTOM_TEST_HEADER_1="test-header-value-1").get("/traced/")
+        spans = self.exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, span.attributes)
+        self.memory_exporter.clear()
+
+    def test_http_custom_response_headers_in_span_attributes(self):
+        expected = {
+            "http.response.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.response.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        Client().get("/traced_custom_header/")
+        spans = self.exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+        self.memory_exporter.clear()
+
+    def test_http_custom_response_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.response.header.custom_test_header_3": (
+                "test-header-value-3",
+            ),
+        }
+        Client().get("/traced_custom_header/")
+        spans = self.exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, span.attributes)
+        self.memory_exporter.clear()