mcp-dbutils 0.2.3__py3-none-any.whl

mcp_dbutils/postgres/server.py

@@ -0,0 +1,205 @@
+"""PostgreSQL MCP server implementation"""
+import psycopg2
+from psycopg2.pool import SimpleConnectionPool
+from typing import Optional, List
+import mcp.types as types
+from ..base import DatabaseServer
+from ..log import create_logger
+from .config import PostgresConfig
+class PostgresServer(DatabaseServer):
+    def __init__(self, config: PostgresConfig, config_path: Optional[str] = None):
+        """初始化PostgreSQL服务器
+        Args:
+            config: 数据库配置
+            config_path: 配置文件路径(可选)
+        """
+        super().__init__("postgres-server", config.debug)
+        self.config = config
+        self.config_path = config_path
+        self.log = create_logger("postgres", config.debug)
+        # 创建连接池
+        try:
+            conn_params = config.get_connection_params()
+            masked_params = config.get_masked_connection_info()
+            self.log("debug", f"正在连接数据库，参数: {masked_params}")
+            # 测试连接
+            test_conn = psycopg2.connect(**conn_params)
+            test_conn.close()
+            self.log("info", "测试连接成功")
+            # 创建连接池
+            self.pool = SimpleConnectionPool(1, 5, **conn_params)
+            self.log("info", "数据库连接池创建成功")
+        except psycopg2.Error as e:
+            self.log("error", f"数据库连接失败: [Code: {e.pgcode}] {e.pgerror or str(e)}")
+            raise
+    async def list_resources(self) -> list[types.Resource]:
+        """列出所有表资源"""
+        try:
+            conn = self.pool.getconn()
+            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.config.host}/{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"获取表列表失败: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            self.log("error", error_msg)
+            raise
+        finally:
+            self.pool.putconn(conn)
+    async def read_resource(self, uri: str) -> str:
+        """读取表结构信息"""
+        try:
+            table_name = uri.split('/')[-2]
+            conn = self.pool.getconn()
+            with conn.cursor() as cur:
+                # 获取列信息
+                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()
+                # 获取约束信息
+                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"读取表结构失败: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            self.log("error", error_msg)
+            raise
+        finally:
+            self.pool.putconn(conn)
+    def get_tools(self) -> list[types.Tool]:
+        """获取可用工具列表"""
+        return [
+            types.Tool(
+                name="query",
+                description="执行只读SQL查询",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "database": {
+                            "type": "string",
+                            "description": "数据库配置名称（可选）"
+                        },
+                        "sql": {
+                            "type": "string",
+                            "description": "SQL查询语句（仅支持SELECT）"
+                        }
+                    },
+                    "required": ["sql"]
+                }
+            )
+        ]
+    async def call_tool(self, name: str, arguments: dict) -> list[types.TextContent]:
+        """执行工具调用"""
+        if name != "query":
+            raise ValueError(f"未知工具: {name}")
+        sql = arguments.get("sql", "").strip()
+        if not sql:
+            raise ValueError("SQL查询不能为空")
+        # 仅允许SELECT语句
+        if not sql.lower().startswith("select"):
+            raise ValueError("仅支持SELECT查询")
+        database = arguments.get("database")
+        use_pool = True
+        conn = None
+        try:
+            if database and self.config_path:
+                # 使用指定的数据库配置
+                config = PostgresConfig.from_yaml(self.config_path, database)
+                conn_params = config.get_connection_params()
+                masked_params = config.get_masked_connection_info()
+                self.log("info", f"使用配置 {database} 连接数据库: {masked_params}")
+                conn = psycopg2.connect(**conn_params)
+                use_pool = False
+            else:
+                # 使用现有连接池
+                conn = self.pool.getconn()
+            self.log("info", f"执行查询: {sql}")
+            with conn.cursor() as cur:
+                # 启动只读事务
+                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({
+                        'type': 'postgres',
+                        'config_name': database or 'default',
+                        'query_result': {
+                            'columns': columns,
+                            'rows': formatted_results,
+                            'row_count': len(results)
+                        }
+                    })
+                    self.log("info", f"查询完成，返回{len(results)}行结果")
+                    return [types.TextContent(type="text", text=result_text)]
+                finally:
+                    cur.execute("ROLLBACK")
+        except Exception as e:
+            if isinstance(e, psycopg2.Error):
+                error = f"查询执行失败: [Code: {e.pgcode}] {e.pgerror or str(e)}"
+            else:
+                error = f"查询执行失败: {str(e)}"
+            error_msg = str({
+                'type': 'postgres',
+                'config_name': database or 'default',
+                'error': error
+            })
+            self.log("error", error_msg)
+            return [types.TextContent(type="text", text=error_msg)]
+        finally:
+            if conn:
+                if use_pool:
+                    self.pool.putconn(conn)
+                else:
+                    conn.close()
+    async def cleanup(self):
+        """清理资源"""
+        if hasattr(self, 'pool'):
+            self.log("info", "关闭数据库连接池")
+            self.pool.closeall()

mcp_dbutils/sqlite/__init__.py

@@ -0,0 +1,6 @@
+"""SQLite module"""
+
+from .handler import SqliteHandler
+from .config import SqliteConfig
+
+__all__ = ['SqliteHandler', 'SqliteConfig']

mcp_dbutils/sqlite/config.py

@@ -0,0 +1,74 @@
+"""SQLite configuration module"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Any, Optional, Literal
+from ..config import DatabaseConfig
+
+@dataclass
+class SqliteConfig(DatabaseConfig):
+    path: str
+    password: Optional[str] = None
+    uri: bool = True  # Enable URI mode to support parameters like password
+    type: Literal['sqlite'] = 'sqlite'
+
+    @property
+    def absolute_path(self) -> str:
+        """Return absolute path to database file"""
+        return str(Path(self.path).expanduser().resolve())
+
+    def get_connection_params(self) -> Dict[str, Any]:
+        """Get sqlite3 connection parameters"""
+        if not self.password:
+            return {'database': self.absolute_path, 'uri': self.uri}
+
+        # Use URI format if password is provided
+        uri = f"file:{self.absolute_path}?mode=rw"
+        if self.password:
+            uri += f"&password={self.password}"
+
+        return {
+            'database': uri,
+            'uri': True
+        }
+
+    def get_masked_connection_info(self) -> Dict[str, Any]:
+        """Return connection information for logging"""
+        info = {
+            'database': self.absolute_path,
+            'uri': self.uri
+        }
+        if self.password:
+            info['password'] = '******'
+        return info
+
+    @classmethod
+    def from_yaml(cls, yaml_path: str, db_name: str, **kwargs) -> 'SqliteConfig':
+        """Create SQLite configuration from YAML
+
+        Args:
+            yaml_path: Path to YAML configuration file
+            db_name: Database configuration name
+        """
+        configs = cls.load_yaml_config(yaml_path)
+
+        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'] != 'sqlite':
+            raise ValueError(f"Configuration is not SQLite type: {db_config['type']}")
+        if 'path' not in db_config:
+            raise ValueError("SQLite configuration must include 'path' field")
+
+        config = cls(
+            path=db_config['path'],
+            password=db_config.get('password'),
+            uri=True
+        )
+        config.debug = cls.get_debug_mode()
+        return config

mcp_dbutils/sqlite/handler.py

@@ -0,0 +1,117 @@
+"""SQLite database handler implementation"""
+
+import sqlite3
+from pathlib import Path
+from contextlib import closing
+import mcp.types as types
+
+from ..base import DatabaseHandler
+from .config import SqliteConfig
+
+class SqliteHandler(DatabaseHandler):
+    def __init__(self, config_path: str, database: str, debug: bool = False):
+        """Initialize SQLite handler
+
+        Args:
+            config_path: Path to configuration file
+            database: Database configuration name
+            debug: Enable debug mode
+        """
+        super().__init__(config_path, database, debug)
+        self.config = SqliteConfig.from_yaml(config_path, database)
+
+        # Ensure database directory exists
+        db_file = Path(self.config.absolute_path)
+        db_file.parent.mkdir(parents=True, exist_ok=True)
+
+        # No connection test during initialization
+        self.log("debug", f"Configuring database: {self.config.get_masked_connection_info()}")
+
+    def _get_connection(self):
+        """Get database connection"""
+        connection_params = self.config.get_connection_params()
+        conn = sqlite3.connect(**connection_params)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    async def get_tables(self) -> list[types.Resource]:
+        """Get all table resources"""
+        try:
+            with closing(self._get_connection()) as conn:
+                cursor = conn.execute(
+                    "SELECT name FROM sqlite_master WHERE type='table'"
+                )
+                tables = cursor.fetchall()
+
+                return [
+                    types.Resource(
+                        uri=f"sqlite://{self.database}/{table[0]}/schema",
+                        name=f"{table[0]} schema",
+                        mimeType="application/json"
+                    ) for table in tables
+                ]
+        except sqlite3.Error as e:
+            error_msg = f"Failed to get table list: {str(e)}"
+            self.log("error", error_msg)
+            raise
+
+    async def get_schema(self, table_name: str) -> str:
+        """Get table schema information"""
+        try:
+            with closing(self._get_connection()) as conn:
+                # Get table structure
+                cursor = conn.execute(f"PRAGMA table_info({table_name})")
+                columns = cursor.fetchall()
+
+                # Get index information
+                cursor = conn.execute(f"PRAGMA index_list({table_name})")
+                indexes = cursor.fetchall()
+
+                schema_info = {
+                    'columns': [{
+                        'name': col['name'],
+                        'type': col['type'],
+                        'nullable': not col['notnull'],
+                        'primary_key': bool(col['pk'])
+                    } for col in columns],
+                    'indexes': [{
+                        'name': idx['name'],
+                        'unique': bool(idx['unique'])
+                    } for idx in indexes]
+                }
+
+                return str(schema_info)
+        except sqlite3.Error as e:
+            error_msg = f"Failed to read table schema: {str(e)}"
+            self.log("error", error_msg)
+            raise
+
+    async def execute_query(self, sql: str) -> str:
+        """Execute SQL query"""
+        try:
+            with closing(self._get_connection()) as conn:
+                self.log("info", f"Executing query: {sql}")
+                cursor = conn.execute(sql)
+                results = cursor.fetchall()
+
+                columns = [desc[0] for desc in cursor.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
+
+        except sqlite3.Error as e:
+            error_msg = f"Query execution failed: {str(e)}"
+            self.log("error", error_msg)
+            raise
+
+    async def cleanup(self):
+        """Cleanup resources"""
+        # No special cleanup needed for SQLite
+        pass

mcp_dbutils/sqlite/server.py

@@ -0,0 +1,202 @@
+"""SQLite MCP server implementation"""
+
+import sqlite3
+from pathlib import Path
+from contextlib import closing
+from typing import Optional, List
+import mcp.types as types
+
+from ..base import DatabaseServer
+from ..log import create_logger
+from .config import SqliteConfig
+
+class SqliteServer(DatabaseServer):
+    def __init__(self, config: SqliteConfig, config_path: Optional[str] = None):
+        """初始化 SQLite 服务器
+
+        Args:
+            config: SQLite 配置
+        """
+        super().__init__("sqlite-server", config.debug)
+        self.config = config
+        self.config_path = config_path
+        self.log = create_logger("sqlite", config.debug)
+
+        # 确保数据库目录存在
+        db_file = Path(self.config.absolute_path)
+        db_file.parent.mkdir(parents=True, exist_ok=True)
+
+        # 测试连接
+        try:
+            self.log("debug", f"正在连接数据库: {self.config.get_masked_connection_info()}")
+            connection_params = self.config.get_connection_params()
+            with closing(sqlite3.connect(**connection_params)) as conn:
+                conn.row_factory = sqlite3.Row
+            self.log("info", "数据库连接测试成功")
+        except sqlite3.Error as e:
+            self.log("error", f"数据库连接失败: {str(e)}")
+            raise
+
+    def _get_connection(self):
+        """获取数据库连接"""
+        connection_params = self.config.get_connection_params()
+        conn = sqlite3.connect(**connection_params)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    async def list_resources(self) -> list[types.Resource]:
+        """列出所有表资源"""
+        use_default = True
+        conn = None
+        try:
+            database = arguments.get("database")
+            if database and self.config_path:
+                # 使用指定的数据库配置
+                config = SqliteConfig.from_yaml(self.config_path, database)
+                connection_params = config.get_connection_params()
+                masked_params = config.get_masked_connection_info()
+                self.log("info", f"使用配置 {database} 连接数据库: {masked_params}")
+                conn = sqlite3.connect(**connection_params)
+                conn.row_factory = sqlite3.Row
+                use_default = False
+            else:
+                # 使用默认连接
+                conn = self._get_connection()
+
+            with closing(conn) as connection:
+                cursor = conn.execute(
+                    "SELECT name FROM sqlite_master WHERE type='table'"
+                )
+                tables = cursor.fetchall()
+
+                return [
+                    types.Resource(
+                        uri=f"sqlite://{table[0]}/schema",
+                        name=f"{table[0]} schema",
+                        mimeType="application/json"
+                    ) for table in tables
+                ]
+        except sqlite3.Error as e:
+            error_msg = f"获取表列表失败: {str(e)}"
+            self.log("error", error_msg)
+            raise
+
+    async def read_resource(self, uri: str) -> str:
+        """读取表结构信息"""
+        try:
+            table_name = uri.split('/')[-2]
+            with closing(self._get_connection()) as conn:
+                # 获取表结构
+                cursor = conn.execute(f"PRAGMA table_info({table_name})")
+                columns = cursor.fetchall()
+
+                # 获取索引信息
+                cursor = conn.execute(f"PRAGMA index_list({table_name})")
+                indexes = cursor.fetchall()
+
+                schema_info = {
+                    'columns': [{
+                        'name': col['name'],
+                        'type': col['type'],
+                        'nullable': not col['notnull'],
+                        'primary_key': bool(col['pk'])
+                    } for col in columns],
+                    'indexes': [{
+                        'name': idx['name'],
+                        'unique': bool(idx['unique'])
+                    } for idx in indexes]
+                }
+
+                return str(schema_info)
+        except sqlite3.Error as e:
+            error_msg = f"读取表结构失败: {str(e)}"
+            self.log("error", error_msg)
+            raise
+
+    def get_tools(self) -> list[types.Tool]:
+        """获取可用工具列表"""
+        return [
+            types.Tool(
+                name="query",
+                description="执行只读SQL查询",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "database": {
+                            "type": "string",
+                            "description": "数据库配置名称（可选）"
+                        },
+                        "sql": {
+                            "type": "string",
+                            "description": "SQL查询语句（仅支持SELECT）"
+                        }
+                    },
+                    "required": ["sql"]
+                }
+            )
+        ]
+
+    async def call_tool(self, name: str, arguments: dict) -> list[types.TextContent]:
+        """执行工具调用"""
+        if name != "query":
+            raise ValueError(f"未知工具: {name}")
+
+        sql = arguments.get("sql", "").strip()
+        if not sql:
+            raise ValueError("SQL查询不能为空")
+
+        # 仅允许SELECT语句
+        if not sql.lower().startswith("select"):
+            raise ValueError("仅支持SELECT查询")
+
+        use_default = True
+        conn = None
+        try:
+            database = arguments.get("database")
+            if database and self.config_path:
+                # 使用指定的数据库配置
+                config = SqliteConfig.from_yaml(self.config_path, database)
+                connection_params = config.get_connection_params()
+                masked_params = config.get_masked_connection_info()
+                self.log("info", f"使用配置 {database} 连接数据库: {masked_params}")
+                conn = sqlite3.connect(**connection_params)
+                conn.row_factory = sqlite3.Row
+                use_default = False
+            else:
+                # 使用默认连接
+                conn = self._get_connection()
+
+            with closing(conn) as connection:
+                self.log("info", f"执行查询: {sql}")
+                cursor = conn.execute(sql)
+                results = cursor.fetchall()
+
+                columns = [desc[0] for desc in cursor.description]
+                formatted_results = [dict(zip(columns, row)) for row in results]
+
+                result_text = str({
+                    'type': 'sqlite',
+                    'config_name': database or 'default',
+                    'query_result': {
+                        'columns': columns,
+                        'rows': formatted_results,
+                        'row_count': len(results)
+                    }
+                })
+
+                self.log("info", f"查询完成，返回{len(results)}行结果")
+                return [types.TextContent(type="text", text=result_text)]
+
+        except sqlite3.Error as e:
+            error_msg = str({
+                'type': 'sqlite',
+                'config_name': database or 'default',
+                'error': f"查询执行失败: {str(e)}"
+            })
+            self.log("error", error_msg)
+            return [types.TextContent(type="text", text=error_msg)]
+
+    async def cleanup(self):
+        """清理资源"""
+        # SQLite不需要特别的清理操作
+        pass