Flask: Capture custom request/response headers as span attributes (#952)

* Capture request/response headers for flask

* Update changelog and fixed lint errors

Author: ashu658

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

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

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/<int:helloid>")(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)

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)