"""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()