unrip/scripts/ops/live_near_intents.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import signal
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import DEFAULT_NAMESPACE, build_env, kafka_brokers, raw_near_topic


DEFAULT_OUTPUT_FORMAT = (
    "%d{go[2006-01-02T15:04:05Z07:00]} topic=%t partition=%p offset=%o\n"
    "%v\n"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Live-inspect raw NEAR intents arriving in Redpanda."
    )
    parser.add_argument(
        "--namespace",
        default=DEFAULT_NAMESPACE,
        help=f"Kubernetes namespace to inspect (default: {DEFAULT_NAMESPACE})",
    )
    parser.add_argument(
        "--brokers",
        default="",
        help="Override Kafka brokers instead of reading them from unrip-config.",
    )
    parser.add_argument(
        "--topic",
        default="",
        help="Topic to consume. Defaults to KAFKA_TOPIC_RAW_NEAR_INTENTS_QUOTE from unrip-config.",
    )
    offset_group = parser.add_mutually_exclusive_group()
    offset_group.add_argument(
        "--offset",
        default="end",
        help=(
            "Kafka offset selector passed to rpk. "
            "The default 'end' means live-tail from now, not replay history."
        ),
    )
    offset_group.add_argument(
        "--last",
        type=int,
        default=0,
        help=(
            "Read the most recent N retained records and exit. "
            "Equivalent to '--offset -N --num N'."
        ),
    )
    parser.add_argument(
        "--num",
        type=int,
        default=0,
        help=(
            "Number of records to consume before exiting, starting from the chosen "
            "offset. 0 means unbounded. With '--offset end' this waits for new records."
        ),
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=0,
        help="Stop after this many seconds even if the stream is still open.",
    )
    parser.add_argument(
        "--value-only",
        action="store_true",
        help="Print only the record value without topic/partition/offset metadata.",
    )
    parser.add_argument(
        "--rpk-json",
        action="store_true",
        help="Use rpk's built-in JSON record formatter instead of the default metadata+value output.",
    )
    parser.add_argument(
        "--output",
        choices=("json", "jsonl", "text"),
        default="json",
        help="Output format (default: json).",
    )
    return parser.parse_args()


def build_command(args: argparse.Namespace, brokers: str, topic: str) -> list[str]:
    offset = args.offset
    num = args.num
    if args.last:
        offset = f"-{args.last}"
        num = args.last

    command = [
        "kubectl",
        "exec",
        "-i",
        "-n",
        args.namespace,
        "deploy/redpanda",
        "--",
        "rpk",
        "topic",
        "consume",
        topic,
        "--brokers",
        brokers,
        "--offset",
        offset,
        "--num",
        str(num),
    ]

    if args.output in {"json", "jsonl"} or args.rpk_json:
        command.extend(["--format", "json", "--pretty-print=false"])
    elif args.rpk_json:
        command.extend(["--format", "json", "--pretty-print=false"])
    elif args.value_only:
        command.extend(["--format", "%v\n"])
    else:
        command.extend(["--format", DEFAULT_OUTPUT_FORMAT])

    return command


def main() -> int:
    args = parse_args()
    brokers = args.brokers or kafka_brokers(namespace=args.namespace)
    topic = args.topic or raw_near_topic(namespace=args.namespace)
    command = build_command(args, brokers, topic)
    effective_offset = args.offset if not args.last else f"-{args.last}"
    effective_num = args.num if not args.last else args.last

    if args.output == "text":
        process = subprocess.Popen(command, env=build_env())
        try:
            process.wait(timeout=args.timeout if args.timeout and args.timeout > 0 else None)
        except subprocess.TimeoutExpired:
            process.send_signal(signal.SIGINT)
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=5)
            return 124
        except KeyboardInterrupt:
            process.send_signal(signal.SIGINT)
            return process.wait()
        return process.returncode

    try:
        result = subprocess.run(
            command,
            env=build_env(),
            text=True,
            capture_output=True,
            timeout=args.timeout if args.timeout and args.timeout > 0 else None,
            check=False,
        )
        stdout = result.stdout
        return_code = result.returncode
    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout or ""
        return_code = 124

    records = parse_rpk_json_lines(stdout, value_only=args.value_only)
    if args.output == "json":
        print(json.dumps(records, indent=2))
    else:
        for record in records:
            print(json.dumps(record, separators=(",", ":")))
    return return_code


def parse_rpk_json_lines(stdout: str, *, value_only: bool) -> list[object]:
    records: list[object] = []
    for raw_line in stdout.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        record = json.loads(line)
        if value_only:
            record = decode_jsonish(record.get("value"))
        else:
            if "value" in record:
                record["value"] = decode_jsonish(record.get("value"))
        records.append(record)
    return records


def decode_jsonish(value: object) -> object:
    if not isinstance(value, str):
        return value
    stripped = value.strip()
    if not stripped:
        return value
    if stripped[0] not in "{[":
        return value
    try:
        return json.loads(stripped)
    except json.JSONDecodeError:
        return value


if __name__ == "__main__":
    raise SystemExit(main())