SkyForces
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
|
|
from psycopg import Connection
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
|
|
class SiteDatabase:
|
|
"""Данный класс представляет собой обертку над пулом соединений
|
|
с базой данных PostgreSQL."""
|
|
|
|
def __init__(self, database_url: str) -> None:
|
|
# Пул может содержать в себе максимум 5 соединений.
|
|
self.pool = ConnectionPool(
|
|
conninfo=database_url,
|
|
min_size=1,
|
|
max_size=5,
|
|
open=False,
|
|
# Фабрика для представления строк БД как словарей.
|
|
kwargs={"row_factory": dict_row},
|
|
)
|
|
|
|
def open(self) -> None:
|
|
self.pool.open(wait=True)
|
|
|
|
def close(self) -> None:
|
|
self.pool.close()
|
|
|
|
@contextmanager
|
|
def transaction(self) -> Generator[Connection]:
|
|
with self.pool.connection() as connection:
|
|
with connection.transaction():
|
|
yield connection
|
|
|
|
def ping(self) -> bool:
|
|
"""Простая проверка подключения к БД."""
|
|
with self.pool.connection() as connection:
|
|
return connection.execute("SELECT 1").fetchone() is not None
|