from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager

from psycopg import AsyncConnection
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool


class BotDatabase:
    """Пул соединений БД для асинхронного процесса бота."""

    def __init__(self, database_url: str) -> None:
        # Пул может содержать в себе максимум 5 соединений.
        self.pool = AsyncConnectionPool(
            conninfo=database_url,
            min_size=1,
            max_size=5,
            open=False,
            # Фабрика для представления строк БД как словарей.
            kwargs={"row_factory": dict_row}
        )

    async def open(self) -> None:
        await self.pool.open(wait=True)

    async def close(self) -> None:
        await self.pool.close()

    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[AsyncConnection]:
        async with self.pool.connection() as connection:
            async with connection.transaction():
                yield connection

    async def ping(self) -> bool:
        async with self.pool.connection() as connection:
            cursor = await connection.execute("SELECT 1")
            return await cursor.fetchone() is not None