app/db.py

"""SQLAlchemy engine/session bound to the external Neon Postgres (pgvector).

DATABASE_URL comes from the environment (HF secret / .env). Paste Neon's raw
connection string; this module normalizes the driver prefix and enforces SSL.
"""
from __future__ import annotations

import os

from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker

load_dotenv()


def _normalize_url(raw: str) -> str:
    """Make a Neon connection string usable by SQLAlchemy + psycopg2.

    - postgres:// / postgresql:// -> postgresql+psycopg2://
    - ensure sslmode=require (Neon requires TLS)
    """
    if not raw:
        raise RuntimeError("DATABASE_URL is not set")

    if raw.startswith("postgresql+psycopg2://"):
        url = raw
    elif raw.startswith("postgresql://"):
        url = "postgresql+psycopg2://" + raw[len("postgresql://"):]
    elif raw.startswith("postgres://"):
        url = "postgresql+psycopg2://" + raw[len("postgres://"):]
    else:
        url = raw

    if "sslmode=" not in url:
        sep = "&" if "?" in url else "?"
        url = f"{url}{sep}sslmode=require"

    return url


DATABASE_URL = _normalize_url(os.environ.get("DATABASE_URL", ""))

# pool_pre_ping: Neon free tier idles/closes connections; this revalidates them.
engine = create_engine(
    DATABASE_URL,
    pool_pre_ping=True,
    pool_recycle=300,
    future=True,
)

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


class Base(DeclarativeBase):
    """Declarative base shared by models and Alembic metadata."""


def get_session():
    """FastAPI dependency: yield a session, always close it."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()