"""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", "alias", "region_host", "slug", "user_data_dir", "debug_port", "password", "note", "last_login_at", } VALID_RUN_LOG_FIELDS = { "status", "done", "success_count", "skipped_count", "failed_count", "finished_at", "summary_json", } 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 @dataclass(frozen=True) class RunLog: id: int run_type: str dry_run: int status: str total: int done: int success_count: int skipped_count: int failed_count: int options_json: Optional[str] summary_json: Optional[str] started_at: str finished_at: Optional[str] @property def options(self) -> dict: return json.loads(self.options_json or "{}") @property def summary(self) -> dict: return json.loads(self.summary_json or "{}") @dataclass(frozen=True) class RunLogEvent: id: int run_id: int task_id: Optional[int] alias: Optional[str] item_id: Optional[str] level: str message: str created_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); CREATE TABLE IF NOT EXISTS run_logs ( id INTEGER PRIMARY KEY, run_type TEXT NOT NULL, dry_run INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL DEFAULT 'running', total INTEGER NOT NULL DEFAULT 0, done INTEGER NOT NULL DEFAULT 0, success_count INTEGER NOT NULL DEFAULT 0, skipped_count INTEGER NOT NULL DEFAULT 0, failed_count INTEGER NOT NULL DEFAULT 0, options_json TEXT, summary_json TEXT, started_at TEXT NOT NULL, finished_at TEXT ); CREATE TABLE IF NOT EXISTS run_log_events ( id INTEGER PRIMARY KEY, run_id INTEGER NOT NULL REFERENCES run_logs(id) ON DELETE CASCADE, task_id INTEGER, alias TEXT, item_id TEXT, level TEXT NOT NULL DEFAULT 'info', message TEXT NOT NULL, created_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_run_logs_started ON run_logs(started_at DESC, id DESC); CREATE INDEX IF NOT EXISTS idx_run_log_events_run ON run_log_events(run_id, 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(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()) + [account_alias] with _connection(conn, path) as database: try: with database: database.execute( f"UPDATE accounts SET {assignments} WHERE alias = ?", params, ) except sqlite3.IntegrityError as exc: raise DbError( f"账号别名或 slug 已存在: {fields.get('alias', account_alias)}" ) from exc 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 get_task(task_id, path=None, conn=None): with _connection(conn, path) as database: return _fetch_one( database, "SELECT * FROM tasks WHERE id = ?", (int(task_id),), 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)), ) def reset_generated(task_id, delete_file=False, path=None, conn=None) -> dict: """Clear local AI result and move the task back to collected state.""" with _connection(conn, path) as database: before = get_task(task_id, conn=database) if before is None: raise DbError(f"任务不存在: {task_id}") new_cover_path = before.new_cover_path deleted_file = None if delete_file: raise DbError("重置生成结果默认不删除本地新封面文件，请手动确认后清理") now = _now() with database: database.execute( """ UPDATE tasks SET new_title = NULL, new_cover_path = NULL, stage = 'collected', status = 'success', last_error = NULL, generated_at = NULL, updated_at = ? WHERE id = ? """, (now, int(task_id)), ) after = get_task(task_id, conn=database) return { "task_id": int(task_id), "before": before, "after": after, "new_cover_path": new_cover_path, "deleted_file": deleted_file, } def reset_apply_status(task_id, path=None, conn=None) -> dict: """Move a generated/applied local task back to pending update state.""" with _connection(conn, path) as database: before = get_task(task_id, conn=database) if before is None: raise DbError(f"任务不存在: {task_id}") now = _now() with database: database.execute( """ UPDATE tasks SET stage = 'generated', status = 'pending', last_error = NULL, updated_at = ? WHERE id = ? """, (now, int(task_id)), ) after = get_task(task_id, conn=database) return { "task_id": int(task_id), "before": before, "after": after, } def create_run_log( run_type, dry_run=False, total=0, options=None, path=None, conn=None, ) -> int: """Create a high-level operation log and return its id.""" now = _now() options_json = json.dumps( appconfig.sanitize_for_log(options or {}), ensure_ascii=False, sort_keys=True, ) with _connection(conn, path) as database: with database: cursor = database.execute( """ INSERT INTO run_logs (run_type, dry_run, status, total, options_json, started_at) VALUES (?, ?, 'running', ?, ?, ?) """, (str(run_type), 1 if dry_run else 0, int(total), options_json, now), ) return int(cursor.lastrowid) def finish_run_log(run_id, path=None, conn=None, **fields) -> None: """Mark a run log complete/blocked/failed with a sanitized summary.""" _validate_fields(fields, VALID_RUN_LOG_FIELDS) if "summary_json" in fields and not isinstance(fields["summary_json"], str): fields["summary_json"] = json.dumps( appconfig.sanitize_for_log(fields["summary_json"] or {}), ensure_ascii=False, sort_keys=True, ) if "finished_at" not in fields: fields["finished_at"] = _now() assignments = ", ".join(f"{field} = ?" for field in fields) params = list(fields.values()) + [int(run_id)] with _connection(conn, path) as database: with database: database.execute( f"UPDATE run_logs SET {assignments} WHERE id = ?", params, ) def add_run_log_event( run_id, message, task_id=None, alias=None, item_id=None, level="info", path=None, conn=None, ) -> int: """Append one sanitized event line to a run log.""" sanitized = appconfig.sanitize_for_log( { "message": str(message), "alias": alias, "item_id": item_id, "level": level, } ) with _connection(conn, path) as database: with database: cursor = database.execute( """ INSERT INTO run_log_events (run_id, task_id, alias, item_id, level, message, created_at) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( int(run_id), int(task_id) if task_id is not None else None, sanitized.get("alias"), sanitized.get("item_id"), sanitized.get("level") or "info", sanitized.get("message") or "", _now(), ), ) return int(cursor.lastrowid) def list_run_logs(limit=50, run_type=None, path=None, conn=None): clauses = [] params = [] if run_type is not None: clauses.append("run_type = ?") params.append(str(run_type)) sql = "SELECT * FROM run_logs" if clauses: sql += " WHERE " + " AND ".join(clauses) sql += " ORDER BY started_at DESC, id DESC LIMIT ?" params.append(max(1, int(limit))) with _connection(conn, path) as database: return _fetch_all(database, sql, params, RunLog) def list_run_log_events(run_id=None, limit=200, path=None, conn=None): clauses = [] params = [] if run_id is not None: clauses.append("run_id = ?") params.append(int(run_id)) sql = "SELECT * FROM run_log_events" if clauses: sql += " WHERE " + " AND ".join(clauses) sql += " ORDER BY id DESC LIMIT ?" params.append(max(1, int(limit))) with _connection(conn, path) as database: events = _fetch_all(database, sql, params, RunLogEvent) return list(reversed(events))