database.py

"""SQLAlchemy database setup — supports SQLite and PostgreSQL."""

import logging

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import DeclarativeBase, sessionmaker

from config import settings

logger = logging.getLogger(__name__)

_is_sqlite = settings.database_url.startswith("sqlite")
_engine_kwargs: dict = {
    "pool_pre_ping": True,  # test connections before reuse (fixes Neon idle drops)
}
if _is_sqlite:
    _engine_kwargs["connect_args"] = {"check_same_thread": False}
else:
    _engine_kwargs["pool_recycle"] = 300  # recycle connections every 5 min

engine = create_engine(settings.database_url, **_engine_kwargs)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


class Base(DeclarativeBase):
    pass


def get_db():
    """FastAPI dependency that provides a database session."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def _auto_migrate():
    """Add any missing columns to existing tables (poor-man's migration)."""
    insp = inspect(engine)
    if not insp.has_table("users"):
        return
    existing = {col["name"] for col in insp.get_columns("users")}
    # Use BOOLEAN DEFAULT FALSE for PostgreSQL, BOOLEAN DEFAULT 0 for SQLite
    default_val = "FALSE" if not _is_sqlite else "0"
    migrations = {
        "is_premium": f"ALTER TABLE users ADD COLUMN is_premium BOOLEAN NOT NULL DEFAULT {default_val}",
        "is_admin": f"ALTER TABLE users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT {default_val}",
        # Nullable with no default: existing rows simply have no reset in flight.
        "reset_token_hash": "ALTER TABLE users ADD COLUMN reset_token_hash VARCHAR(64)",
        "reset_token_expires": "ALTER TABLE users ADD COLUMN reset_token_expires TIMESTAMP",
    }
    with engine.begin() as conn:
        for col_name, ddl in migrations.items():
            if col_name not in existing:
                logger.info("Auto-migrating: adding column %s to users", col_name)
                conn.execute(text(ddl))


def create_tables():
    """Create all tables. Called on app startup."""
    Base.metadata.create_all(bind=engine)
    _auto_migrate()