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