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