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 字符串输入
        pcfg = _maybe_json(data.get("public_cfg"))
        if isinstance(pcfg, str):
            raise ValueError("public_cfg must be a JSON object")
        if isinstance(pcfg, dict):
            data["public_cfg"] = pcfg

        # secret_cfg：若用户输入 JSON 字符串，则自动加密落库；若输入已是 token，则原样保存
        scfg = data.get("secret_cfg", "")
        if scfg is None:
            data["secret_cfg"] = ""
            return
        if isinstance(scfg, str):
            s = scfg.strip()
            if not s:
                data["secret_cfg"] = ""
                return
            parsed = _maybe_json(s)
            if isinstance(parsed, dict):
                data["secret_cfg"] = encrypt_json(parsed)
            else:
                # 非 JSON：视为已加密 token
                data["secret_cfg"] = s
            return
        if isinstance(scfg, dict):
            data["secret_cfg"] = encrypt_json(scfg)
            return
        raise ValueError("secret_cfg must be JSON object or encrypted token string")


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>"
        ),
    }