import json import uuid from datetime import datetime from enum import Enum from typing import Optional, Dict, Any from pydantic import BaseModel, Field, validator, ConfigDict, field_validator, field_serializer from pydantic.alias_generators import to_camel from app.core.security import mask_sensitive_dict from app.db.models.data_collection import CollectionTask, TaskExecution, CollectionTemplate from app.module.dataset.schema import DatasetTypeResponse from app.module.dataset.schema.dataset import DatasetType from app.module.shared.schema import TaskStatus class SyncMode(str, Enum): ONCE = "ONCE" SCHEDULED = "SCHEDULED" class CollectionConfig(BaseModel): parameter: Optional[dict] = Field(None, description="模板参数") reader: Optional[dict] = Field(None, description="reader参数") writer: Optional[dict] = Field(None, description="writer参数") job: Optional[dict] = Field(None, description="任务配置") @field_serializer('parameter', 'reader', 'writer', 'job', when_used='json') def mask_sensitive_fields(self, value: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Mask sensitive fields when serializing to JSON""" if value is None: return value return mask_sensitive_dict(value) class CollectionTaskBase(BaseModel): id: str = Field(..., description="任务id") name: str = Field(..., description="任务名称") description: Optional[str] = Field(None, description="任务描述") target_path: str = Field(..., description="目标存放路径") config: CollectionConfig = Field(..., description="任务配置") template_id: str = Field(..., description="模板ID") template_name: Optional[str] = Field(None, description="模板名称") status: TaskStatus = Field(..., description="任务状态") sync_mode: SyncMode = Field(default=SyncMode.ONCE, description="同步方式") schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)") retry_count: int = Field(default=3, description="重试次数") timeout_seconds: int = Field(default=3600, description="超时时间") last_execution_id: Optional[str] = Field(None, description="最后执行id") created_at: Optional[datetime] = Field(None, description="创建时间") updated_at: Optional[datetime] = Field(None, description="更新时间") created_by: Optional[str] = Field(None, description="创建人") updated_by: Optional[str] = Field(None, description="更新人") model_config = ConfigDict( alias_generator=to_camel, populate_by_name=True ) class CollectionTaskCreate(BaseModel): name: str = Field(..., description="任务名称") description: Optional[str] = Field(None, description="任务描述") sync_mode: SyncMode = Field(default=SyncMode.ONCE, description="同步方式") schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)") config: CollectionConfig = Field(..., description="任务配置") template_id: str = Field(..., description="模板ID") dataset_name: Optional[str] = Field(None, description="数据集名称") dataset_type: Optional[DatasetType] = Field(DatasetType.TEXT, description="数据集类型") model_config = ConfigDict( alias_generator=to_camel, populate_by_name=True ) class CollectionTaskUpdate(BaseModel): """任务更新模型 - 只包含可更新的字段""" description: Optional[str] = Field(None, description="任务描述") schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)- 仅SCHEDULED模式可修改") timeout_seconds: Optional[int] = Field(None, description="超时时间(秒)") config: Optional[CollectionConfig] = Field(None, description="任务配置") @field_validator('schedule_expression') @classmethod def validate_schedule_expression(cls, v, info): """验证调度表达式""" if v is not None: # SCHEDULED 模式必须有调度表达式 if not v or not v.strip(): raise ValueError("schedule_expression is required for SCHEDULED mode") return v @field_validator('timeout_seconds') @classmethod def validate_timeout(cls, v): """验证超时时间""" if v is not None and v <= 0: raise ValueError("timeout_seconds must be greater than 0") return v model_config = ConfigDict( alias_generator=to_camel, populate_by_name=True ) def converter_to_response(task: CollectionTask) -> CollectionTaskBase: config_dict = json.loads(task.config) return CollectionTaskBase( id=task.id, name=task.name, description=task.description, sync_mode=task.sync_mode, template_id=task.template_id, template_name=task.template_name, target_path=task.target_path, config=CollectionConfig(**config_dict), schedule_expression=task.schedule_expression, status=task.status, retry_count=task.retry_count, timeout_seconds=task.timeout_seconds, last_execution_id=task.last_execution_id, created_at=task.created_at, updated_at=task.updated_at, created_by=task.created_by, updated_by=task.updated_by, ) def convert_for_create(task: CollectionTaskCreate, task_id: str) -> CollectionTask: schedule_expression = task.schedule_expression if task.sync_mode == SyncMode.SCHEDULED else None return CollectionTask( id=task_id, name=task.name, description=task.description, sync_mode=task.sync_mode, template_id=task.template_id, target_path=f"/dataset/local/{task_id}", config=json.dumps(task.config.dict()), schedule_expression=schedule_expression, status=TaskStatus.PENDING.name ) def create_execute_record(task: CollectionTask) -> TaskExecution: execution_id = str(uuid.uuid4()) return TaskExecution( id=execution_id, task_id=task.id, task_name=task.name, status=TaskStatus.RUNNING.name, started_at=datetime.now(), log_path=f"/flow/data-collection/{task.id}/{execution_id}.log" ) class TaskExecutionBase(BaseModel): id: str = Field(..., description="执行记录ID") task_id: str = Field(..., description="任务ID") task_name: str = Field(..., description="任务名称") status: Optional[str] = Field(None, description="执行状态") log_path: Optional[str] = Field(None, description="日志文件路径") started_at: Optional[datetime] = Field(None, description="开始时间") completed_at: Optional[datetime] = Field(None, description="完成时间") duration_seconds: Optional[int] = Field(None, description="执行时长(秒)") error_message: Optional[str] = Field(None, description="错误信息") created_at: Optional[datetime] = Field(None, description="创建时间") updated_at: Optional[datetime] = Field(None, description="更新时间") created_by: Optional[str] = Field(None, description="创建者") updated_by: Optional[str] = Field(None, description="更新者") model_config = ConfigDict( alias_generator=to_camel, populate_by_name=True ) def converter_execution_to_response(execution: TaskExecution) -> TaskExecutionBase: return TaskExecutionBase( id=execution.id, task_id=execution.task_id, task_name=execution.task_name, status=execution.status, log_path=execution.log_path, started_at=execution.started_at, completed_at=execution.completed_at, duration_seconds=execution.duration_seconds, error_message=execution.error_message, created_at=execution.created_at, updated_at=execution.updated_at, created_by=execution.created_by, updated_by=execution.updated_by, ) class CollectionTemplateBase(BaseModel): id: str = Field(..., description="模板ID") name: str = Field(..., description="模板名称") description: Optional[str] = Field(None, description="模板描述") source_type: str = Field(..., description="源数据源类型") source_name: str = Field(..., description="源数据源名称") target_type: str = Field(..., description="目标数据源类型") target_name: str = Field(..., description="目标数据源名称") template_content: dict = Field(..., description="模板内容") built_in: Optional[bool] = Field(None, description="是否系统内置模板") created_at: Optional[datetime] = Field(None, description="创建时间") updated_at: Optional[datetime] = Field(None, description="更新时间") created_by: Optional[str] = Field(None, description="创建者") updated_by: Optional[str] = Field(None, description="更新者") model_config = ConfigDict( alias_generator=to_camel, populate_by_name=True ) def converter_template_to_response(template: CollectionTemplate) -> CollectionTemplateBase: return CollectionTemplateBase( id=template.id, name=template.name, description=template.description, source_type=template.source_type, source_name=template.source_name, target_type=template.target_type, target_name=template.target_name, template_content=template.template_content, built_in=template.built_in, created_at=template.created_at, updated_at=template.updated_at, created_by=template.created_by, updated_by=template.updated_by, )