# -*- coding: utf-8 -*-
"""
MySQL 连接测试脚本
====================
测试目标 MySQL server 能否连接、登录，并读取表。

环境要求：
    本脚本依赖 PyMySQL，请先在虚拟环境中安装：
        python -m venv .venv
        .venv\\Scripts\\activate          # Windows
        # source .venv/bin/activate      # Linux/Mac
        pip install pymysql

用法：
    python test_mysql_connect.py                          # 使用脚本内置默认连接信息
    python test_mysql_connect.py -h host -P port -u user  # 自定义连接信息

说明：
    - 该账号为只读账号，仅做 SELECT 测试，不做任何写操作。
"""

import argparse
import sys

try:
    import pymysql
    from pymysql.constants import CLIENT
except ImportError:
    print("缺少 PyMySQL，请先执行: pip install pymysql")
    sys.exit(1)

# 默认连接信息（可用命令行参数覆盖）
DEFAULT_HOST = "zkzh.ykz-ai.com"
DEFAULT_PORT = 60031
DEFAULT_USER = "reader"
DEFAULT_PASSWORD = "Reader@123456!"
DEFAULT_DATABASE = None  # 不指定库，先列出所有库

CONNECT_TIMEOUT = 10   # 秒
READ_TIMEOUT = 30      # 秒
MAX_ROWS = 5           # 每个表最多读取的行数


def connect(host, port, user, password, database=None):
    """尝试建立 MySQL 连接，返回 (connection, error_msg)。"""
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset="utf8mb4",
            connect_timeout=CONNECT_TIMEOUT,
            read_timeout=READ_TIMEOUT,
            # 部分服务端为 MariaDB，开启兼容标志避免认证方式报错
            client_flag=CLIENT.MULTI_STATEMENTS,
        )
        return conn, None
    except pymysql.err.OperationalError as e:
        # 主机名解析失败 / 连接超时 / 拒绝访问 都在这类错误里
        return None, "连接失败: {0}".format(e)
    except Exception as e:
        return None, "未知错误: {0}: {1}".format(type(e).__name__, e)


def safe_query(conn, sql, args=None):
    """执行查询，返回 (rows, columns, error_msg)。"""
    try:
        with conn.cursor() as cur:
            cur.execute(sql, args)
            columns = [d[0] for d in cur.description] if cur.description else None
            rows = cur.fetchmany(MAX_ROWS)
            return rows, columns, None
    except Exception as e:
        return None, None, "{0}: {1}".format(type(e).__name__, e)


def main():
    parser = argparse.ArgumentParser(description="MySQL 连接与读取测试")
    parser.add_argument("-H", "--host", default=DEFAULT_HOST, help="主机地址")
    parser.add_argument("-P", "--port", type=int, default=DEFAULT_PORT, help="端口")
    parser.add_argument("-u", "--user", default=DEFAULT_USER, help="用户名")
    parser.add_argument("-p", "--password", default=DEFAULT_PASSWORD, help="密码")
    parser.add_argument("-d", "--database", default=DEFAULT_DATABASE, help="指定数据库（默认列出所有库）")
    args = parser.parse_args()

    print("=" * 60)
    print("MySQL 连接测试")
    print("  地址: {0}:{1}".format(args.host, args.port))
    print("  用户: {0}".format(args.user))
    print("=" * 60)

    # 1. 建立连接
    conn, err = connect(args.host, args.port, args.user, args.password, args.database)
    if conn is None:
        print("\n[失败] {0}".format(err))
        sys.exit(1)

    print("\n[成功] 已连接到 MySQL 服务器")

    # 2. 服务器信息
    rows, cols, err = safe_query(conn, "SELECT VERSION(), NOW(), CURRENT_USER()")
    if err:
        print("[警告] 获取服务器信息失败: {0}".format(err))
    else:
        print("      服务器版本: {0}".format(rows[0][0]))
        print("      服务器时间: {0}".format(rows[0][1]))
        print("      登录身份:   {0}".format(rows[0][2]))

    # 3. 列出数据库
    rows, cols, err = safe_query(conn, "SHOW DATABASES")
    if err:
        print("[警告] 列出数据库失败: {0}".format(err))
    else:
        print("\n[信息] 可访问的数据库({0} 个):".format(len(rows)))
        for r in rows:
            print("       - {0}".format(r[0]))

    # 4. 列出表并读取部分数据
    if args.database:
        db_list = [args.database]
    else:
        db_list = [r[0] for r in rows] if rows else []

    for db in db_list:
        rows, cols, err = safe_query(conn, "SHOW TABLES FROM `{0}`".format(db))
        if err:
            print("\n[失败] 读取数据库 \"{0}\" 的表列表失败: {1}".format(db, err))
            continue
        if not rows:
            print("\n[信息] 数据库 \"{0}\" 没有可访问的表".format(db))
            continue

        print("\n" + "-" * 60)
        print("数据库: {0}  (表数量: {1})".format(db, len(rows)))
        for tbl in rows:
            table_name = tbl[0]
            trows, tcols, terr = safe_query(conn, "SELECT * FROM `{0}`.`{1}` LIMIT {2}".format(db, table_name, MAX_ROWS))
            if terr:
                print("  [失败] 表 [{0}] 读取失败: {1}".format(table_name, terr))
                continue
            print("  [成功] 表 [{0}]  字段={1}  样例数据行数={2}".format(
                table_name, tcols, len(trows) if trows else 0))
            if trows:
                print("         字段: {0}".format(tcols))
                for r in trows[:3]:
                    print("         行:   {0}".format(r))

    conn.close()
    print("\n" + "=" * 60)
    print("测试结束")


if __name__ == "__main__":
    main()