Pyramid: Capture custom request/response headers

CHANGELOG.md

@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1003])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1003)
 - `opentelemetry-instrumentation-elasticsearch` no longer creates unique span names by including search target, replaces them with `<target>` and puts the value in attribute `elasticsearch.target`
   ([#1018](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1018))
+- `opentelemetry-instrumentation-pyramid` Pyramid: Capture custom request/response headers in span attributes
+  ([#1022])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1022)
 
 ## [1.10.0-0.29b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.10.0-0.29b0) - 2022-03-10
 

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

@@ -104,6 +104,8 @@ def _before_traversal(event):
             ] = request.matched_route.pattern
         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)
 
     activation = trace.use_span(span, end_on_exit=True)
     activation.__enter__()  # pylint: disable=E1101
@@ -163,6 +165,10 @@ def trace_tween(request):
                     response_or_exception.status,
                     response_or_exception.headerlist,
                 )
+                if span.is_recording() and span.kind == trace.SpanKind.SERVER:
+                    otel_wsgi.add_custom_response_headers(
+                        span, response_or_exception.headerlist
+                    )
 
                 propagator = get_global_response_propagator()
                 if propagator:

instrumentation/opentelemetry-instrumentation-pyramid/tests/pyramid_base_test.py

@@ -26,6 +26,16 @@ def _hello_endpoint(request):
             raise exc.HTTPInternalServerError()
         return Response("Hello: " + str(helloid))
 
+    @staticmethod
+    def _custom_response_header_endpoint(request):
+        headers = {
+            "content-type": "text/plain; charset=utf-8",
+            "content-length": "7",
+            "my-custom-header": "my-custom-value-1,my-custom-header-2",
+            "dont-capture-me": "test-value",
+        }
+        return Response("Testing", headers=headers)
+
     def _common_initialization(self, config):
         # pylint: disable=unused-argument
         def excluded_endpoint(request):
@@ -43,6 +53,13 @@ def excluded2_endpoint(request):
         config.add_view(excluded_endpoint, route_name="excluded")
         config.add_route("excluded2", "/excluded_noarg2")
         config.add_view(excluded2_endpoint, route_name="excluded2")
+        config.add_route(
+            "custom_response_headers", "/test_custom_response_headers"
+        )
+        config.add_view(
+            self._custom_response_header_endpoint,
+            route_name="custom_response_headers",
+        )
 
         # pylint: disable=attribute-defined-outside-init
         self.client = Client(config.make_wsgi_app(), BaseResponse)

instrumentation/opentelemetry-instrumentation-pyramid/tests/test_automatic.py

@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import patch
+
 from pyramid.config import Configurator
 
+from opentelemetry import trace
 from opentelemetry.instrumentation.pyramid import PyramidInstrumentor
+from opentelemetry.test.globals_test import reset_trace_globals
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import SpanKind
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+)
 
 # pylint: disable=import-error
 from .pyramid_base_test import InstrumentationTest
@@ -109,3 +117,146 @@ def test_with_existing_span(self):
                 parent_span.get_span_context().span_id,
                 span_list[0].parent.span_id,
             )
+
+
+class TestCustomRequestResponseHeaders(
+    InstrumentationTest, TestBase, WsgiTestBase
+):
+    def setUp(self):
+        super().setUp()
+        PyramidInstrumentor().instrument()
+        self.config = Configurator()
+        self._common_initialization(self.config)
+        self.env_patch = patch.dict(
+            "os.environ",
+            {
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,invalid-header",
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header",
+            },
+        )
+        self.env_patch.start()
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        self.env_patch.stop()
+        with self.disable_logging():
+            PyramidInstrumentor().uninstrument()
+
+    def test_custom_request_header_added_in_server_span(self):
+        headers = {
+            "Custom-Test-Header-1": "Test Value 1",
+            "Custom-Test-Header-2": "TestValue2,TestValue3",
+            "Custom-Test-Header-3": "TestValue4",
+        }
+        resp = self.client.get("/hello/123", headers=headers)
+        self.assertEqual(200, resp.status_code)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.request.header.custom_test_header_1": ("Test Value 1",),
+            "http.request.header.custom_test_header_2": (
+                "TestValue2,TestValue3",
+            ),
+        }
+        not_expected = {
+            "http.request.header.custom_test_header_3": ("TestValue4",),
+        }
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, span.attributes)
+
+    def test_custom_request_header_not_added_in_internal_span(self):
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span("test", kind=SpanKind.SERVER):
+            headers = {
+                "Custom-Test-Header-1": "Test Value 1",
+                "Custom-Test-Header-2": "TestValue2,TestValue3",
+            }
+            resp = self.client.get("/hello/123", headers=headers)
+            self.assertEqual(200, resp.status_code)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.request.header.custom_test_header_1": ("Test Value 1",),
+                "http.request.header.custom_test_header_2": (
+                    "TestValue2,TestValue3",
+                ),
+            }
+            self.assertEqual(span.kind, SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+    def test_custom_response_header_added_in_server_span(self):
+        resp = self.client.get("/test_custom_response_headers")
+        self.assertEqual(200, resp.status_code)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("7",),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        not_expected = {
+            "http.response.header.dont_capture_me": ("test-value",)
+        }
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, span.attributes)
+
+    def test_custom_response_header_not_added_in_internal_span(self):
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span("test", kind=SpanKind.SERVER):
+            resp = self.client.get("/test_custom_response_headers")
+            self.assertEqual(200, resp.status_code)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.response.header.content_type": (
+                    "text/plain; charset=utf-8",
+                ),
+                "http.response.header.content_length": ("7",),
+                "http.response.header.my_custom_header": (
+                    "my-custom-value-1,my-custom-header-2",
+                ),
+            }
+            self.assertEqual(span.kind, SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+
+class TestCustomHeadersNonRecordingSpan(
+    InstrumentationTest, TestBase, WsgiTestBase
+):
+    def setUp(self):
+        super().setUp()
+        # This is done because set_tracer_provider cannot override the
+        # current tracer provider.
+        reset_trace_globals()
+        tracer_provider = trace.NoOpTracerProvider()
+        trace.set_tracer_provider(tracer_provider)
+        PyramidInstrumentor().instrument()
+        self.config = Configurator()
+        self._common_initialization(self.config)
+        self.env_patch = patch.dict(
+            "os.environ",
+            {
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,invalid-header",
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header",
+            },
+        )
+        self.env_patch.start()
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        self.env_patch.stop()
+        with self.disable_logging():
+            PyramidInstrumentor().uninstrument()
+
+    def test_custom_header_non_recording_span(self):
+        try:
+            resp = self.client.get("/hello/123")
+            self.assertEqual(200, resp.status_code)
+        except Exception as exc:  # pylint: disable=W0703
+            self.fail(f"Exception raised with NonRecordingSpan {exc}")