mcp-dbutils 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

mcp_dbutils/__init__.py

@@ -1,4 +1,4 @@
-"""MCP Database Utilities Service"""
+"""MCP Connection Utilities Service"""
 
 import asyncio
 import argparse
@@ -9,7 +9,7 @@ import yaml
 from importlib.metadata import metadata
 
 from .log import create_logger
-from .base import DatabaseServer
+from .base import ConnectionServer
 
 # 获取包信息
 pkg_meta = metadata("mcp-dbutils")
@@ -19,7 +19,7 @@ log = create_logger(pkg_meta["Name"])
 
 async def run_server():
     """服务器运行逻辑"""
-    parser = argparse.ArgumentParser(description='MCP Database Server')
+    parser = argparse.ArgumentParser(description='MCP Connection Server')
     parser.add_argument('--config', required=True, help='YAML配置文件路径')
     parser.add_argument('--local-host', help='本地主机地址')
 
@@ -32,7 +32,7 @@ async def run_server():
     global log
     log = create_logger(pkg_meta["Name"], debug)
 
-    log("info", f"MCP Database Utilities Service v{pkg_meta['Version']}")
+    log("info", f"MCP Connection Utilities Service v{pkg_meta['Version']}")
     if debug:
         log("debug", "Debug模式已开启")
 
@@ -40,11 +40,11 @@ async def run_server():
     try:
         with open(args.config, 'r') as f:
             config = yaml.safe_load(f)
-            if not config or 'databases' not in config:
-                log("error", "配置文件必须包含 databases 配置")
+            if not config or 'connections' not in config:
+                log("error", "配置文件必须包含 connections 配置")
                 sys.exit(1)
-            if not config['databases']:
-                log("error", "配置文件必须包含至少一个数据库配置")
+            if not config['connections']:
+                log("error", "配置文件必须包含至少一个连接配置")
                 sys.exit(1)
     except Exception as e:
         log("error", f"读取配置文件失败: {str(e)}")
@@ -52,7 +52,7 @@ async def run_server():
 
     # 创建并运行服务器
     try:
-        server = DatabaseServer(args.config, debug)
+        server = ConnectionServer(args.config, debug)
         await server.run()
     except KeyboardInterrupt:
         log("info", "服务器已停止")

mcp_dbutils/base.py

@@ -1,14 +1,14 @@
-"""Database server base class"""
+"""Connection server base class"""
 
-class DatabaseError(Exception):
-    """Base exception for database errors"""
+class ConnectionHandlerError(Exception):
+    """Base exception for connection errors"""
     pass
 
-class ConfigurationError(DatabaseError):
+class ConfigurationError(ConnectionHandlerError):
     """Configuration related errors"""
     pass
 
-class ConnectionError(DatabaseError):
+class ConnectionError(ConnectionHandlerError):
     """Connection related errors"""
     pass
 
@@ -17,6 +17,8 @@ from typing import Any, List, Optional, AsyncContextManager
 from contextlib import asynccontextmanager
 import json
 import yaml
+import time
+from datetime import datetime
 from importlib.metadata import metadata
 from mcp.server import Server, NotificationOptions
 import mcp.server.stdio
@@ -29,21 +31,21 @@ from .stats import ResourceStats
 # 获取包信息用于日志命名
 pkg_meta = metadata("mcp-dbutils")
 
-class DatabaseHandler(ABC):
-    """Abstract base class defining common interface for database handlers"""
+class ConnectionHandler(ABC):
+    """Abstract base class defining common interface for connection handlers"""
 
-    def __init__(self, config_path: str, database: str, debug: bool = False):
-        """Initialize database handler
+    def __init__(self, config_path: str, connection: str, debug: bool = False):
+        """Initialize connection handler
 
         Args:
             config_path: Path to configuration file
-            database: Database configuration name
+            connection: Database connection name
             debug: Enable debug mode
         """
         self.config_path = config_path
-        self.database = database
+        self.connection = connection
         self.debug = debug
-        self.log = create_logger(f"{pkg_meta['Name']}.handler.{database}", debug)
+        self.log = create_logger(f"{pkg_meta['Name']}.handler.{connection}", debug)
         self.stats = ResourceStats()
 
     @property
@@ -54,7 +56,7 @@ class DatabaseHandler(ABC):
 
     @abstractmethod
     async def get_tables(self) -> list[types.Resource]:
-        """Get list of table resources from database"""
+        """Get list of table resources from database connection"""
         pass
 
     @abstractmethod
@@ -68,28 +70,144 @@ class DatabaseHandler(ABC):
         pass
 
     async def execute_query(self, sql: str) -> str:
-        """Execute SQL query"""
+        """Execute SQL query with performance tracking"""
+        start_time = datetime.now()
         try:
             self.stats.record_query()
             result = await self._execute_query(sql)
+            duration = (datetime.now() - start_time).total_seconds()
+            self.stats.record_query_duration(sql, duration)
             self.stats.update_memory_usage(result)
-            self.log("info", f"Resource stats: {json.dumps(self.stats.to_dict())}")
+            self.log("info", f"Query executed in {duration*1000:.2f}ms. Resource stats: {json.dumps(self.stats.to_dict())}")
             return result
         except Exception as e:
+            duration = (datetime.now() - start_time).total_seconds()
             self.stats.record_error(e.__class__.__name__)
-            self.log("error", f"Query error - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}")
+            self.log("error", f"Query error after {duration*1000:.2f}ms - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}")
             raise
 
+    @abstractmethod
+    async def get_table_description(self, table_name: str) -> str:
+        """Get detailed table description including columns, types, and comments
+
+        Args:
+            table_name: Name of the table to describe
+
+        Returns:
+            Formatted table description
+        """
+        pass
+
+    @abstractmethod
+    async def get_table_ddl(self, table_name: str) -> str:
+        """Get DDL statement for table including columns, constraints and indexes
+
+        Args:
+            table_name: Name of the table to get DDL for
+
+        Returns:
+            DDL statement as string
+        """
+        pass
+
+    @abstractmethod
+    async def get_table_indexes(self, table_name: str) -> str:
+        """Get index information for table
+
+        Args:
+            table_name: Name of the table to get indexes for
+
+        Returns:
+            Formatted index information
+        """
+        pass
+
+    @abstractmethod
+    async def get_table_stats(self, table_name: str) -> str:
+        """Get table statistics information
+
+        Args:
+            table_name: Name of the table to get statistics for
+
+        Returns:
+            Formatted statistics information including row count, size, etc.
+        """
+        pass
+
+    @abstractmethod
+    async def get_table_constraints(self, table_name: str) -> str:
+        """Get constraint information for table
+
+        Args:
+            table_name: Name of the table to get constraints for
+
+        Returns:
+            Formatted constraint information including primary keys, foreign keys, etc.
+        """
+        pass
+
+    @abstractmethod
+    async def explain_query(self, sql: str) -> str:
+        """Get query execution plan
+
+        Args:
+            sql: SQL query to explain
+
+        Returns:
+            Formatted query execution plan with cost estimates
+        """
+        pass
+
     @abstractmethod
     async def cleanup(self):
         """Cleanup resources"""
         pass
 
-class DatabaseServer:
-    """Unified database server class"""
+    async def execute_tool_query(self, tool_name: str, table_name: str = "", sql: str = "") -> str:
+        """Execute a tool query and return formatted result
+
+        Args:
+            tool_name: Name of the tool to execute
+            table_name: Name of the table to query (for table-related tools)
+            sql: SQL query (for query-related tools)
+
+        Returns:
+            Formatted query result
+        """
+        try:
+            self.stats.record_query()
+            
+            if tool_name == "dbutils-describe-table":
+                result = await self.get_table_description(table_name)
+            elif tool_name == "dbutils-get-ddl":
+                result = await self.get_table_ddl(table_name)
+            elif tool_name == "dbutils-list-indexes":
+                result = await self.get_table_indexes(table_name)
+            elif tool_name == "dbutils-get-stats":
+                result = await self.get_table_stats(table_name)
+            elif tool_name == "dbutils-list-constraints":
+                result = await self.get_table_constraints(table_name)
+            elif tool_name == "dbutils-explain-query":
+                if not sql:
+                    raise ValueError("SQL query required for explain-query tool")
+                result = await self.explain_query(sql)
+            else:
+                raise ValueError(f"Unknown tool: {tool_name}")
+                
+            self.stats.update_memory_usage(result)
+            self.log("info", f"Resource stats: {json.dumps(self.stats.to_dict())}")
+            return f"[{self.db_type}]\n{result}"
+            
+        except Exception as e:
+            self.stats.record_error(e.__class__.__name__)
+            self.log("error", f"Tool error - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}")
+            raise
+
+class ConnectionServer:
+    """Unified connection server class"""
 
     def __init__(self, config_path: str, debug: bool = False):
-        """Initialize database server
+        """Initialize connection server
 
         Args:
             config_path: Path to configuration file
@@ -120,27 +238,27 @@ class DatabaseServer:
                 raise
 
     @asynccontextmanager
-    async def get_handler(self, database: str) -> AsyncContextManager[DatabaseHandler]:
-        """Get database handler
+    async def get_handler(self, connection: str) -> AsyncContextManager[ConnectionHandler]:
+        """Get connection handler
 
-        Get appropriate database handler based on configuration name
+        Get appropriate connection handler based on connection name
 
         Args:
-            database: Database configuration name
+            connection: Database connection name
 
         Returns:
-            AsyncContextManager[DatabaseHandler]: Context manager for database handler
+            AsyncContextManager[ConnectionHandler]: Context manager for connection handler
         """
         # Read configuration file to determine database type
         with open(self.config_path, 'r') as f:
             config = yaml.safe_load(f)
-            if not config or 'databases' not in config:
-                raise ConfigurationError("Configuration file must contain 'databases' section")
-            if database not in config['databases']:
-                available_dbs = list(config['databases'].keys())
-                raise ConfigurationError(f"Database configuration not found: {database}. Available configurations: {available_dbs}")
+            if not config or 'connections' not in config:
+                raise ConfigurationError("Configuration file must contain 'connections' section")
+            if connection not in config['connections']:
+                available_connections = list(config['connections'].keys())
+                raise ConfigurationError(f"Connection not found: {connection}. Available connections: {available_connections}")
 
-            db_config = config['databases'][database]
+            db_config = config['connections'][connection]
 
             handler = None
             try:
@@ -150,16 +268,16 @@ class DatabaseServer:
                 db_type = db_config['type']
                 self.logger("debug", f"Creating handler for database type: {db_type}")
                 if db_type == 'sqlite':
-                    from .sqlite.handler import SqliteHandler
-                    handler = SqliteHandler(self.config_path, database, self.debug)
+                    from .sqlite.handler import SQLiteHandler
+                    handler = SQLiteHandler(self.config_path, connection, self.debug)
                 elif db_type == 'postgres':
-                    from .postgres.handler import PostgresHandler
-                    handler = PostgresHandler(self.config_path, database, self.debug)
+                    from .postgres.handler import PostgreSQLHandler
+                    handler = PostgreSQLHandler(self.config_path, connection, self.debug)
                 else:
                     raise ConfigurationError(f"Unsupported database type: {db_type}")
 
                 handler.stats.record_connection_start()
-                self.logger("debug", f"Handler created successfully for {database}")
+                self.logger("debug", f"Handler created successfully for {connection}")
                 self.logger("info", f"Resource stats: {json.dumps(handler.stats.to_dict())}")
                 yield handler
             except yaml.YAMLError as e:
@@ -168,7 +286,7 @@ class DatabaseServer:
                 raise ConfigurationError(f"Failed to import handler for {db_type}: {str(e)}")
             finally:
                 if handler:
-                    self.logger("debug", f"Cleaning up handler for {database}")
+                    self.logger("debug", f"Cleaning up handler for {connection}")
                     handler.stats.record_connection_end()
                     self.logger("info", f"Final resource stats: {json.dumps(handler.stats.to_dict())}")
                     await handler.cleanup()
@@ -177,75 +295,215 @@ class DatabaseServer:
         """Setup MCP handlers"""
         @self.server.list_resources()
         async def handle_list_resources(arguments: dict | None = None) -> list[types.Resource]:
-            if not arguments or 'database' not in arguments:
-                # Return empty list when no database specified
+            if not arguments or 'connection' not in arguments:
+                # Return empty list when no connection specified
                 return []
 
-            database = arguments['database']
-            async with self.get_handler(database) as handler:
+            connection = arguments['connection']
+            async with self.get_handler(connection) as handler:
                 return await handler.get_tables()
 
         @self.server.read_resource()
         async def handle_read_resource(uri: str, arguments: dict | None = None) -> str:
-            if not arguments or 'database' not in arguments:
-                raise ConfigurationError("Database configuration name must be specified")
+            if not arguments or 'connection' not in arguments:
+                raise ConfigurationError("Connection name must be specified")
 
             parts = uri.split('/')
             if len(parts) < 3:
                 raise ConfigurationError("Invalid resource URI format")
 
-            database = arguments['database']
+            connection = arguments['connection']
             table_name = parts[-2]  # URI format: xxx/table_name/schema
 
-            async with self.get_handler(database) as handler:
+            async with self.get_handler(connection) as handler:
                 return await handler.get_schema(table_name)
 
         @self.server.list_tools()
         async def handle_list_tools() -> list[types.Tool]:
             return [
                 types.Tool(
-                    name="query",
-                    description="Execute read-only SQL query",
+                    name="dbutils-run-query",
+                    description="Execute read-only SQL query on database connection",
                     inputSchema={
                         "type": "object",
                         "properties": {
-                            "database": {
+                            "connection": {
                                 "type": "string",
-                                "description": "Database configuration name"
+                                "description": "Database connection name"
                             },
                             "sql": {
                                 "type": "string",
                                 "description": "SQL query (SELECT only)"
                             }
                         },
-                        "required": ["database", "sql"]
+                        "required": ["connection", "sql"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-list-tables",
+                    description="List all available tables in the specified database connection",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            }
+                        },
+                        "required": ["connection"]
                     }
                 ),
                 types.Tool(
-                    name="list_tables",
-                    description="List all available tables in the specified database",
+                    name="dbutils-describe-table",
+                    description="Get detailed information about a table's structure",
                     inputSchema={
                         "type": "object",
                         "properties": {
-                            "database": {
+                            "connection": {
                                 "type": "string",
-                                "description": "Database configuration name"
+                                "description": "Database connection name"
+                            },
+                            "table": {
+                                "type": "string",
+                                "description": "Table name to describe"
                             }
                         },
-                        "required": ["database"]
+                        "required": ["connection", "table"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-get-ddl",
+                    description="Get DDL statement for creating the table",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "table": {
+                                "type": "string",
+                                "description": "Table name to get DDL for"
+                            }
+                        },
+                        "required": ["connection", "table"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-list-indexes",
+                    description="List all indexes on the specified table",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "table": {
+                                "type": "string",
+                                "description": "Table name to list indexes for"
+                            }
+                        },
+                        "required": ["connection", "table"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-get-stats",
+                    description="Get table statistics like row count and size",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "table": {
+                                "type": "string",
+                                "description": "Table name to get statistics for"
+                            }
+                        },
+                        "required": ["connection", "table"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-list-constraints",
+                    description="List all constraints (primary key, foreign keys, etc) on the table",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "table": {
+                                "type": "string",
+                                "description": "Table name to list constraints for"
+                            }
+                        },
+                        "required": ["connection", "table"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-explain-query",
+                    description="Get execution plan for a SQL query",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "sql": {
+                                "type": "string",
+                                "description": "SQL query to explain"
+                            }
+                        },
+                        "required": ["connection", "sql"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-get-performance",
+                    description="Get database performance statistics",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            }
+                        },
+                        "required": ["connection"]
+                    }
+                ),
+                types.Tool(
+                    name="dbutils-analyze-query",
+                    description="Analyze a SQL query for performance",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "connection": {
+                                "type": "string",
+                                "description": "Database connection name"
+                            },
+                            "sql": {
+                                "type": "string",
+                                "description": "SQL query to analyze"
+                            }
+                        },
+                        "required": ["connection", "sql"]
                     }
                 )
             ]
 
         @self.server.call_tool()
         async def handle_call_tool(name: str, arguments: dict) -> list[types.TextContent]:
-            if "database" not in arguments:
-                raise ConfigurationError("Database configuration name must be specified")
-
-            database = arguments["database"]
+            if "connection" not in arguments:
+                raise ConfigurationError("Connection name must be specified")
+            
+            connection = arguments["connection"]
 
-            if name == "list_tables":
-                async with self.get_handler(database) as handler:
+            if name == "dbutils-list-tables":
+                async with self.get_handler(connection) as handler:
                     tables = await handler.get_tables()
                     if not tables:
                         # 空表列表的情况也返回数据库类型
@@ -260,7 +518,7 @@ class DatabaseServer:
                     ])
                     # 添加数据库类型前缀
                     return [types.TextContent(type="text", text=f"[{handler.db_type}]\n{formatted_tables}")]
-            elif name == "query":
+            elif name == "dbutils-run-query":
                 sql = arguments.get("sql", "").strip()
                 if not sql:
                     raise ConfigurationError("SQL query cannot be empty")
@@ -269,9 +527,76 @@ class DatabaseServer:
                 if not sql.lower().startswith("select"):
                     raise ConfigurationError("Only SELECT queries are supported for security reasons")
 
-                async with self.get_handler(database) as handler:
+                async with self.get_handler(connection) as handler:
                     result = await handler.execute_query(sql)
                     return [types.TextContent(type="text", text=result)]
+            elif name in ["dbutils-describe-table", "dbutils-get-ddl", "dbutils-list-indexes",
+                         "dbutils-get-stats", "dbutils-list-constraints"]:
+                table = arguments.get("table", "").strip()
+                if not table:
+                    raise ConfigurationError("Table name cannot be empty")
+                
+                async with self.get_handler(connection) as handler:
+                    result = await handler.execute_tool_query(name, table_name=table)
+                    return [types.TextContent(type="text", text=result)]
+            elif name == "dbutils-explain-query":
+                sql = arguments.get("sql", "").strip()
+                if not sql:
+                    raise ConfigurationError("SQL query cannot be empty")
+                
+                async with self.get_handler(connection) as handler:
+                    result = await handler.execute_tool_query(name, sql=sql)
+                    return [types.TextContent(type="text", text=result)]
+            elif name == "dbutils-get-performance":
+                async with self.get_handler(connection) as handler:
+                    performance_stats = handler.stats.get_performance_stats()
+                    return [types.TextContent(type="text", text=f"[{handler.db_type}]\n{performance_stats}")]
+            elif name == "dbutils-analyze-query":
+                sql = arguments.get("sql", "").strip()
+                if not sql:
+                    raise ConfigurationError("SQL query cannot be empty")
+                
+                async with self.get_handler(connection) as handler:
+                    # First get the execution plan
+                    explain_result = await handler.explain_query(sql)
+                    
+                    # Then execute the actual query to measure performance
+                    start_time = datetime.now()
+                    if sql.lower().startswith("select"):
+                        try:
+                            await handler.execute_query(sql)
+                        except Exception as e:
+                            # If query fails, we still provide the execution plan
+                            self.logger("error", f"Query execution failed during analysis: {str(e)}")
+                    duration = (datetime.now() - start_time).total_seconds()
+                    
+                    # Combine analysis results
+                    analysis = [
+                        f"[{handler.db_type}] Query Analysis",
+                        f"SQL: {sql}",
+                        f"",
+                        f"Execution Time: {duration*1000:.2f}ms",
+                        f"",
+                        f"Execution Plan:",
+                        explain_result
+                    ]
+                    
+                    # Add optimization suggestions based on execution plan and timing
+                    suggestions = []
+                    if "seq scan" in explain_result.lower() and duration > 0.1:
+                        suggestions.append("- Consider adding an index to avoid sequential scan")
+                    if "hash join" in explain_result.lower() and duration > 0.5:
+                        suggestions.append("- Consider optimizing join conditions")
+                    if duration > 0.5:  # 500ms
+                        suggestions.append("- Query is slow, consider optimizing or adding caching")
+                    if "temporary" in explain_result.lower():
+                        suggestions.append("- Query creates temporary tables, consider restructuring")
+                    
+                    if suggestions:
+                        analysis.append("\nOptimization Suggestions:")
+                        analysis.extend(suggestions)
+                        
+                    return [types.TextContent(type="text", text="\n".join(analysis))]
             else:
                 raise ConfigurationError(f"Unknown tool: {name}")