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)}"
|
||
) |