diff --git a/cli_telemetry/telemetry.py b/cli_telemetry/telemetry.py index b7e7f95..9b53f06 100644 --- a/cli_telemetry/telemetry.py +++ b/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. diff --git a/examples/click_ex.py b/examples/click_ex.py index ec68fb7..2696226 100644 --- a/examples/click_ex.py +++ b/examples/click_ex.py @@ -1,35 +1,29 @@ 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) @@ -37,22 +31,17 @@ def shout(message, 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()