ai-mail-bot/src/database.py

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,))