chatkit/backend/app/db.py

"""Async SQLAlchemy engine + session management for ClassLens.

Supports two providers (DATABASE_PROVIDER):
  - "supabase": Postgres via asyncpg, tables isolated in a dedicated schema
  - "sqlite":   local aiosqlite file (dev/tests; schema flattened to None)
"""

from __future__ import annotations

import os
import ssl
from functools import lru_cache
from pathlib import Path
from typing import AsyncIterator, Optional

from sqlalchemy import event
from sqlalchemy.engine import URL, make_url
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from .config import get_settings


def _is_supabase() -> bool:
    s = get_settings()
    return s.database_provider == "supabase" and bool(s.supabase_db_url)


def active_schema() -> Optional[str]:
    """The Postgres schema for our tables, or None for SQLite (no schemas)."""
    if _is_supabase():
        return get_settings().db_schema or None
    return None


def assert_production_db_config() -> None:
    """Fail loud rather than silently running on ephemeral SQLite.

    The container FS on HF Spaces is ephemeral, so booting on SQLite means all
    data is lost on every sleep/rebuild. Refuse to start in those cases.
    """
    s = get_settings()
    if s.database_provider == "supabase" and not s.supabase_db_url:
        raise RuntimeError(
            "DATABASE_PROVIDER=supabase but SUPABASE_DB_URL is empty — refusing to "
            "start on ephemeral SQLite. Set the SUPABASE_DB_URL secret."
        )
    # On a deployed HF Space (SPACE_ID is injected by the platform), require a
    # real Postgres DB so a missing/typo'd provider can't silently degrade.
    if os.getenv("SPACE_ID") and not _is_supabase():
        raise RuntimeError(
            "Running on a Hugging Face Space without Supabase configured — refusing "
            "to start on ephemeral SQLite. Set DATABASE_PROVIDER=supabase and "
            "SUPABASE_DB_URL."
        )


def _build_url() -> URL:
    s = get_settings()
    if _is_supabase():
        # make_url keeps the raw password (handles '*' etc.); just swap the driver.
        return make_url(s.supabase_db_url).set(drivername="postgresql+asyncpg")
    # SQLite fallback
    path = s.database_path or str(Path(__file__).parent.parent / "classlens.db")
    return make_url(f"sqlite+aiosqlite:///{path}")


def _connect_args() -> dict:
    if _is_supabase():
        # Encrypt but don't verify the CA chain (libpq sslmode=require semantics).
        # Supabase's pooler presents a valid cert in prod, but TLS-intercepting
        # proxies (corp egress) break chain verification — this works in both.
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        return {
            "ssl": ctx,
            # Supabase runs pgbouncer; disable asyncpg's prepared-statement cache
            # so the same code works on both session (5432) and txn (6543) poolers.
            "statement_cache_size": 0,
        }
    return {}


@lru_cache
def get_engine() -> AsyncEngine:
    kwargs: dict = {"pool_pre_ping": True, "future": True}
    if _is_supabase():
        kwargs.update(pool_size=5, max_overflow=5, connect_args=_connect_args())
    engine = create_async_engine(_build_url(), **kwargs)

    if not _is_supabase():
        # SQLite doesn't enforce FK constraints (or ON DELETE CASCADE) unless
        # this pragma is set per-connection — without it, every cascade
        # defined in models.py silently does nothing on SQLite, while still
        # working correctly on Postgres (which always enforces FKs). That
        # gap makes local dev/tests lie about delete behavior, so set it
        # explicitly to keep SQLite semantics matching production.
        @event.listens_for(engine.sync_engine, "connect")
        def _enable_sqlite_fk(dbapi_connection, _record):
            cursor = dbapi_connection.cursor()
            cursor.execute("PRAGMA foreign_keys=ON")
            cursor.close()

    return engine


@lru_cache
def get_sessionmaker() -> async_sessionmaker[AsyncSession]:
    return async_sessionmaker(get_engine(), expire_on_commit=False, class_=AsyncSession)


async def get_session() -> AsyncIterator[AsyncSession]:
    """FastAPI dependency: yields a session, commits on success, rolls back on error."""
    maker = get_sessionmaker()
    async with maker() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


def run_pending_migrations() -> None:
    """Apply any pending Alembic migrations against Postgres/Supabase.

    There's no separate deploy step or CI job that runs `alembic upgrade
    head` against production, and whoever deploys this may not have direct
    database access to run it by hand — so instead this runs once at app
    startup (see main.py's lifespan). `alembic upgrade head` is idempotent
    (a no-op once already current), so this is safe to run on every boot.

    Synchronous (alembic's `command.upgrade` drives env.py's own internal
    `asyncio.run(...)` for the async engine) — call via `asyncio.to_thread`
    from async code, never await it directly.
    """
    from alembic import command
    from alembic.config import Config

    backend_dir = Path(__file__).parent.parent
    cfg = Config(str(backend_dir / "alembic.ini"))
    cfg.set_main_option("script_location", str(backend_dir / "alembic"))
    command.upgrade(cfg, "head")