urllib: add support for capturing request and response headers

Author: xrmx

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

@@ -135,8 +135,15 @@ def response_hook(span: Span, request: Request, response: HTTPResponse):
 )
 from opentelemetry.trace import Span, SpanKind, Tracer, get_tracer
 from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_CLIENT_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_CLIENT_RESPONSE,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
     ExcludeList,
+    get_custom_header_attributes,
+    get_custom_headers,
     get_excluded_urls,
+    normalise_request_header_name,
+    normalise_response_header_name,
     parse_excluded_urls,
     redact_url,
     sanitize_method,
@@ -169,6 +176,9 @@ def _instrument(self, **kwargs: Any):
                 ``response_hook``: An optional callback which is invoked right before the span is finished processing a response
                 ``excluded_urls``: A string containing a comma-delimited
                     list of regexes used to exclude URLs from tracking
+                ``captured_request_headers``: A comma-separated list of regexes to match against request headers to capture
+                ``captured_response_headers``: A comma-separated list of regexes to match against response headers to capture
+                ``sensitive_headers``: A comma-separated list of regexes to match against captured headers to be sanitized
         """
         # initialize semantic conventions opt-in if needed
         _OpenTelemetrySemanticConventionStability._initialize()
@@ -205,6 +215,24 @@ def _instrument(self, **kwargs: Any):
                 else parse_excluded_urls(excluded_urls)
             ),
             sem_conv_opt_in_mode=sem_conv_opt_in_mode,
+            captured_request_headers=kwargs.get(
+                "captured_request_headers",
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_CLIENT_REQUEST
+                ),
+            ),
+            captured_response_headers=kwargs.get(
+                "captured_response_headers",
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_CLIENT_RESPONSE
+                ),
+            ),
+            sensitive_headers=kwargs.get(
+                "sensitive_headers",
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
+                ),
+            ),
         )
 
     def _uninstrument(self, **kwargs: Any):
@@ -223,6 +251,9 @@ def _instrument(
     response_hook: _ResponseHookT = None,
     excluded_urls: ExcludeList | None = None,
     sem_conv_opt_in_mode: _StabilityMode = _StabilityMode.DEFAULT,
+    captured_request_headers: list[str] | None = None,
+    captured_response_headers: list[str] | None = None,
+    sensitive_headers: list[str] | None = None,
 ):
     """Enables tracing of all requests calls that go through
     :code:`urllib.Client._make_request`"""
@@ -275,6 +306,15 @@ def _instrumented_open_call(
         )
         _set_http_url(labels, url, sem_conv_opt_in_mode)
 
+        labels.update(
+            get_custom_header_attributes(
+                request.headers,
+                captured_request_headers,
+                sensitive_headers,
+                normalise_request_header_name,
+            )
+        )
+
         with tracer.start_as_current_span(
             span_name, kind=SpanKind.CLIENT, attributes=labels
         ) as span:
@@ -310,6 +350,16 @@ def _instrumented_open_call(
                         labels, f"{ver_[:1]}.{ver_[:-1]}", sem_conv_opt_in_mode
                     )
 
+                if span.is_recording():
+                    response_headers_to_set = get_custom_header_attributes(
+                        result.headers,
+                        captured_response_headers,
+                        sensitive_headers,
+                        normalise_response_header_name,
+                    )
+                    for header, value in response_headers_to_set.items():
+                        span.set_attribute(header, value)
+
             if exception is not None and _report_new(sem_conv_opt_in_mode):
                 span.set_attribute(ERROR_TYPE, type(exception).__qualname__)
                 labels[ERROR_TYPE] = type(exception).__qualname__