463 lines
16 KiB
Python
463 lines
16 KiB
Python
#!/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()
|