"""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() TRIAL_ELIGIBLE_NOTE = "trial_eligible" TRIAL_INELIGIBLE_NOTE = "trial_ineligible" _TRIAL_EVENT_PATTERNS = ("%免费试用%", "%trial%") 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: conn = sqlite3.connect(DB_PATH) try: conn.executescript(SCHEMA) conn.commit() finally: conn.close() @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, offset: int = 0) -> list[dict]: init_db() with _conn() as c: if status == "trial": rows = c.execute( "SELECT * FROM accounts ORDER BY created_at DESC", ).fetchall() elif status: rows = c.execute( "SELECT * FROM accounts WHERE final_status = ? ORDER BY created_at DESC LIMIT ? OFFSET ?", (status, limit, offset), ).fetchall() else: rows = c.execute( "SELECT * FROM accounts ORDER BY created_at DESC LIMIT ? OFFSET ?", (limit, offset), ).fetchall() accounts = _annotate_accounts([_row_to_dict(r) for r in rows], c) if status == "trial": accounts = [account for account in accounts if account_matches_status(account, "trial")] return accounts[offset:offset + limit] return accounts def count_accounts(status: str | None = None) -> int: init_db() with _conn() as c: if status == "trial": rows = c.execute("SELECT * FROM accounts").fetchall() accounts = _annotate_accounts([_row_to_dict(r) for r in rows], c) return sum(1 for account in accounts if account_matches_status(account, "trial")) if status: row = c.execute("SELECT COUNT(*) FROM accounts WHERE final_status = ?", (status,)).fetchone() else: row = c.execute("SELECT COUNT(*) FROM accounts").fetchone() return row[0] if row else 0 def get_account(email: str) -> dict | None: init_db() with _conn() as c: r = c.execute("SELECT * FROM accounts WHERE email = ?", (email,)).fetchone() if not r: return None return _annotate_accounts([_row_to_dict(r)], c)[0] 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 account_matches_status(account: dict | None, status: str | None) -> bool: if not account: return False if not status: return True if status == "trial": return bool(account.get("is_trial_account") or (account.get("final_status") or "") == "trial") return (account.get("final_status") or "") == status def derive_account_flags(account: dict | None, *, has_trial_event: bool = False) -> dict | None: if account is None: return None annotated = dict(account) final_status = str(annotated.get("final_status") or "") note_flags = _parse_note_flags(annotated.get("notes")) trial_state = "unknown" if TRIAL_ELIGIBLE_NOTE in note_flags: trial_state = "eligible" elif TRIAL_INELIGIBLE_NOTE in note_flags: trial_state = "ineligible" elif has_trial_event: trial_state = "eligible" is_trial_account = trial_state == "eligible" annotated["trial_state"] = trial_state annotated["trial_eligible"] = is_trial_account annotated["is_trial_account"] = is_trial_account annotated["can_retry_payment"] = bool(_session_access_token(annotated) and (is_trial_account or final_status == "trial")) return annotated def set_account_trial_eligibility(email: str, password: str, eligible: bool) -> dict | None: existing = get_account(email) or {} notes = _update_note_flags( existing.get("notes"), add=[TRIAL_ELIGIBLE_NOTE if eligible else TRIAL_INELIGIBLE_NOTE], remove=[TRIAL_INELIGIBLE_NOTE if eligible else TRIAL_ELIGIBLE_NOTE], ) return upsert_account(email, password, fields={"notes": notes}) def _annotate_accounts(accounts: Iterable[dict | None], conn: sqlite3.Connection) -> list[dict]: items = [account for account in accounts if account] if not items: return [] trial_emails = _load_trial_event_emails(conn, [str(account.get("email") or "") for account in items]) out: list[dict] = [] for account in items: annotated = derive_account_flags(account, has_trial_event=str(account.get("email") or "") in trial_emails) if annotated: out.append(annotated) return out def _load_trial_event_emails(conn: sqlite3.Connection, emails: Iterable[str]) -> set[str]: email_list = [email for email in emails if email] if not email_list: return set() placeholders = ",".join("?" for _ in email_list) rows = conn.execute( f""" SELECT DISTINCT email FROM account_events WHERE email IN ({placeholders}) AND stage = 'paypal' AND detail IS NOT NULL AND (detail LIKE ? OR detail LIKE ?) """, [*email_list, *_TRIAL_EVENT_PATTERNS], ).fetchall() return {str(row[0]) for row in rows} def _session_access_token(account: dict) -> str: for key in ("plus_session", "initial_session"): session = account.get(key) if isinstance(session, dict): token = session.get("accessToken") if token: return str(token) return "" def _parse_note_flags(notes: Any) -> set[str]: raw = str(notes or "").replace(",", " ") return {part.strip() for part in raw.split() if part.strip()} def _update_note_flags(notes: Any, *, add: Iterable[str] = (), remove: Iterable[str] = ()) -> str | None: flags = _parse_note_flags(notes) flags.difference_update({flag for flag in remove if flag}) flags.update({flag for flag in add if flag}) if not flags: return None return " ".join(sorted(flags)) 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 count_tasks_by_mode_statuses(mode: str, statuses: tuple[str, ...]) -> int: init_db() status_values = tuple(str(status) for status in statuses if status) if not status_values: return 0 placeholders = ",".join("?" for _ in status_values) with _conn() as c: row = c.execute( f"SELECT COUNT(*) FROM tasks WHERE mode = ? AND status IN ({placeholders})", (mode, *status_values), ).fetchone() return int(row[0]) if row else 0 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