454 lines
16 KiB
Python
454 lines
16 KiB
Python
|
"""MySQL MCP (Model Context Protocol) tool for database operations."""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import pymysql
|
|||
|
|
from typing import List, Dict, Any, Optional
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
from th_agenter.services.agent.base import BaseTool, ToolParameter, ToolParameterType, ToolResult
|
|||
|
|
|
|||
|
|
class MySQLMCPTool(BaseTool):
|
|||
|
|
"""MySQL MCP tool for database operations and intelligent querying."""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
self.connections = {} # 存储用户的数据库连接
|
|||
|
|
|
|||
|
|
def get_name(self) -> str:
|
|||
|
|
return "mysql_mcp"
|
|||
|
|
|
|||
|
|
def get_description(self) -> str:
|
|||
|
|
return "MySQL MCP服务工具，提供数据库连接、表结构查询、SQL执行等功能，支持智能数据问答。"
|
|||
|
|
|
|||
|
|
def get_parameters(self) -> List[ToolParameter]:
|
|||
|
|
return [
|
|||
|
|
ToolParameter(
|
|||
|
|
name="operation",
|
|||
|
|
type=ToolParameterType.STRING,
|
|||
|
|
description="操作类型",
|
|||
|
|
required=True,
|
|||
|
|
enum=["connect", "list_tables", "describe_table", "execute_query", "test_connection", "disconnect"]
|
|||
|
|
),
|
|||
|
|
ToolParameter(
|
|||
|
|
name="connection_config",
|
|||
|
|
type=ToolParameterType.OBJECT,
|
|||
|
|
description="数据库连接配置 {host, port, database, username, password}",
|
|||
|
|
required=False
|
|||
|
|
),
|
|||
|
|
ToolParameter(
|
|||
|
|
name="user_id",
|
|||
|
|
type=ToolParameterType.STRING,
|
|||
|
|
description="用户ID，用于管理连接",
|
|||
|
|
required=False
|
|||
|
|
),
|
|||
|
|
ToolParameter(
|
|||
|
|
name="table_name",
|
|||
|
|
type=ToolParameterType.STRING,
|
|||
|
|
description="表名（用于describe_table操作）",
|
|||
|
|
required=False
|
|||
|
|
),
|
|||
|
|
ToolParameter(
|
|||
|
|
name="sql_query",
|
|||
|
|
type=ToolParameterType.STRING,
|
|||
|
|
description="SQL查询语句（用于execute_query操作）",
|
|||
|
|
required=False
|
|||
|
|
),
|
|||
|
|
ToolParameter(
|
|||
|
|
name="limit",
|
|||
|
|
type=ToolParameterType.INTEGER,
|
|||
|
|
description="查询结果限制数量，默认100",
|
|||
|
|
required=False,
|
|||
|
|
default=100
|
|||
|
|
)
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def _get_tables(self, connection) -> List[Dict[str, Any]]:
|
|||
|
|
"""获取数据库表列表"""
|
|||
|
|
cursor = connection.cursor()
|
|||
|
|
try:
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT
|
|||
|
|
table_name,
|
|||
|
|
table_type,
|
|||
|
|
table_schema
|
|||
|
|
FROM information_schema.tables
|
|||
|
|
WHERE table_schema = DATABASE()
|
|||
|
|
ORDER BY table_name;
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
tables = []
|
|||
|
|
for row in cursor.fetchall():
|
|||
|
|
tables.append({
|
|||
|
|
"table_name": row[0],
|
|||
|
|
"table_type": row[1],
|
|||
|
|
"table_schema": row[2]
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return tables
|
|||
|
|
finally:
|
|||
|
|
cursor.close()
|
|||
|
|
|
|||
|
|
def _describe_table(self, connection, table_name: str) -> Dict[str, Any]:
|
|||
|
|
"""获取表结构信息"""
|
|||
|
|
cursor = connection.cursor()
|
|||
|
|
try:
|
|||
|
|
# 获取列信息
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT
|
|||
|
|
column_name,
|
|||
|
|
data_type,
|
|||
|
|
is_nullable,
|
|||
|
|
column_default,
|
|||
|
|
character_maximum_length,
|
|||
|
|
numeric_precision,
|
|||
|
|
numeric_scale,
|
|||
|
|
column_comment
|
|||
|
|
FROM information_schema.columns
|
|||
|
|
WHERE table_schema = DATABASE() AND table_name = %s
|
|||
|
|
ORDER BY ordinal_position;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
columns = []
|
|||
|
|
for row in cursor.fetchall():
|
|||
|
|
column_info = {
|
|||
|
|
"column_name": row[0],
|
|||
|
|
"data_type": row[1],
|
|||
|
|
"is_nullable": row[2] == 'YES',
|
|||
|
|
"column_default": row[3],
|
|||
|
|
"character_maximum_length": row[4],
|
|||
|
|
"numeric_precision": row[5],
|
|||
|
|
"numeric_scale": row[6],
|
|||
|
|
"column_comment": row[7] or ""
|
|||
|
|
}
|
|||
|
|
columns.append(column_info)
|
|||
|
|
|
|||
|
|
# 获取主键信息
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT column_name
|
|||
|
|
FROM information_schema.key_column_usage
|
|||
|
|
WHERE table_schema = DATABASE()
|
|||
|
|
AND table_name = %s
|
|||
|
|
AND constraint_name = 'PRIMARY'
|
|||
|
|
ORDER BY ordinal_position;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
primary_keys = [row[0] for row in cursor.fetchall()]
|
|||
|
|
|
|||
|
|
# 获取外键信息
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT
|
|||
|
|
column_name,
|
|||
|
|
referenced_table_name,
|
|||
|
|
referenced_column_name
|
|||
|
|
FROM information_schema.key_column_usage
|
|||
|
|
WHERE table_schema = DATABASE()
|
|||
|
|
AND table_name = %s
|
|||
|
|
AND referenced_table_name IS NOT NULL;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
foreign_keys = []
|
|||
|
|
for row in cursor.fetchall():
|
|||
|
|
foreign_keys.append({
|
|||
|
|
"column_name": row[0],
|
|||
|
|
"referenced_table": row[1],
|
|||
|
|
"referenced_column": row[2]
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 获取索引信息
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT
|
|||
|
|
index_name,
|
|||
|
|
column_name,
|
|||
|
|
non_unique
|
|||
|
|
FROM information_schema.statistics
|
|||
|
|
WHERE table_schema = DATABASE()
|
|||
|
|
AND table_name = %s
|
|||
|
|
ORDER BY index_name, seq_in_index;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
indexes = []
|
|||
|
|
for row in cursor.fetchall():
|
|||
|
|
indexes.append({
|
|||
|
|
"index_name": row[0],
|
|||
|
|
"column_name": row[1],
|
|||
|
|
"is_unique": row[2] == 0
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 获取表注释
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT table_comment
|
|||
|
|
FROM information_schema.tables
|
|||
|
|
WHERE table_schema = DATABASE() AND table_name = %s;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
table_comment = ""
|
|||
|
|
result = cursor.fetchone()
|
|||
|
|
if result:
|
|||
|
|
table_comment = result[0] or ""
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"table_name": table_name,
|
|||
|
|
"columns": columns,
|
|||
|
|
"primary_keys": primary_keys,
|
|||
|
|
"foreign_keys": foreign_keys,
|
|||
|
|
"indexes": indexes,
|
|||
|
|
"table_comment": table_comment
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
cursor.close()
|
|||
|
|
|
|||
|
|
def _execute_query(self, connection, sql_query: str, limit: int = 100) -> Dict[str, Any]:
|
|||
|
|
"""执行SQL查询"""
|
|||
|
|
cursor = connection.cursor()
|
|||
|
|
try:
|
|||
|
|
# 添加LIMIT限制（如果查询中没有LIMIT）
|
|||
|
|
if limit and limit > 0 and "LIMIT" not in sql_query.upper():
|
|||
|
|
sql_query = f"{sql_query.rstrip(';')} LIMIT {limit}"
|
|||
|
|
|
|||
|
|
cursor.execute(sql_query)
|
|||
|
|
|
|||
|
|
# 获取列名
|
|||
|
|
columns = [desc[0] for desc in cursor.description] if cursor.description else []
|
|||
|
|
|
|||
|
|
# 获取数据
|
|||
|
|
rows = cursor.fetchall()
|
|||
|
|
|
|||
|
|
# 转换为字典列表
|
|||
|
|
data = []
|
|||
|
|
for row in rows:
|
|||
|
|
row_dict = {}
|
|||
|
|
for i, value in enumerate(row):
|
|||
|
|
if i < len(columns):
|
|||
|
|
# 处理特殊数据类型
|
|||
|
|
if isinstance(value, datetime):
|
|||
|
|
row_dict[columns[i]] = value.isoformat()
|
|||
|
|
else:
|
|||
|
|
row_dict[columns[i]] = value
|
|||
|
|
data.append(row_dict)
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"data": data,
|
|||
|
|
"columns": columns,
|
|||
|
|
"row_count": len(data),
|
|||
|
|
"query": sql_query
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
cursor.close()
|
|||
|
|
|
|||
|
|
def _create_connection(self, config: Dict[str, Any]) -> pymysql.Connection:
|
|||
|
|
"""创建MySQL数据库连接"""
|
|||
|
|
try:
|
|||
|
|
connection = pymysql.connect(
|
|||
|
|
host=config['host'],
|
|||
|
|
port=int(config.get('port', 3306)),
|
|||
|
|
user=config['username'],
|
|||
|
|
password=config['password'],
|
|||
|
|
database=config['database'],
|
|||
|
|
connect_timeout=10,
|
|||
|
|
charset='utf8mb4'
|
|||
|
|
)
|
|||
|
|
return connection
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"MySQL连接失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def _test_connection(self, config: Dict[str, Any]) -> Dict[str, Any]:
|
|||
|
|
"""测试数据库连接"""
|
|||
|
|
try:
|
|||
|
|
conn = self._create_connection(config)
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 获取数据库版本信息
|
|||
|
|
cursor.execute("SELECT VERSION();")
|
|||
|
|
version = cursor.fetchone()[0]
|
|||
|
|
|
|||
|
|
# 获取数据库引擎信息
|
|||
|
|
cursor.execute("SHOW ENGINES;")
|
|||
|
|
engines = cursor.fetchall()
|
|||
|
|
has_innodb = any('InnoDB' in str(engine) for engine in engines)
|
|||
|
|
|
|||
|
|
cursor.close()
|
|||
|
|
conn.close()
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"version": version,
|
|||
|
|
"has_innodb": has_innodb,
|
|||
|
|
"message": "连接测试成功"
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {
|
|||
|
|
"success": False,
|
|||
|
|
"error": str(e),
|
|||
|
|
"message": "连接测试失败"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def execute(self, **kwargs) -> ToolResult:
|
|||
|
|
"""Execute the MySQL MCP tool operation."""
|
|||
|
|
try:
|
|||
|
|
operation = kwargs.get("operation")
|
|||
|
|
connection_config = kwargs.get("connection_config", {})
|
|||
|
|
user_id = kwargs.get("user_id")
|
|||
|
|
table_name = kwargs.get("table_name")
|
|||
|
|
sql_query = kwargs.get("sql_query")
|
|||
|
|
limit = kwargs.get("limit", 100)
|
|||
|
|
|
|||
|
|
logger.info(f"执行MySQL MCP操作: {operation}")
|
|||
|
|
if operation == "test_connection":
|
|||
|
|
if not connection_config:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少连接配置参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
result = self._test_connection(connection_config)
|
|||
|
|
return ToolResult(
|
|||
|
|
success=result["success"],
|
|||
|
|
result=result,
|
|||
|
|
error=result.get("error")
|
|||
|
|
)
|
|||
|
|
elif operation == "connect":
|
|||
|
|
if not connection_config:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少connection_config参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not user_id:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少user_id参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 建立MySQL连接
|
|||
|
|
connection = pymysql.connect(
|
|||
|
|
host=connection_config["host"],
|
|||
|
|
port=int(connection_config["port"]),
|
|||
|
|
user=connection_config["username"],
|
|||
|
|
password=connection_config["password"],
|
|||
|
|
database=connection_config["database"],
|
|||
|
|
charset='utf8mb4',
|
|||
|
|
cursorclass=pymysql.cursors.Cursor
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 存储连接
|
|||
|
|
self.connections[user_id] = {
|
|||
|
|
"connection": connection,
|
|||
|
|
"config": connection_config,
|
|||
|
|
"connected_at": datetime.now().isoformat()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 获取表列表
|
|||
|
|
tables = self._get_tables(connection)
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result={
|
|||
|
|
"message": "数据库连接成功",
|
|||
|
|
"database": connection_config["database"],
|
|||
|
|
"tables": tables,
|
|||
|
|
"table_count": len(tables)
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"连接失败: {str(e)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elif operation == "list_tables":
|
|||
|
|
if not user_id or user_id not in self.connections:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="用户未连接数据库，请先执行connect操作"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
connection = self.connections[user_id]["connection"]
|
|||
|
|
tables = self._get_tables(connection)
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result={
|
|||
|
|
"tables": tables,
|
|||
|
|
"table_count": len(tables)
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elif operation == "describe_table":
|
|||
|
|
if not user_id or user_id not in self.connections:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="用户未连接数据库，请先执行connect操作"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not table_name:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少table_name参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
connection = self.connections[user_id]["connection"]
|
|||
|
|
table_info = self._describe_table(connection, table_name)
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result=table_info
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elif operation == "execute_query":
|
|||
|
|
if not user_id or user_id not in self.connections:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="用户未连接数据库，请先执行connect操作"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not sql_query:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少sql_query参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
connection = self.connections[user_id]["connection"]
|
|||
|
|
query_result = self._execute_query(connection, sql_query, limit)
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result=query_result
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elif operation == "disconnect":
|
|||
|
|
if user_id and user_id in self.connections:
|
|||
|
|
try:
|
|||
|
|
self.connections[user_id]["connection"].close()
|
|||
|
|
del self.connections[user_id]
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result={"message": "数据库连接已断开"}
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"断开连接失败: {str(e)}"
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
result={"message": "用户未连接数据库"}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
result=f"不支持的操作类型: {operation}",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"MySQL MCP工具执行失败: {str(e)}", exc_info=True)
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"工具执行失败: {str(e)}"
|
|||
|
|
)
|