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.execute("DELETE FROM processed_emails") conn.commit() logger.info("数据库初始化完成: %s", DB_PATH) def _get_conn() -> sqlite3.Connection: if not hasattr(_local, "conn") or _local.conn is None: DB_PATH.parent.mkdir(parents=True, exist_ok=True) _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,))