yearning-cli 0.1.5__py3-none-any.whl

yearning_cli/__main__.py

@@ -0,0 +1,3 @@
+from yearning_cli.cli import main
+
+main()

yearning_cli/cli.py

@@ -0,0 +1,459 @@
+"""CLI entry point for yearning-cli."""
+
+import argparse
+import logging
+import os
+import sys
+import tomllib
+from pathlib import Path
+
+from yearning_cli.client import TokenExpiredError, YearningClient
+from mysql_proxy.server import serve
+
+BASE_URL = os.environ.get("YEARNING_URL", "http://192.168.1.135")
+CONFIG_FILE = os.path.expanduser("~/.sqlrc")
+
+
+def load_config() -> dict:
+    """Load config from ~/.sqlrc (key: value format)."""
+    cfg = {}
+    if os.path.exists(CONFIG_FILE):
+        with open(CONFIG_FILE, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if ":" in line:
+                    key, val = line.split(":", 1)
+                    cfg[key.strip()] = val.strip()
+    return cfg
+
+
+CONFIG = load_config()
+DEFAULT_SOURCE = CONFIG.get("source", "")
+DEFAULT_DATABASE = CONFIG.get("database", "")
+USERNAME = (
+    os.environ.get("YEARNING_USER", "")
+    or CONFIG.get("yearning_user", "")
+    or CONFIG.get("username", "")
+    or CONFIG.get("user", "")
+)
+PASSWORD = (
+    os.environ.get("YEARNING_PASS", "")
+    or CONFIG.get("yearning_pass", "")
+    or CONFIG.get("yearning_password", "")
+    or CONFIG.get("password", "")
+    or CONFIG.get("pass", "")
+)
+
+# Check for -v/--verbose before argparse (global flag)
+VERBOSE = "-v" in sys.argv or "--verbose" in sys.argv
+if VERBOSE:
+    sys.argv = [a for a in sys.argv if a not in ("-v", "--verbose")]
+
+
+def get_version() -> str:
+    """Return package version from pyproject.toml."""
+    pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml"
+    if not pyproject.exists():
+        raise RuntimeError(f"缺少版本配置文件: {pyproject}")
+
+    with open(pyproject, "rb") as f:
+        data = tomllib.load(f)
+
+    project_version = data.get("project", {}).get("version", "")
+    if not project_version:
+        raise RuntimeError("pyproject.toml 缺少 project.version")
+
+    return project_version
+
+
+def _get_client() -> YearningClient:
+    client = YearningClient(
+        BASE_URL,
+        verbose=VERBOSE,
+        username=USERNAME,
+        password=PASSWORD,
+    )
+    try:
+        client.ensure_authenticated()
+    except TokenExpiredError as e:
+        print(str(e))
+        sys.exit(1)
+    return client
+
+
+def _format_table(columns: list, rows: list, query_time: int = 0) -> None:
+    """Format query results as a table."""
+    if not columns:
+        if query_time:
+            print(f"(空结果集) [{query_time}ms]")
+        else:
+            print("(空结果集)")
+        return
+
+    # Calculate column widths
+    col_widths = [len(str(c)) for c in columns]
+    for row in rows:
+        for i, val in enumerate(row):
+            col_widths[i] = min(max(col_widths[i], len(str(val))), 60)
+
+    # Print header
+    sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+"
+    print(sep)
+    print("|" + "|".join(f" {str(c).ljust(w)} " for c, w in zip(columns, col_widths)) + "|")
+    print(sep)
+    for row in rows:
+        print("|" + "|".join(f" {str(v).ljust(w)} " for v, w in zip(row, col_widths)) + "|")
+    print(sep)
+    ts = f", {query_time}ms" if query_time else ""
+    print(f"({len(rows)} 行{ts})")
+
+
+def cmd_login(args: argparse.Namespace) -> None:
+    client = YearningClient(BASE_URL, verbose=VERBOSE)
+    user = USERNAME or input("用户名: ")
+    pwd = PASSWORD or input("密码: ")
+    try:
+        token = client.login(user, pwd)
+        print(f"登录成功! Token 已保存")
+    except Exception as e:
+        print(f"登录失败: {e}")
+        sys.exit(1)
+
+
+def cmd_sources(args: argparse.Namespace) -> None:
+    client = _get_client()
+    try:
+        sources = client.list_sources(idc=args.idc)
+        if not sources:
+            print("无数据源")
+            return
+
+        print(f"\n{'IDC':<15} {'名称':<25} {'Source ID'}")
+        print("-" * 80)
+        for src in sources:
+            if isinstance(src, dict):
+                print(
+                    f"{src.get('idc', ''):<15} "
+                    f"{src.get('source', ''):<25} "
+                    f"{src.get('source_id', '')}"
+                )
+            else:
+                print(f"  {src}")
+    except Exception as e:
+        print(f"错误: {e}")
+
+
+def cmd_databases(args: argparse.Namespace) -> None:
+    client = _get_client()
+    try:
+        source_id = _resolve_source_id(client, args.source_id)
+        dbs = client.list_databases(source_id)
+        if not dbs:
+            print("(无数据库)")
+            return
+
+        print("\n数据库列表:")
+        for db in dbs:
+            print(f"  {db}")
+    except Exception as e:
+        print(f"错误: {e}")
+
+
+def cmd_tables(args: argparse.Namespace) -> None:
+    client = _get_client()
+    try:
+        source_id = _resolve_source_id(client, args.source_id)
+        tables = client.list_tables(source_id, args.database)
+        if not tables:
+            print("(无表)")
+            return
+
+        print(f"\n表列表 ({args.database}):")
+        for t in tables:
+            print(f"  {t}")
+    except Exception as e:
+        print(f"错误: {e}")
+
+
+def _export_result(result: dict, database: str, fmt: str) -> None:
+    """Export query result to ~/Downloads/{database}_{timestamp}.{fmt}."""
+    from datetime import datetime
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"{database}_{ts}.{fmt}"
+    filepath = os.path.join(os.path.expanduser("~"), "Downloads", filename)
+
+    if fmt == "json":
+        import json
+        data = {
+            "columns": result["columns"],
+            "rows": [dict(zip(result["columns"], row)) for row in result["rows"]],
+            "query_time_ms": result["query_time"],
+            "row_count": len(result["rows"]),
+        }
+        with open(filepath, "w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+    elif fmt == "csv":
+        import csv
+        with open(filepath, "w", newline="", encoding="utf-8-sig") as f:
+            writer = csv.writer(f)
+            writer.writerow(result["columns"])
+            writer.writerows(result["rows"])
+    elif fmt == "xlsx":
+        try:
+            import openpyxl
+        except ImportError:
+            print("错误: 导出 xlsx 需要 openpyxl，请运行: pip install openpyxl")
+            sys.exit(1)
+        wb = openpyxl.Workbook()
+        ws = wb.active
+        ws.append(result["columns"])
+        for row in result["rows"]:
+            ws.append(row)
+        wb.save(filepath)
+
+    print(f"已导出 {len(result['rows'])} 行到 {filepath}")
+
+
+def cmd_query(args: argparse.Namespace) -> None:
+    if not args.source_id:
+        print("错误: 未指定数据源。请在 ~/.sqlrc 配置 source 或在命令行指定")
+        sys.exit(1)
+    if not args.database:
+        print("错误: 未指定数据库。请在 ~/.sqlrc 配置 database 或在命令行指定")
+        sys.exit(1)
+    client = _get_client()
+    try:
+        source_id = _resolve_source_id(client, args.source_id)
+
+        # Auto-create query order before executing
+        print("正在创建查询工单...", end="", flush=True)
+        if not client.create_query_order(source_id):
+            print(" 失败")
+            print("错误: 无法创建查询工单，请先在 Web 界面确认查询权限")
+            return
+        print(" 成功")
+
+        result = client.execute_query(
+            source_id, args.database, args.sql, timeout=args.timeout
+        )
+        if not result["success"]:
+            print(f"错误: {result['error']}")
+            return
+
+        # Export to file if requested
+        if args.export is not None:
+            fmt = args.export.lstrip(".")
+            if fmt not in ("csv", "json", "xlsx"):
+                print(f"错误: 不支持的格式 '{fmt}'，请使用 csv, json 或 xlsx")
+                return
+            _export_result(result, args.database, fmt)
+            return
+
+        if args.json:
+            import json
+
+            output = {
+                "columns": result["columns"],
+                "rows": [
+                    dict(zip(result["columns"], row)) for row in result["rows"]
+                ],
+                "query_time_ms": result["query_time"],
+                "row_count": len(result["rows"]),
+            }
+            print(json.dumps(output, ensure_ascii=False, indent=2))
+        else:
+            _format_table(result["columns"], result["rows"], result["query_time"])
+    except Exception as e:
+        print(f"错误: {e}")
+
+
+def _resolve_source_id(client: YearningClient, name_or_id: str) -> str:
+    """Resolve source name to source_id, or return as-is if already a UUID."""
+    if len(name_or_id) > 30 and "-" in name_or_id:
+        return name_or_id
+    sources = client.list_sources()
+    for src in sources:
+        if not isinstance(src, dict):
+            continue
+        if src.get("source") == name_or_id or src.get("source_id") == name_or_id:
+            return src["source_id"]
+    return name_or_id
+
+
+def cmd_shell(args: argparse.Namespace) -> None:
+    client = _get_client()
+    source_id = _resolve_source_id(client, args.source_id)
+    current_db = args.database
+
+    # 创建查询工单
+    print("正在创建查询工单...")
+    if not client.create_query_order(source_id):
+        print("警告: 查询工单创建失败，可能无法执行查询")
+
+    print(f"\nYearning SQL Shell")
+    print(f"数据源: {args.source_id} ({source_id[:12]}...) | 数据库: {current_db}")
+    print("输入 SQL 执行 | \\s 切换库 | \\d 查看库 | \\t 查看表 | exit 退出\n")
+
+    while True:
+        try:
+            sql = input(f"sql({current_db})> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nBye!")
+            break
+
+        if not sql:
+            continue
+        if sql.lower() in ("exit", "quit", "\\q"):
+            print("Bye!")
+            break
+
+        if sql.lower().startswith("\\s "):
+            current_db = sql[3:].strip()
+            print(f"  切换到: {current_db}")
+            continue
+
+        if sql.lower() in ("\\d", "\\s"):
+            try:
+                dbs = client.list_databases(source_id)
+                for db in dbs:
+                    print(f"  {db}")
+            except Exception as e:
+                print(f"  错误: {e}")
+            continue
+
+        if sql.lower() == "\\t":
+            try:
+                tables = client.list_tables(source_id, current_db)
+                for t in tables:
+                    print(f"  {t}")
+            except Exception as e:
+                print(f"  错误: {e}")
+            continue
+
+        # multi-line support
+        while not sql.endswith(";"):
+            try:
+                line = input("    ...> ").strip()
+                if line.lower() in ("exit", "quit"):
+                    break
+                sql += " " + line
+            except (EOFError, KeyboardInterrupt):
+                break
+
+        sql = sql.rstrip(";").strip()
+        if not sql:
+            continue
+
+        try:
+            result = client.execute_query(source_id, current_db, sql)
+            if not result["success"]:
+                print(f"错误: {result['error']}")
+                continue
+            _format_table(result["columns"], result["rows"], result["query_time"])
+        except Exception as e:
+            print(f"错误: {e}")
+
+
+def cmd_proxy(args: argparse.Namespace) -> None:
+    """Start MySQL protocol proxy server."""
+    level = logging.DEBUG if VERBOSE else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s [%(levelname)s] %(message)s",
+        datefmt="%H:%M:%S",
+    )
+
+    config = {
+        "base_url": args.url or BASE_URL,
+        "source_name": args.source or DEFAULT_SOURCE,
+        "source_id": "",
+        "database": args.database or DEFAULT_DATABASE,
+        "user": args.user or "",
+        "password": args.password or "",
+    }
+
+    serve(host=args.host, port=args.port, config=config)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="sql",
+        description="Yearning MySQL Audit Platform CLI",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"%(prog)s {get_version()}",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    # login
+    p_login = sub.add_parser("login", help="登录 Yearning")
+
+    # sources
+    p_src = sub.add_parser("sources", help="列出数据源")
+    p_src.add_argument("--idc", default="", help="按 IDC 筛选")
+
+    # databases
+    p_db = sub.add_parser("databases", help="列出数据库")
+    p_db.add_argument("source_id", nargs="?", default=DEFAULT_SOURCE,
+                       help=f"数据源 ID (默认: {DEFAULT_SOURCE or '未配置'})")
+
+    # tables
+    p_tbl = sub.add_parser("tables", help="列出表")
+    p_tbl.add_argument("source_id", nargs="?", default=DEFAULT_SOURCE,
+                        help=f"数据源 ID (默认: {DEFAULT_SOURCE or '未配置'})")
+    p_tbl.add_argument("database", nargs="?", default=DEFAULT_DATABASE,
+                        help=f"数据库名 (默认: {DEFAULT_DATABASE or '未配置'})")
+
+    # query
+    p_q = sub.add_parser("query", help="执行 SQL 查询")
+    p_q.add_argument("sql", help="SQL 语句")
+    p_q.add_argument("source_id", nargs="?", default=DEFAULT_SOURCE,
+                      help=f"数据源 ID 或名称 (默认: {DEFAULT_SOURCE or '未配置'})")
+    p_q.add_argument("database", nargs="?", default=DEFAULT_DATABASE,
+                      help=f"数据库名 (默认: {DEFAULT_DATABASE or '未配置'})")
+    p_q.add_argument("--json", action="store_true", help="以 JSON 格式输出")
+    p_q.add_argument("--timeout", type=int, default=30, help="查询超时秒数 (默认: 30)")
+    p_q.add_argument("-o", "--export", nargs="?", const="xlsx", metavar="FMT",
+                      help="导出结果到 ~/Downloads/ (格式: xlsx/csv/json, 默认 xlsx)")
+
+    # shell
+    p_sh = sub.add_parser("shell", help="交互式 SQL Shell")
+    p_sh.add_argument("source_id", nargs="?", default=DEFAULT_SOURCE,
+                       help=f"数据源 ID (默认: {DEFAULT_SOURCE or '未配置'})")
+    p_sh.add_argument("database", nargs="?", default=DEFAULT_DATABASE,
+                       help=f"数据库名 (默认: {DEFAULT_DATABASE or '未配置'})")
+
+    # proxy
+    p_pr = sub.add_parser("proxy", help="启动 MySQL 协议代理服务器")
+    p_pr.add_argument("--host", default="0.0.0.0", help="绑定地址 (默认: 0.0.0.0；IPv6 可用 ::)")
+    p_pr.add_argument("--port", type=int, default=3307, help="监听端口 (默认: 3307)")
+    p_pr.add_argument("--source", default="", help="数据源名称或 ID")
+    p_pr.add_argument("--database", default="", help="默认数据库名")
+    p_pr.add_argument("--url", default="", help="Yearning URL")
+    p_pr.add_argument("--user", default="", help="要求认证的用户名")
+    p_pr.add_argument("--password", default="", help="要求的认证密码")
+
+    args = parser.parse_args()
+
+    dispatch = {
+        "login": cmd_login,
+        "sources": cmd_sources,
+        "databases": cmd_databases,
+        "tables": cmd_tables,
+        "query": cmd_query,
+        "shell": cmd_shell,
+        "proxy": cmd_proxy,
+    }
+    try:
+        dispatch[args.command](args)
+    except Exception as e:
+        print(f"错误: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()