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 typing import Any, Iterator

import psycopg2
import psycopg2.extras

from core.config import CreationDbConfig

psycopg2.extras.register_uuid()


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


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


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