opsxsql/backend/app/api/v1/endpoints/scheduled_task.py

"""定时任务 API — 基于 SQL 的自动化任务管理。"""
import time
import logging
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy import select, func as sa_func
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import get_db
from app.core.security import get_current_user, require_role
from app.models.user import User
from app.models.instance import Instance
from app.models.scheduled_task import ScheduledTask, ScheduledTaskLog
from app.schemas.scheduled_task import (
    ScheduledTaskCreate, ScheduledTaskUpdate,
    ScheduledTaskResponse, ScheduledTaskLogResponse, TaskRunResponse,
)
from app.services.db_connector import DBConnector

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/scheduled-tasks", tags=["自动化任务"])


def _serialize(task: ScheduledTask, instance_name: str = None, instance_db_type: str = None) -> dict:
    return {
        "id": task.id,
        "name": task.name,
        "description": task.description,
        "task_sql_type": task.task_sql_type,
        "instance_id": task.instance_id,
        "instance_name": instance_name,
        "instance_db_type": instance_db_type,
        "target_db": task.target_db,
        "target_table": task.target_table,
        "sql_content": task.sql_content,
        "cron_expression": task.cron_expression,
        "is_enabled": task.is_enabled,
        "last_run_at": task.last_run_at,
        "last_run_status": task.last_run_status,
        "last_run_message": task.last_run_message,
        "next_run_at": task.next_run_at,
        "run_count": task.run_count or 0,
        "created_by": task.created_by,
        "created_at": task.created_at,
    }


async def _get_instance(db: AsyncSession, instance_id: int) -> Instance:
    result = await db.execute(select(Instance).where(Instance.id == instance_id))
    inst = result.scalar_one_or_none()
    if not inst:
        raise HTTPException(404, "数据源不存在")
    return inst


# ── CRUD ──────────────────────────────────────────────────────

@router.get("", response_model=list[ScheduledTaskResponse])
async def list_tasks(
    task_sql_type: Optional[str] = None,
    is_enabled: Optional[bool] = None,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """列出所有定时任务。"""
    q = select(ScheduledTask).order_by(ScheduledTask.created_at.desc())
    if task_sql_type:
        q = q.where(ScheduledTask.task_sql_type == task_sql_type)
    if is_enabled is not None:
        q = q.where(ScheduledTask.is_enabled == is_enabled)
    result = await db.execute(q)
    tasks = result.scalars().all()

    # batch load instance info
    inst_ids = {t.instance_id for t in tasks}
    inst_map = {}
    if inst_ids:
        ir = await db.execute(select(Instance).where(Instance.id.in_(inst_ids)))
        for inst in ir.scalars().all():
            inst_map[inst.id] = inst

    return [_serialize(t, inst_map.get(t.instance_id).name if inst_map.get(t.instance_id) else None,
                        inst_map.get(t.instance_id).db_type if inst_map.get(t.instance_id) else None)
            for t in tasks]


@router.post("", response_model=ScheduledTaskResponse, status_code=201)
async def create_task(
    body: ScheduledTaskCreate,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_role("admin")),
):
    """创建定时任务。"""
    # validate instance exists
    inst = await _get_instance(db, body.instance_id)

    task = ScheduledTask(
        name=body.name,
        description=body.description,
        task_sql_type=body.task_sql_type.value,
        instance_id=body.instance_id,
        target_db=body.target_db,
        target_table=body.target_table,
        sql_content=body.sql_content,
        cron_expression=body.cron_expression,
        is_enabled=body.is_enabled,
        created_by=user.id,
    )
    db.add(task)
    await db.flush()
    await db.refresh(task)
    return _serialize(task, inst.name, inst.db_type)


@router.get("/{task_id}", response_model=ScheduledTaskResponse)
async def get_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """获取任务详情。"""
    result = await db.execute(select(ScheduledTask).where(ScheduledTask.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(404, "任务不存在")
    inst = await _get_instance(db, task.instance_id)
    return _serialize(task, inst.name, inst.db_type)


@router.put("/{task_id}", response_model=ScheduledTaskResponse)
async def update_task(
    task_id: int,
    body: ScheduledTaskUpdate,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """更新任务。"""
    result = await db.execute(select(ScheduledTask).where(ScheduledTask.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(404, "任务不存在")

    update_data = body.model_dump(exclude_unset=True)
    if "task_sql_type" in update_data and update_data["task_sql_type"] is not None:
        update_data["task_sql_type"] = update_data["task_sql_type"].value
    for k, v in update_data.items():
        setattr(task, k, v)

    await db.flush()
    await db.refresh(task)

    inst = await _get_instance(db, task.instance_id)
    return _serialize(task, inst.name, inst.db_type)


@router.delete("/{task_id}")
async def delete_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """删除任务。"""
    result = await db.execute(select(ScheduledTask).where(ScheduledTask.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(404, "任务不存在")
    await db.delete(task)
    return {"detail": "已删除"}


@router.patch("/{task_id}/toggle")
async def toggle_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """启用/禁用任务。"""
    result = await db.execute(select(ScheduledTask).where(ScheduledTask.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(404, "任务不存在")
    task.is_enabled = not task.is_enabled
    await db.flush()
    return {"is_enabled": task.is_enabled}


# ── 执行 ──────────────────────────────────────────────────────

@router.post("/{task_id}/run", response_model=TaskRunResponse)
async def run_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """手动执行一次任务。"""
    result = await db.execute(select(ScheduledTask).where(ScheduledTask.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(404, "任务不存在")

    inst = await _get_instance(db, task.instance_id)
    target_db = task.target_db or inst.database or ""

    log = ScheduledTaskLog(
        task_id=task.id,
        status="running",
        sql_content=task.sql_content,
    )
    db.add(log)
    await db.flush()

    start = time.time()
    try:
        msg = await DBConnector.execute_dml(inst, target_db, task.sql_content)
        elapsed = int((time.time() - start) * 1000)
        # parse affected rows from message like "执行成功，影响 3 行"
        affected = 0
        import re
        m = re.search(r'(\d+)', msg)
        if m:
            affected = int(m.group(1))

        log.status = "success"
        log.affected_rows = affected
        log.duration_ms = elapsed
        log.finished_at = __import__("datetime").datetime.now()

        task.last_run_at = __import__("datetime").datetime.now()
        task.last_run_status = "success"
        task.last_run_message = msg
        task.run_count = (task.run_count or 0) + 1
        await db.flush()

        return TaskRunResponse(success=True, message=msg, affected_rows=affected, duration_ms=elapsed)

    except Exception as e:
        elapsed = int((time.time() - start) * 1000)
        error_msg = str(e)

        log.status = "failed"
        log.error_message = error_msg
        log.duration_ms = elapsed
        log.finished_at = __import__("datetime").datetime.now()

        task.last_run_at = __import__("datetime").datetime.now()
        task.last_run_status = "failed"
        task.last_run_message = error_msg
        task.run_count = (task.run_count or 0) + 1
        await db.flush()

        return TaskRunResponse(success=False, message=error_msg, duration_ms=elapsed)


# ── 执行日志 ──────────────────────────────────────────────────

@router.get("/{task_id}/logs", response_model=list[ScheduledTaskLogResponse])
async def get_task_logs(
    task_id: int,
    limit: int = Query(20, ge=1, le=100),
    db: AsyncSession = Depends(get_db),
    _user: User = Depends(require_role("admin")),
):
    """获取任务执行日志。"""
    q = (
        select(ScheduledTaskLog)
        .where(ScheduledTaskLog.task_id == task_id)
        .order_by(ScheduledTaskLog.started_at.desc())
        .limit(limit)
    )
    result = await db.execute(q)
    logs = result.scalars().all()
    return [
        {
            "id": l.id,
            "task_id": l.task_id,
            "status": l.status,
            "sql_content": l.sql_content,
            "affected_rows": l.affected_rows or 0,
            "duration_ms": l.duration_ms,
            "error_message": l.error_message,
            "started_at": l.started_at,
            "finished_at": l.finished_at,
        }
        for l in logs
    ]