sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/database/manager.py

@@ -0,0 +1,178 @@
+"""数据库连接管理器 - 统一管理所有数据库连接"""
+
+from typing import Optional
+
+from ..config import get_config_manager
+from ..settings import DatabaseConfig
+from .connectors.base import BaseConnector, QueryResult
+from .connectors.mysql import MySQLConnector
+from .connectors.sqlserver import SQLServerConnector
+from .connectors.postgresql import PostgreSQLConnector
+from .connectors.redis import RedisConnector
+from .connectors.mongodb import MongoDBConnector
+
+
+class DatabaseManager:
+    """数据库连接管理器"""
+
+    def __init__(self):
+        self._connectors: dict[str, BaseConnector] = {}
+        self._schema_cache: dict[str, dict] = {}
+
+    def _create_connector(self, config: DatabaseConfig) -> BaseConnector:
+        """根据配置创建对应的连接器"""
+        connector_classes = {
+            "mysql": MySQLConnector,
+            "sqlserver": SQLServerConnector,
+            "postgresql": PostgreSQLConnector,
+            "redis": RedisConnector,
+            "mongodb": MongoDBConnector,
+        }
+        cls = connector_classes.get(config.db_type)
+        if cls is None:
+            raise ValueError(f"不支持的数据库类型: {config.db_type}")
+
+        return cls(
+            host=config.host,
+            port=config.get_port(),
+            user=config.user,
+            password=config.password,
+            database=config.database,
+        )
+
+    async def get_connector(self) -> BaseConnector:
+        """获取当前激活的数据库连接器（自动连接）"""
+        config = get_config_manager().get_active_database()
+        if not config:
+            raise ValueError("未配置数据库，请先在设置中添加数据库连接")
+
+        key = config.name
+        if key not in self._connectors:
+            connector = self._create_connector(config)
+            await connector.connect()
+            self._connectors[key] = connector
+
+        return self._connectors[key]
+
+    async def execute(self, sql: str) -> QueryResult:
+        """执行 SQL 并返回结果"""
+        connector = await self.get_connector()
+        return await connector.execute(sql)
+
+    async def get_schema(self, force_refresh: bool = False) -> dict:
+        """获取当前激活数据库的 schema（带缓存）"""
+        config = get_config_manager().get_active_database()
+        if not config:
+            return {"db_type": "", "tables": [], "error": "未配置数据库连接"}
+
+        cache_key = config.name
+        if not force_refresh and cache_key in self._schema_cache:
+            return self._schema_cache[cache_key]
+
+        try:
+            connector = await self.get_connector()
+            schema = await connector.get_schema()
+            self._schema_cache[cache_key] = schema
+            return schema
+        except Exception as e:
+            return {"db_type": config.db_type, "tables": [], "error": str(e)}
+
+    async def refresh_schema(self) -> dict:
+        """强制刷新当前数据库 schema"""
+        config = get_config_manager().get_active_database()
+        if config:
+            cache_key = config.name
+            self._schema_cache.pop(cache_key, None)
+            if cache_key in self._connectors:
+                await self._connectors[cache_key].disconnect()
+                del self._connectors[cache_key]
+        return await self.get_schema(force_refresh=True)
+
+    async def get_schema_text(self) -> str:
+        """将 schema 转为给 LLM 的文本描述"""
+        schema = await self.get_schema()
+
+        db_type = schema.get("db_type", "")
+
+        if db_type == "redis":
+            keys = schema.get("keys", [])
+            if not keys:
+                return "(Redis: 当前数据库无 key)"
+            lines = ["## 当前 Redis Key 列表（前50个）:"]
+            for k in keys:
+                lines.append(f"- {k['name']} ({k['type']})")
+            key_count = schema.get("key_count", 0)
+            if key_count > 50:
+                lines.append(f"... 共 {key_count} 个 key")
+            return "\n".join(lines)
+
+        tables = schema.get("tables", [])
+        if not tables:
+            error = schema.get("error", "")
+            if error:
+                return f"(Schema 获取失败: {error})"
+            return "(数据库为空，暂无表)"
+
+        lines = ["## 当前数据库 Schema:", ""]
+        for t in tables:
+            cols = t.get("columns", [])
+            if not cols:
+                lines.append(f"- {t['name']} (无列信息)")
+                continue
+            col_strs = []
+            for c in cols:
+                extras = []
+                if c.get("key") == "PRI":
+                    extras.append("PRIMARY KEY")
+                if not c.get("nullable", True):
+                    extras.append("NOT NULL")
+                suffix = f"  -- {', '.join(extras)}" if extras else ""
+                col_strs.append(f"    {c['name']} {c['type']}{suffix}")
+            lines.append(f"- {t['name']}:")
+            lines.extend(col_strs)
+            lines.append("")
+
+        return "\n".join(lines)
+
+    async def test_active_connection(self) -> dict:
+        """测试当前激活的连接"""
+        config = get_config_manager().get_active_database()
+        if not config:
+            return {"success": False, "error": "未配置数据库连接"}
+
+        connector = self._create_connector(config)
+        try:
+            result = await connector.test_connection()
+            return result
+        except Exception as e:
+            return {"success": False, "error": str(e), "code": "CONNECTION_FAILED"}
+        finally:
+            await connector.disconnect()
+
+    async def test_connection(self, config: DatabaseConfig) -> dict:
+        """测试指定配置的连接"""
+        connector = self._create_connector(config)
+        try:
+            result = await connector.test_connection()
+            return result
+        except Exception as e:
+            return {"success": False, "error": str(e), "code": "CONNECTION_FAILED"}
+        finally:
+            await connector.disconnect()
+
+    async def close_all(self):
+        """关闭所有连接"""
+        for connector in self._connectors.values():
+            await connector.disconnect()
+        self._connectors.clear()
+
+
+# 全局单例
+_db_manager: Optional[DatabaseManager] = None
+
+
+def get_db_manager() -> DatabaseManager:
+    global _db_manager
+    if _db_manager is None:
+        _db_manager = DatabaseManager()
+    return _db_manager

sql_assistant/database/security.py

@@ -0,0 +1,230 @@
+"""
+SQL 安全保护模块 - 防止 SQL 注入攻击
+"""
+import re
+from typing import Optional, Tuple, List
+from dataclasses import dataclass
+
+
+@dataclass
+class SQLCheckResult:
+    """SQL 检查结果"""
+    is_safe: bool
+    warning: str = ""
+    risk_level: str = "none"  # none, low, medium, high
+    blocked_reason: str = ""
+
+
+class SQLSecurityGuard:
+    """SQL 安全守卫类"""
+    
+    # 危险 SQL 关键字和模式
+    DANGEROUS_KEYWORDS = [
+        # 注释相关
+        r"/\*", r"\*/", r"--", r"#",
+        # 联合查询
+        r"\bUNION\b", r"\bUNION\s+ALL\b",
+        # 子查询
+        r"\bSELECT\b.*\bFROM\b.*\bSELECT\b",
+        # 批处理
+        r";",
+        # 信息_schema
+        r"\binformation_schema\b",
+        r"\bmysql\.",
+        r"\bsys\.",
+        # 系统函数
+        r"\bLOAD_FILE\b",
+        r"\bINTO\s+OUTFILE\b",
+        r"\bINTO\s+DUMPFILE\b",
+        r"\bEXEC\b",
+        r"\bEXECUTE\b",
+        r"\bSYSTEM\b",
+        r"\bSHELL\b",
+        r"\bCMD\b",
+        # 数据操作
+        r"\bDROP\b",
+        r"\bTRUNCATE\b",
+        r"\bALTER\b.*\bDROP\b",
+        # 用户权限
+        r"\bGRANT\b",
+        r"\bREVOKE\b",
+        r"\bCREATE\s+USER\b",
+        r"\bALTER\s+USER\b",
+        # 十六进制
+        r"0x[0-9a-fA-F]+",
+        # 字符串拼接
+        r"\|\|",
+        r"CONCAT\s*\(",
+        r"GROUP_CONCAT\s*\(",
+        # 延迟攻击
+        r"\bSLEEP\b",
+        r"\bBENCHMARK\b",
+        # 布尔注入
+        r"\bOR\s+1\s*=\s*1\b",
+        r"\bAND\s+1\s*=\s*1\b",
+        r"\bOR\s+TRUE\b",
+        r"\bAND\s+FALSE\b",
+        # 盲注
+        r"\bIF\s*\(",
+        r"\bCASE\s+WHEN\b",
+        r"\bEXISTS\s*\(",
+        # 其他危险
+        r"\bDELETE\b.*\bWHERE\s*1\s*=\s*1\b",
+        r"\bUPDATE\b.*\bWHERE\s*1\s*=\s*1\b",
+    ]
+    
+    # 允许的 SQL 类型
+    ALLOWED_SQL_TYPES = {
+        "SELECT", "INSERT", "UPDATE", "DELETE",
+        "SHOW", "DESCRIBE", "EXPLAIN",
+        "CREATE", "ALTER", "DROP", "TRUNCATE",
+        "USE", "SET"
+    }
+    
+    def __init__(self):
+        # 编译正则表达式
+        self.dangerous_patterns = [
+            re.compile(pattern, re.IGNORECASE)
+            for pattern in self.DANGEROUS_KEYWORDS
+        ]
+    
+    def check_sql_safety(self, sql: str) -> SQLCheckResult:
+        """
+        检查 SQL 语句安全性
+        
+        Args:
+            sql: SQL 语句
+            
+        Returns:
+            SQLCheckResult 检查结果
+        """
+        sql = sql.strip()
+        
+        if not sql:
+            return SQLCheckResult(
+                is_safe=False,
+                blocked_reason="SQL 语句为空"
+            )
+        
+        # 检查危险模式
+        warnings = []
+        risk_level = "none"
+        
+        for pattern in self.dangerous_patterns:
+            matches = pattern.findall(sql)
+            if matches:
+                for match in matches:
+                    warning_msg = self._get_warning_message(match)
+                    if warning_msg:
+                        warnings.append(warning_msg)
+                        # 升级风险等级
+                        if "high" in warning_msg.lower():
+                            risk_level = "high"
+                        elif risk_level != "high":
+                            risk_level = "medium"
+        
+        if risk_level == "high":
+            return SQLCheckResult(
+                is_safe=False,
+                warning="; ".join(warnings),
+                risk_level=risk_level,
+                blocked_reason="检测到高风险 SQL 注入模式"
+            )
+        
+        if warnings:
+            return SQLCheckResult(
+                is_safe=True,  # 中等风险仍允许执行，但给出警告
+                warning="; ".join(warnings),
+                risk_level=risk_level
+            )
+        
+        return SQLCheckResult(is_safe=True)
+    
+    def _get_warning_message(self, match: str) -> str:
+        """根据匹配内容获取警告信息"""
+        match_lower = match.lower()
+        
+        if any(keyword in match_lower for keyword in ["union", "select.*from.*select"]):
+            return "⚠️ 检测到联合查询模式，可能存在注入风险"
+        elif any(keyword in match_lower for keyword in ["--", "/*", "*/", "#"]):
+            return "⚠️ 检测到注释符，请确认操作意图"
+        elif ";" in match:
+            return "⚠️ 检测到分号分隔符，可能存在多语句执行风险"
+        elif any(keyword in match_lower for keyword in ["information_schema", "mysql.", "sys."]):
+            return "⚠️ 检测到系统表访问，请注意数据安全"
+        elif any(keyword in match_lower for keyword in ["load_file", "into outfile", "into dumpfile", "exec", "execute", "system", "shell"]):
+            return "🔴 检测到高风险系统函数调用"
+        elif any(keyword in match_lower for keyword in ["sleep", "benchmark"]):
+            return "🔴 检测到延迟注入函数"
+        elif any(keyword in match_lower for keyword in ["or 1=1", "and 1=1", "or true", "and false"]):
+            return "🔴 检测到布尔注入模式"
+        elif any(keyword in match_lower for keyword in ["drop", "truncate"]):
+            return "⚠️ 检测到数据删除操作，请谨慎执行"
+        elif any(keyword in match_lower for keyword in ["0x"]):
+            return "⚠️ 检测到十六进制编码"
+        elif any(keyword in match_lower for keyword in ["concat", "group_concat", "||"]):
+            return "⚠️ 检测到字符串拼接，请注意注入风险"
+        
+        return ""
+    
+    def sanitize_sql(self, sql: str) -> Tuple[str, List[str]]:
+        """
+        清理 SQL 语句，移除一些明显的危险内容
+        
+        Args:
+            sql: 原始 SQL
+            
+        Returns:
+            (清理后的 SQL, 警告列表)
+        """
+        warnings = []
+        cleaned = sql
+        
+        # 移除首尾空白
+        cleaned = cleaned.strip()
+        
+        # 检查并警告危险模式（不自动修改）
+        check_result = self.check_sql_safety(cleaned)
+        if check_result.warning:
+            warnings.append(check_result.warning)
+        
+        return cleaned, warnings
+    
+    def requires_confirmation(self, sql: str) -> Tuple[bool, str]:
+        """
+        判断 SQL 是否需要用户确认
+        
+        Args:
+            sql: SQL 语句
+            
+        Returns:
+            (是否需要确认, 确认提示信息)
+        """
+        from .connectors.base import BaseConnector
+        
+        sql_type = BaseConnector.classify_sql(sql)
+        
+        # 增删改操作需要确认
+        if sql_type in ["INSERT", "UPDATE", "DELETE", "DDL"]:
+            reason_map = {
+                "INSERT": "即将执行 INSERT 操作，会添加新数据",
+                "UPDATE": "即将执行 UPDATE 操作，会修改现有数据",
+                "DELETE": "即将执行 DELETE 操作，会删除数据",
+                "DDL": "即将执行 DDL 操作，会修改数据库结构"
+            }
+            return True, reason_map.get(sql_type, "即将执行数据修改操作")
+        
+        # SELECT 不需要确认
+        return False, ""
+
+
+# 全局单例
+_security_guard: Optional[SQLSecurityGuard] = None
+
+
+def get_security_guard() -> SQLSecurityGuard:
+    """获取 SQL 安全守卫单例"""
+    global _security_guard
+    if _security_guard is None:
+        _security_guard = SQLSecurityGuard()
+    return _security_guard

sql_assistant/llm/__init__.py

@@ -0,0 +1 @@
+"""LLM 集成模块"""

sql_assistant/llm/base.py

@@ -0,0 +1,28 @@
+"""LLM Provider 抽象基类"""
+
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator
+
+
+class BaseLLMProvider(ABC):
+    """LLM Provider 抽象基类"""
+
+    def __init__(self, api_key: str, base_url: str, model: str):
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+
+    @abstractmethod
+    async def chat(self, messages: list[dict], temperature: float = 0.1) -> str:
+        """发送对话请求，返回完整响应"""
+        ...
+
+    @abstractmethod
+    async def chat_stream(self, messages: list[dict], temperature: float = 0.1) -> AsyncGenerator[str, None]:
+        """发送流式对话请求"""
+        ...
+
+    @abstractmethod
+    async def test_connection(self) -> dict:
+        """测试连接是否正常"""
+        ...

sql_assistant/llm/exceptions.py

@@ -0,0 +1,96 @@
+"""LLM Provider 统一异常"""
+
+from typing import Optional, Dict, Any
+
+
+class LLMError(Exception):
+    """LLM Provider 基础异常"""
+    def __init__(self, message: str, provider: str = "", code: Optional[str] = None):
+        super().__init__(message)
+        self.provider = provider
+        self.code = code
+
+    def to_dict(self) -> dict:
+        return {
+            "success": False,
+            "error": str(self),
+            "provider": self.provider,
+            "code": self.code,
+        }
+
+
+class LLMConnectionError(LLMError):
+    """LLM 连接失败异常"""
+    def __init__(self, message: str, provider: str = ""):
+        super().__init__(message, provider, code="LLM_CONNECTION_FAILED")
+
+
+class LLMTimeoutError(LLMError):
+    """LLM 请求超时异常"""
+    def __init__(self, message: str, provider: str = ""):
+        super().__init__(message, provider, code="LLM_TIMEOUT")
+
+
+class LLMAuthError(LLMError):
+    """LLM 认证失败异常"""
+    def __init__(self, message: str, provider: str = ""):
+        super().__init__(message, provider, code="LLM_AUTH_FAILED")
+
+
+class LLMRateLimitError(LLMError):
+    """LLM 限流异常"""
+    def __init__(self, message: str, provider: str = "", retry_after: Optional[int] = None):
+        super().__init__(message, provider, code="LLM_RATE_LIMITED")
+        self.retry_after = retry_after
+
+    def to_dict(self) -> dict:
+        result = super().to_dict()
+        if self.retry_after:
+            result["retry_after"] = self.retry_after
+        return result
+
+
+class LLMResponseError(LLMError):
+    """LLM 响应解析异常"""
+    def __init__(self, message: str, provider: str = "", raw_response: Any = None):
+        super().__init__(message, provider, code="LLM_RESPONSE_ERROR")
+        self.raw_response = raw_response
+
+    def to_dict(self) -> dict:
+        result = super().to_dict()
+        if self.raw_response:
+            result["raw_response"] = str(self.raw_response)[:500]
+        return result
+
+
+def format_llm_result(
+    success: bool,
+    data: Any = None,
+    error: Optional[str] = None,
+    provider: str = "",
+    code: Optional[str] = None,
+) -> Dict[str, Any]:
+    """统一格式化 LLM 返回结果
+
+    Args:
+        success: 是否成功
+        data: 成功时返回的数据
+        error: 失败时的错误信息
+        provider: 提供商名称
+        code: 错误码
+
+    Returns:
+        统一格式的字典
+    """
+    result: Dict[str, Any] = {"success": success}
+
+    if success:
+        result["data"] = data
+    else:
+        result["error"] = error or "未知错误"
+        if code:
+            result["code"] = code
+        if provider:
+            result["provider"] = provider
+
+    return result

sql_assistant/llm/manager.py

@@ -0,0 +1,82 @@
+"""LLM Manager - 统一管理所有 Provider"""
+
+from typing import AsyncGenerator, Optional
+
+from ..config import get_config_manager
+from ..settings import LLMProviderConfig
+from .base import BaseLLMProvider
+from .providers.openai_compatible import OpenAICompatibleProvider
+from .providers.gemini import GeminiProvider
+from .providers.claude import ClaudeProvider
+
+
+class LLMManager:
+    """LLM Provider 管理器"""
+
+    def __init__(self):
+        self._providers: dict[str, BaseLLMProvider] = {}
+
+    def _create_provider(self, config: LLMProviderConfig) -> BaseLLMProvider:
+        """根据配置创建对应的 Provider 实例"""
+        if config.provider == "gemini":
+            return GeminiProvider(
+                api_key=config.api_key,
+                base_url=config.get_base_url(),
+                model=config.get_model(),
+            )
+        elif config.provider == "claude" or config.provider == "minimax":
+            return ClaudeProvider(
+                api_key=config.api_key,
+                base_url=config.get_base_url(),
+                model=config.get_model(),
+            )
+        else:
+            # OpenAI 兼容接口
+            return OpenAICompatibleProvider(
+                api_key=config.api_key,
+                base_url=config.get_base_url(),
+                model=config.get_model(),
+                provider_name=config.provider,
+            )
+
+    def get_provider(self, config: LLMProviderConfig) -> BaseLLMProvider:
+        """获取或创建 Provider (带缓存)"""
+        key = config.name
+        if key not in self._providers:
+            self._providers[key] = self._create_provider(config)
+        return self._providers[key]
+
+    async def chat(self, messages: list[dict], temperature: float = 0.1) -> str:
+        """使用当前激活的 LLM 发送对话请求"""
+        config = get_config_manager().get_active_llm()
+        if not config:
+            raise ValueError("未配置 LLM，请先在设置中添加 LLM 提供商")
+        provider = self.get_provider(config)
+        return await provider.chat(messages, temperature)
+
+    async def chat_stream(self, messages: list[dict], temperature: float = 0.1) -> AsyncGenerator[str, None]:
+        """使用当前激活的 LLM 发送流式对话"""
+        config = get_config_manager().get_active_llm()
+        if not config:
+            raise ValueError("未配置 LLM，请先在设置中添加 LLM 提供商")
+        provider = self.get_provider(config)
+        async for chunk in provider.chat_stream(messages, temperature):
+            yield chunk
+
+    async def close_all(self):
+        """关闭所有 Provider 连接"""
+        for provider in self._providers.values():
+            if hasattr(provider, "close"):
+                await provider.close()
+        self._providers.clear()
+
+
+# 全局单例
+_llm_manager: Optional[LLMManager] = None
+
+
+def get_llm_manager() -> LLMManager:
+    global _llm_manager
+    if _llm_manager is None:
+        _llm_manager = LLMManager()
+    return _llm_manager

sql_assistant/llm/prompts.py

@@ -0,0 +1,29 @@
+"""LLM Prompt 模板 - 自然语言转 SQL"""
+
+SYSTEM_PROMPT = """你是一个专业的 SQL 助手。你的任务是将用户的自然语言问题转换为 {db_type} 数据库的 SQL 语句。
+
+{schema_context}
+## 规则
+1. 只返回 SQL 语句，不要返回其他解释性文字
+2. SQL 语句必须符合 {db_type} 的语法规范
+3. **严格使用上面 Schema 中列出的表名和字段名**，不要凭空编造表或字段
+4. 如果用户问的内容在 Schema 中找不到对应的表/字段，返回：-- 数据库中不存在相关表或字段，请确认查询内容
+5. 对于 SELECT 查询，使用合理的 LIMIT 限制（默认 100 条）
+6. 对于 DELETE/UPDATE，必须包含 WHERE 条件，并在语句前加注释提醒危险操作
+7. 如果用户的问题无法转换为 SQL，返回：-- 无法理解的问题
+8. 对于 Redis，使用 Redis 命令格式
+9. 对于 MongoDB，使用 MongoDB 查询语法（JSON 格式）
+
+## 数据库方言注意事项
+- MySQL: 使用反引号包裹标识符，支持 LIMIT
+- PostgreSQL: 使用双引号包裹标识符，支持 LIMIT
+- SQL Server: 使用方括号包裹标识符，使用 TOP 或 OFFSET-FETCH
+- Redis: 返回 Redis 原生命令
+- MongoDB: 返回 MongoDB find/aggregate 查询
+
+现在开始响应。"""
+
+NL_TO_SQL_PROMPT = """将以下自然语言转换为 {db_type} SQL 语句：
+{question}
+
+只返回 SQL 语句："""

sql_assistant/llm/providers/__init__.py

@@ -0,0 +1 @@
+"""LLM Provider 实现"""