from contextlib import contextmanager from typing import Any, Optional import psycopg2 from psycopg2 import pool from psycopg2.extras import RealDictCursor, Json from app.config import settings _connection_pool: Optional[pool.SimpleConnectionPool] = None def init_pool(minconn: int = 1, maxconn: int = 10) -> None: global _connection_pool if _connection_pool is None: _connection_pool = pool.SimpleConnectionPool( minconn, maxconn, dsn=settings.database_dsn, ) def close_pool() -> None: global _connection_pool if _connection_pool is not None: _connection_pool.closeall() _connection_pool = None @contextmanager def get_connection(): if _connection_pool is None: init_pool() conn = _connection_pool.getconn() try: yield conn conn.commit() except Exception: conn.rollback() raise finally: _connection_pool.putconn(conn) def fetch_all(query: str, params: Optional[tuple] = None) -> list[dict[str, Any]]: with get_connection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(query, params) return [dict(row) for row in cur.fetchall()] def fetch_one(query: str, params: Optional[tuple] = None) -> Optional[dict[str, Any]]: with get_connection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(query, params) row = cur.fetchone() return dict(row) if row else None def execute_returning(query: str, params: Optional[tuple] = None) -> Optional[dict[str, Any]]: with get_connection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(query, params) row = cur.fetchone() return dict(row) if row else None def execute(query: str, params: Optional[tuple] = None) -> int: with get_connection() as conn: with conn.cursor() as cur: cur.execute(query, params) return cur.rowcount def json_param(value: Any) -> Json: return Json(value or {})