79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from threading import RLock
|
|
|
|
import psycopg as pg
|
|
|
|
from app.core.config import get_pgconn_string
|
|
|
|
g_conn_dict: dict[str, pg.Connection] = {}
|
|
_registry_lock = RLock()
|
|
_project_locks: dict[str, RLock] = {}
|
|
|
|
|
|
def _is_closed(connection: pg.Connection) -> bool:
|
|
return bool(getattr(connection, "closed", False))
|
|
|
|
|
|
def _close_connection(connection: pg.Connection) -> None:
|
|
if not _is_closed(connection):
|
|
connection.close()
|
|
|
|
|
|
def _is_healthy(connection: pg.Connection) -> bool:
|
|
if _is_closed(connection):
|
|
return False
|
|
try:
|
|
with connection.cursor() as cur:
|
|
cur.execute("SELECT 1")
|
|
except pg.Error:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _get_project_lock(name: str) -> RLock:
|
|
with _registry_lock:
|
|
lock = _project_locks.get(name)
|
|
if lock is None:
|
|
lock = RLock()
|
|
_project_locks[name] = lock
|
|
return lock
|
|
|
|
|
|
def open_connection(name: str) -> pg.Connection:
|
|
with _get_project_lock(name):
|
|
connection = g_conn_dict.get(name)
|
|
if connection is None or not _is_healthy(connection):
|
|
if connection is not None:
|
|
_close_connection(connection)
|
|
connection = pg.connect(
|
|
conninfo=get_pgconn_string(db_name=name), autocommit=True
|
|
)
|
|
g_conn_dict[name] = connection
|
|
return connection
|
|
|
|
|
|
def is_connection_open(name: str) -> bool:
|
|
with _get_project_lock(name):
|
|
connection = g_conn_dict.get(name)
|
|
if connection is None:
|
|
return False
|
|
if not _is_healthy(connection):
|
|
del g_conn_dict[name]
|
|
_close_connection(connection)
|
|
return False
|
|
return True
|
|
|
|
|
|
def close_connection(name: str) -> None:
|
|
with _get_project_lock(name):
|
|
connection = g_conn_dict.pop(name, None)
|
|
if connection is not None:
|
|
_close_connection(connection)
|
|
|
|
|
|
@contextmanager
|
|
def project_connection(name: str) -> Iterator[pg.Connection]:
|
|
with _get_project_lock(name):
|
|
yield open_connection(name)
|