Add skill enable controls and data factory SQL skill

2026-05-08 14:56:32 +08:00
parent 046b0eeab0
commit bd2f260b9f
16 changed files with 1167 additions and 25 deletions
@@ -0,0 +1,246 @@
+"""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: ~/Downloads/data_factory_<timestamp>.csv"""
+    ts = time.strftime("%Y%m%d_%H%M%S")
+    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: ~/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())