SQLite database manager
| 11 | |
| 12 | |
| 13 | class Database: |
| 14 | """SQLite database manager""" |
| 15 | |
| 16 | def __init__(self, db_path: str = None): |
| 17 | if db_path is None: |
| 18 | # Store database in data directory |
| 19 | data_dir = Path(__file__).parent.parent.parent / "data" |
| 20 | data_dir.mkdir(exist_ok=True) |
| 21 | db_path = str(data_dir / "flow.db") |
| 22 | self.db_path = db_path |
| 23 | self._write_lock = asyncio.Lock() |
| 24 | self._connect_timeout = 30 |
| 25 | self._busy_timeout_ms = 30000 |
| 26 | |
| 27 | def db_exists(self) -> bool: |
| 28 | """Check if database file exists""" |
| 29 | return Path(self.db_path).exists() |
| 30 | |
| 31 | async def _configure_connection(self, db): |
| 32 | """Apply SQLite runtime settings for better concurrent behavior.""" |
| 33 | await db.execute(f"PRAGMA busy_timeout = {self._busy_timeout_ms}") |
| 34 | await db.execute("PRAGMA foreign_keys = ON") |
| 35 | |
| 36 | def _current_stats_date(self) -> str: |
| 37 | """Return the logical date used by daily token statistics.""" |
| 38 | return date.today().isoformat() |
| 39 | |
| 40 | @asynccontextmanager |
| 41 | async def _connect(self, *, write: bool = False): |
| 42 | """Open a configured SQLite connection and optionally serialize writes.""" |
| 43 | if write: |
| 44 | async with self._write_lock: |
| 45 | async with aiosqlite.connect(self.db_path, timeout=self._connect_timeout) as db: |
| 46 | await self._configure_connection(db) |
| 47 | yield db |
| 48 | return |
| 49 | |
| 50 | async with aiosqlite.connect(self.db_path, timeout=self._connect_timeout) as db: |
| 51 | await self._configure_connection(db) |
| 52 | yield db |
| 53 | |
| 54 | async def _table_exists(self, db, table_name: str) -> bool: |
| 55 | """Check if a table exists in the database""" |
| 56 | cursor = await db.execute( |
| 57 | "SELECT name FROM sqlite_master WHERE type='table' AND name=?", |
| 58 | (table_name,) |
| 59 | ) |
| 60 | result = await cursor.fetchone() |
| 61 | return result is not None |
| 62 | |
| 63 | async def _column_exists(self, db, table_name: str, column_name: str) -> bool: |
| 64 | """Check if a column exists in a table""" |
| 65 | try: |
| 66 | cursor = await db.execute(f"PRAGMA table_info({table_name})") |
| 67 | columns = await cursor.fetchall() |
| 68 | return any(col[1] == column_name for col in columns) |
| 69 | except: |
| 70 | return False |
no outgoing calls