enhancement: use a common method to capture custom headers

Signed-off-by: Varsha GS <[email protected]>

src/instana/instrumentation/aiohttp/client.py

@@ -12,7 +12,7 @@
 from instana.propagators.format import Format
 from instana.singletons import agent
 from instana.util.secrets import strip_secrets_from_query
-from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+from instana.util.traceutils import get_tracer_tuple, tracing_is_off, extract_custom_headers
 
 try:
     import aiohttp
@@ -21,19 +21,6 @@
         from aiohttp.client import ClientSession
         from instana.span.span import InstanaSpan
 
-    def extract_custom_headers(
-        span: "InstanaSpan", headers: Dict[str, Any]
-    ) -> None:
-        if not agent.options.extra_http_headers or not headers:
-            return
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                if custom_header in headers:
-                    span.set_attribute(
-                        f"http.header.{custom_header}", headers[custom_header]
-                    )
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
 
     async def stan_request_start(
         session: "ClientSession", trace_config_ctx: SimpleNamespace, params

src/instana/instrumentation/aiohttp/server.py

@@ -11,6 +11,7 @@
 from instana.propagators.format import Format
 from instana.singletons import agent, tracer
 from instana.util.secrets import strip_secrets_from_query
+from instana.util.traceutils import extract_custom_headers
 
 if TYPE_CHECKING:
     from instana.span.span import InstanaSpan
@@ -22,20 +23,6 @@
     if TYPE_CHECKING:
         import aiohttp.web
 
-    def extract_custom_headers(
-        span: "InstanaSpan", headers: Dict[str, Any]
-    ) -> None:
-        if not agent.options.extra_http_headers or not headers:
-            return
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                if custom_header in headers:
-                    span.set_attribute(
-                        f"http.header.{custom_header}", headers[custom_header]
-                    )
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
-
     @middleware
     async def stan_middleware(
         request: "aiohttp.web.Request",

src/instana/instrumentation/asgi.py

@@ -5,7 +5,7 @@
 Instana ASGI Middleware
 """
 
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Tuple
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict
 
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import SpanKind
@@ -14,6 +14,7 @@
 from instana.propagators.format import Format
 from instana.singletons import agent, tracer
 from instana.util.secrets import strip_secrets_from_query
+from instana.util.traceutils import extract_custom_headers
 
 if TYPE_CHECKING:
     from starlette.middleware.exceptions import ExceptionMiddleware
@@ -28,23 +29,6 @@ class InstanaASGIMiddleware:
     def __init__(self, app: "ExceptionMiddleware") -> None:
         self.app = app
 
-    def _extract_custom_headers(
-        self, span: "InstanaSpan", headers: List[Tuple[object, ...]]
-    ) -> None:
-        if agent.options.extra_http_headers is None:
-            return
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                # Headers are in the following format: b'x-header-1'
-                for header_pair in headers:
-                    if header_pair[0].decode("utf-8").lower() == custom_header.lower():
-                        span.set_attribute(
-                            f"http.header.{custom_header}",
-                            header_pair[1].decode("utf-8"),
-                        )
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
-
     def _collect_kvs(self, scope: Dict[str, Any], span: "InstanaSpan") -> None:
         try:
             span.set_attribute("span.kind", SpanKind.SERVER)
@@ -94,7 +78,7 @@ async def __call__(
         with tracer.start_as_current_span("asgi", span_context=request_context) as span:
             self._collect_kvs(scope, span)
             if "headers" in scope and agent.options.extra_http_headers:
-                self._extract_custom_headers(span, scope["headers"])
+                extract_custom_headers(span, scope["headers"])
 
             instana_send = self._send_with_instana(
                 span,
@@ -125,7 +109,7 @@ async def send_wrapper(response: Dict[str, Any]) -> Awaitable[None]:
 
                     headers = response.get("headers")
                     if headers:
-                        self._extract_custom_headers(current_span, headers)
+                        extract_custom_headers(current_span, headers)
                         tracer.inject(current_span.context, Format.BINARY, headers)
                 except Exception:
                     logger.debug("ASGI send_wrapper error: ", exc_info=True)

src/instana/instrumentation/django/middleware.py

@@ -13,10 +13,10 @@
     from instana.log import logger
     from instana.singletons import agent, tracer
     from instana.util.secrets import strip_secrets_from_query
+    from instana.util.traceutils import extract_custom_headers
     from instana.propagators.format import Format
 
     if TYPE_CHECKING:
-        from instana.span.span import InstanaSpan
         from django.core.handlers.base import BaseHandler
         from django.http import HttpRequest, HttpResponse
 
@@ -53,29 +53,6 @@ def __init__(
             super(InstanaMiddleware, self).__init__(get_response)
             self.get_response = get_response
 
-        def _extract_custom_headers(
-            self, span: "InstanaSpan", headers: Dict[str, Any], format: bool
-        ) -> None:
-            if agent.options.extra_http_headers is None:
-                return
-
-            try:
-                for custom_header in agent.options.extra_http_headers:
-                    # Headers are available in this format: HTTP_X_CAPTURE_THIS
-                    django_header = (
-                        ("HTTP_" + custom_header.upper()).replace("-", "_")
-                        if format
-                        else custom_header
-                    )
-
-                    if django_header in headers:
-                        span.set_attribute(
-                            f"http.header.{custom_header}", headers[django_header]
-                        )
-
-            except Exception:
-                logger.debug("Instana middleware @ extract_custom_headers: ", exc_info=True)
-
         def process_request(self, request: Type["HttpRequest"]) -> None:
             try:
                 env = request.META
@@ -89,7 +66,7 @@ def process_request(self, request: Type["HttpRequest"]) -> None:
                 token = context.attach(ctx)
                 request.token = token
 
-                self._extract_custom_headers(span, env, format=True)
+                extract_custom_headers(span, env, format=True)
 
                 request.span.set_attribute(SpanAttributes.HTTP_METHOD, request.method)
                 if "PATH_INFO" in env:
@@ -138,7 +115,7 @@ def process_response(
                         SpanAttributes.HTTP_STATUS_CODE, response.status_code
                     )
                     if hasattr(response, "headers"):
-                        self._extract_custom_headers(
+                        extract_custom_headers(
                             request.span, response.headers, format=False
                         )
                     tracer.inject(request.span.context, Format.HTTP_HEADERS, response)

src/instana/instrumentation/flask/common.py

@@ -10,23 +10,15 @@
 from opentelemetry.semconv.trace import SpanAttributes
 
 from instana.log import logger
-from instana.singletons import tracer, agent
+from instana.singletons import tracer
 from instana.propagators.format import Format
-from instana.instrumentation.flask import signals_available
 
 
 if TYPE_CHECKING:
-    from instana.span.span import InstanaSpan
     from werkzeug.exceptions import HTTPException
     from flask.typing import ResponseReturnValue
     from jinja2.environment import Template
 
-    if signals_available:
-        from werkzeug.datastructures.headers import Headers
-    else:
-        from werkzeug.datastructures import Headers
-
-
 @wrapt.patch_function_wrapper('flask', 'templating._render')
 def render_with_instana(
     wrapped: Callable[..., str],
@@ -97,21 +89,3 @@ def handle_user_exception_with_instana(
         logger.debug("handle_user_exception_with_instana:", exc_info=True)
 
     return response
-
-
-def extract_custom_headers(
-    span: "InstanaSpan", headers: Union[Dict[str, Any], "Headers"], format: bool
-) -> None:
-    if agent.options.extra_http_headers is None:
-        return
-    try:
-        for custom_header in agent.options.extra_http_headers:
-            # Headers are available in this format: HTTP_X_CAPTURE_THIS
-            flask_header = ('HTTP_' + custom_header.upper()).replace('-', '_') if format else custom_header
-            if flask_header in headers:
-                span.set_attribute(
-                    "http.header.%s" % custom_header, headers[flask_header]
-                )
-
-    except Exception:
-        logger.debug("extract_custom_headers: ", exc_info=True)

src/instana/instrumentation/flask/vanilla.py

@@ -13,7 +13,7 @@
 from instana.log import logger
 from instana.singletons import agent, tracer
 from instana.util.secrets import strip_secrets_from_query
-from instana.instrumentation.flask.common import extract_custom_headers
+from instana.util.traceutils import extract_custom_headers
 from instana.propagators.format import Format
 
 path_tpl_re = re.compile('<.*>')

src/instana/instrumentation/flask/with_blinker.py

@@ -12,7 +12,7 @@
 from instana.log import logger
 from instana.util.secrets import strip_secrets_from_query
 from instana.singletons import agent, tracer
-from instana.instrumentation.flask.common import extract_custom_headers
+from instana.util.traceutils import extract_custom_headers
 from instana.propagators.format import Format
 
 import flask
@@ -78,7 +78,7 @@ def request_finished_with_instana(
             extract_custom_headers(span, response.headers, format=False)
 
             tracer.inject(span.context, Format.HTTP_HEADERS, response.headers)
-    except:
+    except Exception:
         logger.debug("Flask request_finished_with_instana", exc_info=True)
     finally:
         if span and span.is_recording():

src/instana/instrumentation/pyramid.py

@@ -16,12 +16,12 @@
     from instana.log import logger
     from instana.singletons import tracer, agent
     from instana.util.secrets import strip_secrets_from_query
+    from instana.util.traceutils import extract_custom_headers
     from instana.propagators.format import Format
 
     if TYPE_CHECKING:
         from pyramid.request import Request
         from pyramid.response import Response
-        from instana.span.span import InstanaSpan
         from pyramid.registry import Registry
 
     class InstanaTweenFactory(object):
@@ -32,21 +32,6 @@ def __init__(
         ) -> None:
             self.handler = handler
 
-        def _extract_custom_headers(
-            self, span: "InstanaSpan", headers: Dict[str, Any]
-        ) -> None:
-            if not agent.options.extra_http_headers:
-                return
-            try:
-                for custom_header in agent.options.extra_http_headers:
-                    if custom_header in headers:
-                        span.set_attribute(
-                            f"http.header.{custom_header}", headers[custom_header]
-                        )
-
-            except Exception:
-                logger.debug("extract_custom_headers: ", exc_info=True)
-
         def __call__(self, request: "Request") -> "Response":
             ctx = tracer.extract(Format.HTTP_HEADERS, dict(request.headers))
 
@@ -56,7 +41,7 @@ def __call__(self, request: "Request") -> "Response":
                 span.set_attribute(SpanAttributes.HTTP_METHOD, request.method)
                 span.set_attribute(SpanAttributes.HTTP_URL, request.path)
 
-                self._extract_custom_headers(span, request.headers)
+                extract_custom_headers(span, request.headers)
 
                 if len(request.query_string):
                     scrubbed_params = strip_secrets_from_query(
@@ -74,7 +59,7 @@ def __call__(self, request: "Request") -> "Response":
                             "http.path_tpl", request.matched_route.pattern
                         )
 
-                    self._extract_custom_headers(span, response.headers)
+                    extract_custom_headers(span, response.headers)
 
                     tracer.inject(span.context, Format.HTTP_HEADERS, response.headers)
                 except HTTPException as e:

src/instana/instrumentation/tornado/client.py

@@ -13,25 +13,10 @@
     from instana.log import logger
     from instana.singletons import agent, tracer
     from instana.util.secrets import strip_secrets_from_query
+    from instana.util.traceutils import extract_custom_headers
     from instana.propagators.format import Format
     from instana.span.span import get_current_span
 
-    if TYPE_CHECKING:
-        from instana.span.span import InstanaSpan
-
-    def extract_custom_headers(
-        span: "InstanaSpan", headers: Dict[str, Any]
-    ) -> None:
-        if not agent.options.extra_http_headers or not headers:
-            return
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                if custom_header in headers:
-                    span.set_attribute(
-                        f"http.header.{custom_header}", headers[custom_header]
-                    )
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
 
     @wrapt.patch_function_wrapper('tornado.httpclient', 'AsyncHTTPClient.fetch')
     def fetch_with_instana(wrapped, instance, argv, kwargs):

src/instana/instrumentation/tornado/server.py

@@ -12,18 +12,9 @@
     from instana.log import logger
     from instana.singletons import agent, tracer
     from instana.util.secrets import strip_secrets_from_query
+    from instana.util.traceutils import extract_custom_headers
     from instana.propagators.format import Format
 
-    def extract_custom_headers(span, headers):
-        if not agent.options.extra_http_headers or not headers:
-            return
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                if custom_header in headers:
-                    span.set_attribute("http.header.%s" % custom_header, headers[custom_header])
-
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
 
 
     @wrapt.patch_function_wrapper('tornado.web', 'RequestHandler._execute')

src/instana/instrumentation/urllib3.py

@@ -11,27 +11,14 @@
 from instana.propagators.format import Format
 from instana.singletons import agent
 from instana.util.secrets import strip_secrets_from_query
-from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+from instana.util.traceutils import get_tracer_tuple, tracing_is_off, extract_custom_headers
 
 if TYPE_CHECKING:
     from instana.span.span import InstanaSpan
 
 try:
     import urllib3
 
-    def _extract_custom_headers(span: "InstanaSpan", headers: Dict[str, Any]) -> None:
-        if agent.options.extra_http_headers is None:
-            return
-
-        try:
-            for custom_header in agent.options.extra_http_headers:
-                if custom_header in headers:
-                    span.set_attribute(
-                        f"http.header.{custom_header}", headers[custom_header]
-                    )
-        except Exception:
-            logger.debug("urllib3 _extract_custom_headers error: ", exc_info=True)
-
     def _collect_kvs(
         instance: Union[
             urllib3.connectionpool.HTTPConnectionPool,
@@ -82,7 +69,7 @@ def collect_response(
         try:
             span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, response.status)
 
-            _extract_custom_headers(span, response.headers)
+            extract_custom_headers(span, response.headers)
 
             if 500 <= response.status:
                 span.mark_as_errored()
@@ -121,7 +108,7 @@ def urlopen_with_instana(
                 if "method" in kvs:
                     span.set_attribute(SpanAttributes.HTTP_METHOD, kvs["method"])
                 if "headers" in kwargs:
-                    _extract_custom_headers(span, kwargs["headers"])
+                    extract_custom_headers(span, kwargs["headers"])
                     tracer.inject(span.context, Format.HTTP_HEADERS, kwargs["headers"])
 
                 response = wrapped(*args, **kwargs)