"""SQLite 持久化:注册成功的账号、session 快照、CPA 上传记录。""" from __future__ import annotations import json import os import sqlite3 import threading import time from contextlib import contextmanager from typing import Any, Iterable DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") DB_PATH = os.path.join(DATA_DIR, "accounts.db") _LOCK = threading.Lock() SCHEMA = """ CREATE TABLE IF NOT EXISTS accounts ( email TEXT PRIMARY KEY, password TEXT NOT NULL, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, plan_type TEXT, final_status TEXT, -- registered/paid/plus/cpa_uploaded/failed/stopped last_error TEXT, long_link TEXT, cpa_file_name TEXT, cpa_uploaded_at INTEGER, initial_session_json TEXT, -- 注册成功时拉到的 /api/auth/session plus_session_json TEXT, -- 付款后拉到的 plus session notes TEXT ); CREATE INDEX IF NOT EXISTS idx_accounts_status ON accounts(final_status); CREATE INDEX IF NOT EXISTS idx_accounts_created ON accounts(created_at); CREATE TABLE IF NOT EXISTS account_events ( id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT NOT NULL, ts INTEGER NOT NULL, stage TEXT NOT NULL, status TEXT NOT NULL, -- info/ok/warn/error detail TEXT, payload_json TEXT ); CREATE INDEX IF NOT EXISTS idx_events_email ON account_events(email); CREATE INDEX IF NOT EXISTS idx_events_ts ON account_events(ts); CREATE TABLE IF NOT EXISTS tasks ( task_id TEXT PRIMARY KEY, mode TEXT NOT NULL, -- full | pay_only status TEXT NOT NULL, -- queued | running | success | failed | cancelled stage TEXT, attempts INTEGER DEFAULT 0, max_attempts INTEGER DEFAULT 3, params_json TEXT, result_json TEXT, last_error TEXT, email TEXT, plan_type TEXT, cpa_file_name TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, started_at INTEGER, finished_at INTEGER ); CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status); CREATE INDEX IF NOT EXISTS idx_tasks_created ON tasks(created_at); """ def _now_ms() -> int: return int(time.time() * 1000) def _ensure_dir(): os.makedirs(DATA_DIR, exist_ok=True) def init_db(): _ensure_dir() with _LOCK, sqlite3.connect(DB_PATH) as conn: conn.executescript(SCHEMA) conn.commit() @contextmanager def _conn(): _ensure_dir() with _LOCK: c = sqlite3.connect(DB_PATH) c.row_factory = sqlite3.Row try: yield c c.commit() finally: c.close() def _dump(value: Any) -> str | None: if value is None: return None try: return json.dumps(value, ensure_ascii=False) except Exception: return str(value) def upsert_account(email: str, password: str, *, fields: dict | None = None) -> dict: """Insert or update by email. fields 中只更新非 None 字段。""" init_db() fields = dict(fields or {}) now = _now_ms() with _conn() as c: row = c.execute("SELECT email FROM accounts WHERE email = ?", (email,)).fetchone() if row is None: c.execute( """ INSERT INTO accounts (email, password, created_at, updated_at, plan_type, final_status, last_error, long_link, cpa_file_name, cpa_uploaded_at, initial_session_json, plus_session_json, notes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( email, password, now, now, fields.get("plan_type"), fields.get("final_status") or "registered", fields.get("last_error"), fields.get("long_link"), fields.get("cpa_file_name"), fields.get("cpa_uploaded_at"), _dump(fields.get("initial_session")) if "initial_session" in fields else fields.get("initial_session_json"), _dump(fields.get("plus_session")) if "plus_session" in fields else fields.get("plus_session_json"), fields.get("notes"), ), ) else: sets = ["updated_at = ?"] args: list[Any] = [now] for col in ("plan_type", "final_status", "last_error", "long_link", "cpa_file_name", "cpa_uploaded_at", "notes"): if col in fields and fields[col] is not None: sets.append(f"{col} = ?") args.append(fields[col]) if "initial_session" in fields: sets.append("initial_session_json = ?") args.append(_dump(fields["initial_session"])) elif "initial_session_json" in fields and fields["initial_session_json"] is not None: sets.append("initial_session_json = ?") args.append(fields["initial_session_json"]) if "plus_session" in fields: sets.append("plus_session_json = ?") args.append(_dump(fields["plus_session"])) elif "plus_session_json" in fields and fields["plus_session_json"] is not None: sets.append("plus_session_json = ?") args.append(fields["plus_session_json"]) if password: sets.append("password = ?") args.append(password) args.append(email) c.execute(f"UPDATE accounts SET {', '.join(sets)} WHERE email = ?", args) return _row_to_dict(c.execute("SELECT * FROM accounts WHERE email = ?", (email,)).fetchone()) def add_event(email: str, stage: str, status: str = "info", detail: str | None = None, payload: Any = None) -> int: init_db() with _conn() as c: cur = c.execute( "INSERT INTO account_events (email, ts, stage, status, detail, payload_json) VALUES (?, ?, ?, ?, ?, ?)", (email or "", _now_ms(), stage, status, detail, _dump(payload)), ) return cur.lastrowid def list_accounts(limit: int = 200, status: str | None = None) -> list[dict]: init_db() with _conn() as c: if status: rows = c.execute( "SELECT * FROM accounts WHERE final_status = ? ORDER BY created_at DESC LIMIT ?", (status, limit), ).fetchall() else: rows = c.execute( "SELECT * FROM accounts ORDER BY created_at DESC LIMIT ?", (limit,), ).fetchall() return [_row_to_dict(r) for r in rows] def get_account(email: str) -> dict | None: init_db() with _conn() as c: r = c.execute("SELECT * FROM accounts WHERE email = ?", (email,)).fetchone() return _row_to_dict(r) if r else None def list_events(email: str, limit: int = 100) -> list[dict]: init_db() with _conn() as c: rows = c.execute( "SELECT * FROM account_events WHERE email = ? ORDER BY ts DESC LIMIT ?", (email, limit), ).fetchall() return [_event_row(r) for r in rows] def _row_to_dict(row: sqlite3.Row | None) -> dict | None: if row is None: return None d = {k: row[k] for k in row.keys()} # session JSON 反序列化但保留 raw 副本 for k in ("initial_session_json", "plus_session_json"): raw = d.get(k) if raw: try: d[k.replace("_json", "")] = json.loads(raw) except Exception: d[k.replace("_json", "")] = None return d def _event_row(row: sqlite3.Row) -> dict: d = {k: row[k] for k in row.keys()} raw = d.get("payload_json") if raw: try: d["payload"] = json.loads(raw) except Exception: d["payload"] = None return d # ------------------------- tasks ------------------------- def create_task(task_id: str, mode: str, params: dict, max_attempts: int = 3) -> dict: init_db() now = _now_ms() with _conn() as c: c.execute( """ INSERT INTO tasks (task_id, mode, status, stage, attempts, max_attempts, params_json, result_json, last_error, email, plan_type, cpa_file_name, created_at, updated_at, started_at, finished_at) VALUES (?, ?, 'queued', '', 0, ?, ?, NULL, NULL, NULL, NULL, NULL, ?, ?, NULL, NULL) """, (task_id, mode, max_attempts, _dump(params or {}), now, now), ) return _task_row(c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()) def update_task(task_id: str, fields: dict) -> dict | None: init_db() if not fields: return get_task(task_id) sets = ["updated_at = ?"] args: list[Any] = [_now_ms()] for col in ( "status", "stage", "attempts", "max_attempts", "last_error", "email", "plan_type", "cpa_file_name", "started_at", "finished_at", ): if col in fields and fields[col] is not None: sets.append(f"{col} = ?") args.append(fields[col]) if "result" in fields: sets.append("result_json = ?") args.append(_dump(fields["result"])) elif "result_json" in fields and fields["result_json"] is not None: sets.append("result_json = ?") args.append(fields["result_json"]) if "params" in fields: sets.append("params_json = ?") args.append(_dump(fields["params"])) args.append(task_id) with _conn() as c: c.execute(f"UPDATE tasks SET {', '.join(sets)} WHERE task_id = ?", args) row = c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone() return _task_row(row) if row else None def get_task(task_id: str) -> dict | None: init_db() with _conn() as c: r = c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone() return _task_row(r) if r else None def list_tasks(limit: int = 100, status: str | None = None) -> list[dict]: init_db() with _conn() as c: if status: rows = c.execute( "SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ?", (status, limit), ).fetchall() else: rows = c.execute( "SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?", (limit,), ).fetchall() return [_task_row(r) for r in rows] def get_next_queued_task() -> dict | None: """取一个最早的 queued 任务(FIFO)。""" init_db() with _conn() as c: r = c.execute( "SELECT * FROM tasks WHERE status = 'queued' ORDER BY created_at ASC LIMIT 1" ).fetchone() return _task_row(r) if r else None def _task_row(row: sqlite3.Row | None) -> dict | None: if row is None: return None d = {k: row[k] for k in row.keys()} for k in ("params_json", "result_json"): raw = d.get(k) if raw: try: d[k.replace("_json", "")] = json.loads(raw) except Exception: d[k.replace("_json", "")] = None return d