infra-docker-configs/servers/tools-pve/asr-pipeline/asrservice/db.py

from __future__ import annotations

import sqlite3
import time
from typing import Any

from .config import DB, DEFAULT_PROVIDER, ensure_dirs

STATUSES = {
    "uploading", "queued", "transcribing", "refining", "review",
    "publishing", "complete", "failed",
}


def connect() -> sqlite3.Connection:
    ensure_dirs()
    con = sqlite3.connect(DB, timeout=30)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("PRAGMA foreign_keys=ON")
    con.executescript("""
      CREATE TABLE IF NOT EXISTS intake_job(
        id TEXT PRIMARY KEY, original_name TEXT NOT NULL, stored_name TEXT NOT NULL,
        title TEXT NOT NULL DEFAULT '', context TEXT NOT NULL DEFAULT '',
        size INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL,
        detail TEXT NOT NULL DEFAULT '', transcript TEXT NOT NULL DEFAULT '',
        draft TEXT NOT NULL DEFAULT '', result TEXT NOT NULL DEFAULT '',
        source TEXT NOT NULL DEFAULT 'web', sha256 TEXT NOT NULL DEFAULT '',
        attempts INTEGER NOT NULL DEFAULT 0, error_code TEXT NOT NULL DEFAULT '',
        asr_provider TEXT NOT NULL DEFAULT 'qwen3cam',
        asr_model TEXT NOT NULL DEFAULT '',
        created_at REAL NOT NULL, updated_at REAL NOT NULL
      );
      CREATE TABLE IF NOT EXISTS job_event(
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        job_id TEXT NOT NULL REFERENCES intake_job(id) ON DELETE CASCADE,
        status TEXT NOT NULL, detail TEXT NOT NULL DEFAULT '', created_at REAL NOT NULL
      );
      CREATE INDEX IF NOT EXISTS ix_job_event_job ON job_event(job_id,id);
    """)
    cols = {r[1] for r in con.execute("PRAGMA table_info(intake_job)")}
    additions = {
        "draft": "TEXT NOT NULL DEFAULT ''",
        "result": "TEXT NOT NULL DEFAULT ''",
        "source": "TEXT NOT NULL DEFAULT 'web'",
        "sha256": "TEXT NOT NULL DEFAULT ''",
        "attempts": "INTEGER NOT NULL DEFAULT 0",
        "error_code": "TEXT NOT NULL DEFAULT ''",
        "asr_provider": f"TEXT NOT NULL DEFAULT '{DEFAULT_PROVIDER}'",
        "asr_model": "TEXT NOT NULL DEFAULT ''",
    }
    for name, ddl in additions.items():
        if name not in cols:
            con.execute(f"ALTER TABLE intake_job ADD COLUMN {name} {ddl}")
    con.commit()
    return con


def create_job(
    job_id: str,
    original_name: str,
    stored_name: str,
    *,
    title: str = "",
    context: str = "",
    size: int = 0,
    source: str = "web",
    asr_provider: str = "",
    asr_model: str = "",
) -> dict[str, Any]:
    now = time.time()
    prov = (asr_provider or DEFAULT_PROVIDER or "qwen3cam").strip().lower()
    with connect() as con:
        con.execute(
            "INSERT INTO intake_job(id,original_name,stored_name,title,context,size,"
            "status,source,asr_provider,asr_model,created_at,updated_at) "
            "VALUES(?,?,?,?,?,?,?,?,?,?,?,?)",
            (job_id, original_name, stored_name, title, context, size,
             "queued", source, prov, asr_model, now, now))
        con.execute(
            "INSERT INTO job_event(job_id,status,detail,created_at) VALUES(?,?,?,?)",
            (job_id, "queued", f"录音已进入队列 (ASR: {prov})", now))
    return get_job(job_id) or {}


def update_job(job_id: str, status: str, detail: str = "", **fields: Any) -> bool:
    if status not in STATUSES:
        raise ValueError(f"invalid status: {status}")
    allowed = {
        "transcript", "draft", "result", "sha256", "attempts",
        "error_code", "asr_provider", "asr_model",
    }
    changes = {k: v for k, v in fields.items() if k in allowed and v is not None}
    sets = ["status=?", "updated_at=?", "detail=?"]
    values: list[Any] = [status, time.time(), detail[:1000]]
    for key, value in changes.items():
        sets.append(f"{key}=?")
        values.append(value)
    values.append(job_id)
    with connect() as con:
        cur = con.execute(f"UPDATE intake_job SET {','.join(sets)} WHERE id=?", values)
        if cur.rowcount:
            con.execute(
                "INSERT INTO job_event(job_id,status,detail,created_at) VALUES(?,?,?,?)",
                (job_id, status, detail[:1000], time.time()))
    return bool(cur.rowcount)


def get_job(job_id: str) -> dict[str, Any] | None:
    with connect() as con:
        row = con.execute("SELECT * FROM intake_job WHERE id=?", (job_id,)).fetchone()
        return dict(row) if row else None


def find_by_stored(stored_name: str) -> dict[str, Any] | None:
    with connect() as con:
        row = con.execute(
            "SELECT * FROM intake_job WHERE stored_name=? ORDER BY created_at DESC LIMIT 1",
            (stored_name,)).fetchone()
        return dict(row) if row else None


def list_jobs(limit: int = 30, offset: int = 0, status: str = "") -> list[dict[str, Any]]:
    limit = min(max(limit, 1), 200)
    offset = max(offset, 0)
    sql, args = "SELECT * FROM intake_job", []
    if status:
        sql += " WHERE status=?"
        args.append(status)
    sql += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
    args.extend((limit, offset))
    with connect() as con:
        return [dict(r) for r in con.execute(sql, args).fetchall()]


def events(job_id: str) -> list[dict[str, Any]]:
    with connect() as con:
        return [dict(r) for r in con.execute(
            "SELECT status,detail,created_at FROM job_event WHERE job_id=? ORDER BY id",
            (job_id,)).fetchall()]


def pending_recovery() -> list[dict[str, Any]]:
    with connect() as con:
        return [dict(r) for r in con.execute(
            "SELECT * FROM intake_job WHERE status IN ('queued','transcribing','refining','publishing') "
            "ORDER BY created_at").fetchall()]