src/stage4/database.py

"""
UsTwo SQLite database setup.

Uses SQLAlchemy with SQLite — no external DB server needed.
DB file: data/ustwo.db (created automatically on first run).
"""
import os
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, DeclarativeBase

DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "data")
DATABASE_URL = f"sqlite:///{os.path.join(DATA_DIR, 'ustwo.db')}"


class Base(DeclarativeBase):
    pass


engine = create_engine(DATABASE_URL, echo=False)
SessionLocal = sessionmaker(bind=engine)


def init_db():
    """Create all tables. Safe to call multiple times."""
    os.makedirs(DATA_DIR, exist_ok=True)
    from . import models  # noqa: F401 — ensure models are registered
    Base.metadata.create_all(bind=engine)


def get_db():
    """FastAPI dependency — yields a DB session, closes after request."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()