refactor: replace all in-memory state with SQLite database

2026-07-14 14:49:23 +08:00
parent c3132f746a
commit 54c4bf00fb
4 changed files with 218 additions and 51 deletions
--- a/src/database.py
+++ b/src/database.py
@@ -0,0 +1,173 @@
+import json
+import logging
+import sqlite3
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+DB_PATH = Path(__file__).parent.parent / "data" / "mail_bot.db"
+_local = threading.local()
+
+
+def init_db():
+    """初始化数据库表"""
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+    conn = _get_conn()
+    conn.executescript("""
+        CREATE TABLE IF NOT EXISTS processed_emails (
+            uid TEXT PRIMARY KEY,
+            account TEXT NOT NULL,
+            folder TEXT NOT NULL DEFAULT 'INBOX',
+            subject TEXT,
+            sender TEXT,
+            status TEXT NOT NULL DEFAULT 'fetched',
+            created_at TEXT NOT NULL,
+            sent_at TEXT
+        );
+
+        CREATE TABLE IF NOT EXISTS email_contexts (
+            msg_id INTEGER PRIMARY KEY,
+            data TEXT NOT NULL,
+            created_at TEXT NOT NULL
+        );
+
+        CREATE TABLE IF NOT EXISTS conversations (
+            chat_id INTEGER PRIMARY KEY,
+            data TEXT NOT NULL,
+            updated_at TEXT NOT NULL
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_processed_status ON processed_emails(status);
+        CREATE INDEX IF NOT EXISTS idx_processed_account ON processed_emails(account);
+    """)
+    conn.commit()
+    logger.info("数据库初始化完成: %s", DB_PATH)
+
+
+def _get_conn() -> sqlite3.Connection:
+    if not hasattr(_local, "conn") or _local.conn is None:
+        _local.conn = sqlite3.connect(str(DB_PATH), timeout=10)
+        _local.conn.execute("PRAGMA journal_mode=WAL")
+        _local.conn.execute("PRAGMA busy_timeout=5000")
+        _local.conn.row_factory = sqlite3.Row
+    return _local.conn
+
+
+@contextmanager
+def get_db():
+    conn = _get_conn()
+    try:
+        yield conn
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+
+
+# ── processed_emails ────────────────────────────────────
+
+def is_uid_processed(uid: str) -> bool:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT 1 FROM processed_emails WHERE uid=? AND status != 'failed'",
+            (uid,)
+        ).fetchone()
+        return row is not None
+
+
+def is_uids_processed(uids: list[str]) -> set[str]:
+    """批量检查，返回已处理的 uid 集合"""
+    if not uids:
+        return set()
+    with get_db() as conn:
+        placeholders = ",".join("?" * len(uids))
+        rows = conn.execute(
+            f"SELECT uid FROM processed_emails WHERE uid IN ({placeholders}) AND status != 'failed'",
+            uids
+        ).fetchall()
+        return {row["uid"] for row in rows}
+
+
+def mark_fetched(uid: str, account: str, folder: str = "INBOX",
+                 subject: str = "", sender: str = ""):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR IGNORE INTO processed_emails
+               (uid, account, folder, subject, sender, status, created_at)
+               VALUES (?, ?, ?, ?, ?, 'fetched', ?)""",
+            (uid, account, folder, subject, sender, datetime.now().isoformat())
+        )
+
+
+def mark_processing(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='processing' WHERE uid=?",
+            (uid,)
+        )
+
+
+def mark_sent(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='sent', sent_at=? WHERE uid=?",
+            (datetime.now().isoformat(), uid)
+        )
+
+
+def mark_failed(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='failed' WHERE uid=?",
+            (uid,)
+        )
+
+
+# ── email_contexts ──────────────────────────────────────
+
+def save_email_context(msg_id: int, data: dict):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR REPLACE INTO email_contexts (msg_id, data, created_at)
+               VALUES (?, ?, ?)""",
+            (msg_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
+        )
+
+
+def load_email_context(msg_id: int) -> dict | None:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT data FROM email_contexts WHERE msg_id=?", (msg_id,)
+        ).fetchone()
+        if row:
+            return json.loads(row["data"])
+        return None
+
+
+# ── conversations ───────────────────────────────────────
+
+def save_conversation(chat_id: int, data: dict):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR REPLACE INTO conversations (chat_id, data, updated_at)
+               VALUES (?, ?, ?)""",
+            (chat_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
+        )
+
+
+def load_conversation(chat_id: int) -> dict | None:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT data FROM conversations WHERE chat_id=?", (chat_id,)
+        ).fetchone()
+        if row:
+            return json.loads(row["data"])
+        return None
+
+
+def delete_conversation(chat_id: int):
+    with get_db() as conn:
+        conn.execute("DELETE FROM conversations WHERE chat_id=?", (chat_id,))