cmshoppe/app/db.py

"""SQLite persistence for accounts, import batches, and tasks."""

from __future__ import annotations

import json
import os
import sqlite3
import uuid
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import Iterable, Optional

from . import appconfig
from .config import make_slug


DEFAULT_BUSY_TIMEOUT_MS = 5000
VALID_BATCH_FIELDS = {"source_files_json", "status", "note"}
VALID_ACCOUNT_FIELDS = {
    "account_name",
    "region_host",
    "slug",
    "user_data_dir",
    "debug_port",
    "password",
    "note",
    "last_login_at",
}
PHASE_ATTEMPT_FIELDS = {
    "collect": "collect_attempts",
    "collected": "collect_attempts",
    "generate": "generate_attempts",
    "generated": "generate_attempts",
    "apply": "apply_attempts",
    "applied": "apply_attempts",
    "update": "apply_attempts",
    "updated": "apply_attempts",
}


class DbError(RuntimeError):
    """Raised when SQLite persistence cannot complete an operation."""


@dataclass(frozen=True)
class Batch:
    id: str
    source_files_json: str
    status: str
    note: Optional[str]
    created_at: str
    updated_at: str

    @property
    def source_files(self) -> list[str]:
        return json.loads(self.source_files_json)


@dataclass(frozen=True)
class Account:
    id: int
    account_name: str
    alias: str
    region_host: str
    slug: str
    user_data_dir: str
    debug_port: int
    password: Optional[str]
    note: Optional[str]
    created_at: str
    updated_at: str
    last_login_at: Optional[str]


@dataclass(frozen=True)
class Task:
    id: int
    batch_id: str
    source_file: str
    source_file_abs: str
    source_sheet: str
    source_row: int
    row_key: str
    account_name: Optional[str]
    alias: str
    item_id: str
    old_title: Optional[str]
    old_cover_path: Optional[str]
    new_title: Optional[str]
    new_cover_path: Optional[str]
    committed: int
    stage: str
    status: str
    last_error: Optional[str]
    collect_attempts: int
    generate_attempts: int
    apply_attempts: int
    imported_at: str
    collected_at: Optional[str]
    generated_at: Optional[str]
    applied_at: Optional[str]
    updated_at: str


SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS batches (
  id                TEXT PRIMARY KEY,
  source_files_json TEXT NOT NULL,
  status            TEXT NOT NULL DEFAULT 'active',
  note              TEXT,
  created_at        TEXT NOT NULL,
  updated_at        TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS accounts (
  id            INTEGER PRIMARY KEY,
  account_name  TEXT NOT NULL,
  alias         TEXT UNIQUE NOT NULL,
  region_host   TEXT NOT NULL,
  slug          TEXT UNIQUE NOT NULL,
  user_data_dir TEXT NOT NULL,
  debug_port    INTEGER NOT NULL,
  password      TEXT,
  note          TEXT,
  created_at    TEXT NOT NULL,
  updated_at    TEXT NOT NULL,
  last_login_at TEXT
);

CREATE TABLE IF NOT EXISTS tasks (
  id              INTEGER PRIMARY KEY,
  batch_id        TEXT NOT NULL REFERENCES batches(id),
  source_file     TEXT NOT NULL,
  source_file_abs TEXT NOT NULL,
  source_sheet    TEXT NOT NULL,
  source_row      INTEGER NOT NULL,
  row_key         TEXT NOT NULL UNIQUE,
  account_name    TEXT,
  alias           TEXT NOT NULL,
  item_id         TEXT NOT NULL,
  old_title       TEXT,
  old_cover_path  TEXT,
  new_title       TEXT,
  new_cover_path  TEXT,
  committed       INTEGER NOT NULL DEFAULT 0,
  stage           TEXT NOT NULL DEFAULT 'imported',
  status          TEXT NOT NULL DEFAULT 'pending',
  last_error      TEXT,
  collect_attempts  INTEGER NOT NULL DEFAULT 0,
  generate_attempts INTEGER NOT NULL DEFAULT 0,
  apply_attempts    INTEGER NOT NULL DEFAULT 0,
  imported_at     TEXT NOT NULL,
  collected_at    TEXT,
  generated_at    TEXT,
  applied_at      TEXT,
  updated_at      TEXT NOT NULL,
  UNIQUE(batch_id, source_file_abs, source_sheet, source_row)
);

CREATE INDEX IF NOT EXISTS idx_tasks_batch_stage_status
  ON tasks(batch_id, stage, status);
CREATE INDEX IF NOT EXISTS idx_tasks_alias ON tasks(alias);
CREATE INDEX IF NOT EXISTS idx_tasks_item ON tasks(item_id);
"""


def _now() -> str:
    return datetime.now().isoformat(timespec="seconds")


def _db_path(path=None) -> str:
    return path or appconfig.db_path()


def connect(path=None) -> sqlite3.Connection:
    """Open a SQLite connection with the project concurrency pragmas."""

    db_path = _db_path(path)
    if db_path != ":memory:":
        directory = os.path.dirname(os.path.abspath(db_path))
        if directory:
            os.makedirs(directory, exist_ok=True)
    conn = sqlite3.connect(db_path, timeout=DEFAULT_BUSY_TIMEOUT_MS / 1000)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}")
    conn.execute("PRAGMA synchronous=NORMAL")
    return conn


@contextmanager
def _connection(conn=None, path=None):
    owned = conn is None
    database = connect(path) if owned else conn
    try:
        yield database
    finally:
        if owned:
            database.close()


def _as_dataclass(row, cls):
    return None if row is None else cls(**dict(row))


def _fetch_one(conn, sql, params, cls):
    return _as_dataclass(conn.execute(sql, params).fetchone(), cls)


def _fetch_all(conn, sql, params, cls):
    return [_as_dataclass(row, cls) for row in conn.execute(sql, params).fetchall()]


def _validate_fields(fields, allowed):
    unknown = sorted(set(fields) - allowed)
    if unknown:
        raise DbError(f"不支持更新字段: {', '.join(unknown)}")


def _row_key(batch_id, source_file_abs, source_sheet, source_row) -> str:
    return f"{batch_id}:{source_file_abs}:{source_sheet}:{source_row}"


def _attempt_field(phase: str) -> str:
    key = str(phase).lower()
    if key not in PHASE_ATTEMPT_FIELDS:
        raise DbError(f"未知阶段: {phase}")
    return PHASE_ATTEMPT_FIELDS[key]


def init_db(path=None, conn=None) -> None:
    """Create tables and indexes if they do not already exist."""

    with _connection(conn, path) as database:
        with database:
            database.executescript(SCHEMA_SQL)


def create_batch(file_paths: Iterable[str], note=None, path=None, conn=None) -> str:
    batch_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
    files = [os.path.abspath(file_path) for file_path in file_paths]
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                INSERT INTO batches
                  (id, source_files_json, status, note, created_at, updated_at)
                VALUES (?, ?, 'active', ?, ?, ?)
                """,
                (batch_id, json.dumps(files, ensure_ascii=False), note, now, now),
            )
    return batch_id


def get_batch(batch_id, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM batches WHERE id = ?",
            (batch_id,),
            Batch,
        )


def list_batches(status=None, path=None, conn=None):
    sql = "SELECT * FROM batches"
    params = []
    if status is not None:
        sql += " WHERE status = ?"
        params.append(status)
    sql += " ORDER BY created_at DESC, id DESC"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, Batch)


def update_batch(batch_id, path=None, conn=None, **fields) -> None:
    _validate_fields(fields, VALID_BATCH_FIELDS)
    if not fields:
        return
    fields["updated_at"] = _now()
    assignments = ", ".join(f"{field} = ?" for field in fields)
    params = list(fields.values()) + [batch_id]
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"UPDATE batches SET {assignments} WHERE id = ?",
                params,
            )


def list_accounts(path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_all(
            database,
            "SELECT * FROM accounts ORDER BY alias",
            (),
            Account,
        )


def get_account_by_alias(alias, path=None, conn=None):
    with _connection(conn, path) as database:
        return _fetch_one(
            database,
            "SELECT * FROM accounts WHERE alias = ?",
            (alias,),
            Account,
        )


def add_account(
    account_name,
    alias,
    region_host,
    debug_port,
    password=None,
    note=None,
    slug=None,
    user_data_dir=None,
    path=None,
    conn=None,
):
    slug = slug or make_slug(alias)
    user_data_dir = user_data_dir or os.path.join("chrome_user_data_dir", slug)
    now = _now()
    with _connection(conn, path) as database:
        try:
            with database:
                database.execute(
                    """
                    INSERT INTO accounts
                      (account_name, alias, region_host, slug, user_data_dir,
                       debug_port, password, note, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        account_name,
                        alias,
                        region_host,
                        slug,
                        user_data_dir,
                        int(debug_port),
                        password,
                        note,
                        now,
                        now,
                    ),
                )
        except sqlite3.IntegrityError as exc:
            raise DbError(f"账号别名或 slug 已存在: {alias}") from exc
        return get_account_by_alias(alias, conn=database)


def update_account(alias, path=None, conn=None, **fields) -> None:
    _validate_fields(fields, VALID_ACCOUNT_FIELDS)
    if not fields:
        return
    fields["updated_at"] = _now()
    assignments = ", ".join(f"{field} = ?" for field in fields)
    params = list(fields.values()) + [alias]
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"UPDATE accounts SET {assignments} WHERE alias = ?",
                params,
            )


def delete_account(alias, path=None, conn=None) -> None:
    with _connection(conn, path) as database:
        with database:
            database.execute("DELETE FROM accounts WHERE alias = ?", (alias,))


def insert_tasks(batch_id, rows, path=None, conn=None) -> int:
    now = _now()
    values = []
    for row in rows:
        source_file_abs = os.path.abspath(row["source_file_abs"])
        source_sheet = row["source_sheet"]
        source_row = int(row["source_row"])
        values.append(
            (
                batch_id,
                row.get("source_file") or source_file_abs,
                source_file_abs,
                source_sheet,
                source_row,
                row.get("row_key") or _row_key(batch_id, source_file_abs, source_sheet, source_row),
                row.get("account_name"),
                row["alias"],
                str(row["item_id"]),
                now,
                now,
            )
        )
    with _connection(conn, path) as database:
        try:
            with database:
                database.executemany(
                    """
                    INSERT INTO tasks
                      (batch_id, source_file, source_file_abs, source_sheet,
                       source_row, row_key, account_name, alias, item_id,
                       imported_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    values,
                )
        except sqlite3.IntegrityError as exc:
            raise DbError("任务行重复或批次不存在") from exc
    return len(values)


def list_tasks(batch_id=None, stage=None, status=None, alias=None, path=None, conn=None):
    clauses = []
    params = []
    filters = {
        "batch_id": batch_id,
        "stage": stage,
        "status": status,
        "alias": alias,
    }
    for field, value in filters.items():
        if value is not None:
            clauses.append(f"{field} = ?")
            params.append(value)
    sql = "SELECT * FROM tasks"
    if clauses:
        sql += " WHERE " + " AND ".join(clauses)
    sql += " ORDER BY id"
    with _connection(conn, path) as database:
        return _fetch_all(database, sql, params, Task)


def mark_running(task_id, phase, path=None, conn=None) -> None:
    _attempt_field(phase)
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET status = 'running', last_error = NULL, updated_at = ?
                WHERE id = ?
                """,
                (_now(), int(task_id)),
            )


def mark_failed(task_id, phase, error, path=None, conn=None) -> None:
    attempt_field = _attempt_field(phase)
    with _connection(conn, path) as database:
        with database:
            database.execute(
                f"""
                UPDATE tasks
                SET status = 'failed',
                    last_error = ?,
                    {attempt_field} = {attempt_field} + 1,
                    updated_at = ?
                WHERE id = ?
                """,
                (str(error), _now(), int(task_id)),
            )


def mark_skipped(task_id, reason, path=None, conn=None) -> None:
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET status = 'skipped', last_error = ?, updated_at = ?
                WHERE id = ?
                """,
                (str(reason), _now(), int(task_id)),
            )


def set_collected(task_id, old_title, old_cover_path, path=None, conn=None) -> None:
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET old_title = ?,
                    old_cover_path = ?,
                    stage = 'collected',
                    status = 'success',
                    last_error = NULL,
                    collect_attempts = collect_attempts + 1,
                    collected_at = ?,
                    updated_at = ?
                WHERE id = ?
                """,
                (old_title, old_cover_path, now, now, int(task_id)),
            )


def set_generated(task_id, new_title, new_cover_path, path=None, conn=None) -> None:
    now = _now()
    with _connection(conn, path) as database:
        with database:
            database.execute(
                """
                UPDATE tasks
                SET new_title = ?,
                    new_cover_path = ?,
                    stage = 'generated',
                    status = 'success',
                    last_error = NULL,
                    generate_attempts = generate_attempts + 1,
                    generated_at = ?,
                    updated_at = ?
                WHERE id = ?
                """,
                (new_title, new_cover_path, now, now, int(task_id)),
            )


def set_applied(task_id, committed, error=None, path=None, conn=None) -> None:
    now = _now()
    success = bool(committed) and error is None
    with _connection(conn, path) as database:
        with database:
            if success:
                database.execute(
                    """
                    UPDATE tasks
                    SET committed = 1,
                        stage = 'applied',
                        status = 'success',
                        last_error = NULL,
                        apply_attempts = apply_attempts + 1,
                        applied_at = ?,
                        updated_at = ?
                    WHERE id = ?
                    """,
                    (now, now, int(task_id)),
                )
            else:
                database.execute(
                    """
                    UPDATE tasks
                    SET committed = 0,
                        status = 'failed',
                        last_error = ?,
                        apply_attempts = apply_attempts + 1,
                        updated_at = ?
                    WHERE id = ?
                    """,
                    (str(error or "未提交更新"), now, int(task_id)),
                )