Skip to content

Commit

Permalink
Flask: Capture custom request/response headers as span attributes (op…
Browse files Browse the repository at this point in the history
…en-telemetry#952)

* Capture request/response headers for flask

* Update changelog and fixed lint errors
  • Loading branch information
ashu658 authored Mar 9, 2022
1 parent dbb35a2 commit c60a7e4
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 0 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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"
Expand All @@ -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)
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit c60a7e4

Please sign in to comment.