xiaobaifupan/next/backend/jobs/repository.py

from __future__ import annotations

import ast
import json
import sqlite3
from datetime import datetime, timedelta
from typing import Any


class JobRepository:
    def begin(
        self,
        connection: sqlite3.Connection,
        *,
        kind: str,
        run_key: str,
        requested_date: str,
        trigger: str,
        started_at: datetime,
        stale_after_seconds: int,
    ) -> sqlite3.Row:
        stale_before = (started_at - timedelta(seconds=stale_after_seconds)).isoformat(
            timespec="seconds"
        )
        connection.execute(
            """
            UPDATE job_runs SET status = 'failed', finished_at = ?, duration_ms = ?,
                error_code = 'stale_job', error_message = '任务进程中断或超过最长运行时间'
            WHERE kind = ? AND status = 'running' AND started_at < ?
            """,
            (
                started_at.isoformat(timespec="seconds"),
                stale_after_seconds * 1000,
                kind,
                stale_before,
            ),
        )
        attempt_row = connection.execute(
            """
            SELECT COALESCE(MAX(attempt), 0) + 1 AS attempt
            FROM job_runs WHERE kind = ? AND run_key = ?
            """,
            (kind, run_key),
        ).fetchone()
        attempt = int(attempt_row["attempt"] if attempt_row else 1)
        cursor = connection.execute(
            """
            INSERT INTO job_runs (
                kind, run_key, requested_date, trigger, status, attempt, started_at
            ) VALUES (?, ?, ?, ?, 'running', ?, ?)
            """,
            (
                kind,
                run_key,
                requested_date,
                trigger,
                attempt,
                started_at.isoformat(timespec="seconds"),
            ),
        )
        return connection.execute(
            "SELECT * FROM job_runs WHERE id = ?", (cursor.lastrowid,)
        ).fetchone()

    def finish(
        self,
        connection: sqlite3.Connection,
        run_id: int,
        *,
        finished_at: datetime,
        duration_ms: int,
        coverage: float | None,
        source_set: list[str],
        output_version: str,
        payload: dict[str, Any],
    ) -> None:
        connection.execute(
            """
            UPDATE job_runs SET status = 'completed', finished_at = ?, duration_ms = ?,
                coverage = ?, source_set_json = ?, output_version = ?, payload_json = ?
            WHERE id = ? AND status = 'running'
            """,
            (
                finished_at.isoformat(timespec="seconds"),
                duration_ms,
                coverage,
                _json(source_set),
                output_version,
                _json(payload),
                run_id,
            ),
        )

    def fail(
        self,
        connection: sqlite3.Connection,
        run_id: int,
        *,
        finished_at: datetime,
        duration_ms: int,
        error_code: str,
        error_message: str,
    ) -> None:
        connection.execute(
            """
            UPDATE job_runs SET status = 'failed', finished_at = ?, duration_ms = ?,
                error_code = ?, error_message = ? WHERE id = ? AND status = 'running'
            """,
            (
                finished_at.isoformat(timespec="seconds"),
                duration_ms,
                error_code,
                error_message[:500],
                run_id,
            ),
        )

    def latest(self, connection: sqlite3.Connection, limit: int = 40) -> tuple[sqlite3.Row, ...]:
        return tuple(
            connection.execute(
                "SELECT * FROM job_runs ORDER BY started_at DESC, id DESC LIMIT ?",
                (limit,),
            ).fetchall()
        )

    def latest_for_kind(
        self, connection: sqlite3.Connection, kind: str
    ) -> sqlite3.Row | None:
        return connection.execute(
            "SELECT * FROM job_runs WHERE kind = ? ORDER BY started_at DESC, id DESC LIMIT 1",
            (kind,),
        ).fetchone()

    def latest_success(
        self, connection: sqlite3.Connection, kind: str, requested_date: str = ""
    ) -> sqlite3.Row | None:
        return connection.execute(
            """
            SELECT * FROM job_runs
            WHERE kind = ? AND status = 'completed' AND (? = '' OR requested_date = ?)
            ORDER BY finished_at DESC, id DESC LIMIT 1
            """,
            (kind, requested_date, requested_date),
        ).fetchone()


def public_job(row: sqlite3.Row) -> dict[str, Any]:
    result = dict(row)
    result["source_set"] = json.loads(str(result.pop("source_set_json")))
    result["payload"] = json.loads(str(result.pop("payload_json")))
    result["error_message"] = _public_error_message(str(result["error_message"]))
    return result


def _public_error_message(value: str) -> str:
    """Normalize AppError tuples written before structured error storage."""
    if not value.startswith("("):
        return value
    try:
        legacy = ast.literal_eval(value)
    except (SyntaxError, ValueError):
        return value
    if (
        isinstance(legacy, tuple)
        and len(legacy) == 3
        and isinstance(legacy[0], str)
        and isinstance(legacy[1], str)
        and isinstance(legacy[2], int)
    ):
        return legacy[1]
    return value


def _json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)