157 lines
4.6 KiB
Python
157 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Callable
|
|
from contextlib import AbstractContextManager
|
|
from typing import Any
|
|
|
|
from src.storage.exceptions import PostgresConnectionPoolError
|
|
|
|
|
|
PostgresPoolFactory = Callable[..., Any]
|
|
|
|
|
|
def _default_pool_factory(**kwargs: Any) -> Any:
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
return ConnectionPool(**kwargs)
|
|
|
|
|
|
class PostgresConnectionPool:
|
|
"""Явная обёртка управляемого жизненного цикла пула psycopg."""
|
|
|
|
__slots__ = (
|
|
"_conninfo",
|
|
"_min_size",
|
|
"_max_size",
|
|
"_name",
|
|
"_pool",
|
|
"_pool_factory",
|
|
"_timeout_seconds",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
conninfo: str,
|
|
min_size: int = 1,
|
|
max_size: int = 4,
|
|
timeout_seconds: float = 10.0,
|
|
name: str = "dzentra-storage",
|
|
pool_factory: PostgresPoolFactory = _default_pool_factory,
|
|
) -> None:
|
|
normalized_conninfo = str(conninfo or "").strip()
|
|
normalized_name = str(name or "").strip()
|
|
|
|
if not normalized_conninfo:
|
|
raise ValueError("conninfo must not be empty")
|
|
|
|
if (
|
|
isinstance(min_size, bool)
|
|
or not isinstance(min_size, int)
|
|
or min_size <= 0
|
|
):
|
|
raise ValueError("min_size must be a positive integer")
|
|
|
|
if (
|
|
isinstance(max_size, bool)
|
|
or not isinstance(max_size, int)
|
|
or max_size < min_size
|
|
):
|
|
raise ValueError(
|
|
"max_size must be an integer not smaller than min_size"
|
|
)
|
|
|
|
if (
|
|
isinstance(timeout_seconds, bool)
|
|
or not isinstance(timeout_seconds, (int, float))
|
|
or not math.isfinite(float(timeout_seconds))
|
|
or timeout_seconds <= 0
|
|
):
|
|
raise ValueError("timeout_seconds must be positive and finite")
|
|
|
|
if not normalized_name:
|
|
raise ValueError("name must not be empty")
|
|
|
|
if not callable(pool_factory):
|
|
raise TypeError("pool_factory must be callable")
|
|
|
|
self._conninfo = normalized_conninfo
|
|
self._min_size = min_size
|
|
self._max_size = max_size
|
|
self._timeout_seconds = float(timeout_seconds)
|
|
self._name = normalized_name
|
|
self._pool_factory = pool_factory
|
|
self._pool: Any | None = None
|
|
|
|
@property
|
|
def is_open(self) -> bool:
|
|
return self._pool is not None
|
|
|
|
def open(self) -> None:
|
|
"""Открыть пул и проверить создание минимального числа соединений."""
|
|
if self._pool is not None:
|
|
return
|
|
|
|
pool: Any = None
|
|
|
|
try:
|
|
pool = self._pool_factory(
|
|
conninfo=self._conninfo,
|
|
min_size=self._min_size,
|
|
max_size=self._max_size,
|
|
timeout=self._timeout_seconds,
|
|
kwargs={"autocommit": False},
|
|
name=self._name,
|
|
open=False,
|
|
)
|
|
pool.open(
|
|
wait=True,
|
|
timeout=self._timeout_seconds,
|
|
)
|
|
except BaseException as error:
|
|
if pool is not None:
|
|
try:
|
|
pool.close(timeout=self._timeout_seconds)
|
|
except Exception as cleanup_error:
|
|
error.add_note(
|
|
"PostgreSQL pool cleanup also failed: "
|
|
f"{type(cleanup_error).__name__}."
|
|
)
|
|
|
|
if not isinstance(error, Exception):
|
|
raise
|
|
|
|
raise PostgresConnectionPoolError(
|
|
"Failed to open PostgreSQL connection pool."
|
|
) from error
|
|
|
|
self._pool = pool
|
|
|
|
def connection(self) -> AbstractContextManager[Any]:
|
|
"""Выдать одно транзакционное соединение из открытого пула."""
|
|
pool = self._pool
|
|
|
|
if pool is None:
|
|
raise PostgresConnectionPoolError(
|
|
"PostgreSQL connection pool is not open."
|
|
)
|
|
|
|
return pool.connection(timeout=self._timeout_seconds)
|
|
|
|
def close(self) -> None:
|
|
"""Закрыть пул; повторное закрытие ничего не делает."""
|
|
pool = self._pool
|
|
|
|
if pool is None:
|
|
return
|
|
|
|
self._pool = None
|
|
|
|
try:
|
|
pool.close(timeout=self._timeout_seconds)
|
|
except Exception as error:
|
|
raise PostgresConnectionPoolError(
|
|
"Failed to close PostgreSQL connection pool."
|
|
) from error
|