56 lines
1.2 KiB
Python
56 lines
1.2 KiB
Python
import os
|
|
from contextlib import contextmanager
|
|
|
|
import psycopg2
|
|
from psycopg2 import pool
|
|
|
|
_pool: pool.ThreadedConnectionPool | None = None
|
|
|
|
|
|
def _db_config() -> dict:
|
|
return {
|
|
"host": os.environ.get("POSTGRES_HOST", "postgres"),
|
|
"port": int(os.environ.get("POSTGRES_PORT", 5432)),
|
|
"dbname": os.environ.get("POSTGRES_DB", "novela"),
|
|
"user": os.environ.get("POSTGRES_USER", "novela"),
|
|
"password": os.environ.get("POSTGRES_PASSWORD", ""),
|
|
}
|
|
|
|
|
|
def init_pool(minconn: int = 2, maxconn: int = 10) -> None:
|
|
global _pool
|
|
if _pool is None:
|
|
_pool = pool.ThreadedConnectionPool(minconn=minconn, maxconn=maxconn, **_db_config())
|
|
|
|
|
|
def close_pool() -> None:
|
|
global _pool
|
|
if _pool is not None:
|
|
_pool.closeall()
|
|
_pool = None
|
|
|
|
|
|
def get_conn():
|
|
global _pool
|
|
if _pool is None:
|
|
init_pool()
|
|
return _pool.getconn() # type: ignore[union-attr]
|
|
|
|
|
|
def release_conn(conn) -> None:
|
|
if _pool is not None and conn is not None:
|
|
_pool.putconn(conn)
|
|
|
|
|
|
@contextmanager
|
|
def get_db_conn():
|
|
conn = get_conn()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
release_conn(conn)
|
|
|
|
|
|
def direct_connect():
|
|
return psycopg2.connect(**_db_config())
|