agent-skills/skills/g3fo-db-ops/scripts/db_ops.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
G3FO Database Operations Tool

替代 MCP MySQL，提供完整的数据库操作能力：
- 增删改查 (query, execute)
- 列举所有表 (list_tables)
- 执行存储过程 (call_procedure)
- 批量执行 SQL (batch_execute)
- 读取表结构 (describe_table)
- 读取/修改表和列注释 (get/set comments)

默认连接配置与 MCP user-mysql 一致，可通过参数或环境变量覆盖。
"""

import argparse
import json
import os
import sys
from typing import Any, Optional
from datetime import date, datetime, time
from decimal import Decimal

try:
    import mysql.connector
    from mysql.connector import Error as MySQLError
except ImportError:
    print("Error: mysql-connector-python is required. Run: pip install mysql-connector-python", file=sys.stderr)
    sys.exit(1)

# 默认连接配置（与 MCP user-mysql 一致）
DEFAULT_CONFIG = {
    "host": "192.168.3.233",
    "port": 3306,
    "user": "root",
    "password": "afe123456",
    "database": "g3fo_base",
    "connection_timeout": 10,  # 连接超时（秒），防止长时间挂起
}

ENV_MAPPING = {
    "host": "MYSQL_HOST",
    "port": "MYSQL_PORT",
    "user": "MYSQL_USER",
    "password": "MYSQL_PASSWORD",
    "database": "MYSQL_DATABASE",
    "connection_timeout": "MYSQL_CONNECTION_TIMEOUT",
}


def get_connection_config(
    host: Optional[str] = None,
    port: Optional[int] = None,
    user: Optional[str] = None,
    password: Optional[str] = None,
    database: Optional[str] = None,
    connection_timeout: Optional[int] = None,
) -> dict:
    """从默认配置、环境变量、参数中合并连接配置，参数优先级最高。"""
    config = {}
    params = {"host": host, "port": port, "user": user, "password": password, "database": database, "connection_timeout": connection_timeout}
    for key, default in DEFAULT_CONFIG.items():
        env_key = ENV_MAPPING.get(key)
        env_val = os.environ.get(env_key) if env_key else None
        param_val = params.get(key)
        if param_val is not None:
            config[key] = int(param_val) if key in ("port", "connection_timeout") else param_val
        elif env_val is not None:
            config[key] = int(env_val) if key in ("port", "connection_timeout") else env_val
        else:
            config[key] = default
    return config


def get_connection(config: dict):
    """
    创建数据库连接。
    连接仅在单次命令执行期间存在，命令结束或异常时在 finally 中关闭，不会长期占用。
    """
    return mysql.connector.connect(
        host=config["host"],
        port=config["port"],
        user=config["user"],
        password=config["password"],
        database=config["database"],
        charset="utf8mb4",
        collation="utf8mb4_unicode_ci",
        connection_timeout=config.get("connection_timeout", 10),
    )


def _normalize_value(value: Any) -> Any:
    """将 MySQL 返回的值转换为可 JSON 序列化的类型。"""
    if isinstance(value, (datetime, date, time)):
        return value.isoformat()
    if isinstance(value, Decimal):
        # 大多数场景下用 float 即可，避免 JSON 不支持 Decimal
        return float(value)
    return value


def cmd_query(args: argparse.Namespace, config: dict) -> dict:
    """执行 SELECT 查询。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(args.sql, args.params or [])
        rows = cursor.fetchall()
        cursor.close()
        normalized_rows = [
            {k: _normalize_value(v) for k, v in row.items()} for row in rows
        ]
        return {"success": True, "data": normalized_rows, "rowCount": len(normalized_rows)}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_execute(args: argparse.Namespace, config: dict) -> dict:
    """执行 INSERT/UPDATE/DELETE。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        cursor.execute(args.sql, args.params or [])
        conn.commit()
        affected = cursor.rowcount
        cursor.close()
        return {"success": True, "affectedRows": affected}
    except MySQLError as e:
        conn.rollback()
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_list_tables(args: argparse.Namespace, config: dict) -> dict:
    """列举数据库中的所有表。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE, TABLE_COMMENT "
            "FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s ORDER BY TABLE_NAME",
            (config["database"],),
        )
        rows = cursor.fetchall()
        cursor.close()
        tables = [
            {
                "schema": r[0],
                "name": r[1],
                "type": r[2],
                "comment": r[3] or "",
            }
            for r in rows
        ]
        return {"success": True, "tables": tables, "count": len(tables)}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_describe_table(args: argparse.Namespace, config: dict) -> dict:
    """获取表结构（列信息）。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        table = args.table
        schema = args.schema or config["database"]
        cursor.execute(
            """
            SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY, COLUMN_DEFAULT, COLUMN_COMMENT,
                   EXTRA
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
            ORDER BY ORDINAL_POSITION
            """,
            (schema, table),
        )
        rows = cursor.fetchall()
        cursor.close()
        columns = [
            {
                "name": r[0],
                "type": r[1],
                "nullable": r[2],
                "key": r[3] or "",
                "default": r[4],
                "comment": r[5] or "",
                "extra": r[6] or "",
            }
            for r in rows
        ]
        # 获取表注释
        cursor = conn.cursor()
        cursor.execute(
            "SELECT TABLE_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s",
            (schema, table),
        )
        tbl = cursor.fetchone()
        table_comment = tbl[0] if tbl else ""
        cursor.close()
        return {"success": True, "table": table, "schema": schema, "columns": columns, "tableComment": table_comment}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_call_procedure(args: argparse.Namespace, config: dict) -> dict:
    """执行存储过程。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        placeholders = ", ".join(["%s"] * len(args.params)) if args.params else ""
        sql = f"CALL {args.procedure}({placeholders})" if placeholders else f"CALL {args.procedure}()"
        cursor.execute(sql, args.params or [])
        rows = []
        if cursor.description:
            rows = cursor.fetchall()
        while cursor.nextset():
            if cursor.description:
                rows.extend(cursor.fetchall())
        conn.commit()
        cursor.close()
        return {"success": True, "data": rows, "rowCount": len(rows)}
    except MySQLError as e:
        conn.rollback()
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_batch_execute(args: argparse.Namespace, config: dict) -> dict:
    """批量执行 SQL（支持多条语句、存储过程、DELIMITER）。"""
    sql_text = args.sql or ""
    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            sql_text = f.read()
    if not sql_text.strip():
        return {"success": False, "error": "No SQL provided. Use 'sql' argument or --file"}
    # 处理 DELIMITER（存储过程等）：移除 DELIMITER 行，将 $$ 替换为 ;
    if "DELIMITER" in sql_text.upper():
        lines = []
        for line in sql_text.split("\n"):
            if line.strip().upper().startswith("DELIMITER"):
                continue
            lines.append(line.replace("$$", ";"))
        sql_text = "\n".join(lines)
    conn = get_connection(config)
    results = []
    try:
        cursor = conn.cursor()
        # mysql-connector-python 9.2+ 已移除 multi 参数，直接 execute 即可执行多条语句
        cursor.execute(sql_text)
        # 消费所有结果集（SELECT 返回行、DML 返回 affected rows）
        while True:
            if cursor.description:
                results.append({"type": "query", "rowCount": len(cursor.fetchall())})
            else:
                results.append({"type": "execute", "affectedRows": cursor.rowcount})
            if not cursor.nextset():
                break
        conn.commit()
        cursor.close()
        return {"success": True, "results": results}
    except MySQLError as e:
        conn.rollback()
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_get_table_comment(args: argparse.Namespace, config: dict) -> dict:
    """获取表注释。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        schema = args.schema or config["database"]
        cursor.execute(
            "SELECT TABLE_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s",
            (schema, args.table),
        )
        row = cursor.fetchone()
        cursor.close()
        return {"success": True, "table": args.table, "schema": schema, "comment": row[0] if row else ""}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_set_table_comment(args: argparse.Namespace, config: dict) -> dict:
    """修改表注释。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        schema = args.schema or config["database"]
        sql = f"ALTER TABLE `{schema}`.`{args.table}` COMMENT = %s"
        cursor.execute(sql, (args.comment,))
        conn.commit()
        cursor.close()
        return {"success": True, "table": args.table, "schema": schema}
    except MySQLError as e:
        conn.rollback()
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_get_column_comment(args: argparse.Namespace, config: dict) -> dict:
    """获取列注释。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        schema = args.schema or config["database"]
        cursor.execute(
            "SELECT COLUMN_COMMENT FROM information_schema.COLUMNS "
            "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND COLUMN_NAME = %s",
            (schema, args.table, args.column),
        )
        row = cursor.fetchone()
        cursor.close()
        return {"success": True, "table": args.table, "column": args.column, "comment": row[0] if row else ""}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_set_column_comment(args: argparse.Namespace, config: dict) -> dict:
    """修改列注释（需要 COLUMN_TYPE，可通过 describe_table 获取）。"""
    conn = get_connection(config)
    try:
        cursor = conn.cursor()
        schema = args.schema or config["database"]
        cursor.execute(
            "SELECT COLUMN_TYPE FROM information_schema.COLUMNS "
            "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND COLUMN_NAME = %s",
            (schema, args.table, args.column),
        )
        row = cursor.fetchone()
        if not row:
            cursor.close()
            return {"success": False, "error": f"Column {args.column} not found"}
        col_type = row[0]
        sql = f"ALTER TABLE `{schema}`.`{args.table}` MODIFY COLUMN `{args.column}` {col_type} COMMENT %s"
        cursor.execute(sql, (args.comment,))
        conn.commit()
        cursor.close()
        return {"success": True, "table": args.table, "column": args.column}
    except MySQLError as e:
        conn.rollback()
        return {"success": False, "error": str(e)}
    finally:
        conn.close()


def cmd_test_connection(args: argparse.Namespace, config: dict) -> dict:
    """测试数据库连接。"""
    conn = None
    try:
        conn = get_connection(config)
        cursor = conn.cursor()
        cursor.execute("SELECT 1 AS conn_ok")
        row = cursor.fetchone()
        cursor.close()
        return {"success": True, "message": "Connection OK", "conn_ok": row[0] if row else 1}
    except MySQLError as e:
        return {"success": False, "error": str(e)}
    finally:
        if conn:
            conn.close()


def main():
    parser = argparse.ArgumentParser(description="G3FO Database Operations Tool")
    parser.add_argument("--host", help="MySQL host (default: 192.168.3.233)")
    parser.add_argument("--port", type=int, help="MySQL port (default: 3306)")
    parser.add_argument("--user", help="MySQL user (default: root)")
    parser.add_argument("--password", help="MySQL password")
    parser.add_argument("--database", help="MySQL database (default: g3fo_base)")
    parser.add_argument("--connection-timeout", type=int, dest="connection_timeout", help="Connection timeout in seconds (default: 10)")


    subparsers = parser.add_subparsers(dest="command", required=True)

    # query
    p_query = subparsers.add_parser("query", help="Execute SELECT query")
    p_query.add_argument("sql", help="SQL SELECT statement")
    p_query.add_argument("--params", nargs="*", help="Query parameters")

    # execute
    p_execute = subparsers.add_parser("execute", help="Execute INSERT/UPDATE/DELETE")
    p_execute.add_argument("sql", help="SQL statement")
    p_execute.add_argument("--params", nargs="*", help="Query parameters")

    # list_tables
    subparsers.add_parser("list_tables", help="List all tables in database")

    # describe_table
    p_desc = subparsers.add_parser("describe_table", help="Get table structure")
    p_desc.add_argument("table", help="Table name")
    p_desc.add_argument("--schema", help="Schema/database (default: current database)")

    # call_procedure
    p_proc = subparsers.add_parser("call_procedure", help="Execute stored procedure")
    p_proc.add_argument("procedure", help="Procedure name (e.g. InsertSystemMenu)")
    p_proc.add_argument("--params", nargs="*", help="Procedure parameters")

    # batch_execute
    p_batch = subparsers.add_parser("batch_execute", help="Execute batch SQL")
    p_batch.add_argument("sql", nargs="?", help="SQL text (multiple statements separated by ;)")
    p_batch.add_argument("--file", "-f", help="SQL file path")
    p_batch.add_argument("--continue-on-error", action="store_true", dest="continue_on_error", help="Continue on error")

    # get_table_comment
    p_gtc = subparsers.add_parser("get_table_comment", help="Get table comment")
    p_gtc.add_argument("table", help="Table name")
    p_gtc.add_argument("--schema", help="Schema (default: current database)")

    # set_table_comment
    p_stc = subparsers.add_parser("set_table_comment", help="Set table comment")
    p_stc.add_argument("table", help="Table name")
    p_stc.add_argument("comment", help="Comment text")
    p_stc.add_argument("--schema", help="Schema (default: current database)")

    # get_column_comment
    p_gcc = subparsers.add_parser("get_column_comment", help="Get column comment")
    p_gcc.add_argument("table", help="Table name")
    p_gcc.add_argument("column", help="Column name")
    p_gcc.add_argument("--schema", help="Schema (default: current database)")

    # set_column_comment
    p_scc = subparsers.add_parser("set_column_comment", help="Set column comment")
    p_scc.add_argument("table", help="Table name")
    p_scc.add_argument("column", help="Column name")
    p_scc.add_argument("comment", help="Comment text")
    p_scc.add_argument("--schema", help="Schema (default: current database)")

    # test_connection
    subparsers.add_parser("test_connection", help="Test database connection")

    args = parser.parse_args()
    config = get_connection_config(
        host=args.host,
        port=args.port,
        user=args.user,
        password=args.password,
        database=args.database,
        connection_timeout=getattr(args, "connection_timeout", None),
    )

    commands = {
        "query": cmd_query,
        "execute": cmd_execute,
        "list_tables": cmd_list_tables,
        "describe_table": cmd_describe_table,
        "call_procedure": cmd_call_procedure,
        "batch_execute": cmd_batch_execute,
        "get_table_comment": cmd_get_table_comment,
        "set_table_comment": cmd_set_table_comment,
        "get_column_comment": cmd_get_column_comment,
        "set_column_comment": cmd_set_column_comment,
        "test_connection": cmd_test_connection,
    }
    handler = commands[args.command]
    result = handler(args, config)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    sys.exit(0 if result.get("success", False) else 1)


if __name__ == "__main__":
    main()