mcp-dbutils 0.2.3__py3-none-any.whl

mcp_dbutils/__init__.py

@@ -0,0 +1,62 @@
+"""MCP Database Utilities Service"""
+
+import asyncio
+import argparse
+import os
+import sys
+from pathlib import Path
+import yaml
+
+from .log import create_logger
+from .base import DatabaseServer
+
+# 创建全局logger
+log = create_logger("mcp-dbutils")
+
+async def run_server():
+    """服务器运行逻辑"""
+    parser = argparse.ArgumentParser(description='MCP Database Server')
+    parser.add_argument('--config', required=True, help='YAML配置文件路径')
+    parser.add_argument('--local-host', help='本地主机地址')
+
+    args = parser.parse_args()
+
+    # 检查是否开启debug模式
+    debug = os.getenv('MCP_DEBUG', '').lower() in ('1', 'true', 'yes')
+
+    # 更新logger的debug状态
+    global log
+    log = create_logger("mcp-dbutils", debug)
+
+    if debug:
+        log("debug", "Debug模式已开启")
+
+    # 验证配置文件
+    try:
+        with open(args.config, 'r') as f:
+            config = yaml.safe_load(f)
+            if not config or 'databases' not in config:
+                log("error", "配置文件必须包含 databases 配置")
+                sys.exit(1)
+            if not config['databases']:
+                log("error", "配置文件必须包含至少一个数据库配置")
+                sys.exit(1)
+    except Exception as e:
+        log("error", f"读取配置文件失败: {str(e)}")
+        sys.exit(1)
+
+    # 创建并运行服务器
+    try:
+        server = DatabaseServer(args.config, debug)
+        await server.run()
+    except KeyboardInterrupt:
+        log("info", "服务器已停止")
+    except Exception as e:
+        log("error", str(e))
+        sys.exit(1)
+
+def main():
+    """命令行入口函数"""
+    asyncio.run(run_server())
+
+__all__ = ['main']

mcp_dbutils/base.py

@@ -0,0 +1,182 @@
+"""Database server base class"""
+
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional, AsyncContextManager
+from contextlib import asynccontextmanager
+import yaml
+from mcp.server import Server, NotificationOptions
+import mcp.server.stdio
+import mcp.types as types
+from .log import create_logger
+
+class DatabaseHandler(ABC):
+    """Abstract base class defining common interface for database handlers"""
+
+    def __init__(self, config_path: str, database: str, debug: bool = False):
+        """Initialize database handler
+
+        Args:
+            config_path: Path to configuration file
+            database: Database configuration name
+            debug: Enable debug mode
+        """
+        self.config_path = config_path
+        self.database = database
+        self.debug = debug
+        self.log = create_logger(f"db-handler-{database}", debug)
+
+    @abstractmethod
+    async def get_tables(self) -> list[types.Resource]:
+        """Get list of table resources from database"""
+        pass
+
+    @abstractmethod
+    async def get_schema(self, table_name: str) -> str:
+        """Get schema information for specified table"""
+        pass
+
+    @abstractmethod
+    async def execute_query(self, sql: str) -> str:
+        """Execute SQL query"""
+        pass
+
+    @abstractmethod
+    async def cleanup(self):
+        """Cleanup resources"""
+        pass
+
+class DatabaseServer:
+    """Unified database server class"""
+
+    def __init__(self, config_path: str, debug: bool = False):
+        """Initialize database server
+
+        Args:
+            config_path: Path to configuration file
+            debug: Enable debug mode
+        """
+        self.config_path = config_path
+        self.debug = debug
+        self.log = create_logger("db-server", debug)
+        self.server = Server("database-server")
+        self._setup_handlers()
+
+    @asynccontextmanager
+    async def get_handler(self, database: str) -> AsyncContextManager[DatabaseHandler]:
+        """Get database handler
+
+        Get appropriate database handler based on configuration name
+
+        Args:
+            database: Database configuration name
+
+        Returns:
+            AsyncContextManager[DatabaseHandler]: Context manager for database 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 ValueError("Configuration file must contain 'databases' section")
+            if database not in config['databases']:
+                available_dbs = list(config['databases'].keys())
+                raise ValueError(f"Database configuration not found: {database}. Available configurations: {available_dbs}")
+
+            db_config = config['databases'][database]
+
+            handler = None
+            try:
+                # Create appropriate handler based on configuration
+                if 'path' in db_config:
+                    from .sqlite.handler import SqliteHandler
+                    handler = SqliteHandler(self.config_path, database, self.debug)
+                elif 'dbname' in db_config or 'host' in db_config:
+                    from .postgres.handler import PostgresHandler
+                    handler = PostgresHandler(self.config_path, database, self.debug)
+                else:
+                    raise ValueError("Cannot determine database type, missing required parameters in configuration")
+
+                yield handler
+            finally:
+                if handler:
+                    await handler.cleanup()
+
+    def _setup_handlers(self):
+        """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
+                return []
+
+            database = arguments['database']
+            async with self.get_handler(database) 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 ValueError("Database configuration name must be specified")
+
+            parts = uri.split('/')
+            if len(parts) < 3:
+                raise ValueError("Invalid resource URI")
+
+            database = arguments['database']
+            table_name = parts[-2]  # URI format: xxx/table_name/schema
+
+            async with self.get_handler(database) 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",
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "database": {
+                                "type": "string",
+                                "description": "Database configuration name"
+                            },
+                            "sql": {
+                                "type": "string",
+                                "description": "SQL query (SELECT only)"
+                            }
+                        },
+                        "required": ["database", "sql"]
+                    }
+                )
+            ]
+
+        @self.server.call_tool()
+        async def handle_call_tool(name: str, arguments: dict) -> list[types.TextContent]:
+            if name != "query":
+                raise ValueError(f"Unknown tool: {name}")
+
+            if "database" not in arguments:
+                raise ValueError("Database configuration name must be specified")
+
+            sql = arguments.get("sql", "").strip()
+            if not sql:
+                raise ValueError("SQL query cannot be empty")
+
+            # Only allow SELECT statements
+            if not sql.lower().startswith("select"):
+                raise ValueError("Only SELECT queries are supported")
+
+            database = arguments["database"]
+            async with self.get_handler(database) as handler:
+                result = await handler.execute_query(sql)
+                return [types.TextContent(type="text", text=result)]
+
+    async def run(self):
+        """Run server"""
+        async with mcp.server.stdio.stdio_server() as streams:
+            await self.server.run(
+                streams[0],
+                streams[1],
+                self.server.create_initialization_options()
+            )

mcp_dbutils/config.py

@@ -0,0 +1,59 @@
+"""Common configuration utilities"""
+
+import os
+import yaml
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, Literal
+from pathlib import Path
+
+# Supported database types
+DBType = Literal['sqlite', 'postgres']
+
+class DatabaseConfig(ABC):
+    """Base class for database configuration"""
+
+    debug: bool = False
+    type: DBType  # Database type
+
+    @abstractmethod
+    def get_connection_params(self) -> Dict[str, Any]:
+        """Get connection parameters"""
+        pass
+
+    @abstractmethod
+    def get_masked_connection_info(self) -> Dict[str, Any]:
+        """Get masked connection information for logging"""
+        pass
+
+    @classmethod
+    def load_yaml_config(cls, yaml_path: str) -> Dict[str, Any]:
+        """Load YAML configuration file
+
+        Args:
+            yaml_path: Path to YAML file
+
+        Returns:
+            Parsed configuration dictionary
+        """
+        with open(yaml_path, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+
+        if not config or 'databases' not in config:
+            raise ValueError("Configuration file must contain 'databases' section")
+
+        # Validate type field in each database configuration
+        databases = config['databases']
+        for db_name, db_config in databases.items():
+            if 'type' not in db_config:
+                raise ValueError(f"Database configuration {db_name} missing required 'type' field")
+            db_type = db_config['type']
+            if db_type not in ('sqlite', 'postgres'):
+                raise ValueError(f"Invalid type value in database configuration {db_name}: {db_type}")
+
+        return databases
+
+    @classmethod
+    def get_debug_mode(cls) -> bool:
+        """Get debug mode status"""
+        return os.environ.get('MCP_DEBUG', '').lower() in ('1', 'true')

mcp_dbutils/log.py

@@ -0,0 +1,33 @@
+"""日志处理模块"""
+
+import sys
+from datetime import datetime
+from typing import Callable, Optional
+
+def create_logger(name: str, is_debug: bool = False) -> Callable:
+    """创建日志函数
+    Args:
+        name: 服务名称
+        is_debug: 是否输出debug级别日志
+    """
+    def log(level: str, message: str, notify: Optional[Callable] = None):
+        """输出日志
+        Args:
+            level: 日志级别 (debug/info/warning/error)
+            message: 日志内容
+            notify: MCP通知函数(可选)
+        """
+        if level == "debug" and not is_debug:
+            return
+
+        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
+        log_message = f"[{timestamp}] [{name}] [{level}] {message}"
+
+        # 始终输出到stderr
+        print(log_message, file=sys.stderr, flush=True)
+
+        # 如果提供了notify函数，同时发送MCP通知
+        if notify:
+            notify(level=level, data=message)
+
+    return log

mcp_dbutils/postgres/__init__.py

@@ -0,0 +1,6 @@
+"""PostgreSQL module"""
+
+from .handler import PostgresHandler
+from .config import PostgresConfig
+
+__all__ = ['PostgresHandler', 'PostgresConfig']

mcp_dbutils/postgres/config.py

@@ -0,0 +1,66 @@
+"""PostgreSQL configuration module"""
+from dataclasses import dataclass
+from typing import Optional, Dict, Any, Literal
+from ..config import DatabaseConfig
+
+@dataclass
+class PostgresConfig(DatabaseConfig):
+    dbname: str
+    user: str
+    password: str
+    host: str = 'localhost'
+    port: str = '5432'
+    local_host: Optional[str] = None
+    type: Literal['postgres'] = 'postgres'
+
+    @classmethod
+    def from_yaml(cls, yaml_path: str, db_name: str, local_host: Optional[str] = None) -> 'PostgresConfig':
+        """Create configuration from YAML file
+
+        Args:
+            yaml_path: Path to YAML configuration file
+            db_name: Database configuration name to use
+            local_host: Optional local host address
+        """
+        configs = cls.load_yaml_config(yaml_path)
+        if not db_name:
+            raise ValueError("Database name must be specified")
+        if db_name not in configs:
+            available_dbs = list(configs.keys())
+            raise ValueError(f"Database configuration not found: {db_name}. Available configurations: {available_dbs}")
+
+        db_config = configs[db_name]
+        if 'type' not in db_config:
+            raise ValueError("Database configuration must include 'type' field")
+        if db_config['type'] != 'postgres':
+            raise ValueError(f"Configuration is not PostgreSQL type: {db_config['type']}")
+
+        config = cls(
+            dbname=db_config.get('dbname', ''),
+            user=db_config.get('user', ''),
+            password=db_config.get('password', ''),
+            host=db_config.get('host', 'localhost'),
+            port=str(db_config.get('port', 5432)),
+            local_host=local_host,
+        )
+        config.debug = cls.get_debug_mode()
+        return config
+
+    def get_connection_params(self) -> Dict[str, Any]:
+        """Get psycopg2 connection parameters"""
+        params = {
+            'dbname': self.dbname,
+            'user': self.user,
+            'password': self.password,
+            'host': self.local_host or self.host,
+            'port': self.port
+        }
+        return {k: v for k, v in params.items() if v}
+
+    def get_masked_connection_info(self) -> Dict[str, Any]:
+        """Return masked connection information for logging"""
+        return {
+            'dbname': self.dbname,
+            'host': self.local_host or self.host,
+            'port': self.port
+        }

mcp_dbutils/postgres/handler.py

@@ -0,0 +1,150 @@
+"""PostgreSQL database handler implementation"""
+
+import psycopg2
+from psycopg2.pool import SimpleConnectionPool
+import mcp.types as types
+
+from ..base import DatabaseHandler
+from .config import PostgresConfig
+
+class PostgresHandler(DatabaseHandler):
+    def __init__(self, config_path: str, database: str, debug: bool = False):
+        """Initialize PostgreSQL handler
+
+        Args:
+            config_path: Path to configuration file
+            database: Database configuration name
+            debug: Enable debug mode
+        """
+        super().__init__(config_path, database, debug)
+        self.config = PostgresConfig.from_yaml(config_path, database)
+
+        # No connection pool creation during initialization
+        masked_params = self.config.get_masked_connection_info()
+        self.log("debug", f"Configuring database with parameters: {masked_params}")
+        self.pool = None
+
+    async def get_tables(self) -> list[types.Resource]:
+        """Get all table resources"""
+        try:
+            conn_params = self.config.get_connection_params()
+            conn = psycopg2.connect(**conn_params)
+            with conn.cursor() as cur:
+                cur.execute("""
+                    SELECT
+                        table_name,
+                        obj_description(
+                            (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
+                            'pg_class'
+                        ) as description
+                    FROM information_schema.tables
+                    WHERE table_schema = 'public'
+                """)
+                tables = cur.fetchall()
+                return [
+                    types.Resource(
+                        uri=f"postgres://{self.database}/{table[0]}/schema",
+                        name=f"{table[0]} schema",
+                        description=table[1] if table[1] else None,
+                        mimeType="application/json"
+                    ) for table in tables
+                ]
+        except psycopg2.Error as e:
+            error_msg = f"Failed to get table list: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            self.log("error", error_msg)
+            raise
+        finally:
+            if conn:
+                conn.close()
+
+    async def get_schema(self, table_name: str) -> str:
+        """Get table schema information"""
+        try:
+            conn_params = self.config.get_connection_params()
+            conn = psycopg2.connect(**conn_params)
+            with conn.cursor() as cur:
+                # Get column information
+                cur.execute("""
+                    SELECT
+                        column_name,
+                        data_type,
+                        is_nullable,
+                        col_description(
+                            (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
+                            ordinal_position
+                        ) as description
+                    FROM information_schema.columns
+                    WHERE table_name = %s
+                    ORDER BY ordinal_position
+                """, (table_name,))
+                columns = cur.fetchall()
+
+                # Get constraint information
+                cur.execute("""
+                    SELECT
+                        conname as constraint_name,
+                        contype as constraint_type
+                    FROM pg_constraint c
+                    JOIN pg_class t ON c.conrelid = t.oid
+                    WHERE t.relname = %s
+                """, (table_name,))
+                constraints = cur.fetchall()
+
+                return str({
+                    'columns': [{
+                        'name': col[0],
+                        'type': col[1],
+                        'nullable': col[2] == 'YES',
+                        'description': col[3]
+                    } for col in columns],
+                    'constraints': [{
+                        'name': con[0],
+                        'type': con[1]
+                    } for con in constraints]
+                })
+        except psycopg2.Error as e:
+            error_msg = f"Failed to read table schema: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            self.log("error", error_msg)
+            raise
+        finally:
+            if conn:
+                conn.close()
+
+    async def execute_query(self, sql: str) -> str:
+        """Execute SQL query"""
+        try:
+            conn_params = self.config.get_connection_params()
+            conn = psycopg2.connect(**conn_params)
+            self.log("info", f"Executing query: {sql}")
+
+            with conn.cursor() as cur:
+                # Start read-only transaction
+                cur.execute("BEGIN TRANSACTION READ ONLY")
+                try:
+                    cur.execute(sql)
+                    results = cur.fetchall()
+                    columns = [desc[0] for desc in cur.description]
+                    formatted_results = [dict(zip(columns, row)) for row in results]
+
+                    result_text = str({
+                        'columns': columns,
+                        'rows': formatted_results,
+                        'row_count': len(results)
+                    })
+
+                    self.log("info", f"Query completed, returned {len(results)} rows")
+                    return result_text
+                finally:
+                    cur.execute("ROLLBACK")
+        except psycopg2.Error as e:
+            error_msg = f"Query execution failed: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            self.log("error", error_msg)
+            raise
+        finally:
+            if conn:
+                conn.close()
+
+    async def cleanup(self):
+        """Cleanup resources"""
+        # No special cleanup needed since we're not using connection pool
+        pass