enhancement: use a common method to capture custom headers

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

Author: GSVarsha

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)
@@ -93,8 +77,8 @@ 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"])
+            if "headers" in scope:
+                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/boto3_inst.py

@@ -10,7 +10,7 @@
 
 from instana.log import logger
 from instana.singletons import tracer, agent
-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
 from instana.propagators.format import Format
 from instana.span.span import get_current_span
 
@@ -23,21 +23,6 @@
     import boto3
     from boto3.s3 import inject
 
-    def extract_custom_headers(
-        span: "InstanaSpan", headers: Optional[Dict[str, Any]] = None
-    ) -> 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(
-                        "http.header.%s" % custom_header, headers[custom_header]
-                    )
-
-        except Exception:
-            logger.debug("extract_custom_headers: ", exc_info=True)
-
     def lambda_inject_context(payload: Dict[str, Any], span: "InstanaSpan") -> None:
         """
         When boto3 lambda client 'Invoke' is called, we want to inject the tracing context.

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/sanic_inst.py

@@ -75,8 +75,7 @@ def request_with_instana(request: Request) -> None:
                     )
                     span.set_attribute("http.params", scrubbed_params)
 
-                if agent.options.extra_http_headers:
-                    extract_custom_headers(span, headers)
+                extract_custom_headers(span, headers)
                 if hasattr(request, "uri_template") and request.uri_template:
                     span.set_attribute("http.path_tpl", request.uri_template)
             except Exception:
@@ -113,8 +112,7 @@ def response_with_instana(request: Request, response: HTTPResponse) -> None:
                     span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, status_code)
 
                 if hasattr(response, "headers"):
-                    if agent.options.extra_http_headers:
-                        extract_custom_headers(span, response.headers)
+                    extract_custom_headers(span, response.headers)
                     tracer.inject(span.context, Format.HTTP_HEADERS, response.headers)
 
                 if span.is_recording():

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):