Add support for programmatic instrumentation in Flask

Author: ocelotl

docs/examples/opentelemetry-example-app/src/opentelemetry_example_app/flask_example.py

@@ -33,8 +33,8 @@
     SimpleExportSpanProcessor(ConsoleSpanExporter())
 )
 
-FlaskInstrumentor().instrument()
-app = flask.Flask(__name__)
+Flask = FlaskInstrumentor().instrument(flask_class=flask.Flask)
+app = Flask(__name__)
 opentelemetry.ext.http_requests.enable(trace.get_tracer_provider())
 
 

docs/getting-started.rst

@@ -184,9 +184,6 @@ And let's write a small Flask application that sends an HTTP request, activating
 .. code-block:: python
 
     # flask_example.py
-    from opentelemetry.ext.flask import FlaskInstrumentor
-    FlaskInstrumentor().instrument()  # This needs to be executed before importing Flask
-
     import flask
     import requests
 
@@ -195,13 +192,16 @@ And let's write a small Flask application that sends an HTTP request, activating
     from opentelemetry.sdk.trace import TracerProvider
     from opentelemetry.sdk.trace.export import ConsoleSpanExporter
     from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor
+    from opentelemetry.ext.flask import FlaskInstrumentor
+
+    Flask = FlaskInstrumentor().instrument(flask_class=flask.Flask)
 
     trace.set_tracer_provider(TracerProvider())
     trace.get_tracer_provider().add_span_processor(
         SimpleExportSpanProcessor(ConsoleSpanExporter())
     )
 
-    app = flask.Flask(__name__)
+    app = Flask(__name__)
     opentelemetry.ext.http_requests.enable(trace.get_tracer_provider())
 
     @app.route("/")

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

@@ -29,9 +29,10 @@
 
 .. code-block:: python
 
-    from opentelemetry.ext.flask import FlaskInstrumentor
-    FlaskInstrumentor().instrument()  # This needs to be executed before importing Flask
     from flask import Flask
+    from opentelemetry.ext.flask import FlaskInstrumentor
+
+    Flask = FlaskInstrumentor().instrument(flask_class=Flask)
 
     app = Flask(__name__)
 
@@ -158,11 +159,20 @@ class FlaskInstrumentor(BaseInstrumentor):
 
     def __init__(self):
         super().__init__()
-        self._original_flask = None
+        self._original_flask_class = None
 
-    def _instrument(self, **kwargs):
-        self._original_flask = flask.Flask
+    def _instrument(
+        self, flask_class=None
+    ):  # pylint: disable=arguments-differ
+        if flask_class is not None:
+            self._original_flask_class = flask_class
+            return _InstrumentedFlask
+
+        self._original_flask_class = flask.Flask
         flask.Flask = _InstrumentedFlask
 
-    def _uninstrument(self, **kwargs):
-        flask.Flask = self._original_flask
+        return None
+
+    def _uninstrument(self):  # pylint: disable=arguments-differ
+        flask.Flask = self._original_flask_class
+        return self._original_flask_class

ext/opentelemetry-ext-flask/tests/conftest.py

@@ -19,8 +19,11 @@
 from werkzeug.wrappers import BaseResponse
 
 from opentelemetry import trace as trace_api
+from opentelemetry.ext.flask import FlaskInstrumentor
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 
+Flask = FlaskInstrumentor().instrument(flask_class=Flask)
+
 
 def expected_attributes(override_attributes):
     default_attributes = {

ext/opentelemetry-ext-flask/tests/test_flask_integration.py

@@ -52,7 +52,9 @@ def instrument(self, **kwargs):
             self._is_instrumented = True
             return result
 
-        _LOG.warning("Attempting to instrument while already instrumented")
+        _LOG.warning(
+            "Attempting to automatically instrument while already instrumented"
+        )
 
         return None
 
@@ -64,7 +66,10 @@ def uninstrument(self, **kwargs):
             self._is_instrumented = False
             return result
 
-        _LOG.warning("Attempting to uninstrument while already uninstrumented")
+        _LOG.warning(
+            "Attempting to automatically uninstrument while already"
+            " uninstrumented"
+        )
 
         return None
 

opentelemetry-auto-instrumentation/src/opentelemetry/auto_instrumentation/instrumentor.py

@@ -34,6 +34,7 @@ def test_protect(self):
             self.assertIs(instrumentor.uninstrument(), None)
 
         self.assertEqual(instrumentor.instrument(), "instrumented")
+
         with self.assertLogs(level=WARNING):
             self.assertIs(instrumentor.instrument(), None)