# -*- coding: utf-8 -*- """异步 SQLite 访问层(aiosqlite + WAL)。 核心目标:**读写同时进行**。 - WAL 模式:读事务与写事务互不阻塞(读快照隔离),这是 SQLite 层并发的根本; - 读路径:每次操作独立连接、无锁、可并发(`fetch_all`/`fetch_one`/`fetch_val`); - 写路径:SQLite 全库同一时刻只允许一个写者,用 `asyncio.Lock` 串行化写事务, 但写者不阻塞读者(WAL); - `busy_timeout`:多进程/多连接下遇到短暂写锁自旋等待而非立刻 SQLITE_BUSY。 与旧实现的差异:旧 `auth_system` 用 `threading.RLock` 把**全部**读写串行化, 读也被写挡住;本层只串行化写,读放行并发,即「读写同时进行」。 """ import asyncio import os from collections.abc import AsyncIterator, Iterable, Sequence from contextlib import asynccontextmanager from pathlib import Path from typing import Any import aiosqlite # 站点数据目录(可经环境变量覆盖,便于测试隔离) DATA_DIR = Path(os.environ.get('SITE_DATA_DIR', 'data')) DB_PATH = DATA_DIR / 'site.db' # 锁按事件循环隔离:asyncio.Lock 绑定创建它的 loop,跨 loop 复用会抛 # "bound to a different event loop"(迁移期 FastAPI loop + 同步桥 loop + 测试各自 # 的 loop 并存,必须分开持有)。 _LOOP_LOCKS: dict[tuple[int, str], asyncio.Lock] = {} def _loop_lock(kind: str) -> asyncio.Lock: key = (id(asyncio.get_running_loop()), kind) lock = _LOOP_LOCKS.get(key) if lock is None: lock = asyncio.Lock() _LOOP_LOCKS[key] = lock return lock def _write_lock() -> asyncio.Lock: """写事务串行锁(每事件循环一个)。""" return _loop_lock('write') def _wal_lock() -> asyncio.Lock: """WAL 切换专用锁(不能用 _write_lock:execute() 已持它再调 connect→_configure 会自锁死)。""" return _loop_lock('wal') # 已确认切到 WAL 的库路径(首连设置一次;后续连接不再重复执行 PRAGMA 争排他锁) _WAL_READY_FOR: str | None = None async def _configure(conn: aiosqlite.Connection) -> None: global _WAL_READY_FOR conn.row_factory = aiosqlite.Row # busy_timeout 先设:让后续语句遇锁自旋等待,而不是立刻报 database is locked await conn.execute('PRAGMA busy_timeout=15000') key = str(DB_PATH) if _WAL_READY_FOR != key: # 切换日志模式需要排他锁,且 `_WAL_READY_FOR` 的检查不是原子的—— # 并发连接须在专用锁内串行完成「切换 + 标记」,并二次复查,否则批量并发 # 连接会同时执行 PRAGMA journal_mode=WAL 争锁(实测 database is locked) async with _wal_lock(): if _WAL_READY_FOR != key: await conn.execute('PRAGMA journal_mode=WAL') _WAL_READY_FOR = key await conn.execute('PRAGMA synchronous=NORMAL') # WAL 下兼顾安全与吞吐 await conn.execute('PRAGMA foreign_keys=ON') async def connect() -> aiosqlite.Connection: """新建一个已配置好的连接(调用方负责 close)。""" DATA_DIR.mkdir(parents=True, exist_ok=True) # noqa: ASYNC240 - 幂等的微秒级一次目录创建,无需线程池 conn = await aiosqlite.connect(str(DB_PATH), timeout=15) await _configure(conn) return conn # ============ 读路径(并发,无锁) ============ async def fetch_all(sql: str, params: Sequence = ()) -> list[aiosqlite.Row]: conn = await connect() try: cur = await conn.execute(sql, params) return list(await cur.fetchall()) finally: await conn.close() async def fetch_one(sql: str, params: Sequence = ()) -> aiosqlite.Row | None: conn = await connect() try: cur = await conn.execute(sql, params) return await cur.fetchone() finally: await conn.close() async def fetch_val(sql: str, params: Sequence = (), default: Any = None) -> Any: row = await fetch_one(sql, params) return row[0] if row is not None else default # ============ 写路径(串行,不阻塞读) ============ async def execute(sql: str, params: Sequence = ()) -> int: """执行单条写语句,返回 lastrowid(INSERT)或 rowcount(UPDATE/DELETE)。""" async with _write_lock(): conn = await connect() try: cur = await conn.execute(sql, params) await conn.commit() return cur.lastrowid if cur.lastrowid else cur.rowcount finally: await conn.close() async def execute_many(sql: str, seq: Iterable[Sequence]) -> int: """批量写,返回受影响行数。""" async with _write_lock(): conn = await connect() try: cur = await conn.executemany(sql, seq) await conn.commit() return cur.rowcount finally: await conn.close() @asynccontextmanager async def write_transaction() -> AsyncIterator[aiosqlite.Connection]: """写事务上下文:多条写语句原子提交(串行锁内)。 用法: async with write_transaction() as conn: await conn.execute(...) await conn.execute(...) # 退出即 commit;异常则 rollback """ async with _write_lock(): conn = await connect() try: yield conn await conn.commit() except Exception: await conn.rollback() raise finally: await conn.close() @asynccontextmanager async def read_connection() -> AsyncIterator[aiosqlite.Connection]: """只读连接上下文(无锁、可并发);需要在一个连接上做多次读时使用。""" conn = await connect() try: yield conn finally: await conn.close()