Merge pull request #1 from amjith/amjith/auto-instrument

Add auto-instrumentation for Click.

Author: amjith

cli_telemetry/telemetry.py

@@ -81,6 +81,10 @@ def init_telemetry(service_name: str, db_path: Optional[str] = None, user_id_fil
         if _initialized:
             return
 
+        # monkeypatch click if allowed
+        if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH"):
+            _monkeypatch_click()
+
         # determine base path under XDG_DATA_HOME
         xdg = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
         base = os.path.join(xdg, "cli-telemetry", service_name)
@@ -266,3 +270,39 @@ def end_session() -> None:
             _conn.close()
         except Exception:
             pass
+
+
+def _monkeypatch_click():
+    """Monkeypatch click.Command and click.Group to auto-wrap in telemetry spans."""
+    import os
+
+    if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH") == "1":
+        return  # User has opted out of monkeypatching
+
+    try:
+        import click
+
+        def telemetry_wrapper(original_invoke):
+            def invoke_with_span(self, ctx):
+                # Assumes init_telemetry() already called by app
+                with Span(self.name, attributes={"cli.command": ctx.command_path}) as span:
+                    try:
+                        for param in self.params:
+                            if param.name in ctx.params:
+                                add_tag(f"args.{param.name}", ctx.params[param.name])
+                    except Exception:
+                        pass
+                    return original_invoke(self, ctx)
+            return invoke_with_span
+
+        # Avoid double-patching
+        if not getattr(click.Command, "_telemetry_patched", False):
+            click.Command.invoke = telemetry_wrapper(click.Command.invoke)
+            click.Command._telemetry_patched = True
+
+        if not getattr(click.Group, "_telemetry_patched", False):
+            click.Group.invoke = telemetry_wrapper(click.Group.invoke)
+            click.Group._telemetry_patched = True
+
+    except ImportError:
+        pass  # No click? No telemetry sadness.

examples/click_ex.py

@@ -1,58 +1,47 @@
 from time import sleep
 import click
-from cli_telemetry.telemetry import start_session, end_session, profile, add_tag, profile_block
 
+# import this first so it can patch click before any commands are defined
+import cli_telemetry.telemetry as telemetry
+from cli_telemetry.telemetry import profile, add_tag, profile_block
 
-@click.group()
-@click.pass_context
-def cli(ctx):
-    """
-    Example CLI with telemetry instrumentation.
-    """
-    # Start the root span for this invocation
-    cmd = ctx.invoked_subcommand or "cli"
-    start_session(command_name=cmd, service_name="example-cli")
-    # Ensure we end the span when the CLI exits
-    ctx.call_on_close(end_session)
+telemetry.init_telemetry("example-cli")
 
+@click.group()
+def cli():
+    """Example CLI with telemetry instrumentation."""
+    # no manual start_session() or call_on_close() needed any more
+    pass
 
 @cli.command()
 @click.argument("message")
-@profile
 def echo(message):
     """Prints the message as-is."""
     # Tag the argument so it shows up on this span
     add_tag("args.message", message)
     click.echo(message)
 
-
 @cli.command()
 @click.argument("message")
 @click.option("--times", "-n", default=1, show_default=True, help="How many times to shout")
-@profile
 def shout(message, times):
     """Prints the message uppercased with exclamation."""
     add_tag("args.message", message)
     add_tag("args.times", times)
     for _ in range(times):
         click.echo(f"{message.upper()}!")
 
-
 @cli.command()
-@profile
 def work():
     """Simulate some nested work using a profile_block."""
     add_tag("phase", "start_work")
     with profile_block("step_1", tags={"step": 1}):
         click.echo("step1")
         sleep(0.1)
-        pass
     with profile_block("step_2", tags={"step": 2}):
         click.echo("step2")
         sleep(0.2)
-        pass
     click.echo("Work done!")
 
-
 if __name__ == "__main__":
     cli()