58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
"""
|
|
Database engine, session management, and WAL mode setup.
|
|
|
|
SQLite with WAL mode for single-user concurrent read/write safety.
|
|
"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from sqlalchemy import event, create_engine, Engine
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
|
|
from app.config import settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""SQLAlchemy declarative base for all models."""
|
|
pass
|
|
|
|
|
|
def _set_sqlite_wal_mode(dbapi_conn: Any, _connection_record: Any) -> None:
|
|
"""Enable WAL mode and other SQLite performance pragmas on every connection."""
|
|
cursor = dbapi_conn.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA synchronous=NORMAL")
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.close()
|
|
|
|
|
|
def create_db_engine() -> Engine:
|
|
"""Create the SQLAlchemy engine with SQLite and WAL mode."""
|
|
# Ensure the database directory exists
|
|
settings.DATABASE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
engine = create_engine(
|
|
f"sqlite:///{settings.DATABASE_PATH}",
|
|
echo=settings.DEBUG,
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
|
|
# Register WAL mode on every new connection
|
|
event.listen(engine, "connect", _set_sqlite_wal_mode)
|
|
|
|
return engine
|
|
|
|
|
|
engine = create_db_engine()
|
|
SessionLocal = sessionmaker(bind=engine, class_=Session, expire_on_commit=False)
|
|
|
|
|
|
def get_db() -> Session:
|
|
"""FastAPI dependency that yields a database session."""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db # type: ignore[misc]
|
|
finally:
|
|
db.close()
|