core/database.py

"""
AWS RDS (Postgres + pgvector) connection management.

Uses a threaded psycopg2 connection pool. A single shared connection breaks
under concurrent Streamlit/Gradio requests and doesn't survive RDS idle
timeouts well — a pool hands out healthy connections and recycles them.

Usage:
    from legal_ai.core.database import get_cursor

    with get_cursor() as cur:
        cur.execute("SELECT 1;")
        row = cur.fetchone()
    # connection is committed and returned to the pool automatically
"""

from __future__ import annotations

from contextlib import contextmanager
from functools import lru_cache

import psycopg2
from psycopg2.pool import ThreadedConnectionPool

from .config import get_settings


@lru_cache(maxsize=1)
def get_pool() -> ThreadedConnectionPool:
    """Create (once) and return the shared connection pool."""
    settings = get_settings()
    return ThreadedConnectionPool(
        minconn=settings.db_pool_min,
        maxconn=settings.db_pool_max,
        dsn=settings.rds_dsn,
    )


@contextmanager
def get_connection():
    """
    Borrow a connection from the pool. Commits on success, rolls back on
    error, and always returns the connection to the pool.
    """
    pool = get_pool()
    conn = pool.getconn()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        pool.putconn(conn)


@contextmanager
def get_cursor():
    """
    Borrow a connection and yield a cursor. Handles commit/rollback and
    returns the connection to the pool when done.
    """
    with get_connection() as conn:
        cur = conn.cursor()
        try:
            yield cur
        finally:
            cur.close()


def close_pool() -> None:
    """Close all pooled connections. Call on application shutdown."""
    try:
        pool = get_pool()
        pool.closeall()
    except Exception:
        pass
    finally:
        get_pool.cache_clear()


def healthcheck() -> bool:
    """Return True if the database is reachable."""
    try:
        with get_cursor() as cur:
            cur.execute("SELECT 1;")
            return cur.fetchone() is not None
    except Exception:
        return False