SkyForces
39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
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
|