WIP

Author: xrmx

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

@@ -627,12 +627,11 @@ 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 = {}
+def _cli_invoke_wrapper(wrapped, instance, args, kwargs, tracer):
+    callback = args[0] if args else kwargs.get("__callback")
 
-    print("AAAAAAAAAAA", args, kwargs)
-    print("AAAAAAAAAAA", wrapped, instance)
+    span_name = callback.__name__
+    span_attributes = {}
 
     with tracer.start_as_current_span(
         name=span_name,
@@ -696,16 +695,20 @@ def _instrument(self, **kwargs):
             tracer_provider,
             schema_url=_get_schema_url(sem_conv_opt_in_mode),
         )
+        import click
+
         wrap_function_wrapper(
-            "flask.cli",
-            "AppGroup.command",
-            partial(_cli_command_wrapper, tracer=tracer),
+            click.core.Context,
+            "invoke",
+            partial(_cli_invoke_wrapper, tracer=tracer),
         )
 
     def _uninstrument(self, **kwargs):
         flask.Flask = self._original_flask
 
-        unwrap(flask.cli.AppGroup, "command")
+        import click
+
+        unwrap(click.core.Context, "invoke")
 
     # pylint: disable=too-many-locals
     @staticmethod

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

@@ -20,6 +20,7 @@
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.flask import FlaskInstrumentor
 from opentelemetry.test.wsgitestutil import WsgiTestBase
+from opentelemetry.trace.status import StatusCode
 
 # pylint: disable=import-error
 from .base_test import InstrumentationTest
@@ -96,26 +97,58 @@ def test_no_op_tracer_provider(self):
         span_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(span_list), 0)
 
+    def test_cli_command_wrapping(self):
+        @self.app.cli.command()
+        def flask_command():
+            print("flask")
+            pass
+
+        runner = self.app.test_cli_runner()
+        result = runner.invoke(args=["flask-command"])
+        (span,) = self.memory_exporter.get_finished_spans()
+
+        self.assertEqual(span.status.status_code, StatusCode.UNSET)
+        self.assertEqual(span.name, "flask_command")
+
+    def test_cli_command_wrapping_with_name(self):
+        @self.app.cli.command("mycommand")
+        def flask_command():
+            print("my")
+            pass
+
+        runner = self.app.test_cli_runner()
+        print(runner)
+        result = runner.invoke(args=["mycommand"])
+        (span,) = self.memory_exporter.get_finished_spans()
+
+        self.assertEqual(span.status.status_code, StatusCode.UNSET)
+        self.assertEqual(span.name, "mycommand")
+
+    def test_cli_command_wrapping_with_options(self):
+        @self.app.cli.command()
+        @click.option("--option", default="default")
+        def my_command_with_opts(option):
+            print("opts")
+            pass
+
+        runner = self.app.test_cli_runner()
+        result = runner.invoke(
+            args=["my-command-with-opts", "--option", "option"],
+        )
+        (span,) = self.memory_exporter.get_finished_spans()
 
-@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()
+        self.assertEqual(span.status.status_code, StatusCode.UNSET)
+        self.assertEqual(span.name, "my_command_with_opts")
 
-    FlaskInstrumentor().uninstrument()
+    def test_cli_command_raises_error(self):
+        @self.app.cli.command()
+        def command_raises():
+            print("raises")
+            raise ValueError()
 
+        runner = self.app.test_cli_runner()
+        result = runner.invoke(args=["command-raises"])
+        (span,) = self.memory_exporter.get_finished_spans()
 
-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()
+        self.assertEqual(span.status.status_code, StatusCode.ERROR)
+        self.assertEqual(span.name, "command_raises")