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