File size: 1,527 Bytes
788eab4 | 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 | """
Database connection and session management.
Ensures SQLite WAL mode and foreign key constraint enforcement.
"""
import os
from collections.abc import Generator
from sqlalchemy import create_engine, event
from sqlalchemy.engine import Engine
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./placement_engine.db")
# For SQLite, ensure connection arguments for multithreading
connect_args = {}
if DATABASE_URL.startswith("sqlite"):
connect_args = {"check_same_thread": False}
engine = create_engine(
DATABASE_URL,
connect_args=connect_args,
echo=os.getenv("DEBUG", "false").lower() == "true",
)
# SQLite PRAGMA configuration: WAL mode, foreign keys, synchronous normal
@event.listens_for(Engine, "connect")
def set_sqlite_pragma(dbapi_connection: object, connection_record: object) -> None:
# Check if SQLite connection
if hasattr(dbapi_connection, "cursor"):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA synchronous=NORMAL")
cursor.close()
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
class Base(DeclarativeBase):
pass
def get_db() -> Generator[Session, None, None]:
"""Dependency for API route handlers."""
db = SessionLocal()
try:
yield db
finally:
db.close()
|