amazingdata_platform/backend/models/schemas.py

"""
AmazingData 数据服务平台 - Pydantic 数据模型
"""

from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from datetime import datetime
from enum import Enum


# ==================== 通用响应模型 ====================

class ResponseCode(int, Enum):
    SUCCESS = 200
    ERROR = 500
    UNAUTHORIZED = 401
    FORBIDDEN = 403
    NOT_FOUND = 404


class BaseResponse(BaseModel):
    """基础响应模型"""
    code: int = ResponseCode.SUCCESS
    message: str = "success"
    data: Optional[Any] = None


# ==================== 认证模型 ====================

class LoginRequest(BaseModel):
    """登录请求"""
    username: str = Field(..., min_length=3, max_length=50)
    password: str = Field(..., min_length=6)


class LoginResponse(BaseModel):
    """登录响应"""
    access_token: str
    token_type: str = "bearer"
    user_info: Dict[str, Any]


class UserInfo(BaseModel):
    """用户信息"""
    id: int
    username: str
    role: str
    is_active: bool
    last_login: Optional[datetime]


# ==================== 历史数据模型 ====================

class PeriodEnum(str, Enum):
    """K线周期枚举"""
    DAILY = "day"
    MIN1 = "min1"
    MIN5 = "min5"
    MIN15 = "min15"
    MIN30 = "min30"
    MIN60 = "min60"


class SingleKlineRequest(BaseModel):
    """单只K线请求"""
    code: str = Field(..., description="代码，如 000001.SZ 或 ag2605.SHF")
    trading_day: Optional[str] = Field(None, description="交易日，格式 YYYYMMDD，默认最近交易日")
    period: PeriodEnum = Field(PeriodEnum.DAILY, description="K线周期")
    save_path: Optional[str] = Field("./data/single", description="保存路径")


class BatchStockRequest(BaseModel):
    """批量股票请求"""
    codes: Optional[List[str]] = Field(None, description="股票代码列表，None表示全部A股")
    trading_days: Optional[List[str]] = Field(None, description="交易日列表")
    save_path: Optional[str] = Field("./data/stock", description="保存路径")
    batch_size: int = Field(100, description="批次大小")


class BatchFutureRequest(BaseModel):
    """批量期货请求"""
    underlying_codes: Optional[List[str]] = Field(None, description="品种代码列表")
    use_main_contract: bool = Field(True, description="是否使用主力合约")
    trading_days: Optional[List[str]] = Field(None, description="交易日列表")
    save_path: Optional[str] = Field("./data/future", description="保存路径")


class KlineDataResponse(BaseModel):
    """K线数据响应"""
    code: str
    trading_day: str
    period: str
    count: int
    data: List[Dict[str, Any]]
    file_path: str


# ==================== 实时订阅模型 ====================

class SubscribePeriodEnum(str, Enum):
    """订阅周期枚举"""
    MIN1 = "min1"
    MIN5 = "min5"
    MIN15 = "min15"
    MIN30 = "min30"
    MIN60 = "min60"


class SubscribeRequest(BaseModel):
    """订阅请求"""
    codes: List[str] = Field(..., description="品种代码列表")
    periods: List[SubscribePeriodEnum] = Field([SubscribePeriodEnum.MIN5], description="订阅周期")
    save_path: Optional[str] = Field("./data/realtime", description="保存路径")
    duration: int = Field(0, description="运行时长(秒)，0=无限")
    save_interval: int = Field(60, description="保存间隔(秒)")
    task_name: Optional[str] = Field(None, description="任务名称")


class SubscribeResponse(BaseModel):
    """订阅响应"""
    task_id: int
    task_name: str
    status: str
    message: str


class TaskStatus(BaseModel):
    """任务状态"""
    id: int
    task_name: str
    codes: List[str]
    periods: List[str]
    status: str
    started_at: Optional[datetime]
    stopped_at: Optional[datetime]
    created_at: datetime


# ==================== 批量任务模型 ====================

class BatchTaskRequest(BaseModel):
    """批量任务请求"""
    task_type: str = Field(..., description="任务类型: stock/future")
    codes: Optional[List[str]] = Field(None, description="代码列表")
    use_main_contract: bool = Field(True, description="是否主力合约")
    trading_days: Optional[List[str]] = Field(None, description="交易日列表")
    save_path: Optional[str] = Field("./data/batch", description="保存路径")
    batch_size: int = Field(100, description="批次大小")


class BatchTaskStatus(BaseModel):
    """批量任务状态"""
    id: int
    task_type: str
    total_count: int
    processed_count: int
    success_count: int
    failed_count: int
    status: str
    output_path: Optional[str]
    error_message: Optional[str]
    started_at: Optional[datetime]
    completed_at: Optional[datetime]


# ==================== 缓存管理模型 ====================

class CacheFileItem(BaseModel):
    """缓存文件项"""
    id: int
    filename: str
    file_type: str
    trading_day: Optional[str]
    code: Optional[str]
    period: Optional[str]
    record_count: int
    file_size: int
    file_path: str
    created_at: datetime


class CacheStats(BaseModel):
    """缓存统计"""
    total_files: int
    total_size: int
    by_type: Dict[str, int]
    by_day: Dict[str, int]


# ==================== 系统配置模型 ====================

class ConfigItem(BaseModel):
    """配置项"""
    id: int
    config_key: str
    config_value: str
    config_type: str
    description: Optional[str]


class ConfigUpdateRequest(BaseModel):
    """配置更新请求"""
    config_value: str


class TestConnectionResponse(BaseModel):
    """连接测试响应"""
    success: bool
    message: str
    details: Optional[Dict[str, Any]] = None