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"


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.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
    }