Create-knowledge-find-decode/core/db_session.py

"""Formal PostgreSQL session helpers for the creation-knowledge state store."""
from __future__ import annotations

from contextlib import contextmanager
from threading import Lock
from typing import Any, Iterator

import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool

from core.config import CreationDbConfig

psycopg2.extras.register_uuid()

_pool: ThreadedConnectionPool | None = None
_pool_key: tuple[Any, ...] | None = None
_pool_lock = Lock()


def connect(config: CreationDbConfig) -> Any:
    """Open a PostgreSQL connection scoped to the formal state schema."""
    return psycopg2.connect(
        host=config.host,
        port=config.port,
        user=config.user,
        password=config.password,
        dbname=config.database,
        connect_timeout=config.timeout,
        application_name=config.application_name,
        options=f"-c search_path={config.schema},public",
    )


def _connection_kwargs(config: CreationDbConfig) -> dict[str, Any]:
    return {
        "host": config.host,
        "port": config.port,
        "user": config.user,
        "password": config.password,
        "dbname": config.database,
        "connect_timeout": config.timeout,
        "application_name": config.application_name,
        "options": f"-c search_path={config.schema},public",
    }


def get_pool(config: CreationDbConfig) -> ThreadedConnectionPool:
    """Return a process-local pool for the current creation DB config."""
    global _pool, _pool_key
    key = (
        config.host,
        config.port,
        config.user,
        config.database,
        config.schema,
        config.application_name,
        config.pool_min,
        config.pool_max,
    )
    with _pool_lock:
        if _pool is not None and _pool_key == key:
            return _pool
        if _pool is not None:
            _pool.closeall()
        _pool = ThreadedConnectionPool(
            minconn=max(1, config.pool_min),
            maxconn=max(config.pool_min, config.pool_max),
            **_connection_kwargs(config),
        )
        _pool_key = key
        return _pool


@contextmanager
def transaction(config: CreationDbConfig) -> Iterator[Any]:
    """Yield a connection and commit or roll back around the caller's work."""
    conn = connect(config)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


@contextmanager
def pooled_transaction(config: CreationDbConfig) -> Iterator[Any]:
    """Yield a pooled connection and return it to the pool after the request."""
    pool = get_pool(config)
    conn = pool.getconn()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        pool.putconn(conn, close=bool(getattr(conn, "closed", False)))


def fetch_all(
    config: CreationDbConfig,
    sql: str,
    params: tuple[Any, ...] = (),
) -> list[dict[str, Any]]:
    """Small read helper for diagnostics and repository smoke checks."""
    conn = connect(config)
    try:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            return [dict(row) for row in cur.fetchall()]
    finally:
        conn.close()