multimymcp 1.0.0__py3-none-any.whl

multimymcp/__init__.py

@@ -0,0 +1,24 @@
+"""
+MultiMyMCP - 生产级 MySQL 多数据源 MCP 工具
+
+版本: 1.0.0
+作者: TraeMCP Team
+适配环境: TRAE CN IDE
+
+核心功能:
+- 增强型连接池管理
+- 数据源配置加密存储
+- SQL 执行超时控制
+- 连接池状态监控
+- SQL 白名单/黑名单机制
+"""
+
+__version__ = "1.0.0"
+__author__ = "TraeMCP Team"
+__all__ = ["MultiMyMCP", "DataSourceConfig", "ConnectionPool", "SQLExecutor", "Monitor"]
+
+from multimymcp.core import MultiMyMCP
+from multimymcp.config import DataSourceConfig
+from multimymcp.pool import ConnectionPool
+from multimymcp.executor import SQLExecutor
+from multimymcp.monitor import Monitor

multimymcp/cli.py

@@ -0,0 +1,138 @@
+"""
+命令行接口模块
+
+提供命令行工具，用于管理数据源、查看状态等
+"""
+
+import argparse
+import json
+from multimymcp import MultiMyMCP
+from multimymcp.mcp_config_loader import MCPConfigLoader
+
+
+def resolve_config_path(config_path: str = None) -> str:
+    """
+    解析 CLI 应使用的配置文件路径。
+
+    Args:
+        config_path: 命令行显式传入的配置文件路径
+
+    Returns:
+        str: 可用的配置文件路径；若未找到则返回原始入参
+    """
+    if config_path:
+        return config_path
+
+    loader = MCPConfigLoader()
+    default_config = loader.get_default_config_path()
+
+    # 保持与 MCP Server 一致的默认查找顺序。
+    if default_config.exists():
+        return str(default_config)
+
+    if loader.multimymcp_config_file.exists():
+        return str(loader.multimymcp_config_file)
+
+    return config_path
+
+
+def main():
+    """
+    主函数
+    """
+    parser = argparse.ArgumentParser(description='MultiMyMCP 命令行工具')
+    
+    # 全局参数
+    parser.add_argument('--config', '-c', type=str, help='配置文件路径')
+    parser.add_argument('--key', '-k', type=str, help='加密密钥')
+    
+    # 子命令
+    subparsers = parser.add_subparsers(dest='command', help='子命令')
+    
+    # 连接数据源
+    connect_parser = subparsers.add_parser('connect', help='连接数据源')
+    connect_parser.add_argument('datasource', type=str, help='数据源名称')
+    
+    # 执行SQL
+    execute_parser = subparsers.add_parser('execute', help='执行SQL')
+    execute_parser.add_argument('sql', type=str, help='SQL语句')
+    execute_parser.add_argument('--datasource', '-d', type=str, default='default', help='数据源名称')
+    execute_parser.add_argument('--params', '-p', type=str, help='参数(JSON格式)')
+    
+    # 查看状态
+    status_parser = subparsers.add_parser('status', help='查看连接池状态')
+    status_parser.add_argument('--datasource', '-d', type=str, default='default', help='数据源名称')
+    
+    # 健康检查
+    health_parser = subparsers.add_parser('health', help='健康检查')
+    health_parser.add_argument('--datasource', '-d', type=str, default='default', help='数据源名称')
+    
+    # 性能报告
+    perf_parser = subparsers.add_parser('performance', help='性能报告')
+    perf_parser.add_argument('--datasource', '-d', type=str, default='default', help='数据源名称')
+    
+    # 列出数据源
+    list_parser = subparsers.add_parser('list', help='列出所有数据源')
+    
+    # 调整连接池大小
+    resize_parser = subparsers.add_parser('resize', help='调整连接池大小')
+    resize_parser.add_argument('min_size', type=int, help='最小连接数')
+    resize_parser.add_argument('max_size', type=int, help='最大连接数')
+    resize_parser.add_argument('--datasource', '-d', type=str, default='default', help='数据源名称')
+    
+    # 保存配置
+    save_parser = subparsers.add_parser('save', help='保存配置')
+    save_parser.add_argument('file', type=str, help='保存文件路径')
+    
+    # 加载配置
+    load_parser = subparsers.add_parser('load', help='加载配置')
+    load_parser.add_argument('file', type=str, help='配置文件路径')
+    
+    args = parser.parse_args()
+    resolved_config_path = resolve_config_path(args.config)
+
+    with MultiMyMCP(resolved_config_path, args.key) as mcp:
+        if args.command == 'connect':
+            success = mcp.connect(args.datasource)
+            print(f"连接成功: {success}")
+            
+        elif args.command == 'execute':
+            params = None
+            if args.params:
+                params = json.loads(args.params)
+            result = mcp.execute(args.sql, params, args.datasource)
+            print(json.dumps(result, indent=2, ensure_ascii=False))
+            
+        elif args.command == 'status':
+            status = mcp.get_pool_status(args.datasource)
+            print(json.dumps(status, indent=2, ensure_ascii=False))
+            
+        elif args.command == 'health':
+            health = mcp.get_health_status(args.datasource)
+            print(json.dumps(health, indent=2, ensure_ascii=False))
+            
+        elif args.command == 'performance':
+            report = mcp.get_performance_report(args.datasource)
+            print(json.dumps(report, indent=2, ensure_ascii=False))
+            
+        elif args.command == 'list':
+            datasources = mcp.list_data_sources()
+            print("数据源列表:")
+            for ds in datasources:
+                print(f"  - {ds}")
+            
+        elif args.command == 'resize':
+            mcp.resize_pool(args.min_size, args.max_size, args.datasource)
+            print(f"连接池大小已调整: 最小={args.min_size}, 最大={args.max_size}")
+            
+        elif args.command == 'save':
+            mcp.save_config(args.file)
+            print(f"配置已保存到: {args.file}")
+            
+        elif args.command == 'load':
+            mcp.load_config(args.file)
+            print(f"配置已加载: {args.file}")
+
+
+if __name__ == '__main__':
+    main()

multimymcp/config.py

@@ -0,0 +1,300 @@
+"""
+配置管理模块
+
+支持多数据源配置、环境变量配置、JSON5配置文件
+配置加密存储,运行时解密
+"""
+
+import os
+import json5
+from typing import Dict, Optional, Any
+from dataclasses import dataclass, field
+from pathlib import Path
+from dotenv import load_dotenv
+from multimymcp.encryption import EncryptionUtil
+from multimymcp.exceptions import ConfigurationError
+
+
+@dataclass
+class DataSourceConfig:
+    """
+    数据源配置类
+    
+    Attributes:
+        name: 数据源名称
+        host: 数据库主机地址
+        port: 数据库端口
+        user: 数据库用户名
+        password: 数据库密码(加密存储)
+        database: 数据库名称
+        pool_min_size: 连接池最小连接数
+        pool_max_size: 连接池最大连接数
+        pool_timeout: 连接池超时时间(秒)
+        sql_timeout: SQL执行超时时间(秒)
+        charset: 字符集
+        autocommit: 是否自动提交
+        encrypted: 密码是否已加密
+    """
+    name: str
+    host: str
+    port: int
+    user: str
+    password: str
+    database: str
+    pool_min_size: int = 2
+    pool_max_size: int = 10
+    pool_timeout: int = 30
+    sql_timeout: int = 60
+    charset: str = "utf8mb4"
+    autocommit: bool = False
+    encrypted: bool = False
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        转换为字典格式
+        
+        Returns:
+            Dict[str, Any]: 配置字典
+        """
+        return {
+            "name": self.name,
+            "host": self.host,
+            "port": self.port,
+            "user": self.user,
+            "password": self.password,
+            "database": self.database,
+            "pool_min_size": self.pool_min_size,
+            "pool_max_size": self.pool_max_size,
+            "pool_timeout": self.pool_timeout,
+            "sql_timeout": self.sql_timeout,
+            "charset": self.charset,
+            "autocommit": self.autocommit,
+            "encrypted": self.encrypted,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DataSourceConfig":
+        """
+        从字典创建配置对象
+        
+        Args:
+            data: 配置字典
+            
+        Returns:
+            DataSourceConfig: 配置对象
+        """
+        return cls(**data)
+
+
+@dataclass
+class SecurityConfig:
+    """
+    安全配置类
+    
+    Attributes:
+        whitelist_enabled: 是否启用白名单
+        blacklist_enabled: 是否启用黑名单
+        whitelist: SQL白名单列表
+        blacklist: SQL黑名单列表
+    """
+    whitelist_enabled: bool = False
+    blacklist_enabled: bool = True
+    whitelist: list = field(default_factory=list)
+    blacklist: list = field(default_factory=lambda: ["DROP", "TRUNCATE", "ALTER", "CREATE"])
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        转换为字典格式
+        
+        Returns:
+            Dict[str, Any]: 配置字典
+        """
+        return {
+            "whitelist_enabled": self.whitelist_enabled,
+            "blacklist_enabled": self.blacklist_enabled,
+            "whitelist": self.whitelist,
+            "blacklist": self.blacklist,
+        }
+
+
+class ConfigManager:
+    """
+    配置管理器
+    
+    管理多数据源配置、安全配置、加密配置等
+    支持从环境变量、JSON5文件加载配置
+    """
+    
+    def __init__(self, encryption_key: Optional[str] = None):
+        """
+        初始化配置管理器
+        
+        Args:
+            encryption_key: 加密密钥(可选,默认从环境变量读取)
+        """
+        load_dotenv()
+        
+        self.encryption_key = encryption_key or os.getenv("MYSQL_ENCRYPTION_KEY")
+        if not self.encryption_key:
+            self.encryption_key = EncryptionUtil.generate_key()
+        
+        self.encryption = EncryptionUtil(self.encryption_key)
+        self.data_sources: Dict[str, DataSourceConfig] = {}
+        self.security_config = SecurityConfig()
+        self._load_from_env()
+    
+    def _load_from_env(self):
+        """
+        从环境变量加载配置
+        """
+        host = os.getenv("MYSQL_HOST")
+        if host:
+            config = DataSourceConfig(
+                name="default",
+                host=host,
+                port=int(os.getenv("MYSQL_PORT", "3306")),
+                user=os.getenv("MYSQL_USER", "root"),
+                password=os.getenv("MYSQL_PASSWORD", ""),
+                database=os.getenv("MYSQL_DATABASE", ""),
+                pool_min_size=int(os.getenv("MYSQL_POOL_MIN_SIZE", "2")),
+                pool_max_size=int(os.getenv("MYSQL_POOL_MAX_SIZE", "10")),
+                pool_timeout=int(os.getenv("MYSQL_POOL_TIMEOUT", "30")),
+                sql_timeout=int(os.getenv("MYSQL_SQL_TIMEOUT", "60")),
+            )
+            self.add_data_source(config)
+        
+        self.security_config.whitelist_enabled = (
+            os.getenv("MYSQL_WHITELIST_ENABLED", "false").lower() == "true"
+        )
+        self.security_config.blacklist_enabled = (
+            os.getenv("MYSQL_BLACKLIST_ENABLED", "true").lower() == "true"
+        )
+        
+        blacklist_str = os.getenv("MYSQL_BLACKLIST", "")
+        if blacklist_str:
+            self.security_config.blacklist = blacklist_str.split(",")
+    
+    def load_from_file(self, file_path: str):
+        """
+        从JSON5文件加载配置
+        
+        Args:
+            file_path: 配置文件路径
+            
+        Raises:
+            ConfigurationError: 配置文件不存在或格式错误
+        """
+        path = Path(file_path)
+        if not path.exists():
+            raise ConfigurationError(f"配置文件不存在: {file_path}")
+        
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                config_data = json5.load(f)
+            
+            if "datasources" in config_data:
+                for ds_name, ds_config in config_data["datasources"].items():
+                    ds_config["name"] = ds_name
+                    config = DataSourceConfig.from_dict(ds_config)
+                    self.add_data_source(config)
+            
+            if "security" in config_data:
+                security_data = config_data["security"]
+                self.security_config = SecurityConfig(**security_data)
+                
+        except Exception as e:
+            raise ConfigurationError(f"配置文件解析失败: {str(e)}")
+    
+    def add_data_source(self, config: DataSourceConfig, encrypt: bool = True):
+        """
+        添加数据源配置
+        
+        Args:
+            config: 数据源配置对象
+            encrypt: 是否加密密码
+        """
+        if encrypt and not config.encrypted:
+            config.password = self.encryption.encrypt(config.password)
+            config.encrypted = True
+        
+        self.data_sources[config.name] = config
+    
+    def get_data_source(self, name: str) -> Optional[DataSourceConfig]:
+        """
+        获取数据源配置
+        
+        Args:
+            name: 数据源名称
+            
+        Returns:
+            Optional[DataSourceConfig]: 数据源配置对象
+        """
+        config = self.data_sources.get(name)
+        if config and config.encrypted:
+            config.password = self.encryption.decrypt(config.password)
+            config.encrypted = False
+        return config
+    
+    def remove_data_source(self, name: str) -> bool:
+        """
+        移除数据源配置
+        
+        Args:
+            name: 数据源名称
+            
+        Returns:
+            bool: 是否成功移除
+        """
+        if name in self.data_sources:
+            del self.data_sources[name]
+            return True
+        return False
+    
+    def list_data_sources(self) -> list:
+        """
+        列出所有数据源名称
+        
+        Returns:
+            list: 数据源名称列表
+        """
+        return list(self.data_sources.keys())
+    
+    def save_to_file(self, file_path: str):
+        """
+        保存配置到JSON5文件
+        
+        Args:
+            file_path: 配置文件路径
+        """
+        config_data = {
+            "datasources": {
+                name: config.to_dict() for name, config in self.datasources.items()
+            },
+            "security": self.security_config.to_dict(),
+        }
+        
+        path = Path(file_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(path, "w", encoding="utf-8") as f:
+            json5.dump(config_data, f, indent=2, ensure_ascii=False)
+    
+    def get_security_config(self) -> SecurityConfig:
+        """
+        获取安全配置
+        
+        Returns:
+            SecurityConfig: 安全配置对象
+        """
+        return self.security_config
+    
+    def update_security_config(self, **kwargs):
+        """
+        更新安全配置
+        
+        Args:
+            **kwargs: 安全配置参数
+        """
+        for key, value in kwargs.items():
+            if hasattr(self.security_config, key):
+                setattr(self.security_config, key, value)