Falcon: Capture request/response headers as span attributes (#1003)

Author: ashu658

CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#999])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/999)
 - `opentelemetry-instrumentation-tornado` Fix non-recording span bug
   ([#999])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/999)
+- `opentelemetry-instrumentation-falcon` Falcon: Capture custom request/response headers in span attributes
+  ([#1003])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1003)
 
 ## [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-falcon/src/opentelemetry/instrumentation/falcon/__init__.py

@@ -190,6 +190,8 @@ def __call__(self, env, start_response):
             attributes = otel_wsgi.collect_request_attributes(env)
             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)
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()
@@ -295,6 +297,10 @@ def process_response(
                     description=reason,
                 )
             )
+            if span.is_recording() and span.kind == trace.SpanKind.SERVER:
+                otel_wsgi.add_custom_response_headers(
+                    span, resp.headers.items()
+                )
         except ValueError:
             pass
 

instrumentation/opentelemetry-instrumentation-falcon/tests/app.py

@@ -33,6 +33,18 @@ def on_get(self, req, resp):
         print(non_existent_var)  # noqa
 
 
+class CustomResponseHeaderResource:
+    def on_get(self, _, resp):
+        # pylint: disable=no-member
+        resp.status = falcon.HTTP_201
+        resp.set_header("content-type", "text/plain; charset=utf-8")
+        resp.set_header("content-length", "0")
+        resp.set_header(
+            "my-custom-header", "my-custom-value-1,my-custom-header-2"
+        )
+        resp.set_header("dont-capture-me", "test-value")
+
+
 def make_app():
     if hasattr(falcon, "App"):
         # Falcon 3
@@ -43,4 +55,7 @@ def make_app():
     app.add_route("/hello", HelloWorldResource())
     app.add_route("/ping", HelloWorldResource())
     app.add_route("/error", ErrorResource())
+    app.add_route(
+        "/test_custom_response_headers", CustomResponseHeaderResource()
+    )
     return app

instrumentation/opentelemetry-instrumentation-falcon/tests/test_falcon.py

@@ -28,6 +28,10 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import StatusCode
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+)
 
 from .app import make_app
 
@@ -280,3 +284,105 @@ def test_mark_span_internal_in_presence_of_span_from_other_framework(self):
             self.assertEqual(
                 span.parent.span_id, parent_span.get_span_context().span_id
             )
+
+
+@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",
+    },
+)
+class TestCustomRequestResponseHeaders(TestFalconBase):
+    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",
+        }
+        self.client().simulate_request(
+            method="GET", path="/hello", headers=headers
+        )
+        span = self.memory_exporter.get_finished_spans()[0]
+        assert span.status.is_ok
+
+        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, trace.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=trace.SpanKind.SERVER):
+            headers = {
+                "Custom-Test-Header-1": "Test Value 1",
+                "Custom-Test-Header-2": "TestValue2,TestValue3",
+            }
+            self.client().simulate_request(
+                method="GET", path="/hello", headers=headers
+            )
+            span = self.memory_exporter.get_finished_spans()[0]
+            assert span.status.is_ok
+            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, trace.SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+    def test_custom_response_header_added_in_server_span(self):
+        self.client().simulate_request(
+            method="GET", path="/test_custom_response_headers"
+        )
+        span = self.memory_exporter.get_finished_spans()[0]
+        assert span.status.is_ok
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("0",),
+            "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, trace.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=trace.SpanKind.SERVER):
+            self.client().simulate_request(
+                method="GET", path="/test_custom_response_headers"
+            )
+            span = self.memory_exporter.get_finished_spans()[0]
+            assert span.status.is_ok
+            not_expected = {
+                "http.response.header.content_type": (
+                    "text/plain; charset=utf-8",
+                ),
+                "http.response.header.content_length": ("0",),
+                "http.response.header.my_custom_header": (
+                    "my-custom-value-1,my-custom-header-2",
+                ),
+            }
+            self.assertEqual(span.kind, trace.SpanKind.INTERNAL)
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)