"""SQLite persistence for accounts, import batches, and tasks.""" from __future__ import annotations import hashlib import json import os import re import sqlite3 import uuid from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime from typing import Iterable, Optional from . import appconfig 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 _slug_from_alias(alias: str) -> str: base = re.sub(r"[^a-z0-9_]+", "_", alias.lower()).strip("_") suffix = hashlib.sha1(alias.encode("utf-8")).hexdigest()[:8] return f"{base}_{suffix}" if base else f"account_{suffix}" 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 _slug_from_alias(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)), )