InboundVerify/state_store.py

# state_store.py
# 阶段0：站点状态持久化（SQLite）。记录各站登录态 + 应到/实到数据文件生成时间，
# 供后台心跳刷新、菜单状态盘展示，以及未来 FastAPI / Web 前端读取。
# 重启不丢——程序重启后状态从本库恢复（登录态会随心跳重新探测校正）。
#
# 设计：纯 Python（sqlite3 标准库，无新依赖），每次读写开短连接，主线程使用。

import os
import sqlite3
from datetime import datetime

from paths import STATE_DB_PATH

# 登录态枚举
LOGIN_UNKNOWN = "unknown"  # 尚未探测过
LOGIN_IN = "logged_in"
LOGIN_OUT = "logged_out"

# 任务状态枚举（task_history.status）
TASK_PENDING = "pending"  # 已入队，待执行
TASK_RUNNING = "running"  # 正在执行
TASK_SUCCESS = "success"  # 成功（有数据）
TASK_NO_DATA = "no_data"  # 成功但本站本次无数据
TASK_FAILED = "failed"  # 失败（重试耗尽 / 未登录 / 异常）


def _now():
    """本地时间的字符串（到秒），用于时间戳列。"""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def init_db():
    """建库建表（幂等）。确保 state 目录存在。"""
    os.makedirs(os.path.dirname(STATE_DB_PATH), exist_ok=True)
    with sqlite3.connect(STATE_DB_PATH) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS site_status (
                site                  TEXT PRIMARY KEY,
                login_state           TEXT,
                login_checked_at      TEXT,
                expected_ready        INTEGER,
                expected_generated_at TEXT,
                actual_ready          INTEGER,
                actual_generated_at   TEXT,
                updated_at            TEXT
            )
            """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS task_history (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                site        TEXT,
                kind        TEXT,
                status      TEXT,
                started_at  TEXT,
                finished_at TEXT,
                error       TEXT
            )
            """)
        conn.commit()


def _upsert(conn, site, **fields):
    """更新（或插入）单站：保留未传字段，刷新 updated_at。"""
    row = conn.execute(
        "SELECT login_state, login_checked_at, expected_ready, expected_generated_at, "
        "actual_ready, actual_generated_at FROM site_status WHERE site = ?",
        (site,),
    ).fetchone()
    cur = {
        "login_state": LOGIN_UNKNOWN,
        "login_checked_at": "",
        "expected_ready": 0,
        "expected_generated_at": "",
        "actual_ready": 0,
        "actual_generated_at": "",
    }
    if row:
        (
            cur["login_state"],
            cur["login_checked_at"],
            cur["expected_ready"],
            cur["expected_generated_at"],
            cur["actual_ready"],
            cur["actual_generated_at"],
        ) = row
    cur.update(fields)
    cur["updated_at"] = _now()

    values = (
        site,
        cur["login_state"],
        cur["login_checked_at"],
        cur["expected_ready"],
        cur["expected_generated_at"],
        cur["actual_ready"],
        cur["actual_generated_at"],
        cur["updated_at"],
    )
    if row:
        conn.execute(
            "UPDATE site_status SET login_state=?, login_checked_at=?, expected_ready=?, "
            "expected_generated_at=?, actual_ready=?, actual_generated_at=?, updated_at=? "
            "WHERE site=?",
            values[1:] + (site,),
        )
    else:
        conn.execute(
            "INSERT INTO site_status (site, login_state, login_checked_at, expected_ready, "
            "expected_generated_at, actual_ready, actual_generated_at, updated_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            values,
        )
    conn.commit()


def set_login_state(site, logged_in):
    """更新单站登录态。logged_in: bool。"""
    state = LOGIN_IN if logged_in else LOGIN_OUT
    with sqlite3.connect(STATE_DB_PATH) as conn:
        _upsert(conn, site, login_state=state, login_checked_at=_now())


def set_data_state(site, kind, ready, generated_at):
    """更新单站数据态。kind: 'expected'/'actual'；ready: bool；generated_at: str 或 ''。"""
    key_ready = "expected_ready" if kind == "expected" else "actual_ready"
    key_at = "expected_generated_at" if kind == "expected" else "actual_generated_at"
    with sqlite3.connect(STATE_DB_PATH) as conn:
        _upsert(
            conn, site, **{key_ready: 1 if ready else 0, key_at: generated_at or ""}
        )


def get_all_status():
    """返回 {site: {各字段}}；库不存在则返回 {}。"""
    if not os.path.exists(STATE_DB_PATH):
        return {}
    with sqlite3.connect(STATE_DB_PATH) as conn:
        rows = conn.execute(
            "SELECT site, login_state, login_checked_at, expected_ready, "
            "expected_generated_at, actual_ready, actual_generated_at, updated_at "
            "FROM site_status"
        ).fetchall()
    return {
        r[0]: {
            "login_state": r[1],
            "login_checked_at": r[2],
            "expected_ready": bool(r[3]),
            "expected_generated_at": r[4],
            "actual_ready": bool(r[5]),
            "actual_generated_at": r[6],
            "updated_at": r[7],
        }
        for r in rows
    }


# ============================ 任务历史 ============================


def create_task(site, kind):
    """新建一条 pending 任务，返回其 id。"""
    now = _now()
    with sqlite3.connect(STATE_DB_PATH) as conn:
        cur = conn.execute(
            "INSERT INTO task_history (site, kind, status, started_at, finished_at, error) "
            "VALUES (?, ?, ?, ?, '', '')",
            (site, kind, TASK_PENDING, now),
        )
        conn.commit()
        return cur.lastrowid


def update_task(task_id, status, error=None):
    """更新任务状态。终态(success/no_data/failed)写入 finished_at。"""
    finished = _now() if status in (TASK_SUCCESS, TASK_NO_DATA, TASK_FAILED) else ""
    with sqlite3.connect(STATE_DB_PATH) as conn:
        if finished:
            conn.execute(
                "UPDATE task_history SET status=?, finished_at=?, error=? WHERE id=?",
                (status, finished, error or "", task_id),
            )
        else:
            conn.execute(
                "UPDATE task_history SET status=?, error=? WHERE id=?",
                (status, error or "", task_id),
            )
        conn.commit()


def get_task(task_id):
    """返回单条任务 dict，不存在返回 None。"""
    with sqlite3.connect(STATE_DB_PATH) as conn:
        row = conn.execute(
            "SELECT id, site, kind, status, started_at, finished_at, error "
            "FROM task_history WHERE id=?",
            (task_id,),
        ).fetchone()
    if not row:
        return None
    return {
        "id": row[0],
        "site": row[1],
        "kind": row[2],
        "status": row[3],
        "started_at": row[4],
        "finished_at": row[5],
        "error": row[6],
    }


def list_tasks(limit=20):
    """返回最近 limit 条任务（按 id 倒序）。"""
    with sqlite3.connect(STATE_DB_PATH) as conn:
        rows = conn.execute(
            "SELECT id, site, kind, status, started_at, finished_at, error "
            "FROM task_history ORDER BY id DESC LIMIT ?",
            (limit,),
        ).fetchall()
    return [
        {
            "id": r[0],
            "site": r[1],
            "kind": r[2],
            "status": r[3],
            "started_at": r[4],
            "finished_at": r[5],
            "error": r[6],
        }
        for r in rows
    ]