sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/api/query.py

@@ -0,0 +1,275 @@
+"""查询相关 API 路由"""
+
+import hashlib
+from typing import Any
+
+from fastapi import APIRouter, HTTPException
+
+from ..llm.prompts import SYSTEM_PROMPT, NL_TO_SQL_PROMPT
+from ..database.connectors.base import QueryResult
+from ..database.security import get_security_guard
+
+from .dependencies import get_config, get_llm, get_db, get_history
+from .models import QueryRequest, QueryResponse, SQLPreviewResponse
+
+router = APIRouter()
+
+
+def _extract_sql(text: str) -> str:
+    """从 LLM 输出中提取 SQL 语句。
+
+    处理多种常见格式：
+    - Markdown 代码块包裹：```sql ... ```
+    - 带前导注释：-- 注释\\nDELETE FROM ...
+    - 带解释性前缀：以下是 SQL：\\nSELECT ...
+    """
+    text = text.strip()
+
+    # 1. 提取 markdown 代码块内容
+    if text.startswith("```"):
+        lines = text.split("\n")
+        if lines[0].strip().startswith("```"):
+            lines = lines[1:]
+        if lines and lines[-1].strip() == "```":
+            lines = lines[:-1]
+        text = "\n".join(lines).strip()
+
+    # 2. 如果提取后内容仍不以 SQL 关键字开头，尝试定位 SQL 起始行
+    sql_keywords = (
+        "SELECT", "INSERT", "UPDATE", "DELETE", "WITH", "CREATE", "ALTER",
+        "DROP", "TRUNCATE", "SHOW", "DESCRIBE", "EXPLAIN", "USE", "SET",
+        "GRANT", "REVOKE", "BEGIN", "COMMIT", "ROLLBACK",
+    )
+    lines = text.split("\n")
+    # 跳过前导的注释行和空行，找到第一条实际语句
+    start_idx = 0
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if not stripped or stripped.startswith("--") or stripped.startswith("#") or stripped.startswith("/*"):
+            continue
+        # 检查是否以 SQL 关键字开头
+        upper = stripped.upper()
+        if any(upper.startswith(kw + " ") or upper == kw or upper.startswith(kw + "\t")
+               for kw in sql_keywords):
+            start_idx = i
+            break
+        else:
+            # 不是 SQL 关键字也不是注释，可能是解释性文字，继续往后找
+            continue
+
+    if start_idx > 0:
+        text = "\n".join(lines[start_idx:]).strip()
+
+    return text
+
+
+def _result_to_dict(result: Any) -> dict:
+    if isinstance(result, QueryResult):
+        return {
+            "columns": result.columns,
+            "rows": result.rows,
+            "row_count": result.row_count,
+            "affected_rows": result.affected_rows,
+            "sql_type": result.sql_type,
+        }
+    return {"value": str(result)}
+
+
+def _hash_sql(sql: str) -> str:
+    return hashlib.sha256(sql.strip().encode('utf-8')).hexdigest()
+
+
+def _validate_request() -> tuple:
+    config = get_config()
+    llm = get_llm()
+    db = get_db()
+
+    active_llm = config.get_active_llm()
+    active_db = config.get_active_database()
+
+    if not active_llm:
+        raise HTTPException(status_code=400, detail="请先在设置中配置并选择 LLM 提供商")
+    if not active_db:
+        raise HTTPException(status_code=400, detail="请先在设置中配置并选择数据库")
+
+    return config, llm, db, active_llm, active_db
+
+
+@router.post("/query/preview", response_model=SQLPreviewResponse)
+async def preview_query(request: QueryRequest):
+    config, llm, db, active_llm, active_db = _validate_request()
+    db_type = request.db_type_override or active_db.db_type
+    schema_text = await db.get_schema_text()
+
+    try:
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT.format(
+                db_type=db_type,
+                schema_context=schema_text,
+            )},
+            {"role": "user", "content": NL_TO_SQL_PROMPT.format(
+                db_type=db_type,
+                question=request.question,
+            )},
+        ]
+        sql_text = await llm.chat(messages, temperature=0.1)
+        sql_text = _extract_sql(sql_text)
+
+        security_guard = get_security_guard()
+        check_result = security_guard.check_sql_safety(sql_text)
+
+        if not check_result.is_safe:
+            return SQLPreviewResponse(
+                success=False,
+                error=f"SQL 安全检查失败: {check_result.blocked_reason}"
+            )
+
+        requires_confirmation, confirmation_reason = security_guard.requires_confirmation(sql_text)
+
+        return SQLPreviewResponse(
+            success=True,
+            sql=sql_text,
+            sql_hash=_hash_sql(sql_text),
+            requires_confirmation=requires_confirmation,
+            confirmation_reason=confirmation_reason,
+            warning=check_result.warning,
+            risk_level=check_result.risk_level
+        )
+
+    except Exception as e:
+        return SQLPreviewResponse(
+            success=False,
+            error=f"生成 SQL 失败: {e}"
+        )
+
+
+@router.post("/query", response_model=QueryResponse)
+async def execute_query(request: QueryRequest):
+    config, llm, db, active_llm, active_db = _validate_request()
+    history = get_history()
+    db_type = request.db_type_override or active_db.db_type
+    schema_text = await db.get_schema_text()
+
+    sql_text = ""
+    
+    # 如果是确认执行（confirmed=True），直接使用前端传递的 SQL
+    # 避免 LLM 非确定性导致 hash 不匹配
+    if request.confirmed:
+        if not request.sql:
+            return QueryResponse(
+                success=False,
+                question=request.question,
+                error="确认执行时需要提供 SQL 语句",
+            )
+        
+        sql_text = request.sql
+        
+        # 如果提供了 sql_hash，验证 SQL 的完整性
+        if request.sql_hash:
+            expected_hash = _hash_sql(sql_text)
+            if request.sql_hash != expected_hash:
+                return QueryResponse(
+                    success=False,
+                    question=request.question,
+                    sql=sql_text,
+                    error="SQL 验证失败，请重新预览并确认",
+                )
+    else:
+        # 非确认执行，正常调用 LLM 生成 SQL
+        try:
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT.format(
+                    db_type=db_type,
+                    schema_context=schema_text,
+                )},
+                {"role": "user", "content": NL_TO_SQL_PROMPT.format(
+                    db_type=db_type,
+                    question=request.question,
+                )},
+            ]
+            sql_text = await llm.chat(messages, temperature=0.1)
+            sql_text = _extract_sql(sql_text)
+        except Exception as e:
+            return QueryResponse(
+                success=False,
+                question=request.question,
+                error=f"LLM 调用失败: {e}",
+            )
+
+    security_guard = get_security_guard()
+    check_result = security_guard.check_sql_safety(sql_text)
+    if not check_result.is_safe:
+        return QueryResponse(
+            success=False,
+            question=request.question,
+            sql=sql_text,
+            error=f"SQL 安全检查失败: {check_result.blocked_reason}",
+        )
+
+    requires_confirmation, _ = security_guard.requires_confirmation(sql_text)
+
+    if requires_confirmation and not request.confirmed:
+        return QueryResponse(
+            success=False,
+            question=request.question,
+            sql=sql_text,
+            error="此操作需要用户确认，请预览后确认执行",
+        )
+
+    try:
+        result = await db.execute(sql_text)
+        result_dict = _result_to_dict(result)
+
+        total_rows = result_dict.get("row_count", 0)
+        total_pages = (total_rows + request.page_size - 1) // request.page_size if total_rows > 0 else 1
+
+        if total_rows > 0 and request.page > 1:
+            start_idx = (request.page - 1) * request.page_size
+            end_idx = start_idx + request.page_size
+            result_dict["rows"] = result_dict["rows"][start_idx:end_idx]
+            result_dict["row_count"] = len(result_dict["rows"])
+
+        pagination = {
+            "page": request.page,
+            "page_size": request.page_size,
+            "total_rows": total_rows,
+            "total_pages": total_pages,
+        }
+
+        history_id = await history.add_record(
+            question=request.question,
+            sql=sql_text,
+            result=result_dict,
+            db_type=db_type,
+            llm_provider=active_llm.provider,
+            success=True,
+            conversation_id=request.conversation_id,
+        )
+
+        return QueryResponse(
+            success=True,
+            question=request.question,
+            sql=sql_text,
+            result=result_dict,
+            history_id=history_id,
+            conversation_id=request.conversation_id,
+            pagination=pagination,
+        )
+    except Exception as e:
+        history_id = await history.add_record(
+            question=request.question,
+            sql=sql_text,
+            db_type=db_type,
+            llm_provider=active_llm.provider,
+            success=False,
+            error_message=str(e),
+            conversation_id=request.conversation_id,
+        )
+        return QueryResponse(
+            success=False,
+            question=request.question,
+            sql=sql_text,
+            error=f"SQL 执行失败: {e}",
+            history_id=history_id,
+            conversation_id=request.conversation_id,
+        )

sql_assistant/api/routes.py

@@ -0,0 +1,19 @@
+"""API 路由"""
+
+from fastapi import APIRouter
+
+from .query import router as query_router
+from .config import router as config_router
+from .conversation import router as conversation_router
+from .backup import router as backup_router
+from .schema import router as schema_router
+from .history import router as history_router
+
+router = APIRouter(prefix="/api")
+
+router.include_router(query_router)
+router.include_router(config_router)
+router.include_router(conversation_router)
+router.include_router(backup_router)
+router.include_router(schema_router)
+router.include_router(history_router)

sql_assistant/api/schema.py

@@ -0,0 +1,21 @@
+"""Schema 相关 API 路由"""
+
+from fastapi import APIRouter
+
+from .dependencies import get_db
+
+router = APIRouter()
+
+
+@router.post("/schema/refresh")
+async def refresh_schema():
+    db = get_db()
+    schema = await db.refresh_schema()
+    return schema
+
+
+@router.get("/schema")
+async def get_schema():
+    db = get_db()
+    schema = await db.get_schema()
+    return schema

sql_assistant/config.py

@@ -0,0 +1,144 @@
+"""配置管理 - YAML 文件读写"""
+
+import os
+import yaml
+from pathlib import Path
+from typing import Optional
+
+from .settings import AppSettings, LLMProviderConfig, DatabaseConfig
+
+# 配置文件放在项目目录下的 .data 文件夹
+PROJECT_DIR = Path(__file__).resolve().parent.parent.parent
+DEFAULT_CONFIG_DIR = PROJECT_DIR / ".data"
+DEFAULT_CONFIG_FILE = DEFAULT_CONFIG_DIR / "config.yaml"
+
+
+class ConfigManager:
+    """统一配置管理器"""
+
+    def __init__(self, config_path: Optional[Path] = None):
+        self.config_path = Path(config_path) if config_path else DEFAULT_CONFIG_FILE
+        self._settings: AppSettings = AppSettings()
+        self._ensure_config_exists()
+
+    def _ensure_config_exists(self):
+        """确保配置目录存在"""
+        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # ---- 加载/保存 ----
+
+    def load(self) -> AppSettings:
+        """从 YAML 文件加载配置"""
+        try:
+            with open(self.config_path, "r", encoding="utf-8") as f:
+                data = yaml.safe_load(f) or {}
+            self._settings = AppSettings(**data)
+        except Exception:
+            self._settings = AppSettings()
+        return self._settings
+
+    def save(self) -> None:
+        """保存配置到 YAML 文件"""
+        self._ensure_config_exists()
+        data = self._settings.model_dump(
+            exclude={"llm_providers": {"__all__": {"DEFAULT_BASE_URLS", "DEFAULT_MODELS"}},
+                     "databases": {"__all__": {"DEFAULT_PORTS"}}}
+        )
+        with open(self.config_path, "w", encoding="utf-8") as f:
+            yaml.safe_dump(data, f, allow_unicode=True, default_flow_style=False, sort_keys=False)
+
+    def get_settings(self) -> AppSettings:
+        return self._settings
+
+    # ---- LLM Provider CRUD ----
+
+    def get_llm_providers(self) -> list[LLMProviderConfig]:
+        return self._settings.llm_providers
+
+    def get_llm_provider(self, name: str) -> Optional[LLMProviderConfig]:
+        for p in self._settings.llm_providers:
+            if p.name == name:
+                return p
+        return None
+
+    def add_llm_provider(self, config: LLMProviderConfig) -> None:
+        existing = self.get_llm_provider(config.name)
+        if existing:
+            idx = self._settings.llm_providers.index(existing)
+            self._settings.llm_providers[idx] = config
+        else:
+            self._settings.llm_providers.append(config)
+        self.save()
+
+    def remove_llm_provider(self, name: str) -> bool:
+        provider = self.get_llm_provider(name)
+        if provider:
+            self._settings.llm_providers.remove(provider)
+            if self._settings.active_llm == name:
+                self._settings.active_llm = ""
+            self.save()
+            return True
+        return False
+
+    def get_active_llm(self) -> Optional[LLMProviderConfig]:
+        return self.get_llm_provider(self._settings.active_llm)
+
+    def set_active_llm(self, name: str) -> bool:
+        if self.get_llm_provider(name):
+            self._settings.active_llm = name
+            self.save()
+            return True
+        return False
+
+    # ---- Database CRUD ----
+
+    def get_databases(self) -> list[DatabaseConfig]:
+        return self._settings.databases
+
+    def get_database(self, name: str) -> Optional[DatabaseConfig]:
+        for db in self._settings.databases:
+            if db.name == name:
+                return db
+        return None
+
+    def add_database(self, config: DatabaseConfig) -> None:
+        existing = self.get_database(config.name)
+        if existing:
+            idx = self._settings.databases.index(existing)
+            self._settings.databases[idx] = config
+        else:
+            self._settings.databases.append(config)
+        self.save()
+
+    def remove_database(self, name: str) -> bool:
+        db = self.get_database(name)
+        if db:
+            self._settings.databases.remove(db)
+            if self._settings.active_database == name:
+                self._settings.active_database = ""
+            self.save()
+            return True
+        return False
+
+    def get_active_database(self) -> Optional[DatabaseConfig]:
+        return self.get_database(self._settings.active_database)
+
+    def set_active_database(self, name: str) -> bool:
+        if self.get_database(name):
+            self._settings.active_database = name
+            self.save()
+            return True
+        return False
+
+
+# 全局单例
+_config_instance: Optional[ConfigManager] = None
+
+
+def get_config_manager() -> ConfigManager:
+    """获取全局 ConfigManager 单例"""
+    global _config_instance
+    if _config_instance is None:
+        _config_instance = ConfigManager()
+        _config_instance.load()
+    return _config_instance

sql_assistant/database/__init__.py

@@ -0,0 +1 @@
+"""数据库连接器模块"""