# -*- coding: utf-8 -*- """SQLite 连接、迁移入口与单写队列。 所有写操作经单个后台线程串行执行,避免并发任务出现 ``database is locked``; 读操作使用短生命周期连接。网络请求绝不在写事务内执行。 """ from __future__ import annotations 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 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