# 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 ]