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.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())
    return create_async_engine(_build_url(), **kwargs)


@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