diff --git a/CHANGELOG.md b/CHANGELOG.md index e68f5bd0c0..780e1888d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `opentelemetry-instrumentation-urllib3` Add urllib3 instrumentation ([#299](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/299)) + +- `opentelemetry-instrumentation-flask` Added `request_hook` and `response_hook` callbacks. + ([#416](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/416)) + - `opentelemetry-instrumenation-django` now supports request and response hooks. ([#407](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/407)) - `opentelemetry-instrumentation-falcon` FalconInstrumentor now supports request/response hooks. 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 9ae40b9fa5..f37859ac30 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py @@ -85,7 +85,7 @@ def get_default_span_name(): return span_name -def _rewrapped_app(wsgi_app): +def _rewrapped_app(wsgi_app, response_hook=None): def _wrapped_app(wrapped_app_environ, start_response): # We want to measure the time for route matching, etc. # In theory, we could start the span here and use @@ -114,7 +114,8 @@ def _start_response(status, response_headers, *args, **kwargs): "missing at _start_response(%s)", status, ) - + if response_hook is not None: + response_hook(span, status, response_headers) return start_response(status, response_headers, *args, **kwargs) return wsgi_app(wrapped_app_environ, _start_response) @@ -122,13 +123,12 @@ def _start_response(status, response_headers, *args, **kwargs): return _wrapped_app -def _wrapped_before_request(name_callback, tracer): +def _wrapped_before_request(request_hook=None, tracer=None): def _before_request(): if _excluded_urls.url_disabled(flask.request.url): return - flask_request_environ = flask.request.environ - span_name = name_callback() + span_name = get_default_span_name() token = context.attach( extract(flask_request_environ, getter=otel_wsgi.wsgi_getter) ) @@ -138,6 +138,9 @@ def _before_request(): kind=trace.SpanKind.SERVER, start_time=flask_request_environ.get(_ENVIRON_STARTTIME_KEY), ) + if request_hook: + request_hook(span, flask_request_environ) + if span.is_recording(): attributes = otel_wsgi.collect_request_attributes( flask_request_environ @@ -183,21 +186,25 @@ def _teardown_request(exc): class _InstrumentedFlask(flask.Flask): - name_callback = get_default_span_name _tracer_provider = None + _request_hook = None + _response_hook = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._original_wsgi_ = self.wsgi_app - self.wsgi_app = _rewrapped_app(self.wsgi_app) + + self.wsgi_app = _rewrapped_app( + self.wsgi_app, _InstrumentedFlask._response_hook + ) tracer = trace.get_tracer( __name__, __version__, _InstrumentedFlask._tracer_provider ) _before_request = _wrapped_before_request( - _InstrumentedFlask.name_callback, tracer, + _InstrumentedFlask._request_hook, tracer, ) self._before_request = _before_request self.before_request(_before_request) @@ -216,26 +223,30 @@ def instrumentation_dependencies(self) -> Collection[str]: def _instrument(self, **kwargs): self._original_flask = flask.Flask - name_callback = kwargs.get("name_callback") + request_hook = kwargs.get("request_hook") + response_hook = kwargs.get("response_hook") + if callable(request_hook): + _InstrumentedFlask._request_hook = request_hook + if callable(response_hook): + _InstrumentedFlask._response_hook = response_hook + flask.Flask = _InstrumentedFlask tracer_provider = kwargs.get("tracer_provider") - if callable(name_callback): - _InstrumentedFlask.name_callback = name_callback _InstrumentedFlask._tracer_provider = tracer_provider flask.Flask = _InstrumentedFlask def instrument_app( - self, app, name_callback=get_default_span_name, tracer_provider=None + self, app, request_hook=None, response_hook=None, tracer_provider=None ): # pylint: disable=no-self-use if not hasattr(app, "_is_instrumented"): app._is_instrumented = False if not app._is_instrumented: app._original_wsgi_app = app.wsgi_app - app.wsgi_app = _rewrapped_app(app.wsgi_app) + app.wsgi_app = _rewrapped_app(app.wsgi_app, response_hook) tracer = trace.get_tracer(__name__, __version__, tracer_provider) - _before_request = _wrapped_before_request(name_callback, tracer) + _before_request = _wrapped_before_request(request_hook, tracer) app._before_request = _before_request app.before_request(_before_request) app.teardown_request(_teardown_request) diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py index 56cbc7de85..8eb916b3d0 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py @@ -220,19 +220,28 @@ def test_exclude_lists(self): self.assertEqual(len(span_list), 1) -class TestProgrammaticCustomSpanName( - InstrumentationTest, TestBase, WsgiTestBase -): +class TestProgrammaticHooks(InstrumentationTest, TestBase, WsgiTestBase): def setUp(self): super().setUp() - def custom_span_name(): - return "flask-custom-span-name" + hook_headers = ( + "hook_attr", + "hello otel", + ) + + def request_hook_test(span, environ): + span.update_name("name from hook") + + def response_hook_test(span, environ, response_headers): + span.set_attribute("hook_attr", "hello world") + response_headers.append(hook_headers) self.app = Flask(__name__) FlaskInstrumentor().instrument_app( - self.app, name_callback=custom_span_name + self.app, + request_hook=request_hook_test, + response_hook=response_hook_test, ) self._common_initialization() @@ -242,24 +251,44 @@ def tearDown(self): with self.disable_logging(): FlaskInstrumentor().uninstrument_app(self.app) - def test_custom_span_name(self): - self.client.get("/hello/123") + def test_hooks(self): + expected_attrs = expected_attributes( + { + "http.target": "/hello/123", + "http.route": "/hello/", + "hook_attr": "hello world", + } + ) + resp = self.client.get("/hello/123") span_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(span_list), 1) - self.assertEqual(span_list[0].name, "flask-custom-span-name") + self.assertEqual(span_list[0].name, "name from hook") + self.assertEqual(span_list[0].attributes, expected_attrs) + self.assertEqual(resp.headers["hook_attr"], "hello otel") -class TestProgrammaticCustomSpanNameCallbackWithoutApp( +class TestProgrammaticHooksWithoutApp( InstrumentationTest, TestBase, WsgiTestBase ): def setUp(self): super().setUp() - def custom_span_name(): - return "instrument-without-app" + hook_headers = ( + "hook_attr", + "hello otel without app", + ) + + def request_hook_test(span, environ): + span.update_name("without app") + + def response_hook_test(span, environ, response_headers): + span.set_attribute("hook_attr", "hello world without app") + response_headers.append(hook_headers) - FlaskInstrumentor().instrument(name_callback=custom_span_name) + FlaskInstrumentor().instrument( + request_hook=request_hook_test, response_hook=response_hook_test + ) # pylint: disable=import-outside-toplevel,reimported,redefined-outer-name from flask import Flask @@ -272,12 +301,21 @@ def tearDown(self): with self.disable_logging(): FlaskInstrumentor().uninstrument() - def test_custom_span_name(self): - self.client.get("/hello/123") + def test_no_app_hooks(self): + expected_attrs = expected_attributes( + { + "http.target": "/hello/123", + "http.route": "/hello/", + "hook_attr": "hello world without app", + } + ) + resp = self.client.get("/hello/123") span_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(span_list), 1) - self.assertEqual(span_list[0].name, "instrument-without-app") + self.assertEqual(span_list[0].name, "without app") + self.assertEqual(span_list[0].attributes, expected_attrs) + self.assertEqual(resp.headers["hook_attr"], "hello otel without app") class TestProgrammaticCustomTracerProvider(