Mutual_Fund/nl2sql/jobs.py

"""NL2SQL 可由外部调度器调用的运维任务封装。"""
from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import Awaitable, Callable

logger = logging.getLogger("nl2sql.jobs")


@dataclass(frozen=True)
class JobResult:
    """运维任务的稳定结果结构。"""

    name: str
    status: str
    attempts: int
    detail: dict | None = None
    error_type: str | None = None


async def run_metadata_sync(*, redis=None, worker=None, retries: int = 2) -> JobResult:
    """执行数据字典增量同步。"""
    if worker is None:
        from scripts.sync_nl2sql_metadata import synchronize

        worker = synchronize
    return await run_job("metadata_sync", worker, redis=redis, retries=retries)


async def run_vector_cleanup(*, redis=None, worker=None, retries: int = 2) -> JobResult:
    """执行失效向量清理。"""
    if worker is None:
        from config.database.milvus import client
        from nl2sql.operations import cleanup_invalid_vectors

        worker = lambda: _cleanup_vectors(cleanup_invalid_vectors, client)
    return await run_job("vector_cleanup", worker, redis=redis, retries=retries)


async def run_consistency_check(*, redis=None, worker, retries: int = 1, backoff: float = 0.5) -> JobResult:
    """执行元数据与向量一致性巡检。"""
    return await run_job(
        "consistency_check", worker, redis=redis, retries=retries, backoff=backoff
    )


async def run_history_cleanup(
    db,
    before: datetime,
    *,
    redis=None,
    repo_factory=None,
    retries: int = 1,
    backoff: float = 0.5,
) -> JobResult:
    """删除保留期之前的查询历史。"""
    if repo_factory is None:
        from repositories.nl2sql_permission import Nl2SqlPermissionRepo

        repo_factory = Nl2SqlPermissionRepo

    async def worker():
        deleted = await repo_factory(db).delete_history_before(before)
        return {"deleted": deleted}

    return await run_job(
        "history_cleanup", worker, redis=redis, retries=retries, backoff=backoff
    )


async def _cleanup_vectors(cleanup, client_factory):
    return {"deleted": await cleanup(client_factory())}


async def run_job(
    name: str,
    worker: Callable[[], Awaitable[dict] | dict],
    *,
    redis=None,
    retries: int = 2,
    backoff: float = 0.5,
    lock_ttl: int = 900,
) -> JobResult:
    """以 Redis 锁和有限重试执行一个幂等任务。"""
    lock_key = f"nl2sql:job:{name}:lock"
    token = uuid.uuid4().hex
    locked = True
    if redis is not None:
        try:
            locked = bool(await redis.set(lock_key, token, ex=lock_ttl, nx=True))
        except Exception:  # noqa: BLE001 Redis 故障不应伪造锁成功
            logger.warning("NL2SQL 任务锁不可用，继续执行任务：%s", name)
    if not locked:
        return JobResult(name=name, status="skipped", attempts=0, detail={"reason": "lock_held"})

    try:
        for attempt in range(1, max(0, retries) + 2):
            try:
                result = worker()
                if asyncio.iscoroutine(result):
                    result = await result
                return JobResult(name=name, status="success", attempts=attempt, detail=result or {})
            except Exception as exc:  # noqa: BLE001 任务失败返回类型化摘要
                if attempt > retries:
                    logger.warning("NL2SQL 任务失败：%s", name, exc_info=True)
                    return JobResult(
                        name=name,
                        status="failed",
                        attempts=attempt,
                        error_type=type(exc).__name__,
                    )
                if backoff > 0:
                    await asyncio.sleep(backoff * (2 ** (attempt - 1)))
    finally:
        if redis is not None and locked:
            try:
                await redis.delete(lock_key)
            except Exception:  # noqa: BLE001 释放锁失败仅记录日志
                logger.warning("NL2SQL 任务锁释放失败：%s", name, exc_info=True)
    raise RuntimeError("NL2SQL 任务执行流程异常")