77 lines
2.1 KiB
Python
77 lines
2.1 KiB
Python
|
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 {})
|