68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""Async-пул соединений Postgres (psycopg v3).
|
||
|
||
Все репозитории берут соединение отсюда. Соединения возвращаются в пул и
|
||
коммитятся при успешном выходе из контекстного менеджера.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
|
||
from psycopg.rows import dict_row
|
||
from psycopg_pool import AsyncConnectionPool
|
||
|
||
from common.config import Settings, get_settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_pool: AsyncConnectionPool | None = None
|
||
|
||
|
||
async def init_pool(settings: Settings | None = None) -> AsyncConnectionPool:
|
||
"""Создать и открыть пул. Идемпотентно.
|
||
|
||
``settings`` передаёт админка (свои настройки); по умолчанию — настройки бота.
|
||
"""
|
||
global _pool
|
||
if _pool is not None:
|
||
return _pool
|
||
|
||
settings = settings or get_settings()
|
||
_pool = AsyncConnectionPool(
|
||
conninfo=settings.dsn,
|
||
min_size=settings.db_pool_min,
|
||
max_size=settings.db_pool_max,
|
||
open=False,
|
||
kwargs={
|
||
"row_factory": dict_row,
|
||
"options": f"-c statement_timeout={settings.db_statement_timeout_ms}",
|
||
},
|
||
)
|
||
await _pool.open(wait=True)
|
||
logger.info(
|
||
"Postgres pool opened: %s@%s:%s/%s (min=%s, max=%s)",
|
||
settings.db_user,
|
||
settings.db_host,
|
||
settings.db_port,
|
||
settings.db_name,
|
||
settings.db_pool_min,
|
||
settings.db_pool_max,
|
||
)
|
||
return _pool
|
||
|
||
|
||
async def close_pool() -> None:
|
||
"""Закрыть пул при остановке приложения."""
|
||
global _pool
|
||
if _pool is not None:
|
||
await _pool.close()
|
||
_pool = None
|
||
logger.info("Postgres pool closed")
|
||
|
||
|
||
def get_pool() -> AsyncConnectionPool:
|
||
"""Вернуть открытый пул или упасть с понятной ошибкой."""
|
||
if _pool is None:
|
||
raise RuntimeError("Postgres pool не инициализирован: вызовите init_pool() при старте")
|
||
return _pool
|