query_runner.py

"""
查询执行工具 — SQL保存 → 执行 → 结果保存 一体化

遵循 AGENTS.md 命名规范：
  最终查询模式（默认）:
    SQL 文件   → sql/queryYYYYMMDD_HHMMSS.sql
    结果文件   → export_dir/queryYYYYMMDD_HHMMSS.out
    Excel 文件 → export_dir/查询内容_YYYYMMDD_HHMMSS.xlsx
  探索模式 (--dry-run):
    仅输出到控制台，不保存任何文件

用法：
    python query_runner.py dm8 local USER_ZWWW_BQGL "SELECT * FROM ZCZQ_TAG_INFO"
    python query_runner.py dm8 local USER_ZWWW_BQGL "SELECT * FROM tab" --dry-run
"""

import sys
import os
import argparse

# 确保能导入项目模块
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import config as cfg
from db_utils import execute_sql, export_to_excel, close_connection, validate_sql_safe


def save_sql_file(sql: str, timestamp: str) -> str:
    """保存 SQL 到 sql/ 目录，写入时间戳和连接目标信息"""
    os.makedirs(cfg.SQL_DIR, exist_ok=True)
    filename = f"query{timestamp}.sql"
    filepath = os.path.join(cfg.SQL_DIR, filename)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(f"-- {filename}\n")
        f.write(f"-- 生成时间: {timestamp}\n")
        f.write(f"-- 目标: [{cfg.DEFAULT_DB_TYPE}:{cfg.DEFAULT_INSTANCE}:{cfg.DEFAULT_USER}]\n\n")
        f.write(sql.strip() + "\n")
    return filepath


def save_result_file(content: str, timestamp: str) -> str:
    """保存查询结果到 export_dir/ 目录"""
    os.makedirs(cfg.EXPORT_DIR, exist_ok=True)
    filename = f"query{timestamp}.out"
    filepath = os.path.join(cfg.EXPORT_DIR, filename)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)
    return filepath


def run_query(db_type: str, instance: str, user: str, sql: str,
              show: bool = True, export_excel: bool = False, dry_run: bool = False):
    """执行查询并保存结果"""

    ts = cfg.timestamp_str()

    # 1. 打印规范
    cfg.print_rules()

    print("=" * 70)
    mode_label = "🔍 探索查询 (--dry-run)" if dry_run else "🔍 查询执行"
    print(f"{mode_label}: [{db_type}:{instance}:{user}]")
    print(f"   时间戳: {ts}")
    print("=" * 70)

    # 2. SQL 安全校验（在保存前拒绝危险 SQL）
    try:
        validate_sql_safe(sql)
    except ValueError as e:
        print(f"   ❌ 安全校验失败: {e}")
        sys.exit(1)

    # 3. 探索模式不保存 SQL
    if not dry_run:
        sql_path = save_sql_file(sql, ts)
        print(f"\n📄 SQL 已保存: {sql_path}")
    else:
        print(f"\n📄 SQL (探索模式，不保存): {sql[:80]}{'...' if len(sql) > 80 else ''}")

    # 4. 连接并执行
    print(f"\n⚡ 正在执行 SQL...")
    try:
        df = execute_sql(sql)
        print(f"   ✅ 查询成功！共 {len(df)} 条记录，{len(df.columns)} 列")
    except Exception as e:
        print(f"   ❌ 执行失败: {e}")
        # 最终模式才记录错误文件
        if not dry_run:
            err_content = f"查询执行失败\n时间戳: {ts}\nSQL: {sql}\n错误: {e}\n"
            result_path = save_result_file(err_content, ts)
            print(f"   📄 错误已记录: {result_path}")
        sys.exit(1)
    finally:
        close_connection()

    # 4. 构建结果内容
    result_lines = []
    result_lines.append("=" * 70)
    result_lines.append(f"查询结果")
    result_lines.append(f"时间戳: {ts}")
    result_lines.append(f"目标: [{db_type}:{instance}:{user}]")
    result_lines.append(f"记录数: {len(df)} 行, {len(df.columns)} 列")
    result_lines.append("=" * 70)
    result_lines.append("")

    # 表头
    header = " | ".join(str(c) for c in df.columns)
    result_lines.append(header)
    result_lines.append("-" * len(header))

    # 数据
    for _, row in df.iterrows():
        result_lines.append(" | ".join(str(v) for v in row.values))

    result_lines.append("")
    result_lines.append(f"共 {len(df)} 条记录")

    result_content = "\n".join(result_lines)

    # 5. 显示到控制台
    if show:
        print(f"\n📊 结果:")
        print("-" * 70)
        print(result_content)
        print("-" * 70)

    # 6. 最终模式才保存结果文件和 Excel
    if not dry_run:
        result_path = save_result_file(result_content, ts)
        print(f"\n📄 结果已保存: {result_path}")

        if export_excel:
            excel_path = export_to_excel(df, filename=f"查询结果_{ts}.xlsx")
            print(f"📊 Excel 已导出: {excel_path}")
    else:
        print(f"   (探索模式，不保存文件)")

    print(f"\n{'=' * 70}")
    print("✅ 执行完成" if not dry_run else "✅ 探索完成 (未保存文件)")
    print("=" * 70)


def main():
    parser = argparse.ArgumentParser(
        description="查询执行工具 — SQL保存→执行→结果保存 一体化",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python query_runner.py dm8 local USER_ZWWW_BQGL "SELECT * FROM ZCZQ_TAG_INFO"
  python query_runner.py dm8 local SYSDBA "SELECT COUNT(*) FROM DUAL" --excel
  python query_runner.py dm8 local USER_ZWWW_BQGL "SELECT * FROM tab" --dry-run
        """,
    )
    parser.add_argument("db_type", help="数据库类型（如 dm8, mysql, pgsql）")
    parser.add_argument("instance", help="实例名")
    parser.add_argument("user", help="用户名")
    parser.add_argument("sql", help="要执行的 SQL 语句")
    parser.add_argument("--excel", action="store_true", help="同时导出 Excel")
    parser.add_argument("--quiet", action="store_true", help="安静模式，不打印结果")
    parser.add_argument("--dry-run", action="store_true", help="探索模式：仅输出到控制台，不保存文件")

    args = parser.parse_args()
    run_query(args.db_type, args.instance, args.user, args.sql,
              show=not args.quiet, export_excel=args.excel, dry_run=args.dry_run)


if __name__ == "__main__":
    main()