test_conn.py

"""
数据库连接测试工具 - 测试各类数据库的连接是否正常

用法：
    # 交互模式：选择数据库类型 → 实例 → 用户 → 测试连接
    python test_conn.py

    # 直接指定参数
    python test_conn.py dm8 local SYSDBA
    python test_conn.py dm8 local USER_ZWWW_BQGL
    python test_conn.py dm8 local USER_ZWWW_BQGL --sql "SELECT 1 FROM DUAL"
    python test_conn.py dm8 local USER_ZWWW_BQGL --sql "SELECT * FROM ZCZQ_TAG_INFO" --show
"""

import sys
import argparse

import config as cfg
from db_utils import validate_sql_safe


def print_header(title: str):
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def select_from_list(items: list, title: str) -> str:
    """交互式菜单选择：展示编号列表，用户输入序号返回对应项"""
    if not items:
        print(f"  ⚠️  没有可选的{title}")
        return None
    print(f"\n可用的{title}:")
    for i, item in enumerate(items, 1):
        print(f"  [{i}] {item}")
    while True:
        try:
            choice = input(f"\n请选择{title} (1-{len(items)}): ").strip()
            idx = int(choice) - 1
            if 0 <= idx < len(items):
                return items[idx]
            print(f"  ⚠️  请输入 1-{len(items)} 之间的数字")
        except ValueError:
            print("  ⚠️  请输入有效的数字")


def interactive_mode():
    """交互模式：逐步选择并测试"""
    print_header("🔌 数据库连接测试工具（交互模式）")

    # 1. 选择数据库类型
    db_types = cfg.list_db_types()
    if not db_types:
        print("\n⚠️  未在 .env 中配置任何数据库")
        return

    db_type = select_from_list(db_types, "数据库类型")
    if not db_type:
        return

    # 2. 选择实例
    instances = cfg.list_instances(db_type)
    instance = select_from_list(instances, f"{db_type} 实例")
    if not instance:
        return

    # 3. 选择用户
    users = cfg.list_users(db_type, instance)
    user = select_from_list(users, f"{db_type}:{instance} 用户")
    if not user:
        return

    # 4. 测试连接
    test_connection(db_type, instance, user)


def test_connection(db_type: str, instance: str, user: str, sql: str = None, show: bool = False):
    """测试指定目标的数据库连接"""
    print_header(f"🔌 测试连接: [{db_type}:{instance}:{user}]")

    # 显示连接参数
    try:
        inst_cfg = cfg.get_instance_config(db_type, instance)
        user_cred = cfg.get_user_credential(db_type, instance, user)
    except ValueError as e:
        print(f"  ❌ 配置错误: {e}")
        return

    print(f"  数据库类型: {db_type}")
    print(f"  实例:       {instance}")
    print(f"  主机:       {inst_cfg.get('host', '-')}")
    print(f"  端口:       {inst_cfg.get('port', '-')}")
    print(f"  用户名:     {user_cred['user']}")
    print(f"  密码:       {'*' * 8}")

    # 测试连接
    print(f"\n  正在连接...")
    conn = None
    try:
        conn = cfg.create_connection(db_type, instance, user)
        cursor = conn.cursor()

        # 不同数据库的测试查询
        test_sqls = {
            "dm8": "SELECT '连接成功' AS result FROM DUAL",
            "mysql": "SELECT '连接成功' AS result",
            "mariadb": "SELECT '连接成功' AS result",
            "pgsql": "SELECT '连接成功' AS result",
            "oracle": "SELECT '连接成功' AS result FROM DUAL",
            "sqlserver": "SELECT '连接成功' AS result",
            "sqlite": "SELECT '连接成功' AS result",
            "kingbase": "SELECT '连接成功' AS result",
            "gaussdb": "SELECT '连接成功' AS result",
            "opengauss": "SELECT '连接成功' AS result",
            "polardb": "SELECT '连接成功' AS result",
            "oceanbase": "SELECT '连接成功' AS result",
            "tidb": "SELECT '连接成功' AS result",
            "gbase": "SELECT '连接成功' AS result",
        }

        test_sql = test_sqls.get(db_type, "SELECT 1")
        cursor.execute(test_sql)
        result = cursor.fetchone()
        print(f"  ✅ 连接成功！测试查询: {result}")

        # 如果指定了 SQL，执行查询
        if sql:
            execute_query(cursor, sql, show)

        cursor.close()

    except Exception as e:
        print(f"  ❌ 连接失败: {e}")
    finally:
        if conn:
            try:
                conn.close()
                print(f"  ℹ️  连接已关闭")
            except Exception:
                pass


def execute_query(cursor, sql: str, show: bool):
    """在已有游标上执行 SQL 并显示结果（可选预览）"""
    import pandas as pd

    # SQL 安全校验
    try:
        validate_sql_safe(sql)
    except ValueError as e:
        print(f"  ❌ 安全校验失败: {e}")
        return

    print(f"\n  📝 执行 SQL: {sql[:80]}{'...' if len(sql) > 80 else ''}")
    try:
        cursor.execute(sql)
        columns = [desc[0] for desc in cursor.description]
        data = cursor.fetchall()
        df = pd.DataFrame(data, columns=columns)

        print(f"  ✅ 查询成功！共 {len(df)} 条记录，{len(df.columns)} 列")

        if show and len(df) > 0:
            print(f"\n  结果预览（前10条）:")
            print(f"  {'─' * 66}")
            for _, row in df.head(10).iterrows():
                print(f"  {list(row.values)}")
            print(f"  {'─' * 66}")
            if len(df) > 10:
                print(f"  ... 共 {len(df)} 条，仅显示前10条")

    except Exception as e:
        print(f"  ❌ 查询失败: {e}")


def cli_mode(db_type: str, instance: str, user: str, sql: str = None, show: bool = False):
    """命令行模式：直接测试指定目标"""
    test_connection(db_type, instance, user, sql, show)


def main():
    parser = argparse.ArgumentParser(
        description="数据库连接测试工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python test_conn.py                    # 交互模式
  python test_conn.py dm8 local SYSDBA   # 直接指定参数
        """,
    )
    parser.add_argument("db_type", nargs="?", help="数据库类型（如 dm8, mysql, pgsql）")
    parser.add_argument("instance", nargs="?", help="实例名")
    parser.add_argument("user", nargs="?", help="用户名")
    parser.add_argument("--sql", help="要执行的 SQL 语句（可选）")
    parser.add_argument("--show", action="store_true", help="显示查询结果预览")

    args = parser.parse_args()

    if args.db_type and args.instance and args.user:
        # 命令行模式
        cli_mode(args.db_type, args.instance, args.user, args.sql, args.show)
    else:
        # 交互模式
        interactive_mode()


if __name__ == "__main__":
    main()