diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e38858d3a..3c5e82bab8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes ([#925])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/925) +- `opentelemetry-instrumentation-flask` Flask: Capture custom request/response headers in span attributes + ([#952])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/952) + ### Added - `opentelemetry-instrumentation-sqlalchemy` added experimental sql commenter capability diff --git a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py index 6a46460be7..1db768a2c0 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py @@ -153,6 +153,10 @@ def _start_response(status, response_headers, *args, **kwargs): otel_wsgi.add_response_attributes( span, status, response_headers ) + if span.kind == trace.SpanKind.SERVER: + otel_wsgi.add_custom_response_headers( + span, response_headers + ) else: _logger.warning( "Flask environ's OpenTelemetry span " @@ -200,6 +204,10 @@ def _before_request(): ] = flask.request.url_rule.rule for key, value in attributes.items(): span.set_attribute(key, value) + if span.kind == trace.SpanKind.SERVER: + otel_wsgi.add_custom_request_headers( + span, flask_request_environ + ) activation = trace.use_span(span, end_on_exit=True) activation.__enter__() # pylint: disable=E1101 diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py b/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py index d989c66474..d5424f9079 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from flask import Response from werkzeug.test import Client from werkzeug.wrappers import BaseResponse @@ -23,6 +24,16 @@ def _hello_endpoint(helloid): raise ValueError(":-(") return "Hello: " + str(helloid) + @staticmethod + def _custom_response_headers(): + resp = Response("test response") + resp.headers["content-type"] = "text/plain; charset=utf-8" + resp.headers["content-length"] = "13" + resp.headers[ + "my-custom-header" + ] = "my-custom-value-1,my-custom-header-2" + return resp + def _common_initialization(self): def excluded_endpoint(): return "excluded" @@ -35,6 +46,9 @@ def excluded2_endpoint(): self.app.route("/excluded/")(self._hello_endpoint) self.app.route("/excluded")(excluded_endpoint) self.app.route("/excluded2")(excluded2_endpoint) + self.app.route("/test_custom_response_headers")( + self._custom_response_headers + ) # pylint: disable=attribute-defined-outside-init self.client = Client(self.app, BaseResponse) diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py index a3064b52e5..6329bf1d30 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py @@ -442,3 +442,101 @@ def test_mark_span_internal_in_presence_of_span_from_other_framework(self): self.assertEqual( span_list[0].parent.span_id, span_list[1].context.span_id ) + + +class TestCustomRequestResponseHeaders( + InstrumentationTest, TestBase, WsgiTestBase +): + 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": "content-type,content-length,my-custom-header,invalid-header", + }, + ) + self.env_patch.start() + self.app = Flask(__name__) + FlaskInstrumentor().instrument_app(self.app) + + self._common_initialization() + + def tearDown(self): + super().tearDown() + self.env_patch.stop() + with self.disable_logging(): + FlaskInstrumentor().uninstrument_app(self.app) + + def test_custom_request_header_added_in_server_span(self): + 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] + 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.SERVER) + self.assertSpanHasAttributes(span, expected) + + 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", + } + 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, trace.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(resp.status_code, 200) + span = self.memory_exporter.get_finished_spans()[0] + expected = { + "http.response.header.content_type": ( + "text/plain; charset=utf-8", + ), + "http.response.header.content_length": ("13",), + "http.response.header.my_custom_header": ( + "my-custom-value-1,my-custom-header-2", + ), + } + self.assertEqual(span.kind, trace.SpanKind.SERVER) + self.assertSpanHasAttributes(span, expected) + + 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): + resp = self.client.get("/test_custom_response_headers") + self.assertEqual(resp.status_code, 200) + 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": ("13",), + "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)