from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field


class QueryRequest(BaseModel):
    """统一查询请求模型。

    支持 MySQL、PostgreSQL、Redis、MongoDB 四种数据库的查询请求。

    Attributes:
        instance_id: 数据源实例ID
        database: 目标数据库名/Redis DB号/MongoDB database
        sql: 查询语句（SQL/Redis命令/MongoDB查询JSON）
        options: 额外查询选项
    """
    instance_id: int
    database: Optional[str] = None
    sql: str = Field(..., min_length=1, description="查询语句（SQL/Redis命令/MongoDB查询JSON）")
    options: Optional[Dict[str, Any]] = Field(default=None, description="额外选项（如 limit, offset 等）")


class UnifiedQueryResult(BaseModel):
    """统一查询结果模型。

    兼容 MySQL、PostgreSQL、Redis、MongoDB 四种数据库的返回格式。
    SQL 数据库返回标准表格格式，Redis/MongoDB 同时返回格式化表格和原始结果。

    Attributes:
        columns: SQL表格列名列表
        rows: SQL表格行数据列表（字典格式）
        raw_result: Redis/MongoDB原始结果
        total: 总行数/总文档数
        elapsed_ms: 执行耗时（毫秒）
        truncated: 是否因超过行数限制而截断
        db_type: 数据库类型（mysql/postgres/redis/mongo）
        query_type: 查询类型（SELECT/GET/find等）
    """
    columns: List[str] = Field(default=[], description="SQL表格列名")
    rows: List[Dict[str, Any]] = Field(default=[], description="SQL表格行数据")
    raw_result: Any = Field(default=None, description="Redis/MongoDB原始结果")
    total: int = Field(default=0, description="总行数")
    elapsed_ms: int = Field(default=0, description="执行耗时（毫秒）")
    truncated: bool = Field(default=False, description="是否截断")
    db_type: str = Field(default="", description="数据库类型")
    query_type: str = Field(default="", description="查询类型")


# 保留旧名兼容
QueryResult = UnifiedQueryResult