Vastai-ConnectHub/app/admin/views.py

from __future__ import annotations

import json
from datetime import datetime
from typing import Any
from zoneinfo import ZoneInfo

from croniter import croniter
from markupsafe import Markup
from sqladmin import ModelView, action
from sqladmin.models import Request
from starlette.responses import RedirectResponse

from app.db.models import Job, JobLog
from app.plugins.manager import load_job_class
from app.security.fernet import encrypt_json
from app.tasks.execute import execute_job


def _maybe_json(value: Any) -> Any:
    if isinstance(value, str):
        s = value.strip()
        if not s:
            return value
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            return value
    return value


def _fmt_dt_seconds(dt: datetime | None) -> str:
    if not dt:
        return ""
    # DB 中保存的时间多为 naive；按 UTC 解释后转换为 Asia/Shanghai 展示
    tz = ZoneInfo("Asia/Shanghai")
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=ZoneInfo("UTC"))
    return dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S")


def _truncate(s: str, n: int = 120) -> str:
    s = s or ""
    return (s[: n - 3] + "...") if len(s) > n else s


class JobAdmin(ModelView, model=Job):
    name = "Job"
    name_plural = "Jobs"
    icon = "fa fa-cogs"

    column_list = [Job.id, Job.enabled, Job.cron_expr, Job.handler_path, Job.updated_at]
    column_details_list = [
        Job.id,
        Job.enabled,
        Job.cron_expr,
        Job.handler_path,
        Job.public_cfg,
        Job.secret_cfg,
        Job.last_run_at,
        Job.created_at,
        Job.updated_at,
    ]

    # 允许在表单中编辑主键（创建 Job 必填）
    form_include_pk = True
    form_columns = [Job.id, Job.enabled, Job.cron_expr, Job.handler_path, Job.public_cfg, Job.secret_cfg]

    # 为 Job 详情页指定模板（用于调整按钮间距）
    details_template = "job_details.html"

    # 列表页模板：加入每行 Run Now
    list_template = "job_list.html"

    @action(
        name="run_now",
        label="Run Now",
        confirmation_message="Trigger this job now?",
        add_in_list=True,
        add_in_detail=True,
    )
    async def run_now(self, request: Request):  # type: ignore[override]
        pks = request.query_params.get("pks", "").split(",")
        for pk in [p for p in pks if p]:
            execute_job.delay(job_id=pk)
        referer = request.headers.get("Referer")
        return RedirectResponse(referer or request.url_for("admin:list", identity=self.identity), status_code=303)

    async def on_model_change(self, data: dict, model: Job, is_created: bool, request) -> None:  # type: ignore[override]
        # id 必填（避免插入时触发 NOT NULL）
        raw_id = data.get("id") if is_created else (data.get("id") or getattr(model, "id", None))
        if raw_id is None or not str(raw_id).strip():
            raise ValueError("id is required")

        # handler_path 强校验：必须可 import 且继承 BaseJob
        handler_path = data.get("handler_path") if is_created else (data.get("handler_path") or model.handler_path)
        if handler_path is None or not str(handler_path).strip():
            raise ValueError("handler_path is required")
        load_job_class(str(handler_path).strip())

        # cron_expr 校验：必须是合法 cron 表达式
        cron_expr = data.get("cron_expr") if is_created else (data.get("cron_expr") or model.cron_expr)
        if cron_expr is None or not str(cron_expr).strip():
            raise ValueError("cron_expr is required")
        base = datetime.now(ZoneInfo("Asia/Shanghai"))
        itr = croniter(str(cron_expr).strip(), base)
        _ = itr.get_next(datetime)

        # public_cfg：必须是合法 JSON 对象（dict），否则直接报错阻止落库
        pcfg = data.get("public_cfg")
        if isinstance(pcfg, str):
            try:
                pcfg = json.loads(pcfg)
            except json.JSONDecodeError as e:
                raise ValueError("public_cfg must be a JSON object") from e
        if not isinstance(pcfg, dict):
            raise ValueError("public_cfg must be a JSON object")
        data["public_cfg"] = pcfg

        # secret_cfg：必须是合法 JSON 对象（dict），并且保存时必须加密落库
        scfg = data.get("secret_cfg")
        if isinstance(scfg, str):
            try:
                scfg = json.loads(scfg)
            except json.JSONDecodeError as e:
                raise ValueError("secret_cfg must be a JSON object") from e
        if not isinstance(scfg, dict):
            raise ValueError("secret_cfg must be a JSON object")
        data["secret_cfg"] = encrypt_json(scfg)


class JobLogAdmin(ModelView, model=JobLog):
    name = "JobLog"
    name_plural = "JobLogs"
    icon = "fa fa-list"

    can_create = False
    can_edit = False
    can_delete = False

    # 列表更适合扫读：保留关键字段 + message（截断）
    column_list = [JobLog.id, JobLog.job_id, JobLog.status, JobLog.started_at, JobLog.finished_at, JobLog.message]
    # 默认按 started_at 倒序（最新在前）
    column_default_sort = [(JobLog.started_at, True)]
    column_details_list = [
        JobLog.id,
        JobLog.job_id,
        JobLog.status,
        JobLog.snapshot_params,
        JobLog.message,
        JobLog.traceback,
        JobLog.run_log,
        JobLog.celery_task_id,
        JobLog.attempt,
        JobLog.started_at,
        JobLog.finished_at,
    ]

    # 列表页模板：加入每行 Retry
    list_template = "joblog_list.html"
    # 为 JobLog 详情页单独指定模板（用于加入 Retry 按钮）
    details_template = "joblog_details.html"

    column_formatters = {
        JobLog.started_at: lambda m, a: _fmt_dt_seconds(m.started_at),
        JobLog.finished_at: lambda m, a: _fmt_dt_seconds(m.finished_at),
        JobLog.message: lambda m, a: _truncate(m.message, 120),
    }

    column_formatters_detail = {
        JobLog.started_at: lambda m, a: _fmt_dt_seconds(m.started_at),
        JobLog.finished_at: lambda m, a: _fmt_dt_seconds(m.finished_at),
        JobLog.traceback: lambda m, a: Markup(f"<pre style='white-space:pre-wrap'>{m.traceback or ''}</pre>"),
        JobLog.run_log: lambda m, a: Markup(
            "<pre style='max-height:480px;overflow:auto;white-space:pre-wrap'>"
            + (m.run_log or "")
            + "</pre>"
        ),
        JobLog.snapshot_params: lambda m, a: Markup(
            "<pre style='white-space:pre-wrap'>"
            + json.dumps(m.snapshot_params or {}, ensure_ascii=False, indent=2, sort_keys=True)
            + "</pre>"
        ),
    }