From d5d0d9df9f45015d9a0570f9ee116a7fd6fae298 Mon Sep 17 00:00:00 2001 From: Ashutosh Goel Date: Wed, 9 Mar 2022 16:02:19 +0530 Subject: [PATCH 1/2] Capture request/response headers for flask --- .../instrumentation/flask/__init__.py | 4 + .../tests/base_test.py | 13 ++- .../tests/test_programmatic.py | 96 +++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) 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..f518d3e94d 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,8 @@ 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 +202,8 @@ 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..26b40dd295 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py @@ -14,7 +14,7 @@ from werkzeug.test import Client from werkzeug.wrappers import BaseResponse - +from flask import Response class InstrumentationTest: @staticmethod @@ -23,18 +23,27 @@ 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" def excluded2_endpoint(): return "excluded2" - + # pylint: disable=no-member self.app.route("/hello/")(self._hello_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..4556ecdc2d 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py @@ -442,3 +442,99 @@ 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) \ No newline at end of file From 8c9b6c996a23ad263548e0a1755c0506017d7fed Mon Sep 17 00:00:00 2001 From: Ashutosh Goel Date: Wed, 9 Mar 2022 16:09:50 +0530 Subject: [PATCH 2/2] Update changelog and fixed lint errors --- CHANGELOG.md | 3 +++ .../instrumentation/flask/__init__.py | 8 ++++++-- .../tests/base_test.py | 15 ++++++++++----- .../tests/test_programmatic.py | 14 ++++++++------ 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f8967010a7..41a489ec22 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-dbapi` add 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 f518d3e94d..1db768a2c0 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py @@ -154,7 +154,9 @@ def _start_response(status, response_headers, *args, **kwargs): span, status, response_headers ) if span.kind == trace.SpanKind.SERVER: - otel_wsgi.add_custom_response_headers(span, response_headers) + otel_wsgi.add_custom_response_headers( + span, response_headers + ) else: _logger.warning( "Flask environ's OpenTelemetry span " @@ -203,7 +205,9 @@ def _before_request(): 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) + 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 26b40dd295..d5424f9079 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py @@ -12,9 +12,10 @@ # 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 -from flask import Response + class InstrumentationTest: @staticmethod @@ -28,22 +29,26 @@ 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" + resp.headers[ + "my-custom-header" + ] = "my-custom-value-1,my-custom-header-2" return resp - + def _common_initialization(self): def excluded_endpoint(): return "excluded" def excluded2_endpoint(): return "excluded2" - + # pylint: disable=no-member self.app.route("/hello/")(self._hello_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) + 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 4556ecdc2d..6329bf1d30 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py @@ -444,7 +444,9 @@ def test_mark_span_internal_in_presence_of_span_from_other_framework(self): ) -class TestCustomRequestResponseHeaders(InstrumentationTest, TestBase, WsgiTestBase): +class TestCustomRequestResponseHeaders( + InstrumentationTest, TestBase, WsgiTestBase +): def setUp(self): super().setUp() @@ -452,7 +454,7 @@ def setUp(self): "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" + "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE": "content-type,content-length,my-custom-header,invalid-header", }, ) self.env_patch.start() @@ -466,7 +468,7 @@ def tearDown(self): 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", @@ -503,7 +505,7 @@ def test_custom_request_header_not_added_in_internal_span(self): 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) @@ -519,7 +521,7 @@ def test_custom_response_header_added_in_server_span(self): } 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): @@ -537,4 +539,4 @@ def test_custom_response_header_not_added_in_internal_span(self): } self.assertEqual(span.kind, trace.SpanKind.INTERNAL) for key, _ in not_expected.items(): - self.assertNotIn(key, span.attributes) \ No newline at end of file + self.assertNotIn(key, span.attributes)