Starlette: capture custom request response headers in span attributes

CHANGELOG.md

@@ -22,6 +22,9 @@ 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)
 
 ### Added
+
+- `opentelemetry-instrumentation-starlette` Capture custom request/response headers in span attributes
+  ([#1046])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1046)
 - `opentelemetry-instrumentation-fastapi` Capture custom request/response headers in span attributes
   ([#1032])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1032)
 - `opentelemetry-instrumentation-django` Capture custom request/response headers in span attributes

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

@@ -74,7 +74,6 @@ def client_response_hook(span: Span, message: dict):
 
    FastAPIInstrumentor().instrument(server_request_hook=server_request_hook, client_request_hook=client_request_hook, client_response_hook=client_response_hook)
 
-
 Capture HTTP request and response headers
 *****************************************
 You can configure the agent to capture predefined HTTP headers as span attributes, according to the `semantic convention <https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers>`_.
@@ -93,7 +92,7 @@ def client_response_hook(span: Span, message: dict):
 will extract ``content-type`` and ``custom_request_header`` from request headers and add them as span attributes.
 
 It is recommended that you should give the correct names of the headers to be captured in the environment variable.
-Request header names in FastAPI are case insensitive. So, giving header name as ``CUStom-Header`` in environment variable will be able capture header with name ``custom-header``.
+Request header names in fastapi are case insensitive. So, giving header name as ``CUStom-Header`` in environment variable will be able capture header with name ``custom-header``.
 
 The name of the added span attribute will follow the format ``http.request.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
 The value of the attribute will be single item list containing all the header values.
@@ -115,7 +114,7 @@ def client_response_hook(span: Span, message: dict):
 will extract ``content-type`` and ``custom_response_header`` from response headers and add them as span attributes.
 
 It is recommended that you should give the correct names of the headers to be captured in the environment variable.
-Response header names captured in FastAPI are case insensitive. So, giving header name as ``CUStomHeader`` in environment variable will be able capture header with name ``customheader``.
+Response header names captured in fastapi are case insensitive. So, giving header name as ``CUStomHeader`` in environment variable will be able capture header with name ``customheader``.
 
 The name of the added span attribute will follow the format ``http.response.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
 The value of the attribute will be single item list containing all the header values.

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

@@ -68,6 +68,57 @@ def client_response_hook(span: Span, message: dict):
 
    StarletteInstrumentor().instrument(server_request_hook=server_request_hook, client_request_hook=client_request_hook, client_response_hook=client_response_hook)
 
+Capture HTTP request and response headers
+*****************************************
+You can configure the agent to capture predefined HTTP headers as span attributes, according to the `semantic convention <https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers>`_.
+
+Request headers
+***************
+To capture predefined HTTP request headers as span attributes, set the environment variable ``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST``
+to a comma-separated list of HTTP header names.
+
+For example,
+
+::
+
+    export OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST="content-type,custom_request_header"
+
+will extract ``content-type`` and ``custom_request_header`` from request headers and add them as span attributes.
+
+It is recommended that you should give the correct names of the headers to be captured in the environment variable.
+Request header names in starlette are case insensitive. So, giving header name as ``CUStom-Header`` in environment variable will be able capture header with name ``custom-header``.
+
+The name of the added span attribute will follow the format ``http.request.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
+The value of the attribute will be single item list containing all the header values.
+
+Example of the added span attribute,
+``http.request.header.custom_request_header = ["<value1>,<value2>"]``
+
+Response headers
+****************
+To capture predefined HTTP response headers as span attributes, set the environment variable ``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE``
+to a comma-separated list of HTTP header names.
+
+For example,
+
+::
+
+    export OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE="content-type,custom_response_header"
+
+will extract ``content-type`` and ``custom_response_header`` from response headers and add them as span attributes.
+
+It is recommended that you should give the correct names of the headers to be captured in the environment variable.
+Response header names captured in starlette are case insensitive. So, giving header name as ``CUStomHeader`` in environment variable will be able capture header with name ``customheader``.
+
+The name of the added span attribute will follow the format ``http.response.header.<header_name>`` where ``<header_name>`` being the normalized HTTP header name (lowercase, with - characters replaced by _ ).
+The value of the attribute will be single item list containing all the header values.
+
+Example of the added span attribute,
+``http.response.header.custom_response_header = ["<value1>,<value2>"]``
+
+Note:
+    Environment variable names to caputre http headers are still experimental, and thus are subject to change.
+
 API
 ---
 """

instrumentation/opentelemetry-instrumentation-starlette/tests/test_starlette_instrumentation.py

@@ -19,13 +19,24 @@
 from starlette.responses import PlainTextResponse
 from starlette.routing import Route
 from starlette.testclient import TestClient
+from starlette.websockets import WebSocket
 
 import opentelemetry.instrumentation.starlette as otel_starlette
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.test.globals_test import reset_trace_globals
 from opentelemetry.test.test_base import TestBase
-from opentelemetry.trace import SpanKind, get_tracer
-from opentelemetry.util.http import get_excluded_urls
+from opentelemetry.trace import (
+    NoOpTracerProvider,
+    SpanKind,
+    get_tracer,
+    set_tracer_provider,
+)
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_excluded_urls,
+)
 
 
 class TestStarletteManualInstrumentation(TestBase):
@@ -244,3 +255,272 @@ def test_mark_span_internal_in_presence_of_another_span(self):
             self.assertEqual(
                 parent_span.context.span_id, starlette_span.parent.span_id
             )
+
+
+class TestBaseWithCustomHeaders(TestBase):
+    def create_app(self):
+        app = self.create_starlette_app()
+        self._instrumentor.instrument_app(app=app)
+        return app
+
+    def setUp(self):
+        super().setUp()
+        self.env_patch = patch.dict(
+            "os.environ",
+            {
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+            },
+        )
+        self.env_patch.start()
+        self._instrumentor = otel_starlette.StarletteInstrumentor()
+        self._app = self.create_app()
+        self._client = TestClient(self._app)
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        self.env_patch.stop()
+        with self.disable_logging():
+            self._instrumentor.uninstrument()
+
+    @staticmethod
+    def create_starlette_app():
+        app = applications.Starlette()
+
+        @app.route("/foobar")
+        def _(request):
+            return PlainTextResponse(
+                content="hi",
+                headers={
+                    "custom-test-header-1": "test-header-value-1",
+                    "custom-test-header-2": "test-header-value-2",
+                },
+            )
+
+        @app.websocket_route("/foobar_web")
+        async def _(websocket: WebSocket) -> None:
+            message = await websocket.receive()
+            if message.get("type") == "websocket.connect":
+                await websocket.send(
+                    {
+                        "type": "websocket.accept",
+                        "headers": [
+                            (b"custom-test-header-1", b"test-header-value-1"),
+                            (b"custom-test-header-2", b"test-header-value-2"),
+                        ],
+                    }
+                )
+                await websocket.send_json({"message": "hello world"})
+                await websocket.close()
+            if message.get("type") == "websocket.disconnect":
+                pass
+
+        return app
+
+
+class TestHTTPAppWithCustomHeaders(TestBaseWithCustomHeaders):
+    def test_custom_request_headers_in_span_attributes(self):
+        expected = {
+            "http.request.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.request.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        resp = self._client.get(
+            "/foobar",
+            headers={
+                "custom-test-header-1": "test-header-value-1",
+                "custom-test-header-2": "test-header-value-2",
+            },
+        )
+        self.assertEqual(200, resp.status_code)
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 3)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        self.assertSpanHasAttributes(server_span, expected)
+
+    def test_custom_request_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.request.header.custom_test_header_3": (
+                "test-header-value-3",
+            ),
+        }
+        resp = self._client.get(
+            "/foobar",
+            headers={
+                "custom-test-header-1": "test-header-value-1",
+                "custom-test-header-2": "test-header-value-2",
+            },
+        )
+        self.assertEqual(200, resp.status_code)
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 3)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        for key in not_expected:
+            self.assertNotIn(key, server_span.attributes)
+
+    def test_custom_response_headers_in_span_attributes(self):
+        expected = {
+            "http.response.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.response.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        resp = self._client.get("/foobar")
+        self.assertEqual(200, resp.status_code)
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 3)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        self.assertSpanHasAttributes(server_span, expected)
+
+    def test_custom_response_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.response.header.custom_test_header_3": (
+                "test-header-value-3",
+            ),
+        }
+        resp = self._client.get("/foobar")
+        self.assertEqual(200, resp.status_code)
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 3)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        for key in not_expected:
+            self.assertNotIn(key, server_span.attributes)
+
+
+class TestWebSocketAppWithCustomHeaders(TestBaseWithCustomHeaders):
+    def test_custom_request_headers_in_span_attributes(self):
+        expected = {
+            "http.request.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.request.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        with self._client.websocket_connect(
+            "/foobar_web",
+            headers={
+                "custom-test-header-1": "test-header-value-1",
+                "custom-test-header-2": "test-header-value-2",
+            },
+        ) as websocket:
+            data = websocket.receive_json()
+            self.assertEqual(data, {"message": "hello world"})
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 5)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+        self.assertSpanHasAttributes(server_span, expected)
+
+    def test_custom_request_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.request.header.custom_test_header_3": (
+                "test-header-value-3",
+            ),
+        }
+        with self._client.websocket_connect(
+            "/foobar_web",
+            headers={
+                "custom-test-header-1": "test-header-value-1",
+                "custom-test-header-2": "test-header-value-2",
+            },
+        ) as websocket:
+            data = websocket.receive_json()
+            self.assertEqual(data, {"message": "hello world"})
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 5)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, server_span.attributes)
+
+    def test_custom_response_headers_in_span_attributes(self):
+        expected = {
+            "http.response.header.custom_test_header_1": (
+                "test-header-value-1",
+            ),
+            "http.response.header.custom_test_header_2": (
+                "test-header-value-2",
+            ),
+        }
+        with self._client.websocket_connect("/foobar_web") as websocket:
+            data = websocket.receive_json()
+            self.assertEqual(data, {"message": "hello world"})
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 5)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        self.assertSpanHasAttributes(server_span, expected)
+
+    def test_custom_response_headers_not_in_span_attributes(self):
+        not_expected = {
+            "http.response.header.custom_test_header_3": (
+                "test-header-value-3",
+            ),
+        }
+        with self._client.websocket_connect("/foobar_web") as websocket:
+            data = websocket.receive_json()
+            self.assertEqual(data, {"message": "hello world"})
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 5)
+
+        server_span = [
+            span for span in span_list if span.kind == SpanKind.SERVER
+        ][0]
+
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, server_span.attributes)
+
+
+class TestNonRecordingSpanWithCustomHeaders(TestBaseWithCustomHeaders):
+    def setUp(self):
+        super().setUp()
+        reset_trace_globals()
+        set_tracer_provider(tracer_provider=NoOpTracerProvider())
+
+        self._app = self.create_app()
+        self._client = TestClient(self._app)
+
+    def test_custom_header_not_present_in_non_recording_span(self):
+        resp = self._client.get(
+            "/foobar",
+            headers={
+                "custom-test-header-1": "test-header-value-1",
+            },
+        )
+        self.assertEqual(200, resp.status_code)
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 0)