Capture custom request/response headers for wsgi and change in passing response_headers in django, pyramid

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.9.1-0.28b1...HEAD)
 
+- `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes
+  ([#925])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/925)
+
 ### Added
 
 - `opentelemetry-instrumentation-dbapi` add experimental sql commenter capability

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

@@ -275,7 +275,7 @@ def process_response(self, request, response):
                 add_response_attributes(
                     span,
                     f"{response.status_code} {response.reason_phrase}",
-                    response,
+                    response.items(),
                 )
 
             propagator = get_global_response_propagator()

instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py

@@ -161,7 +161,7 @@ def trace_tween(request):
                 otel_wsgi.add_response_attributes(
                     span,
                     response_or_exception.status,
-                    response_or_exception.headers,
+                    response_or_exception.headerlist,
                 )
 
                 propagator = get_global_response_propagator()

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

@@ -117,7 +117,14 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
 from opentelemetry.propagators.textmap import Getter
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
-from opentelemetry.util.http import remove_url_credentials
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_custom_headers,
+    normalise_request_header_name,
+    normalise_response_header_name,
+    remove_url_credentials,
+)
 
 _HTTP_VERSION_PREFIX = "HTTP/"
 _CARRIER_KEY_PREFIX = "HTTP_"
@@ -205,9 +212,48 @@ def collect_request_attributes(environ):
     if flavor:
         result[SpanAttributes.HTTP_FLAVOR] = flavor
 
+    result.update(collect_custom_request_headers(environ))
     return result
 
 
+def collect_custom_request_headers(environ):
+    """Collects and returns custom HTTP request headers configured by the user from the PEP3333-conforming
+    WSGI environ to be used as span creation attributes as described
+    in the specification https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers"""
+    attributes = {}
+    custom_request_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+    )
+    for header_name in custom_request_headers_name:
+        wsgi_env_var = header_name.upper().replace("-", "_")
+        header_values = environ.get(f"HTTP_{wsgi_env_var}")
+        if header_values:
+            key = normalise_request_header_name(header_name)
+            attributes[key] = [header_values]
+    return attributes
+
+
+def collect_custom_response_headers(response_headers):
+    """Collects and returns custom HTTP response headers configured by the user from the
+    PEP3333-conforming WSGI environ as described in the specification
+    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers"""
+    attributes = {}
+    custom_response_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+    )
+    response_headers_dict = {}
+    if response_headers:
+        for header_name, header_value in response_headers:
+            response_headers_dict[header_name.lower()] = header_value
+
+    for header_name in custom_response_headers_name:
+        header_values = response_headers_dict.get(header_name.lower())
+        if header_values:
+            key = normalise_response_header_name(header_name)
+            attributes[key] = [header_values]
+    return attributes
+
+
 def add_response_attributes(
     span, start_response_status, response_headers
 ):  # pylint: disable=unused-argument
@@ -232,6 +278,8 @@ def add_response_attributes(
             Status(http_status_to_status_code(status_code, server_span=True))
         )
 
+    span.set_attributes(collect_custom_response_headers(response_headers))
+
 
 def get_default_span_name(environ):
     """Default implementation for name_callback, returns HTTP {METHOD_NAME}."""

instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py

@@ -25,6 +25,10 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import StatusCode
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+)
 
 
 class Response:
@@ -377,6 +381,53 @@ def test_credential_removal(self):
             expected.items(),
         )
 
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3"
+        },
+    )
+    def test_collect_request_attributes_collect_custom_request_headers(self):
+        attributes = {}
+        self.environ.update(
+            {
+                "HTTP_CUSTOM_TEST_HEADER_1": "Test Value 1",
+                "HTTP_CUSTOM_TEST_HEADER_2": "TestValue2,TestValue3",
+            }
+        )
+        attributes = otel_wsgi.collect_request_attributes(self.environ)
+        expected = {
+            "http.request.header.custom_test_header_1": ["Test Value 1"],
+            "http.request.header.custom_test_header_2": [
+                "TestValue2,TestValue3",
+            ],
+        }
+        self.assertGreaterEqual(attributes.items(), expected.items())
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header"
+        },
+    )
+    def test_add_response_attributes_collects_custom_response_headers(self):
+        response_headers = [
+            ("content-type", "text/plain; charset=utf-8"),
+            ("content-length", "100"),
+            ("my-custom-header", "my-custom-value-1,my-custom-header-2"),
+        ]
+        otel_wsgi.add_response_attributes(
+            self.span, "200 OK", response_headers
+        )
+        expected = {
+            "http.response.header.content_type": ["text/plain; charset=utf-8"],
+            "http.response.header.content_length": ["100"],
+            "http.response.header.my_custom_header": [
+                "my-custom-value-1,my-custom-header-2"
+            ],
+        }
+        self.span.set_attributes.assert_called_once_with(expected)
+
 
 class TestWsgiMiddlewareWithTracerProvider(WsgiTestBase):
     def validate_response(

util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py

@@ -15,9 +15,16 @@
 from os import environ
 from re import compile as re_compile
 from re import search
-from typing import Iterable
+from typing import Iterable, List
 from urllib.parse import urlparse, urlunparse
 
+OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST = (
+    "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST"
+)
+OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE = (
+    "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE"
+)
+
 
 class ExcludeList:
     """Class to exclude certain paths (given as a list of regexes) from tracing requests"""
@@ -98,3 +105,23 @@ def remove_url_credentials(url: str) -> str:
     except ValueError:  # an unparseable url was passed
         pass
     return url
+
+
+def normalise_request_header_name(header: str) -> str:
+    key = header.lower().replace("-", "_")
+    return f"http.request.header.{key}"
+
+
+def normalise_response_header_name(header: str) -> str:
+    key = header.lower().replace("-", "_")
+    return f"http.response.header.{key}"
+
+
+def get_custom_headers(env_var: str) -> List[str]:
+    custom_headers = environ.get(env_var, [])
+    if custom_headers:
+        custom_headers = [
+            custom_headers.strip()
+            for custom_headers in custom_headers.split(",")
+        ]
+    return custom_headers

util/opentelemetry-util-http/tests/test_capture_custom_headers.py

@@ -0,0 +1,67 @@
+# Copyright The OpenTelemetry 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.
+
+from unittest.mock import patch
+
+from opentelemetry.test.test_base import TestBase
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_custom_headers,
+    normalise_request_header_name,
+    normalise_response_header_name,
+)
+
+
+class TestCaptureCustomHeaders(TestBase):
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "User-Agent,Test-Header"
+        },
+    )
+    def test_get_custom_request_header(self):
+        custom_headers_to_capture = get_custom_headers(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+        )
+        self.assertEqual(
+            custom_headers_to_capture, ["User-Agent", "Test-Header"]
+        )
+
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,test-header"
+        },
+    )
+    def test_get_custom_response_header(self):
+        custom_headers_to_capture = get_custom_headers(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+        )
+        self.assertEqual(
+            custom_headers_to_capture,
+            [
+                "content-type",
+                "content-length",
+                "test-header",
+            ],
+        )
+
+    def test_normalise_request_header_name(self):
+        key = normalise_request_header_name("Test-Header")
+        self.assertEqual(key, "http.request.header.test_header")
+
+    def test_normalise_response_header_name(self):
+        key = normalise_response_header_name("Test-Header")
+        self.assertEqual(key, "http.response.header.test_header")