feat(g3fo-db-ops): replace MCP MySQL with Python db_ops tool, add connection timeout

Made-with: Cursor
2026-03-11 09:27:19 +08:00
parent 8e2afda170
commit 23df8ffd75
5 changed files with 634 additions and 26 deletions
@@ -0,0 +1,462 @@
+#!/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
+
+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 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()
+        return {"success": True, "data": rows, "rowCount": len(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()