Refactoring custom header collection API for consistency (#1064)

Author: ashu658

CHANGELOG.md

@@ -8,9 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.11.1-0.30b1...HEAD)
 
 ### Fixed
-- `opentelemetry-instrumentation-aws-lambda` Fixed an issue - in some rare cases (API GW proxy integration test) 
+- `opentelemetry-instrumentation-aws-lambda` Fixed an issue - in some rare cases (API GW proxy integration test)
   headers are set to None, breaking context propagators.
   ([#1055](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1055))
+- Refactoring custom header collection API for consistency
+  ([#1064](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1064))
 
 ## [1.11.1-0.30b1](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.11.1-0.30b1) - 2022-04-21
 

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

@@ -28,13 +28,13 @@
     _start_internal_or_server_span,
     extract_attributes_from_object,
 )
+from opentelemetry.instrumentation.wsgi import add_response_attributes
 from opentelemetry.instrumentation.wsgi import (
-    add_custom_request_headers as wsgi_add_custom_request_headers,
+    collect_custom_request_headers_attributes as wsgi_collect_custom_request_headers_attributes,
 )
 from opentelemetry.instrumentation.wsgi import (
-    add_custom_response_headers as wsgi_add_custom_response_headers,
+    collect_custom_response_headers_attributes as wsgi_collect_custom_response_headers_attributes,
 )
-from opentelemetry.instrumentation.wsgi import add_response_attributes
 from opentelemetry.instrumentation.wsgi import (
     collect_request_attributes as wsgi_collect_request_attributes,
 )
@@ -231,7 +231,11 @@ def process_request(self, request):
                     )
             else:
                 if span.is_recording() and span.kind == SpanKind.SERVER:
-                    wsgi_add_custom_request_headers(span, carrier)
+                    custom_attributes = (
+                        wsgi_collect_custom_request_headers_attributes(carrier)
+                    )
+                    if len(custom_attributes) > 0:
+                        span.set_attributes(custom_attributes)
 
             for key, value in attributes.items():
                 span.set_attribute(key, value)
@@ -309,7 +313,13 @@ def process_response(self, request, response):
                     response.items(),
                 )
                 if span.is_recording() and span.kind == SpanKind.SERVER:
-                    wsgi_add_custom_response_headers(span, response.items())
+                    custom_attributes = (
+                        wsgi_collect_custom_response_headers_attributes(
+                            response.items()
+                        )
+                    )
+                    if len(custom_attributes) > 0:
+                        span.set_attributes(custom_attributes)
 
             propagator = get_global_response_propagator()
             if propagator:

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

@@ -268,7 +268,11 @@ def __call__(self, env, start_response):
             for key, value in attributes.items():
                 span.set_attribute(key, value)
             if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-                otel_wsgi.add_custom_request_headers(span, env)
+                custom_attributes = (
+                    otel_wsgi.collect_custom_request_headers_attributes(env)
+                )
+                if len(custom_attributes) > 0:
+                    span.set_attributes(custom_attributes)
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()
@@ -382,9 +386,13 @@ def process_response(
                 response_headers = resp.headers
 
             if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-                otel_wsgi.add_custom_response_headers(
-                    span, response_headers.items()
+                custom_attributes = (
+                    otel_wsgi.collect_custom_response_headers_attributes(
+                        response_headers.items()
+                    )
                 )
+                if len(custom_attributes) > 0:
+                    span.set_attributes(custom_attributes)
         except ValueError:
             pass
 

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

@@ -208,9 +208,11 @@ def _start_response(status, response_headers, *args, **kwargs):
                         span.is_recording()
                         and span.kind == trace.SpanKind.SERVER
                     ):
-                        otel_wsgi.add_custom_response_headers(
-                            span, response_headers
+                        custom_attributes = otel_wsgi.collect_custom_response_headers_attributes(
+                            response_headers
                         )
+                        if len(custom_attributes) > 0:
+                            span.set_attributes(custom_attributes)
                 else:
                     _logger.warning(
                         "Flask environ's OpenTelemetry span "
@@ -259,9 +261,13 @@ def _before_request():
             for key, value in attributes.items():
                 span.set_attribute(key, value)
             if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-                otel_wsgi.add_custom_request_headers(
-                    span, flask_request_environ
+                custom_attributes = (
+                    otel_wsgi.collect_custom_request_headers_attributes(
+                        flask_request_environ
+                    )
                 )
+                if len(custom_attributes) > 0:
+                    span.set_attributes(custom_attributes)
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=E1101

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

@@ -105,7 +105,13 @@ def _before_traversal(event):
         for key, value in attributes.items():
             span.set_attribute(key, value)
         if span.kind == trace.SpanKind.SERVER:
-            otel_wsgi.add_custom_request_headers(span, request_environ)
+            custom_attributes = (
+                otel_wsgi.collect_custom_request_headers_attributes(
+                    request_environ
+                )
+            )
+            if len(custom_attributes) > 0:
+                span.set_attributes(custom_attributes)
 
     activation = trace.use_span(span, end_on_exit=True)
     activation.__enter__()  # pylint: disable=E1101
@@ -178,9 +184,13 @@ def trace_tween(request):
                     )
 
                 if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-                    otel_wsgi.add_custom_response_headers(
-                        span, getattr(response, "headerlist", None)
+                    custom_attributes = (
+                        otel_wsgi.collect_custom_response_headers_attributes(
+                            getattr(response, "headerlist", None)
+                        )
                     )
+                    if len(custom_attributes) > 0:
+                        span.set_attributes(custom_attributes)
 
                 propagator = get_global_response_propagator()
                 if propagator and hasattr(response, "headers"):

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

@@ -315,7 +315,7 @@ def _log_exception(tracer, func, handler, args, kwargs):
     return func(*args, **kwargs)
 
 
-def _add_custom_request_headers(span, request_headers):
+def _collect_custom_request_headers_attributes(request_headers):
     custom_request_headers_name = get_custom_headers(
         OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
     )
@@ -325,10 +325,10 @@ def _add_custom_request_headers(span, request_headers):
         if header_values:
             key = normalise_request_header_name(header_name.lower())
             attributes[key] = [header_values]
-    span.set_attributes(attributes)
+    return attributes
 
 
-def _add_custom_response_headers(span, response_headers):
+def _collect_custom_response_headers_attributes(response_headers):
     custom_response_headers_name = get_custom_headers(
         OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
     )
@@ -338,7 +338,7 @@ def _add_custom_response_headers(span, response_headers):
         if header_values:
             key = normalise_response_header_name(header_name.lower())
             attributes[key] = [header_values]
-    span.set_attributes(attributes)
+    return attributes
 
 
 def _get_attributes_from_request(request):
@@ -392,7 +392,11 @@ def _start_span(tracer, handler, start_time) -> _TraceContext:
             span.set_attribute(key, value)
         span.set_attribute("tornado.handler", _get_full_handler_name(handler))
         if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-            _add_custom_request_headers(span, handler.request.headers)
+            custom_attributes = _collect_custom_request_headers_attributes(
+                handler.request.headers
+            )
+            if len(custom_attributes) > 0:
+                span.set_attributes(custom_attributes)
 
     activation = trace.use_span(span, end_on_exit=True)
     activation.__enter__()  # pylint: disable=E1101
@@ -448,7 +452,11 @@ def _finish_span(tracer, handler, error=None):
             )
         )
         if ctx.span.is_recording() and ctx.span.kind == trace.SpanKind.SERVER:
-            _add_custom_response_headers(ctx.span, handler._headers)
+            custom_attributes = _collect_custom_response_headers_attributes(
+                handler._headers
+            )
+            if len(custom_attributes) > 0:
+                ctx.span.set_attributes(custom_attributes)
 
     ctx.activation.__exit__(*finish_args)  # pylint: disable=E1101
     if ctx.token:

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

@@ -266,8 +266,8 @@ def collect_request_attributes(environ):
     return result
 
 
-def add_custom_request_headers(span, environ):
-    """Adds custom HTTP request headers into the span which are configured by the user
+def collect_custom_request_headers_attributes(environ):
+    """Returns custom HTTP request headers which are 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 = {}
@@ -280,11 +280,11 @@ def add_custom_request_headers(span, environ):
         if header_values:
             key = normalise_request_header_name(header_name)
             attributes[key] = [header_values]
-    span.set_attributes(attributes)
+    return attributes
 
 
-def add_custom_response_headers(span, response_headers):
-    """Adds custom HTTP response headers into the sapn which are configured by the user from the
+def collect_custom_response_headers_attributes(response_headers):
+    """Returns custom HTTP response headers which are 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 = {}
@@ -301,7 +301,7 @@ def add_custom_response_headers(span, response_headers):
         if header_values:
             key = normalise_response_header_name(header_name)
             attributes[key] = [header_values]
-    span.set_attributes(attributes)
+    return attributes
 
 
 def add_response_attributes(
@@ -365,7 +365,11 @@ def _create_start_response(span, start_response, response_hook):
         def _start_response(status, response_headers, *args, **kwargs):
             add_response_attributes(span, status, response_headers)
             if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-                add_custom_response_headers(span, response_headers)
+                custom_attributes = collect_custom_response_headers_attributes(
+                    response_headers
+                )
+                if len(custom_attributes) > 0:
+                    span.set_attributes(custom_attributes)
             if response_hook:
                 response_hook(status, response_headers)
             return start_response(status, response_headers, *args, **kwargs)
@@ -388,7 +392,11 @@ def __call__(self, environ, start_response):
             attributes=collect_request_attributes(environ),
         )
         if span.is_recording() and span.kind == trace.SpanKind.SERVER:
-            add_custom_request_headers(span, environ)
+            custom_attributes = collect_custom_request_headers_attributes(
+                environ
+            )
+            if len(custom_attributes) > 0:
+                span.set_attributes(custom_attributes)
 
         if self.request_hook:
             self.request_hook(span, environ)