zk-data-agent/skills/data-factory-sql/run_sql.py

"""CLI runner for data-factory-sql skill.

Usage:
    python run_sql.py "SELECT 1"                          # auto engine, save CSV
    python run_sql.py -f query.sql                        # read SQL from file
    python run_sql.py "SELECT ..." --engine spark
    python run_sql.py "SELECT ..." --output ./out.csv
    python run_sql.py "SELECT ..." --no-save              # only print to stdout
    python run_sql.py "SELECT ..." --preview-rows 0       # no preview
    python run_sql.py --print-config                      # show effective config

Token resolution order:
    1. --token <value>
    2. KYUUBI_TOKEN env var
    3. C:\\workspace\\data-factory-token.txt
    4. ~/.config/data-factory/token

Output format:
    Always prints a JSON summary to stdout (last line) so the agent can parse:
    {"queryId": "...", "engine": "TRINO", "rows": 42, "cols": 5,
     "elapsed_ms": 3210, "output": "/path/to.csv"}

    Errors → JSON with "error" key, exit code != 0.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from pathlib import Path

# Allow running as a script: python run_sql.py
sys.path.insert(0, str(Path(__file__).parent))
from kyuubi_client import (  # noqa: E402
    AuthError,
    CLUSTER_BASE_URLS,
    KyuubiClient,
    KyuubiError,
    QueryError,
)


TOKEN_PATHS = [
    Path("C:/workspace/data-factory-token.txt"),
    Path.home() / ".config" / "data-factory" / "token",
]


def resolve_token(arg_token: str | None) -> str:
    if arg_token:
        return arg_token.strip()
    env = os.environ.get("KYUUBI_TOKEN")
    if env:
        return env.strip()
    for p in TOKEN_PATHS:
        if p.exists():
            return p.read_text(encoding="utf-8").strip()
    raise SystemExit(
        "no token found. set --token, $KYUUBI_TOKEN, or write token to "
        f"{TOKEN_PATHS[0]}"
    )


def read_sql(args: argparse.Namespace) -> str:
    if args.file:
        return Path(args.file).read_text(encoding="utf-8")
    if args.sql:
        return args.sql
    if not sys.stdin.isatty():
        return sys.stdin.read()
    raise SystemExit("provide SQL as positional arg, with -f FILE, or via stdin")


def default_output_path(query_id: str) -> Path:
    """Default to current session output when python_exec exposes it."""
    ts = time.strftime("%Y%m%d_%H%M%S")
    scratchpad = os.environ.get("PYTHON_EXEC_SCRATCHPAD")
    if scratchpad:
        output_dir = Path(scratchpad).resolve().parent / "output"
        output_dir.mkdir(parents=True, exist_ok=True)
        return output_dir / f"data_factory_{ts}.csv"
    downloads = Path.home() / "Downloads"
    downloads.mkdir(parents=True, exist_ok=True)
    return downloads / f"data_factory_{ts}.csv"


def save_csv(path: Path, columns: list[dict], rows: list[list]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    headers = [c.get("name", f"col{i}") for i, c in enumerate(columns)]
    with path.open("w", encoding="utf-8-sig", newline="") as f:
        w = csv.writer(f)
        w.writerow(headers)
        for row in rows:
            w.writerow(["" if v is None else v for v in row])


def print_preview(columns: list[dict], rows: list[list], n: int) -> None:
    if n <= 0 or not rows:
        return
    headers = [c.get("name", f"col{i}") for i, c in enumerate(columns)]
    types = [c.get("type", "?") for c in columns]
    widths = [max(len(h), len(t)) for h, t in zip(headers, types)]
    for row in rows[:n]:
        widths = [
            max(w, len(str(v)) if v is not None else 4)
            for w, v in zip(widths, row)
        ]
    sep = "  "
    print(sep.join(h.ljust(w) for h, w in zip(headers, widths)), file=sys.stderr)
    print(sep.join(t.ljust(w) for t, w in zip(types, widths)), file=sys.stderr)
    print(sep.join("-" * w for w in widths), file=sys.stderr)
    for row in rows[:n]:
        cells = [str(v) if v is not None else "NULL" for v in row]
        print(sep.join(c.ljust(w) for c, w in zip(cells, widths)), file=sys.stderr)
    if len(rows) > n:
        print(f"... ({len(rows) - n} more rows)", file=sys.stderr)


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Run SQL on Xiaomi Data Factory via Kyuubi HTTP API"
    )
    p.add_argument("sql", nargs="?", help="SQL string (or use -f / stdin)")
    p.add_argument("-f", "--file", help="read SQL from file")
    p.add_argument("--token", help="override token (default: file/env)")
    p.add_argument(
        "--cluster",
        choices=sorted(CLUSTER_BASE_URLS.keys()),
        default="cnbj1",
        help="cluster (default: cnbj1)",
    )
    p.add_argument(
        "--engine",
        default="auto",
        help="auto/presto/spark/doris/hologres (default: auto)",
    )
    p.add_argument("--catalog", help="default catalog")
    p.add_argument("--schema", help="default schema")
    p.add_argument(
        "--output",
        help=(
            "output CSV path (default: session output when PYTHON_EXEC_SCRATCHPAD "
            "is set, otherwise ~/Downloads/data_factory_<ts>.csv)"
        ),
    )
    p.add_argument(
        "--no-save", action="store_true", help="don't save CSV, only print to stdout"
    )
    p.add_argument(
        "--preview-rows",
        type=int,
        default=10,
        help="preview rows shown to stderr (default: 10, 0 to disable)",
    )
    p.add_argument(
        "--query-timeout",
        type=int,
        default=600,
        help="overall timeout in seconds (default: 600)",
    )
    p.add_argument(
        "--poll-interval",
        type=float,
        default=2.0,
        help="status poll interval seconds (default: 2.0)",
    )
    p.add_argument(
        "--print-config",
        action="store_true",
        help="print effective config (token redacted) and exit",
    )
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    try:
        token = resolve_token(args.token)
    except SystemExit:
        print(json.dumps({"error": "no token configured"}), flush=True)
        raise

    base_url = CLUSTER_BASE_URLS[args.cluster]

    if args.print_config:
        cfg = {
            "cluster": args.cluster,
            "base_url": base_url,
            "engine": args.engine,
            "catalog": args.catalog,
            "schema": args.schema,
            "token": f"{token[:6]}...{token[-4:]}",
            "token_len": len(token),
        }
        print(json.dumps(cfg, indent=2))
        return 0

    sql = read_sql(args).strip()
    if not sql:
        print(json.dumps({"error": "empty SQL"}), flush=True)
        return 2

    client = KyuubiClient(
        token=token,
        base_url=base_url,
        engine=args.engine,
        catalog=args.catalog,
        schema=args.schema,
        poll_interval=args.poll_interval,
        query_timeout_seconds=args.query_timeout,
    )

    try:
        result = client.execute(sql)
    except AuthError as e:
        print(json.dumps({"error": f"auth: {e}"}), flush=True)
        return 3
    except QueryError as e:
        print(json.dumps({"error": f"query: {e}"}), flush=True)
        return 4
    except KyuubiError as e:
        print(json.dumps({"error": f"kyuubi: {e}"}), flush=True)
        return 5

    print_preview(result.columns, result.rows, args.preview_rows)

    summary: dict = {
        "queryId": result.query_id,
        "engine": result.engine,
        "rows": len(result.rows),
        "cols": len(result.columns),
        "elapsed_ms": result.elapsed_ms,
        "columns": [
            {"name": c.get("name"), "type": c.get("type")} for c in result.columns
        ],
    }

    if not args.no_save:
        out_path = Path(args.output) if args.output else default_output_path(
            result.query_id or "x"
        )
        save_csv(out_path, result.columns, result.rows)
        summary["output"] = str(out_path)

    print(json.dumps(summary, ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())