opsxsql/backend/app/services/ticket_scheduler.py

"""工单定时执行调度器 — 基于 asyncio 的轻量级定时任务。"""
import asyncio
import logging
from datetime import datetime, timezone
from typing import Dict

logger = logging.getLogger(__name__)

# 内存中保存活跃的定时任务句柄
_active_tasks: Dict[int, asyncio.Task] = {}


async def schedule_ticket_execution(ticket_id: int, scheduled_at: datetime) -> None:
    """注册一个定时执行任务。

    使用 asyncio.create_task 创建后台等待任务，到达 scheduled_at 后自动执行工单。

    Args:
        ticket_id: 工单 ID
        scheduled_at: 定时执行时间 (UTC)
    """
    now = datetime.now(timezone.utc)
    delay = (scheduled_at - now).total_seconds()

    if delay <= 0:
        # 已到时间，立即执行
        logger.info("工单 %d 定时时间已到，立即执行", ticket_id)
        task = asyncio.create_task(_execute_scheduled_ticket(ticket_id))
    else:
        logger.info("工单 %d 定时执行: %s (%.1f 秒后)", ticket_id, scheduled_at, delay)
        task = asyncio.create_task(_wait_and_execute(ticket_id, delay))

    _active_tasks[ticket_id] = task


async def _wait_and_execute(ticket_id: int, delay: float) -> None:
    """等待 delay 秒后执行工单。"""
    try:
        await asyncio.sleep(delay)
        await _execute_scheduled_ticket(ticket_id)
    except asyncio.CancelledError:
        logger.info("工单 %d 定时执行已取消", ticket_id)
    except Exception:
        logger.exception("工单 %d 定时执行异常", ticket_id)
    finally:
        _active_tasks.pop(ticket_id, None)


async def _execute_scheduled_ticket(ticket_id: int) -> None:
    """执行定时工单 — 创建独立数据库会话执行。"""
    from app.core.database import async_session
    from app.models.ticket import Ticket as TicketModel
    from app.models.instance import Instance
    from app.services.ticket_executor import TicketExecutor
    from app.services.audit_log import AuditLogService
    from app.core.security import decrypt_password
    from sqlalchemy import select
    import json

    logger.info("开始执行定时工单 %d", ticket_id)

    async with async_session() as db:
        try:
            # 查询工单
            result = await db.execute(select(TicketModel).where(TicketModel.id == ticket_id))
            ticket = result.scalar_one_or_none()
            if not ticket:
                logger.error("定时工单 %d 不存在", ticket_id)
                return

            # 状态检查
            if ticket.status != "scheduled":
                logger.warning("工单 %d 状态已变更 (%s)，跳过执行", ticket_id, ticket.status)
                return

            # 更新状态为 executing
            ticket.status = "executing"
            await db.commit()
            await db.refresh(ticket)

            # 查询实例
            inst_result = await db.execute(select(Instance).where(Instance.id == ticket.instance_id))
            instance = inst_result.scalar_one_or_none()
            if not instance:
                ticket.status = "failed"
                ticket.execute_result = json.dumps({"status": "failed", "errors": ["数据源实例不存在"]}, ensure_ascii=False)
                ticket.execute_time = datetime.now(timezone.utc)
                await db.commit()
                return

            # 解密密码
            instance.password = decrypt_password(instance.password)

            # 执行
            executor = TicketExecutor()
            report = await executor.execute(ticket, instance)

            # 更新结果
            ticket.status = "success" if report.status == "success" else "failed"
            ticket.execute_result = report.to_json()
            ticket.execute_time = datetime.now(timezone.utc)
            await db.commit()

            logger.info("定时工单 %d 执行完成: %s", ticket_id, report.status)

        except Exception as e:
            logger.exception("定时工单 %d 执行异常", ticket_id)
            try:
                result = await db.execute(select(TicketModel).where(TicketModel.id == ticket_id))
                ticket = result.scalar_one_or_none()
                if ticket:
                    ticket.status = "failed"
                    ticket.execute_result = json.dumps({"status": "failed", "errors": [f"{type(e).__name__}: {e}"]}, ensure_ascii=False)
                    ticket.execute_time = datetime.now(timezone.utc)
                    await db.commit()
            except Exception:
                logger.exception("更新定时工单 %d 失败状态也失败", ticket_id)
        finally:
            _active_tasks.pop(ticket_id, None)


def cancel_ticket_execution(ticket_id: int) -> bool:
    """取消定时执行任务。

    Args:
        ticket_id: 工单 ID

    Returns:
        是否成功取消
    """
    task = _active_tasks.get(ticket_id)
    if task and not task.done():
        task.cancel()
        _active_tasks.pop(ticket_id, None)
        logger.info("已取消工单 %d 的定时执行", ticket_id)
        return True
    return False


def get_active_scheduled_count() -> int:
    """返回当前活跃的定时任务数量。"""
    return len(_active_tasks)


def get_active_scheduled_ids() -> list[int]:
    """返回当前活跃的定时工单 ID 列表。"""
    return list(_active_tasks.keys())