autoChatGPTPlus-Paypal/storage.py

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