sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/__init__.py

@@ -0,0 +1,3 @@
+"""SQL 智能助手 - 自然语言转 SQL 查询工具"""
+
+__version__ = "1.0.0"

sql_assistant/api/__init__.py

@@ -0,0 +1 @@
+"""API 路由模块"""

sql_assistant/api/backup.py

@@ -0,0 +1,116 @@
+"""备份相关 API 路由"""
+
+from fastapi import APIRouter, HTTPException
+
+from .dependencies import get_config, get_db
+from .models import (
+    BackupRequest, BackupResponse, BackupInfoResponse, BackupListResponse,
+    RestoreRequest, RestoreResponse,
+)
+
+router = APIRouter()
+
+
+@router.post("/backup", response_model=BackupResponse)
+async def create_backup(request: BackupRequest):
+    from ..database.backup import get_backup_manager, BackupConfig
+
+    backup_manager = get_backup_manager()
+
+    config = BackupConfig(
+        backup_type=request.backup_type,
+        tables=request.tables,
+        include_schema=request.include_schema,
+        include_data=request.include_data
+    )
+
+    result = await backup_manager.backup(config)
+
+    return BackupResponse(
+        success=result.success,
+        message=result.message,
+        backup_id=result.backup_id,
+        backup_path=result.backup_path,
+        tables_backed_up=result.tables_backed_up,
+        total_records=result.total_records,
+        backup_size=result.backup_size,
+        backup_time=result.backup_time
+    )
+
+
+@router.get("/backup/list", response_model=BackupListResponse)
+async def list_backups():
+    from ..database.backup import get_backup_manager
+
+    backup_manager = get_backup_manager()
+    backups = backup_manager.list_backups()
+
+    return BackupListResponse(
+        backups=[BackupInfoResponse(
+            backup_id=b.backup_id,
+            backup_type=b.backup_type,
+            db_type=b.db_type,
+            db_name=b.db_name,
+            tables=b.tables,
+            record_count=b.record_count,
+            backup_time=b.backup_time,
+            file_size=b.file_size
+        ) for b in backups],
+        total=len(backups)
+    )
+
+
+@router.get("/backup/{backup_id}", response_model=BackupInfoResponse)
+async def get_backup_info(backup_id: str):
+    from ..database.backup import get_backup_manager
+
+    backup_manager = get_backup_manager()
+    backup_info = backup_manager.get_backup_info(backup_id)
+
+    if not backup_info:
+        raise HTTPException(status_code=404, detail="备份不存在")
+
+    return BackupInfoResponse(
+        backup_id=backup_info.backup_id,
+        backup_type=backup_info.backup_type,
+        db_type=backup_info.db_type,
+        db_name=backup_info.db_name,
+        tables=backup_info.tables,
+        record_count=backup_info.record_count,
+        backup_time=backup_info.backup_time,
+        file_size=backup_info.file_size
+    )
+
+
+@router.delete("/backup/{backup_id}")
+async def delete_backup(backup_id: str):
+    from ..database.backup import get_backup_manager
+
+    backup_manager = get_backup_manager()
+
+    if not backup_manager.delete_backup(backup_id):
+        raise HTTPException(status_code=404, detail="备份不存在")
+
+    return {"ok": True, "message": "备份删除成功"}
+
+
+@router.post("/backup/restore", response_model=RestoreResponse)
+async def restore_backup(request: RestoreRequest):
+    from ..database.backup import get_backup_manager
+
+    backup_manager = get_backup_manager()
+
+    result = await backup_manager.restore(
+        backup_id=request.backup_id,
+        restore_schema=request.restore_schema,
+        restore_data=request.restore_data,
+        tables=request.tables
+    )
+
+    return RestoreResponse(
+        success=result.success,
+        message=result.message,
+        backup_id=result.backup_id,
+        tables_restored=result.tables_restored,
+        total_records=result.total_records
+    )

sql_assistant/api/config.py

@@ -0,0 +1,183 @@
+"""配置相关 API 路由"""
+
+from typing import Optional
+
+from fastapi import APIRouter, HTTPException
+
+from ..settings import LLMProviderConfig, DatabaseConfig
+
+from .dependencies import get_config, get_llm, get_db
+from .models import (
+    LLMConfigRequest, LLMConfigResponse,
+    DatabaseConfigRequest, DatabaseConfigResponse,
+    TestConnectionRequest, TestConnectionResponse, TestLLMConnectionRequest,
+    AppSettingsResponse,
+)
+
+router = APIRouter()
+
+
+def _mask_api_key(key: str) -> str:
+    if len(key) <= 4:
+        return "****"
+    return "*" * (len(key) - 4) + key[-4:]
+
+
+@router.get("/config/llm/models")
+async def get_llm_models(provider: Optional[str] = None):
+    from ..settings import PROVIDER_MODELS
+
+    if provider:
+        models = PROVIDER_MODELS.get(provider, [])
+        return {"provider": provider, "models": models}
+
+    return {"models": PROVIDER_MODELS}
+
+
+@router.get("/config/llm", response_model=list[LLMConfigResponse])
+async def list_llm_configs():
+    config = get_config()
+    return [
+        LLMConfigResponse(
+            name=p.name, provider=p.provider, model=p.get_model(),
+            enabled=p.enabled, api_key_masked=_mask_api_key(p.api_key),
+        )
+        for p in config.get_llm_providers()
+    ]
+
+
+@router.post("/config/llm")
+async def save_llm_config(req: LLMConfigRequest):
+    config = get_config()
+    llm_config = LLMProviderConfig(**req.model_dump())
+    config.add_llm_provider(llm_config)
+    return {"ok": True}
+
+
+@router.delete("/config/llm/{name}")
+async def delete_llm_config(name: str):
+    config = get_config()
+    if not config.remove_llm_provider(name):
+        raise HTTPException(status_code=404, detail="LLM 配置不存在")
+    return {"ok": True}
+
+
+@router.put("/config/llm/active/{name}")
+async def set_active_llm(name: str):
+    config = get_config()
+    if not config.set_active_llm(name):
+        raise HTTPException(status_code=404, detail="LLM 配置不存在")
+    return {"ok": True}
+
+
+@router.post("/config/llm/test", response_model=TestConnectionResponse)
+async def test_llm_connection(req: TestLLMConnectionRequest):
+    llm = get_llm()
+    config_mgr = get_config()
+
+    if req.config:
+        llm_config = LLMProviderConfig(**req.config.model_dump())
+        from ..llm.manager import LLMManager
+        llm_manager = LLMManager()
+        provider = llm_manager.get_provider(llm_config)
+        result = await provider.test_connection()
+    elif req.name:
+        llm_config = config_mgr.get_llm_provider(req.name)
+        if not llm_config:
+            raise HTTPException(status_code=404, detail="LLM 配置不存在")
+        from ..llm.manager import LLMManager
+        llm_manager = LLMManager()
+        provider = llm_manager.get_provider(llm_config)
+        result = await provider.test_connection()
+    else:
+        result = {"success": False, "error": "请提供 LLM 配置名称或完整配置"}
+
+    return TestConnectionResponse(
+        success=result.get("success", False),
+        message=result.get("message") or (result.get("data") or {}).get("message") or result.get("error") or "未知错误",
+    )
+
+
+@router.get("/config/database", response_model=list[DatabaseConfigResponse])
+async def list_database_configs():
+    config = get_config()
+    return [
+        DatabaseConfigResponse(
+            name=db.name, db_type=db.db_type, host=db.host,
+            port=db.get_port(), user=db.user,
+            database=db.database, enabled=db.enabled,
+        )
+        for db in config.get_databases()
+    ]
+
+
+@router.post("/config/database")
+async def save_database_config(req: DatabaseConfigRequest):
+    config = get_config()
+    db_config = DatabaseConfig(**req.model_dump())
+    config.add_database(db_config)
+    return {"ok": True}
+
+
+@router.delete("/config/database/{name}")
+async def delete_database_config(name: str):
+    config = get_config()
+    if not config.remove_database(name):
+        raise HTTPException(status_code=404, detail="数据库配置不存在")
+    return {"ok": True}
+
+
+@router.put("/config/database/active/{name}")
+async def set_active_database(name: str):
+    config = get_config()
+    if not config.set_active_database(name):
+        raise HTTPException(status_code=404, detail="数据库配置不存在")
+    return {"ok": True}
+
+
+@router.post("/config/database/test", response_model=TestConnectionResponse)
+async def test_database_connection(req: TestConnectionRequest):
+    db = get_db()
+    config_mgr = get_config()
+
+    if req.config:
+        db_config = DatabaseConfig(**req.config.model_dump())
+        result = await db.test_connection(db_config)
+    elif req.name:
+        db_config = config_mgr.get_database(req.name)
+        if not db_config:
+            raise HTTPException(status_code=404, detail="数据库配置不存在")
+        result = await db.test_connection(db_config)
+    else:
+        result = {"success": False, "error": "请提供数据库名称或完整配置"}
+
+    return TestConnectionResponse(
+        success=result.get("success", False),
+        message=result.get("message") or (result.get("data") or {}).get("message") or result.get("error") or "未知错误",
+    )
+
+
+@router.get("/config/settings", response_model=AppSettingsResponse)
+async def get_settings():
+    config = get_config()
+    settings = config.get_settings()
+    return AppSettingsResponse(
+        active_llm=settings.active_llm,
+        active_database=settings.active_database,
+        max_history_rows=settings.max_history_rows,
+        llm_providers=[
+            LLMConfigResponse(
+                name=p.name, provider=p.provider, model=p.get_model(),
+                enabled=p.enabled, api_key_masked=_mask_api_key(p.api_key),
+            )
+            for p in settings.llm_providers
+        ],
+        databases=[
+            DatabaseConfigResponse(
+                name=db.name, db_type=db.db_type, host=db.host,
+                port=db.get_port(), user=db.user,
+                database=db.database, enabled=db.enabled,
+            )
+            for db in settings.databases
+        ],
+    )

sql_assistant/api/conversation.py

@@ -0,0 +1,71 @@
+"""对话相关 API 路由"""
+
+from fastapi import APIRouter, HTTPException
+
+from .dependencies import get_history
+from .models import (
+    ConversationRequest, ConversationResponse, ConversationDetailResponse,
+    ConversationListResponse, ConversationUpdateRequest, HistoryRecord,
+)
+
+router = APIRouter()
+
+
+@router.post("/conversations", response_model=ConversationResponse, status_code=201)
+async def create_conversation(req: ConversationRequest):
+    history = get_history()
+    conversation_id = await history.create_conversation(title=req.title)
+    conversation = await history.get_conversation(conversation_id)
+    if not conversation:
+        raise HTTPException(status_code=500, detail="创建对话失败")
+    return ConversationResponse(**conversation)
+
+
+@router.get("/conversations", response_model=ConversationListResponse)
+async def list_conversations(limit: int = 50, offset: int = 0):
+    history = get_history()
+    conversations = await history.get_conversations(limit=limit, offset=offset)
+    total = await history.get_conversation_count()
+    return ConversationListResponse(
+        conversations=[ConversationResponse(**c) for c in conversations],
+        total=total,
+        limit=limit,
+        offset=offset,
+    )
+
+
+@router.get("/conversations/{conversation_id}", response_model=ConversationDetailResponse)
+async def get_conversation(conversation_id: int):
+    history = get_history()
+    conversation = await history.get_conversation(conversation_id)
+    if not conversation:
+        raise HTTPException(status_code=404, detail="对话不存在")
+
+    messages = await history.get_conversation_messages(conversation_id)
+    return ConversationDetailResponse(
+        id=conversation["id"],
+        title=conversation["title"],
+        created_at=conversation.get("created_at", ""),
+        updated_at=conversation.get("updated_at", ""),
+        messages=[HistoryRecord(**m) for m in messages],
+    )
+
+
+@router.put("/conversations/{conversation_id}", response_model=ConversationResponse)
+async def update_conversation(conversation_id: int, req: ConversationUpdateRequest):
+    history = get_history()
+    if not await history.update_conversation(conversation_id, req.title):
+        raise HTTPException(status_code=404, detail="对话不存在")
+
+    conversation = await history.get_conversation(conversation_id)
+    if not conversation:
+        raise HTTPException(status_code=500, detail="获取对话失败")
+    return ConversationResponse(**conversation)
+
+
+@router.delete("/conversations/{conversation_id}")
+async def delete_conversation(conversation_id: int):
+    history = get_history()
+    if not await history.delete_conversation(conversation_id):
+        raise HTTPException(status_code=404, detail="对话不存在")
+    return {"ok": True}

sql_assistant/api/dependencies.py

@@ -0,0 +1,22 @@
+"""FastAPI 依赖注入"""
+
+from ..config import get_config_manager
+from ..llm.manager import get_llm_manager
+from ..database.manager import get_db_manager
+from ..database.history import get_history_manager
+
+
+def get_config():
+    return get_config_manager()
+
+
+def get_llm():
+    return get_llm_manager()
+
+
+def get_db():
+    return get_db_manager()
+
+
+def get_history():
+    return get_history_manager()

sql_assistant/api/history.py

@@ -0,0 +1,61 @@
+"""历史记录相关 API 路由"""
+
+from typing import Optional
+
+from fastapi import APIRouter
+
+from .dependencies import get_history
+from .models import HistoryRecord, HistoryListResponse
+
+router = APIRouter()
+
+
+@router.get("/history", response_model=HistoryListResponse)
+async def list_history(limit: int = 50, offset: int = 0, conversation_id: Optional[int] = None):
+    history = get_history()
+    records = await history.get_records(limit=limit, offset=offset, conversation_id=conversation_id)
+    total = await history.get_count(conversation_id=conversation_id)
+    return HistoryListResponse(
+        records=[
+            HistoryRecord(
+                id=r["id"],
+                conversation_id=r.get("conversation_id"),
+                question=r["question"],
+                sql=r["sql"],
+                result_json=r.get("result_json"),
+                db_type=r.get("db_type", ""),
+                llm_provider=r.get("llm_provider", ""),
+                success=bool(r.get("success", 1)),
+                error_message=r.get("error_message", ""),
+                created_at=r.get("created_at", ""),
+            )
+            for r in records
+        ],
+        total=total,
+        limit=limit,
+        offset=offset,
+    )
+
+
+@router.delete("/history/{record_id}")
+async def delete_history_record(record_id: int):
+    history = get_history()
+    if not await history.delete_record(record_id):
+        return {"ok": False, "error": "记录不存在"}
+    return {"ok": True}
+
+
+@router.delete("/history")
+async def clear_history():
+    history = get_history()
+    count = await history.clear_history()
+    return {"ok": True, "deleted": count}
+
+
+@router.get("/history/{record_id}")
+async def get_history_record(record_id: int):
+    history = get_history()
+    record = await history.get_record(record_id)
+    if not record:
+        return {"error": "记录不存在"}
+    return record

sql_assistant/api/models.py

@@ -0,0 +1,221 @@
+"""API 数据模型"""
+
+from pydantic import BaseModel, Field
+from typing import Optional, Any, List
+
+
+# ---- Query ----
+
+class QueryRequest(BaseModel):
+    question: str = Field(description="自然语言查询问题")
+    db_type_override: Optional[str] = Field(default=None, description="临时切换数据库类型")
+    conversation_id: Optional[int] = Field(default=None, description="对话ID，用于在已有对话中添加消息")
+    page: int = Field(default=1, ge=1, description="结果页码")
+    page_size: int = Field(default=100, ge=1, le=1000, description="每页记录数")
+    confirmed: bool = Field(default=False, description="用户是否已确认执行")
+    sql_hash: Optional[str] = Field(default=None, description="要确认执行的 SQL 哈希，用于验证")
+    sql: Optional[str] = Field(default=None, description="用户确认后要执行的 SQL 语句（confirmed=True 时使用）")
+
+
+class SQLPreviewResponse(BaseModel):
+    """SQL 预览响应"""
+    success: bool
+    sql: str = ""
+    sql_hash: str = ""
+    requires_confirmation: bool = False
+    confirmation_reason: str = ""
+    warning: str = ""
+    risk_level: str = "none"
+    error: Optional[str] = None
+
+
+class QueryResponse(BaseModel):
+    success: bool
+    question: str = ""
+    sql: str = ""
+    result: Optional[dict] = None
+    error: Optional[str] = None
+    history_id: Optional[int] = None
+    conversation_id: Optional[int] = None
+    pagination: Optional[dict] = None
+
+
+# ---- Conversation ----
+
+class ConversationRequest(BaseModel):
+    title: Optional[str] = Field(default="", description="对话标题")
+
+
+class ConversationUpdateRequest(BaseModel):
+    title: str = Field(description="新的对话标题")
+
+
+class ConversationResponse(BaseModel):
+    id: int
+    title: str
+    created_at: str = ""
+    updated_at: str = ""
+    message_count: int = 0
+
+
+class ConversationDetailResponse(BaseModel):
+    id: int
+    title: str
+    created_at: str = ""
+    updated_at: str = ""
+    messages: list["HistoryRecord"] = Field(default_factory=list)
+
+
+class ConversationListResponse(BaseModel):
+    conversations: list[ConversationResponse]
+    total: int
+    limit: int
+    offset: int
+
+
+# ---- Config: LLM ----
+
+class LLMConfigRequest(BaseModel):
+    name: str
+    provider: str  # deepseek, doubao, kimi, qwen, gemini, claude, openai
+    api_key: str = ""
+    base_url: str = ""
+    model: str = ""
+    enabled: bool = True
+
+
+class LLMConfigResponse(BaseModel):
+    name: str
+    provider: str
+    model: str
+    enabled: bool
+    api_key_masked: str = ""  # 只显示后4位
+
+
+# ---- Config: Database ----
+
+class DatabaseConfigRequest(BaseModel):
+    name: str
+    db_type: str  # mysql, sqlserver, postgresql, redis, mongodb
+    host: str = "localhost"
+    port: int = 0
+    user: str = ""
+    password: str = ""
+    database: str = ""
+    enabled: bool = True
+
+
+class DatabaseConfigResponse(BaseModel):
+    name: str
+    db_type: str
+    host: str
+    port: int
+    user: str
+    database: str
+    enabled: bool
+
+
+class TestConnectionRequest(BaseModel):
+    """测试连接请求 - 可传完整配置或只传名称"""
+    name: Optional[str] = None
+    config: Optional[DatabaseConfigRequest] = None
+
+
+class TestLLMConnectionRequest(BaseModel):
+    """测试 LLM 连接请求"""
+    name: Optional[str] = None
+    config: Optional[LLMConfigRequest] = None
+
+
+class TestConnectionResponse(BaseModel):
+    success: bool
+    message: str
+
+
+# ---- Config: Settings ----
+
+class AppSettingsResponse(BaseModel):
+    active_llm: str = ""
+    active_database: str = ""
+    max_history_rows: int = 1000
+    llm_providers: list[LLMConfigResponse] = Field(default_factory=list)
+    databases: list[DatabaseConfigResponse] = Field(default_factory=list)
+
+
+# ---- History ----
+
+class HistoryRecord(BaseModel):
+    id: int
+    conversation_id: Optional[int] = None
+    question: str
+    sql: str
+    result_json: Optional[str] = None
+    db_type: str = ""
+    llm_provider: str = ""
+    success: bool = True
+    error_message: str = ""
+    created_at: str = ""
+
+
+class HistoryListResponse(BaseModel):
+    records: list[HistoryRecord]
+    total: int
+    limit: int
+    offset: int
+
+
+# ---- Backup ----
+
+class BackupRequest(BaseModel):
+    """备份请求模型"""
+    backup_type: str = Field(default="full", description="备份类型: full(全量备份) / incremental(增量备份)")
+    tables: Optional[List[str]] = Field(default=None, description="要备份的表名列表，为空则备份所有表")
+    include_schema: bool = Field(default=True, description="是否包含表结构")
+    include_data: bool = Field(default=True, description="是否包含数据")
+
+
+class BackupResponse(BaseModel):
+    """备份响应模型"""
+    success: bool
+    message: str
+    backup_id: str = ""
+    backup_path: str = ""
+    tables_backed_up: List[str] = Field(default_factory=list)
+    total_records: int = 0
+    backup_size: int = 0  # bytes
+    backup_time: str = ""
+
+
+class BackupInfoResponse(BaseModel):
+    """备份信息响应模型"""
+    backup_id: str
+    backup_type: str
+    db_type: str = ""
+    db_name: str = ""
+    tables: List[str] = Field(default_factory=list)
+    record_count: int = 0
+    backup_time: str = ""
+    file_size: int = 0  # bytes
+
+
+class BackupListResponse(BaseModel):
+    """备份列表响应模型"""
+    backups: List[BackupInfoResponse]
+    total: int
+
+
+class RestoreRequest(BaseModel):
+    """恢复备份请求模型"""
+    backup_id: str = Field(description="要恢复的备份ID")
+    restore_schema: bool = Field(default=False, description="是否重建表结构（谨慎使用，会删除现有表）")
+    restore_data: bool = Field(default=True, description="是否恢复数据")
+    tables: Optional[List[str]] = Field(default=None, description="指定要恢复的表，为空则恢复所有表")
+
+
+class RestoreResponse(BaseModel):
+    """恢复备份响应模型"""
+    success: bool
+    message: str
+    backup_id: str = ""
+    tables_restored: List[str] = Field(default_factory=list)
+    total_records: int = 0