385 lines
14 KiB
Python
385 lines
14 KiB
Python
|
"""PostgreSQL MCP (Model Context Protocol) tool for database operations."""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import psycopg2
|
|||
|
|
from typing import List, Dict, Any, Optional
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
from th_agenter.services.agent.base import BaseTool, ToolParameter, ToolParameterType, ToolResult
|
|||
|
|
|
|||
|
|
class PostgreSQLMCPTool(BaseTool):
|
|||
|
|
"""PostgreSQL MCP tool for database operations and intelligent querying."""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
self.connections = {} # 存储用户的数据库连接
|
|||
|
|
|
|||
|
|
def get_name(self) -> str:
|
|||
|
|
return "postgresql_mcp"
|
|||
|
|
|
|||
|
|
def get_description(self) -> str:
|
|||
|
|
return "PostgreSQL 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 _create_connection(self, config: Dict[str, Any]) -> psycopg2.extensions.connection:
|
|||
|
|
"""创建PostgreSQL数据库连接"""
|
|||
|
|
try:
|
|||
|
|
connection = psycopg2.connect(
|
|||
|
|
host=config['host'],
|
|||
|
|
port=int(config.get('port', 5432)),
|
|||
|
|
user=config['username'],
|
|||
|
|
password=config['password'],
|
|||
|
|
database=config['database'],
|
|||
|
|
connect_timeout=10
|
|||
|
|
)
|
|||
|
|
return connection
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"PostgreSQL连接失败: {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]
|
|||
|
|
|
|||
|
|
# 检查pgvector扩展
|
|||
|
|
cursor.execute("SELECT * FROM pg_extension WHERE extname = 'vector';")
|
|||
|
|
has_vector = bool(cursor.fetchall())
|
|||
|
|
|
|||
|
|
cursor.close()
|
|||
|
|
conn.close()
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"version": version,
|
|||
|
|
"has_pgvector": has_vector,
|
|||
|
|
"message": "连接测试成功"
|
|||
|
|
}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {
|
|||
|
|
"success": False,
|
|||
|
|
"error": str(e),
|
|||
|
|
"message": "连接测试失败"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
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 = 'public'
|
|||
|
|
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
|
|||
|
|
FROM information_schema.columns
|
|||
|
|
WHERE table_name = %s AND table_schema = 'public'
|
|||
|
|
ORDER BY ordinal_position;
|
|||
|
|
""", (table_name,))
|
|||
|
|
|
|||
|
|
columns = []
|
|||
|
|
for row in cursor.fetchall():
|
|||
|
|
columns.append({
|
|||
|
|
"column_name": row[0],
|
|||
|
|
"data_type": row[1],
|
|||
|
|
"is_nullable": row[2],
|
|||
|
|
"column_default": row[3],
|
|||
|
|
"character_maximum_length": row[4],
|
|||
|
|
"numeric_precision": row[5],
|
|||
|
|
"numeric_scale": row[6]
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 获取主键信息
|
|||
|
|
cursor.execute("""
|
|||
|
|
SELECT column_name
|
|||
|
|
FROM information_schema.key_column_usage
|
|||
|
|
WHERE table_name = %s AND table_schema = 'public'
|
|||
|
|
AND constraint_name IN (
|
|||
|
|
SELECT constraint_name
|
|||
|
|
FROM information_schema.table_constraints
|
|||
|
|
WHERE table_name = %s AND constraint_type = 'PRIMARY KEY'
|
|||
|
|
);
|
|||
|
|
""", (table_name, table_name))
|
|||
|
|
|
|||
|
|
primary_keys = [row[0] for row in cursor.fetchall()]
|
|||
|
|
|
|||
|
|
# 获取表行数
|
|||
|
|
cursor.execute(f"SELECT COUNT(*) FROM {table_name};")
|
|||
|
|
row_count = cursor.fetchone()[0]
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"table_name": table_name,
|
|||
|
|
"columns": columns,
|
|||
|
|
"primary_keys": primary_keys,
|
|||
|
|
"row_count": row_count
|
|||
|
|
}
|
|||
|
|
finally:
|
|||
|
|
cursor.close()
|
|||
|
|
|
|||
|
|
def _execute_query(self, connection, sql_query: str, limit: int = 100) -> Dict[str, Any]:
|
|||
|
|
"""执行SQL查询"""
|
|||
|
|
cursor = connection.cursor()
|
|||
|
|
try:
|
|||
|
|
# 添加LIMIT限制（如果查询中没有）
|
|||
|
|
if limit 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 []
|
|||
|
|
|
|||
|
|
# 获取结果
|
|||
|
|
if cursor.description: # SELECT查询
|
|||
|
|
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
|
|||
|
|
}
|
|||
|
|
else: # INSERT/UPDATE/DELETE查询
|
|||
|
|
affected_rows = cursor.rowcount
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"affected_rows": affected_rows,
|
|||
|
|
"query": sql_query,
|
|||
|
|
"message": f"查询执行成功，影响 {affected_rows} 行"
|
|||
|
|
}
|
|||
|
|
finally:
|
|||
|
|
cursor.close()
|
|||
|
|
|
|||
|
|
async def execute(self, operation: str, connection_config: Optional[Dict[str, Any]] = None,
|
|||
|
|
user_id: Optional[str] = None, table_name: Optional[str] = None,
|
|||
|
|
sql_query: Optional[str] = None, limit: int = 100) -> ToolResult:
|
|||
|
|
"""执行PostgreSQL MCP操作"""
|
|||
|
|
try:
|
|||
|
|
logger.info(f"执行PostgreSQL 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 or not user_id:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="缺少连接配置或用户ID参数"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
connection = self._create_connection(connection_config)
|
|||
|
|
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,
|
|||
|
|
error=f"不支持的操作类型: {operation}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"PostgreSQL MCP工具执行失败: {str(e)}", exc_info=True)
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"工具执行失败: {str(e)}"
|
|||
|
|
)
|