[FLASK] added request and response hook (#416)

Author: NickSulistio

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.

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,21 +114,21 @@ 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)
 
     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)

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