WIP flask cli

Author: xrmx

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

@@ -240,13 +240,15 @@ def response_hook(span: Span, status: str, response_headers: List):
 """
 
 import weakref
+from functools import partial
 from logging import getLogger
 from time import time_ns
 from timeit import default_timer
 from typing import Collection
 
 import flask
 from packaging import version as package_version
+from wrapt import wrap_function_wrapper
 
 import opentelemetry.instrumentation.wsgi as otel_wsgi
 from opentelemetry import context, trace
@@ -264,14 +266,19 @@ def response_hook(span: Span, status: str, response_headers: List):
 from opentelemetry.instrumentation.propagators import (
     get_global_response_propagator,
 )
-from opentelemetry.instrumentation.utils import _start_internal_or_server_span
+from opentelemetry.instrumentation.utils import (
+    _start_internal_or_server_span,
+    unwrap,
+)
 from opentelemetry.metrics import get_meter
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv.attributes.http_attributes import HTTP_ROUTE
 from opentelemetry.semconv.metrics import MetricInstruments
 from opentelemetry.semconv.metrics.http_metrics import (
     HTTP_SERVER_REQUEST_DURATION,
 )
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace.status import StatusCode
 from opentelemetry.util._importlib_metadata import version
 from opentelemetry.util.http import (
     get_excluded_urls,
@@ -620,6 +627,26 @@ def __init__(self, *args, **kwargs):
         self.teardown_request(_teardown_request)
 
 
+def _cli_command_wrapper(wrapped, instance, args, kwargs, tracer):
+    span_name = "cli span name"
+    span_attributes = {}
+
+    print("AAAAAAAAAAA", args, kwargs)
+    print("AAAAAAAAAAA", wrapped, instance)
+
+    with tracer.start_as_current_span(
+        name=span_name,
+        kind=trace.SpanKind.INTERNAL,
+        attributes=span_attributes,
+    ) as span:
+        try:
+            return wrapped(*args, **kwargs)
+        except Exception as exc:
+            span.set_status(StatusCode.ERROR, str(exc))
+            span.set_attribute(ERROR_TYPE, exc.__class__.__qualname__)
+            raise
+
+
 class FlaskInstrumentor(BaseInstrumentor):
     # pylint: disable=protected-access,attribute-defined-outside-init
     """An instrumentor for flask.Flask
@@ -662,9 +689,24 @@ def _instrument(self, **kwargs):
 
         flask.Flask = _InstrumentedFlask
 
+        # TODO: this is duplicated from _InstrumentedFlask, hopefully will be resolved once moving out of subclassing
+        tracer = trace.get_tracer(
+            __name__,
+            __version__,
+            tracer_provider,
+            schema_url=_get_schema_url(sem_conv_opt_in_mode),
+        )
+        wrap_function_wrapper(
+            "flask.cli",
+            "AppGroup.command",
+            partial(_cli_command_wrapper, tracer=tracer),
+        )
+
     def _uninstrument(self, **kwargs):
         flask.Flask = self._original_flask
 
+        unwrap(flask.cli.AppGroup, "command")
+
     # pylint: disable=too-many-locals
     @staticmethod
     def instrument_app(

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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import click
 import flask
 from werkzeug.test import Client
 from werkzeug.wrappers import Response
@@ -94,3 +95,27 @@ def test_no_op_tracer_provider(self):
 
         span_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(span_list), 0)
+
+
+@app.cli.command("mycommand")
+@click.option("--option", default="default")
+def flask_command(option):
+    pass
+
+
+def test_cli_command_wrapping(runner):
+    FlaskInstrumentor().instrument()
+
+    result = runner.invoke(args=["mycommand"])
+    (span,) = self.memory_exporter.get_finished_spans()
+
+    FlaskInstrumentor().uninstrument()
+
+
+def test_cli_command_wrapping_with_options(runner):
+    FlaskInstrumentor().instrument()
+    result_with_option = runner.invoke(
+        args=["mycommand", "--option", "option"]
+    )
+    (span,) = self.memory_exporter.get_finished_spans()
+    FlaskInstrumentor().uninstrument()