re-xxm/reaction_board/database.py

# -*- coding: utf-8 -*-
"""SQLite 连接、迁移入口与单写队列。

所有写操作经单个后台线程串行执行，避免并发任务出现 ``database is locked``；
读操作使用短生命周期连接。网络请求绝不在写事务内执行。
"""

from __future__ import annotations

import os
import queue
import sqlite3
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, TypeVar

from .migrations import current_version, run_migrations

T = TypeVar("T")


class DatabaseClosed(RuntimeError):
    pass


@dataclass
class _Job:
    fn: Callable[[sqlite3.Connection], Any]
    done: threading.Event = field(default_factory=threading.Event)
    result: Any = None
    error: BaseException | None = None


def _configure(conn: sqlite3.Connection) -> sqlite3.Connection:
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    conn.execute("PRAGMA busy_timeout = 5000")
    conn.execute("PRAGMA synchronous = NORMAL")
    conn.execute("PRAGMA journal_mode = WAL")
    return conn


class Database:
    """线程安全的 SQLite 访问入口。"""

    def __init__(
        self,
        path: str | Path,
        *,
        backup_dir: str | Path | None = None,
        backup_keep_days: int = 14,
        logger=None,
    ) -> None:
        self.path = Path(path)
        self.backup_dir = Path(backup_dir) if backup_dir is not None else None
        self.backup_keep_days = backup_keep_days
        self.logger = logger
        self.schema_version: int | None = None
        self._queue: queue.Queue[_Job | None] = queue.Queue()
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._writer_error: BaseException | None = None
        self._closed = False

    # ------------------------------------------------------------ setup

    def _connect(self) -> sqlite3.Connection:
        # isolation_level=None：单条 DML 立即提交；显式事务由 transaction() 管理。
        return _configure(
            sqlite3.connect(self.path, timeout=5.0, isolation_level=None)
        )

    def initialize(self) -> list[int]:
        """创建目录、执行迁移并启动单写线程。"""
        self.path.parent.mkdir(parents=True, exist_ok=True)
        conn = self._connect()
        try:
            applied = run_migrations(
                conn,
                backup_dir=self.backup_dir,
                backup_keep_days=self.backup_keep_days,
                logger=self.logger,
            )
            self.schema_version = current_version(conn)
        finally:
            conn.close()
        self._start_writer()
        if self.logger is not None:
            self.logger.info(
                "db.ready", path=str(self.path), schema_version=self.schema_version
            )
        return applied

    def _start_writer(self) -> None:
        if self._thread is not None:
            return
        self._thread = threading.Thread(target=self._writer_loop, name="sqlite-writer", daemon=True)
        self._thread.start()
        self._ready.wait(timeout=10)
        if self._writer_error is not None:
            raise RuntimeError(f"SQLite 写线程启动失败: {self._writer_error}") from self._writer_error

    def _writer_loop(self) -> None:
        conn: sqlite3.Connection | None = None
        try:
            conn = self._connect()
        except BaseException as exc:  # noqa: BLE001
            self._writer_error = exc
            self._ready.set()
            return
        self._ready.set()
        try:
            while True:
                job = self._queue.get()
                if job is None:
                    break
                try:
                    job.result = job.fn(conn)
                except BaseException as exc:  # noqa: BLE001  原样返回给调用方
                    job.error = exc
                finally:
                    job.done.set()
        finally:
            conn.close()

    # ------------------------------------------------------------ access

    def write(self, fn: Callable[[sqlite3.Connection], T]) -> T:
        """把函数交给单写线程执行，并同步等待结果。"""
        if self._closed:
            raise DatabaseClosed("数据库已关闭")
        if self._thread is None:
            raise RuntimeError("请先调用 initialize()")
        job = _Job(fn=fn)
        self._queue.put(job)
        job.done.wait()
        if job.error is not None:
            raise job.error
        return job.result

    def transaction(self, fn: Callable[[sqlite3.Connection], T]) -> T:
        """在单写线程内执行一个显式事务。"""
        def run(conn: sqlite3.Connection) -> T:
            conn.execute("BEGIN IMMEDIATE")
            try:
                result = fn(conn)
                conn.execute("COMMIT")
                return result
            except BaseException:
                try:
                    conn.execute("ROLLBACK")
                except sqlite3.Error:
                    pass
                raise

        return self.write(run)

    def read(self, fn: Callable[[sqlite3.Connection], T]) -> T:
        """使用独立的短生命周期连接执行只读查询。"""
        if self._closed:
            raise DatabaseClosed("数据库已关闭")
        conn = self._connect()
        try:
            return fn(conn)
        finally:
            conn.close()

    def backup_to(self, target: str | Path) -> Path:
        """用 SQLite backup API 生成一致性备份；先写临时文件再原子替换。"""
        target = Path(target)
        target.parent.mkdir(parents=True, exist_ok=True)
        tmp = target.with_name(target.name + ".tmp")
        if tmp.exists():
            tmp.unlink()

        def run(conn: sqlite3.Connection) -> None:
            dst = sqlite3.connect(tmp)
            try:
                conn.backup(dst)
            finally:
                dst.close()

        self.write(run)
        os.replace(tmp, target)
        return target

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self._thread is not None:
            self._queue.put(None)
            self._thread.join(timeout=10)
            self._thread = None