[FLASK] added request and response hook

CHANGELOG.md

@@ -19,6 +19,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.
@@ -36,6 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Remove `http.status_text` from span attributes
   ([#406](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/406))
 
+
 ## [0.19b0](https://github.com/open-telemetry/opentelemetry-python-contrib/releases/tag/v0.19b0) - 2021-03-26
 
 - Implement context methods for `_InterceptorChannel`

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -59,6 +59,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"
@@ -78,8 +79,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
@@ -101,21 +101,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):
+def _wrapped_before_request(request_hook=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)
         )
@@ -127,6 +129,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
@@ -143,7 +148,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
 
 
@@ -170,16 +176,15 @@ def _teardown_request(exc):
 
 class _InstrumentedFlask(flask.Flask):
 
-    name_callback = get_default_span_name
 
     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)
 
         _before_request = _wrapped_before_request(
-            _InstrumentedFlask.name_callback
+            _InstrumentedFlask.request_hook
         )
         self._before_request = _before_request
         self.before_request(_before_request)
@@ -192,25 +197,25 @@ class FlaskInstrumentor(BaseInstrumentor):
 
     See `BaseInstrumentor`
     """
-
     def _instrument(self, **kwargs):
         self._original_flask = flask.Flask
-        name_callback = kwargs.get("name_callback")
-        if callable(name_callback):
-            _InstrumentedFlask.name_callback = name_callback
+        request_hook = kwargs.get("request_hook")
+        response_hook = kwargs.get("response_hook")
+        _InstrumentedFlask.request_hook = request_hook
+        _InstrumentedFlask.response_hook = response_hook
         flask.Flask = _InstrumentedFlask
 
     def instrument_app(
-        self, app, name_callback=get_default_span_name
+        self, app, request_hook=None, response_hook=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)
 
-            _before_request = _wrapped_before_request(name_callback)
+            _before_request = _wrapped_before_request(request_hook)
             app._before_request = _before_request
             app.before_request(_before_request)
             app.teardown_request(_teardown_request)

instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py

@@ -185,46 +185,127 @@ def test_exclude_lists(self):
         self.assertEqual(len(span_list), 1)
 
 
-class TestProgrammaticCustomSpanName(
+# class TestProgrammaticCustomSpanName(
+#     InstrumentationTest, TestBase, WsgiTestBase
+# ):
+#     def setUp(self):
+#         super().setUp()
+
+#         def custom_span_name():
+#             return "flask-custom-span-name"
+
+#         self.app = Flask(__name__)
+
+#         FlaskInstrumentor().instrument_app(
+#             self.app, name_callback=custom_span_name
+#         )
+
+#         self._common_initialization()
+
+#     def tearDown(self):
+#         super().tearDown()
+#         with self.disable_logging():
+#             FlaskInstrumentor().uninstrument_app(self.app)
+
+#     def test_custom_span_name(self):
+#         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")
+
+
+
+# class TestProgrammaticCustomSpanNameCallbackWithoutApp(
+#     InstrumentationTest, TestBase, WsgiTestBase
+# ):
+#     def setUp(self):
+#         super().setUp()
+
+#         def custom_span_name():
+#             return "instrument-without-app"
+
+#         FlaskInstrumentor().instrument(name_callback=custom_span_name, request_hook=None)
+#         # pylint: disable=import-outside-toplevel,reimported,redefined-outer-name
+#         from flask import Flask
+
+#         self.app = Flask(__name__)
+
+#         self._common_initialization()
+
+#     def tearDown(self):
+#         super().tearDown()
+#         with self.disable_logging():
+#             FlaskInstrumentor().uninstrument()
+
+#     def test_custom_span_name(self):
+#         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")
+
+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()
-
+    
     def tearDown(self):
         super().tearDown()
         with self.disable_logging():
             FlaskInstrumentor().uninstrument_app(self.app)
+
+    def test_hooks(self):
+        expected_attrs = expected_attributes(
+            {"http.target": "/hello/123", "http.route": "/hello/<int:helloid>", "hook_attr":"hello world"}
+        )
 
-    def test_custom_span_name(self):
         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)
 
-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")
 
-        FlaskInstrumentor().instrument(name_callback=custom_span_name)
+        def response_hook_test(span, environ, response_headers):
+            span.set_attribute("hook_attr", "hello world without app")
+            response_headers.append(hook_headers)
+
+        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
 
@@ -237,9 +318,13 @@ def tearDown(self):
         with self.disable_logging():
             FlaskInstrumentor().uninstrument()
 
-    def test_custom_span_name(self):
+    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"}
+        )
         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)