File size: 1,871 Bytes
de1e3fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """Database engine, session factory, and FastAPI dependency.
All DB access goes through the ORM session so a future swap to Postgres+pgvector is mechanical
(spec §5, A2/A8).
"""
from __future__ import annotations
from collections.abc import Iterator
from pathlib import Path
from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session, sessionmaker
from .config import settings
# Ensure the sqlite directory exists before the engine touches the file.
if settings.database_url.startswith("sqlite"):
db_path = settings.database_url.split("///")[-1]
Path(db_path).resolve().parent.mkdir(parents=True, exist_ok=True)
connect_args = (
{"check_same_thread": False} if settings.database_url.startswith("sqlite") else {}
)
engine = create_engine(settings.database_url, connect_args=connect_args, future=True)
@event.listens_for(engine, "connect")
def _sqlite_pragmas(dbapi_connection, _):
"""Per-connection SQLite tuning.
- ``foreign_keys=ON``: enforce FKs (off by default on SQLite).
- ``journal_mode=WAL``: let readers proceed concurrently with a writer, so status polls during a
long batch load don't block on / error against the load's write transaction.
- ``busy_timeout``: wait for a briefly-held lock instead of failing immediately with
"database is locked" (which previously surfaced as a frozen load progress bar).
"""
if settings.database_url.startswith("sqlite"):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA busy_timeout=10000") # 10s
cursor.close()
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
def get_db() -> Iterator[Session]:
db = SessionLocal()
try:
yield db
finally:
db.close()
|