diff --git a/src/warnet/control.py b/src/warnet/control.py index 782764cd9..346f83367 100644 --- a/src/warnet/control.py +++ b/src/warnet/control.py @@ -231,17 +231,18 @@ def run(scenario_file: str, additional_args: tuple[str]): @click.command() -@click.argument("pod_name", type=str, default="") -@click.option("--follow", "-f", is_flag=True, default=False, help="Follow logs") -def logs(pod_name: str, follow: bool): - """Show the logs of a pod""" +@click.argument("pod_name", type=str, required=False) +@click.option("--follow", "-f", is_flag=True, default=False, help="Follow logs in real-time") +@click.option("--grep", "-g", type=str, required=False, help="Pattern to grep for in logs") +def logs(pod_name: str = "", follow: bool = False, grep: str = ""): + """Show the logs of a Kubernetes pod, optionally filtering by a grep pattern.""" follow_flag = "--follow" if follow else "" namespace = get_default_namespace() if pod_name: try: command = f"kubectl logs pod/{pod_name} -n {namespace} {follow_flag}" - stream_command(command) + stream_command(command, grep) return except Exception as e: print(f"Could not find the pod {pod_name}: {e}") @@ -270,7 +271,7 @@ def logs(pod_name: str, follow: bool): pod_name = selected["pod"] try: command = f"kubectl logs pod/{pod_name} -n {namespace} {follow_flag}" - stream_command(command) + stream_command(command, grep) except Exception as e: print(f"Please consider waiting for the pod to become available. Encountered: {e}") else: diff --git a/src/warnet/process.py b/src/warnet/process.py index 626124b71..bfe477d6d 100644 --- a/src/warnet/process.py +++ b/src/warnet/process.py @@ -1,3 +1,4 @@ +import re import subprocess @@ -8,7 +9,8 @@ def run_command(command: str) -> str: return result.stdout -def stream_command(command: str) -> bool: +def stream_command(command: str, grep_pattern: str = "") -> bool: + """Stream output and apply an optional pattern filter.""" process = subprocess.Popen( ["bash", "-c", command], stdout=subprocess.PIPE, @@ -18,10 +20,16 @@ def stream_command(command: str) -> bool: universal_newlines=True, ) + pattern = re.compile(grep_pattern) if grep_pattern else None message = "" + # Only display lines matching the pattern if grep is specified for line in iter(process.stdout.readline, ""): message += line - print(line, end="") + if pattern: + if pattern.search(line): + print(line, end="") + else: + print(line, end="") process.stdout.close() return_code = process.wait()