[FLASK] added request and response hook

CHANGELOG.md

@@ -61,6 +61,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

@@ -63,6 +63,7 @@ def hello():
 from opentelemetry.util._time import _time_ns
 from opentelemetry.util.http import get_excluded_urls
 
+
 _logger = getLogger(__name__)
 
 _ENVIRON_STARTTIME_KEY = "opentelemetry-flask.starttime_key"
@@ -82,8 +83,7 @@ def get_default_span_name():
         span_name = otel_wsgi.get_default_span_name(flask.request.environ)
     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
@@ -112,21 +112,23 @@ def _start_response(status, response_headers, *args, **kwargs):
                         "missing at _start_response(%s)",
                         status,
                     )
-
+                if response_hook:
+                    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)
         )
@@ -136,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
@@ -154,7 +159,8 @@ def _before_request():
         flask_request_environ[_ENVIRON_ACTIVATION_KEY] = activation
         flask_request_environ[_ENVIRON_SPAN_KEY] = span
         flask_request_environ[_ENVIRON_TOKEN] = token
-
+
+
     return _before_request
 
 
@@ -181,21 +187,23 @@ 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)
@@ -208,29 +216,32 @@ class FlaskInstrumentor(BaseInstrumentor):
 
     See `BaseInstrumentor`
     """
-
     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

@@ -142,7 +142,8 @@ def test_trace_response(self):
 
         self.assertIn("traceresponse", headers)
         self.assertEqual(
-            headers["access-control-expose-headers"], "traceresponse",
+            headers["access-control-expose-headers"],
+            "traceresponse",
         )
         self.assertEqual(
             headers["traceresponse"],
@@ -184,7 +185,7 @@ def test_404(self):
         self.assertEqual(span_list[0].name, "HTTP POST")
         self.assertEqual(span_list[0].kind, trace.SpanKind.SERVER)
         self.assertEqual(span_list[0].attributes, expected_attrs)
-
+        
     def test_internal_error(self):
         expected_attrs = expected_attributes(
             {
@@ -220,19 +221,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 +252,45 @@ 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")
+            # environ.headers.set("apple", "cat")
+            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 +303,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(