From b22c3890d29e7bd2cf0f81d0f85c35637992ebd0 Mon Sep 17 00:00:00 2001 From: nadin-Starkware Date: Tue, 18 Nov 2025 17:38:43 +0200 Subject: [PATCH] scripts: add set log level script --- scripts/set_log_level.py | 203 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) create mode 100644 scripts/set_log_level.py diff --git a/scripts/set_log_level.py b/scripts/set_log_level.py new file mode 100644 index 00000000000..63785449d05 --- /dev/null +++ b/scripts/set_log_level.py @@ -0,0 +1,203 @@ +import argparse +import subprocess +import sys +import time +from typing import List + +import requests +import signal +import socket + +SLEEP_INTERVAL = 0.4 + + +def parse_args(args: List[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Set the log level for a module or crate", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + # Add port-forwarding arguments + add_port_forward_args(parser) + + parser.add_argument( + "--target", + type=str, + help="Crate or module name whose log level should be inspected or updated", + ) + parser.add_argument("--log_level", type=str, help="The log level to set for the module/crate") + parser.add_argument( + "--method", + type=str, + choices=["get", "post"], + default="post", + help="HTTP method to use: 'get' to read current log level, 'post' to set a log level", + ) + return parser.parse_args(args) + + +def add_port_forward_args(parser: argparse.ArgumentParser) -> None: + """Add port-forwarding related CLI options to the parser.""" + + pf_group = parser.add_argument_group("port-forwarding options") + + pf_group.add_argument( + "--pod_name", + type=str, + help="Pod to port-forward to; omit when no port forwarding is needed", + ) + + pf_group.add_argument( + "--local_port", + type=int, + default=8082, + help="Local port to bind the port-forward tunnel to", + ) + + pf_group.add_argument( + "--cluster_name", + type=str, + help="Optional k8s cluster name; if supplied, the script runs 'kubectx ' before port-forwarding", + ) + + pf_group.add_argument( + "--namespace", + type=str, + help="Optional k8s namespace; if supplied, the script runs 'kubens ' before port-forwarding", + ) + + pf_group.add_argument( + "--monitoring_port", + type=int, + default=8082, + help="Monitoring endpoint port", + ) + + +def port_forward( + pod_name: str, + local_port: int, + remote_port: int, + max_attempts: int = 5, +) -> subprocess.Popen: + """Start a kubectl port-forward and wait until it is ready. + + Returns the Popen handle so the caller can terminate it later. + Raises RuntimeError if the local port is still unreachable after + `max_attempts` connection checks. + """ + + cmd = ["kubectl", "port-forward", pod_name, f"{local_port}:{remote_port}"] + print("Starting port-forward:", " ".join(cmd)) + proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + for _attempt in range(1, max_attempts + 1): + try: + with socket.create_connection(("localhost", local_port), timeout=1): + print( + f"✅ Port-forward to {pod_name}:{remote_port} is ready on localhost:{local_port}" + ) + return proc + except OSError: + time.sleep(SLEEP_INTERVAL) + + proc.terminate() + proc.wait(timeout=5) + raise RuntimeError( + f"❌ Port-forward to {pod_name}:{remote_port} failed after {max_attempts} attempts." + ) + + +def format_log_levels(raw: str) -> str: + """Convert comma-separated log-level string to human-readable form. + + If the last token has no "=", treat it as the global default level. + Returns a newline-joined string ready for printing. + """ + + raw_entries = [item.strip() for item in raw.split(",") if item.strip()] + + default_level = None + if raw_entries and "=" not in raw_entries[-1]: + default_level = raw_entries.pop() + + lines: list[str] = [] + if default_level: + lines.append(f"default={default_level}") + + lines.extend(sorted(raw_entries)) + return "\n".join(lines) + + +def main(): + args = parse_args(sys.argv[1:]) + + # If a pod name is supplied, establish a port-forward before making the request + port_forward_proc = None + + target_port = args.monitoring_port + setup_port_forwarding = bool(args.pod_name) + base_port = args.local_port if setup_port_forwarding else target_port + + if setup_port_forwarding: + # Switch context/namespace if requested + try: + if args.cluster_name: + subprocess.check_call(["kubectx", args.cluster_name]) + if args.namespace: + subprocess.check_call(["kubens", args.namespace]) + except subprocess.CalledProcessError as err: + print(f"Failed to switch kubectl context/namespace: {err}") + sys.exit(1) + + try: + port_forward_proc = port_forward(args.pod_name, args.local_port, args.monitoring_port) + except RuntimeError as err: + print(err) + sys.exit(1) + + try: + if args.method == "get": + full_url = f"http://localhost:{base_port}/monitoring/logLevel" + print(f"Fetching current log level from {full_url}") + response = requests.get(full_url, timeout=5) + + if response.status_code != 200: + print(f"Failed to fetch log level: {response.status_code} {response.text}") + sys.exit(1) + + print("Current log levels:\n", format_log_levels(response.text)) + elif args.method == "post": + # Validate required arguments + if not args.target or not args.log_level: + print("--target and --log_level are required when --method=post") + sys.exit(1) + + base_url = f"http://localhost:{base_port}/monitoring/setLogLevel" + full_url = f"{base_url}/{args.target}/{args.log_level}" + + print(f"Setting log level for {args.target} to {args.log_level} at {full_url}") + + response = requests.post(full_url, timeout=5) + + if response.status_code != 200: + print( + f"❌ Failed to set log level for {args.target} to {args.log_level}: {response.text}" + ) + sys.exit(1) + + print(f"✅ Successfully set log level for {args.target} to {args.log_level}") + + finally: + # Clean up the port-forward process if we started one + if port_forward_proc: + port_forward_proc.send_signal(signal.SIGINT) + try: + port_forward_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + port_forward_proc.kill() + port_forward_proc.wait() + + +if __name__ == "__main__": + main()