db.py (3378B)
1 """Database connection and migration runner for SQLite.""" 2 3 import sqlite3 4 from concurrent.futures import ThreadPoolExecutor 5 from pathlib import Path 6 from typing import Any 7 8 9 class Database: 10 """Async-friendly SQLite wrapper with migration support.""" 11 12 def __init__(self, db_path: str) -> None: # pragma: no mutate: block 13 self.db_path = db_path 14 self.connection: sqlite3.Connection | None = None 15 self._executor = ThreadPoolExecutor(max_workers=1) 16 17 async def connect(self) -> None: # pragma: no mutate: block 18 """Open a connection and enable WAL mode.""" 19 import asyncio 20 21 loop = asyncio.get_running_loop() 22 self.connection = await loop.run_in_executor(self._executor, self._create_connection) 23 24 async def close(self) -> None: # pragma: no mutate: block 25 """Close the connection and release the executor.""" 26 import asyncio 27 28 if self.connection: 29 loop = asyncio.get_running_loop() 30 await loop.run_in_executor(self._executor, self.connection.close) 31 self.connection = None 32 self._executor.shutdown(wait=False) 33 34 async def execute(self, sql: str, params: tuple = ()) -> None: # pragma: no mutate: block 35 """Execute a statement (no return).""" 36 import asyncio 37 38 self._require_connection() 39 loop = asyncio.get_running_loop() 40 await loop.run_in_executor(self._executor, lambda: self._exec(sql, params)) 41 42 async def execute_fetchall(self, sql: str, params: tuple = ()) -> list: # pragma: no mutate: block 43 """Execute and return all rows.""" 44 import asyncio 45 46 self._require_connection() 47 loop = asyncio.get_running_loop() 48 return await loop.run_in_executor(self._executor, lambda: self._fetchall(sql, params)) 49 50 async def fetch_all(self, sql: str, params: tuple = ()) -> list: # pragma: no mutate: block 51 """Alias for execute_fetchall.""" 52 return await self.execute_fetchall(sql, params) 53 54 async def run_migrations(self, migrations_dir: Path) -> None: # pragma: no mutate: block 55 """Apply SQL files from migrations_dir in sorted order.""" 56 import asyncio 57 58 self._require_connection() 59 sql_files = sorted(migrations_dir.glob("*.sql")) 60 loop = asyncio.get_running_loop() 61 for sql_file in sql_files: 62 sql = sql_file.read_text() 63 await loop.run_in_executor(self._executor, lambda s=sql: self.connection.executescript(s)) 64 65 def _create_connection(self) -> sqlite3.Connection: # pragma: no mutate: block 66 conn = sqlite3.connect(self.db_path, check_same_thread=False) 67 conn.row_factory = sqlite3.Row 68 conn.execute("PRAGMA journal_mode=WAL") 69 conn.execute("PRAGMA foreign_keys=ON") 70 return conn 71 72 def _exec(self, sql: str, params: tuple) -> None: # pragma: no mutate: block 73 assert self.connection is not None 74 self.connection.execute(sql, params) 75 self.connection.commit() 76 77 def _fetchall(self, sql: str, params: tuple) -> list[Any]: # pragma: no mutate: block 78 assert self.connection is not None 79 return self.connection.execute(sql, params).fetchall() 80 81 def _require_connection(self) -> None: # pragma: no mutate: block 82 if self.connection is None: 83 raise RuntimeError("Database not connected")