diff --git a/hasql/abc.py b/hasql/abc.py new file mode 100644 index 0000000..9a4369c --- /dev/null +++ b/hasql/abc.py @@ -0,0 +1,86 @@ +import warnings +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import Any, Generic, TypeVar + +from .acquire import AcquireContext +from .metrics import DriverMetrics, PoolStats +from .utils import Dsn + +PoolT = TypeVar("PoolT") +ConnT = TypeVar("ConnT") + + +class PoolDriver(ABC, Generic[PoolT, ConnT]): + """Database driver interface for pool operations.""" + + @abstractmethod + def get_pool_freesize(self, pool: PoolT) -> int: ... + + @abstractmethod + def acquire_from_pool( + self, + pool: PoolT, + *, + timeout: float | None = None, + **kwargs, + ) -> AcquireContext[ConnT]: ... + + @abstractmethod + async def release_to_pool( + self, + connection: ConnT, + pool: PoolT, + **kwargs, + ) -> None: ... + + @abstractmethod + async def is_master(self, connection: ConnT) -> bool: ... + + @abstractmethod + async def fetch_scalar(self, connection: ConnT, query: str) -> Any: + """Execute a query and return a single scalar value.""" + ... + + @abstractmethod + async def pool_factory(self, dsn: Dsn, **kwargs) -> PoolT: ... + + @abstractmethod + async def close_pool(self, pool: PoolT) -> None: ... + + @abstractmethod + async def terminate_pool(self, pool: PoolT) -> None: ... + + @abstractmethod + def is_connection_closed(self, connection: ConnT) -> bool: ... + + @abstractmethod + def host(self, pool: PoolT) -> str: ... + + @abstractmethod + def pool_stats(self, pool: PoolT) -> PoolStats: ... + + def driver_metrics( + self, + pools: Sequence[PoolT | None], + ) -> Sequence[DriverMetrics]: + warnings.warn( + "driver_metrics() is deprecated, implement pool_stats() instead", + DeprecationWarning, + stacklevel=2, + ) + return [ + DriverMetrics( + min=s.min, max=s.max, idle=s.idle, used=s.used, + host=self.host(p), + ) + for p in pools if p + for s in [self.pool_stats(p)] + ] + + def prepare_pool_factory_kwargs(self, kwargs: dict) -> dict: + """Hook for drivers to adjust pool factory kwargs.""" + return kwargs + + +__all__ = ("PoolDriver",) diff --git a/hasql/acquire.py b/hasql/acquire.py new file mode 100644 index 0000000..3350cc8 --- /dev/null +++ b/hasql/acquire.py @@ -0,0 +1,184 @@ +import asyncio +from collections.abc import Callable, Generator +from contextlib import AbstractAsyncContextManager +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Protocol, + TypeVar, +) + +from .exceptions import NoAvailablePoolError +from .metrics import CalculateMetrics + +if TYPE_CHECKING: + from .balancer_policy.base import AbstractBalancerPolicy + from .pool_state import PoolState + +PoolT = TypeVar("PoolT") +ConnT = TypeVar("ConnT") +ConnT_co = TypeVar("ConnT_co", covariant=True) + + +class AcquireContext(Protocol[ConnT_co]): + async def __aenter__(self) -> ConnT_co: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: ... + def __await__(self) -> Generator[Any, None, ConnT_co]: ... + + +class TimeoutAcquireContext(Generic[ConnT]): + __slots__ = ("_context", "_timeout") + + def __init__(self, context: AcquireContext[ConnT], timeout: float): + self._context = context + self._timeout = timeout + + async def __aenter__(self) -> ConnT: + return await asyncio.wait_for( + self._context.__aenter__(), + timeout=self._timeout, + ) + + async def __aexit__(self, *exc): + # TODO: consider adding a bounded timeout here. Currently if the + # underlying driver hangs during connection release this will block + # indefinitely. A timeout risks leaking the connection (not returned + # to pool), so this needs careful design. + return await self._context.__aexit__(*exc) + + def __await__(self) -> Generator[Any, None, ConnT]: + return asyncio.wait_for( + self._context.__aenter__(), + timeout=self._timeout, + ).__await__() + + +class PoolAcquireContext( + AbstractAsyncContextManager[ConnT], + Generic[PoolT, ConnT], +): + def __init__( + self, + pool_state: "PoolState[PoolT, ConnT]", + balancer: "AbstractBalancerPolicy[PoolT]", + register_connection: Callable[[ConnT, PoolT], None], + unregister_connection: Callable[[ConnT], None], + read_only: bool, + master_as_replica_weight: float | None, + timeout: float, + metrics: CalculateMetrics, + fallback_master: bool = False, + **kwargs, + ): + self._pool_state = pool_state + self._balancer = balancer + self._register_connection = register_connection + self._unregister_connection = unregister_connection + self._read_only = read_only + self._fallback_master = fallback_master + self._master_as_replica_weight = master_as_replica_weight + self._timeout = timeout + self._kwargs = kwargs + self._metrics = metrics + self._pool: PoolT | None = None + self._conn: ConnT | None = None + self._context: AcquireContext[ConnT] | None = None + + def _deadline(self) -> float: + return asyncio.get_running_loop().time() + self._timeout + + def _remaining_timeout(self, deadline: float) -> float: + remaining_timeout = deadline - asyncio.get_running_loop().time() + if remaining_timeout <= 0: + raise asyncio.TimeoutError + return remaining_timeout + + async def _get_pool(self, deadline: float) -> PoolT: + async def get_pool() -> PoolT: + with self._metrics.with_get_pool(): + pool = await self._balancer.get_pool( + read_only=self._read_only, + fallback_master=self._fallback_master, + master_as_replica_weight=self._master_as_replica_weight, + ) + if pool is None: + raise NoAvailablePoolError("No available pool") + return pool + + return await asyncio.wait_for( + get_pool(), + timeout=self._remaining_timeout(deadline), + ) + + async def _resolve_pool_and_acquire_context( + self, + ) -> tuple[PoolT, AcquireContext[ConnT]]: + deadline = self._deadline() + pool = await self._get_pool(deadline) + remaining = self._remaining_timeout(deadline) + driver_ctx = self._pool_state.acquire_from_pool( + pool, + timeout=remaining, + **self._kwargs, + ) + return pool, driver_ctx + + async def _acquire_connection(self) -> ConnT: + pool, driver_ctx = await self._resolve_pool_and_acquire_context() + + host = self._pool_state.host(pool) + with self._metrics.with_acquire(host): + conn: ConnT = await driver_ctx + + try: + self._metrics.add_connection(host) + self._register_connection(conn, pool) + except BaseException: + await self._pool_state.release_to_pool(conn, pool) + raise + return conn + + async def __aenter__(self) -> ConnT: + pool, driver_ctx = await self._resolve_pool_and_acquire_context() + + host = self._pool_state.host(pool) + with self._metrics.with_acquire(host): + conn: ConnT = await driver_ctx.__aenter__() + + try: + self._metrics.add_connection(host) + self._register_connection(conn, pool) + except BaseException: + await driver_ctx.__aexit__(None, None, None) + raise + + self._pool = pool + self._conn = conn + self._context = driver_ctx + return conn + + async def __aexit__(self, *exc): + if self._conn is None or self._pool is None or self._context is None: + return + self._unregister_connection(self._conn) + self._metrics.remove_connection( + self._pool_state.host(self._pool), + ) + await self._context.__aexit__(*exc) + + def __await__(self): + return self._acquire_connection().__await__() + + +__all__ = ( + "AcquireContext", + "TimeoutAcquireContext", + "PoolAcquireContext", +) diff --git a/hasql/balancer_policy/__init__.py b/hasql/balancer_policy/__init__.py index 6de1e02..a32b956 100644 --- a/hasql/balancer_policy/__init__.py +++ b/hasql/balancer_policy/__init__.py @@ -1,9 +1,10 @@ +from .base import AbstractBalancerPolicy from .greedy import GreedyBalancerPolicy from .random_weighted import RandomWeightedBalancerPolicy from .round_robin import RoundRobinBalancerPolicy - __all__ = ( + "AbstractBalancerPolicy", "GreedyBalancerPolicy", "RandomWeightedBalancerPolicy", "RoundRobinBalancerPolicy", diff --git a/hasql/balancer_policy/base.py b/hasql/balancer_policy/base.py index 69b33de..581f3bb 100644 --- a/hasql/balancer_policy/base.py +++ b/hasql/balancer_policy/base.py @@ -1,20 +1,22 @@ import random -from abc import abstractmethod -from typing import Any, Optional +from abc import ABC, abstractmethod +from typing import Generic, TypeVar -from ..base import AbstractBalancerPolicy, BasePoolManager +from ..pool_state import PoolStateProvider +PoolT = TypeVar("PoolT") -class BaseBalancerPolicy(AbstractBalancerPolicy): - def __init__(self, pool_manager: BasePoolManager): - self._pool_manager = pool_manager + +class AbstractBalancerPolicy(ABC, Generic[PoolT]): + def __init__(self, pool_state: PoolStateProvider[PoolT]): + self._pool_state = pool_state async def get_pool( self, read_only: bool, fallback_master: bool = False, - master_as_replica_weight: Optional[float] = None, - ) -> Any: + master_as_replica_weight: float | None = None, + ) -> PoolT | None: if not read_only and master_as_replica_weight is not None: raise ValueError( "Field master_as_replica_weight is used only when " @@ -23,8 +25,9 @@ async def get_pool( choose_master_as_replica = False if master_as_replica_weight is not None: - rand = random.random() - choose_master_as_replica = 0 < rand <= master_as_replica_weight + choose_master_as_replica = ( + random.random() < master_as_replica_weight + ) return await self._get_pool( read_only=read_only, @@ -32,14 +35,41 @@ async def get_pool( choose_master_as_replica=choose_master_as_replica, ) + async def _get_candidates( + self, + read_only: bool, + fallback_master: bool = False, + choose_master_as_replica: bool = False, + ) -> list[PoolT]: + candidates: list[PoolT] = [] + + if read_only: + candidates.extend( + await self._pool_state.get_replica_pools( + fallback_master=fallback_master, + ), + ) + + if not read_only or ( + choose_master_as_replica + and self._pool_state.master_pool_count > 0 + and self._pool_state.replica_pool_count > 0 + ): + candidates.extend(await self._pool_state.get_master_pools()) + + return candidates + @abstractmethod async def _get_pool( self, read_only: bool, fallback_master: bool = False, choose_master_as_replica: bool = False, - ): + ) -> PoolT | None: pass -__all__ = ["BaseBalancerPolicy"] +# Backward-compatible alias +BaseBalancerPolicy = AbstractBalancerPolicy + +__all__ = ["AbstractBalancerPolicy", "BaseBalancerPolicy"] diff --git a/hasql/balancer_policy/greedy.py b/hasql/balancer_policy/greedy.py index fbbf9d3..5083cf2 100644 --- a/hasql/balancer_policy/greedy.py +++ b/hasql/balancer_policy/greedy.py @@ -1,40 +1,37 @@ import random -from hasql.balancer_policy.base import BaseBalancerPolicy +from .base import AbstractBalancerPolicy, PoolT -class GreedyBalancerPolicy(BaseBalancerPolicy): +class GreedyBalancerPolicy(AbstractBalancerPolicy[PoolT]): async def _get_pool( self, read_only: bool, fallback_master: bool = False, choose_master_as_replica: bool = False, - ): - candidates = [] - - if read_only: - candidates.extend( - await self._pool_manager.get_replica_pools( - fallback_master=fallback_master, - ), - ) - - if ( - not read_only or - ( - choose_master_as_replica and - self._pool_manager.master_pool_count > 0 - ) - ): - candidates.extend(await self._pool_manager.get_master_pools()) - - fat_pool = max(candidates, key=self._pool_manager.get_pool_freesize) - max_freesize = self._pool_manager.get_pool_freesize(fat_pool) - return random.choice([ - candidate + ) -> PoolT | None: + candidates = await self._get_candidates( + read_only=read_only, + fallback_master=fallback_master, + choose_master_as_replica=choose_master_as_replica, + ) + + if not candidates: + return None + + freesizes = [ + (candidate, self._pool_state.get_pool_freesize(candidate)) for candidate in candidates - if self._pool_manager.get_pool_freesize(candidate) == max_freesize - ]) + ] + + max_freesize = max(freesize for _, freesize in freesizes) + best = [ + candidate + for candidate, freesize in freesizes + if freesize == max_freesize + ] + + return random.choice(best) __all__ = ("GreedyBalancerPolicy",) diff --git a/hasql/balancer_policy/random_weighted.py b/hasql/balancer_policy/random_weighted.py index db20852..f314cf6 100644 --- a/hasql/balancer_policy/random_weighted.py +++ b/hasql/balancer_policy/random_weighted.py @@ -1,76 +1,41 @@ import random -from typing import Iterable, Optional +from collections.abc import Iterable -from hasql.balancer_policy.base import BaseBalancerPolicy +from .base import AbstractBalancerPolicy, PoolT -MACHINE_EPSILON: float = 1e-16 - - -class RandomWeightedBalancerPolicy(BaseBalancerPolicy): +class RandomWeightedBalancerPolicy(AbstractBalancerPolicy[PoolT]): async def _get_pool( self, read_only: bool, fallback_master: bool = False, choose_master_as_replica: bool = False, - ): - candidates = [] - - if read_only: - candidates.extend( - await self._pool_manager.get_replica_pools( - fallback_master=fallback_master, - ), - ) - if ( - not read_only or - ( - choose_master_as_replica and - self._pool_manager.master_pool_count > 0 - ) - ): - candidates.extend(await self._pool_manager.get_master_pools()) - - choiced_index = self._weighted_choice( - self._normalize_times( - self._reflect_times( - self._get_response_times(candidates), - ), - ), + ) -> PoolT | None: + candidates = await self._get_candidates( + read_only=read_only, + fallback_master=fallback_master, + choose_master_as_replica=choose_master_as_replica, ) - return candidates[choiced_index] - - def _get_response_times(self, pools: list) -> Iterable[Optional[float]]: - for pool in pools: - yield self._pool_manager.get_last_response_time(pool) + if not candidates: + return None - @staticmethod - def _reflect_times( - times: Iterable[Optional[float]], - ) -> Iterable[float]: - list_times = [value or 0 for value in times] - sum_time = sum(list_times) - yield from map(lambda x: sum_time - x + MACHINE_EPSILON, list_times) - - @staticmethod - def _normalize_times(times: Iterable[float]) -> Iterable[float]: - list_times = list(times) - sum_time = sum(list_times) - yield from map(lambda x: sum_time / x, list_times) + weights = self._compute_weights( + self._pool_state.get_last_response_time(pool) + for pool in candidates + ) + return random.choices(candidates, weights=weights)[0] @staticmethod - def _weighted_choice(probability_distribution: Iterable[float]) -> int: - rand = random.random() - prefix_sum: float = 0.0 - - length = 0 - for i, p in enumerate(probability_distribution): - length += 1 - prefix_sum += p - if rand <= prefix_sum: - return i - return length - 1 + def _compute_weights( + times: Iterable[float | None], + ) -> list[float]: + values = [0 if t is None else t for t in times] + max_time = max(values) if values else 0 + # Reflect: faster (lower time) gets higher weight. + # +1 ensures all-zero case produces uniform weights + # rather than all-zero weights (which random.choices rejects). + return [max_time - v + 1 for v in values] __all__ = ["RandomWeightedBalancerPolicy"] diff --git a/hasql/balancer_policy/round_robin.py b/hasql/balancer_policy/round_robin.py index f41e927..88353d5 100644 --- a/hasql/balancer_policy/round_robin.py +++ b/hasql/balancer_policy/round_robin.py @@ -1,8 +1,8 @@ from collections import defaultdict -from types import MappingProxyType -from typing import NamedTuple, Optional +from typing import NamedTuple -from hasql.balancer_policy.base import BaseBalancerPolicy +from .base import AbstractBalancerPolicy, PoolT +from ..pool_state import PoolStateProvider class PoolOptions(NamedTuple): @@ -10,57 +10,31 @@ class PoolOptions(NamedTuple): choose_master_as_replica: bool -class RoundRobinBalancerPolicy(BaseBalancerPolicy): - def __init__(self, pool_manager): - super().__init__(pool_manager) - self._indexes = defaultdict(lambda: 0) - self._choose_predicates = MappingProxyType({ - PoolOptions(True, False): self._replica_predicate, - PoolOptions(True, True): self._master_as_replica_predicate, - PoolOptions(False, False): self._master_predicate, - }) +class RoundRobinBalancerPolicy(AbstractBalancerPolicy[PoolT]): + def __init__(self, pool_state: PoolStateProvider[PoolT]): + super().__init__(pool_state) + self._indexes: defaultdict[PoolOptions, int] = defaultdict(lambda: 0) async def _get_pool( - self, - read_only: bool, - fallback_master: Optional[bool] = None, - choose_master_as_replica: bool = False, - ): - if read_only: - if self._pool_manager.replica_pool_count == 0: - if fallback_master: - read_only = False - choose_master_as_replica = False - if self._pool_manager.master_pool_count == 0: - await self._pool_manager.wait_masters_ready(1) - else: - await self._pool_manager.wait_replicas_ready(1) - else: - if self._pool_manager.master_pool_count == 0: - await self._pool_manager.wait_masters_ready(1) + self, + read_only: bool, + fallback_master: bool = False, + choose_master_as_replica: bool = False, + ) -> PoolT | None: + candidates = await self._get_candidates( + read_only=read_only, + fallback_master=fallback_master, + choose_master_as_replica=choose_master_as_replica, + ) + + if not candidates: + return None pool_options = PoolOptions(read_only, choose_master_as_replica) - assert pool_options in self._choose_predicates - - predicate = self._choose_predicates[pool_options] start_index = self._indexes[pool_options] - - pools = self._pool_manager.pools - for offset in range(len(pools)): - index = (start_index + offset) % len(pools) - current_pool = pools[index] - if current_pool is not None and predicate(current_pool): - self._indexes[pool_options] = (index + 1) % len(pools) - return current_pool - - def _master_predicate(self, pool) -> bool: - return self._pool_manager.pool_is_master(pool) - - def _replica_predicate(self, pool) -> bool: - return self._pool_manager.pool_is_replica(pool) - - def _master_as_replica_predicate(self, pool) -> bool: - return self._master_predicate(pool) or self._replica_predicate(pool) + index = start_index % len(candidates) + self._indexes[pool_options] = (start_index + 1) % len(candidates) + return candidates[index] __all__ = ("RoundRobinBalancerPolicy",) diff --git a/hasql/pool_state.py b/hasql/pool_state.py new file mode 100644 index 0000000..1c4d5a6 --- /dev/null +++ b/hasql/pool_state.py @@ -0,0 +1,28 @@ +# Minimal stub for PR2 — full PoolState implementation is in PR3. +from typing import Protocol, TypeVar, runtime_checkable + +PoolT = TypeVar("PoolT") + + +@runtime_checkable +class PoolStateProvider(Protocol[PoolT]): + """Protocol for read access to pool sets — used by balancer policies.""" + + @property + def master_pool_count(self) -> int: ... + + @property + def replica_pool_count(self) -> int: ... + + async def get_master_pools(self) -> list[PoolT]: ... + + async def get_replica_pools( + self, fallback_master: bool = False, + ) -> list[PoolT]: ... + + def get_pool_freesize(self, pool: PoolT) -> int: ... + + def get_last_response_time(self, pool: PoolT) -> float | None: ... + + +__all__ = ("PoolStateProvider",) diff --git a/tests/test_abc.py b/tests/test_abc.py new file mode 100644 index 0000000..54a361f --- /dev/null +++ b/tests/test_abc.py @@ -0,0 +1,48 @@ +"""Tests for PoolDriver ABC.""" +import pytest + +from hasql.abc import PoolDriver + + +def test_pool_driver_cannot_be_instantiated(): + with pytest.raises(TypeError): + PoolDriver() + + +def test_pool_driver_requires_all_abstract_methods(): + class IncompleteDriver(PoolDriver): + pass + + with pytest.raises(TypeError): + IncompleteDriver() + + +def test_pool_driver_prepare_pool_factory_kwargs_default(): + """Default prepare_pool_factory_kwargs returns kwargs unchanged.""" + from tests.mocks.pool_manager import TestDriver + + driver = TestDriver() + kwargs = {"minsize": 5, "maxsize": 20} + result = driver.prepare_pool_factory_kwargs(kwargs) + assert result is kwargs + assert result == {"minsize": 5, "maxsize": 20} + + +def test_test_driver_is_pool_driver(): + from tests.mocks.pool_manager import TestDriver + + driver = TestDriver() + assert isinstance(driver, PoolDriver) + + +async def test_pool_manager_has_driver_via_pool_state(): + """BasePoolManager exposes driver via _pool_state.""" + from tests.mocks import TestPoolManager + + manager = TestPoolManager( + "postgresql://test:test@master:5432/test", + ) + try: + assert isinstance(manager._pool_state.driver, PoolDriver) + finally: + await manager.close() diff --git a/tests/test_acquire.py b/tests/test_acquire.py new file mode 100644 index 0000000..55123d0 --- /dev/null +++ b/tests/test_acquire.py @@ -0,0 +1,256 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from hasql.acquire import TimeoutAcquireContext, PoolAcquireContext +from hasql.exceptions import NoAvailablePoolError +from hasql.metrics import CalculateMetrics + + +class FakeAcquireContext: + def __init__(self, conn=None): + self.conn = conn or object() + self.entered = False + self.exited = False + self.exit_args = None + + async def __aenter__(self): + self.entered = True + return self.conn + + async def __aexit__(self, *exc): + self.exited = True + self.exit_args = exc + + def __await__(self): + return self.__aenter__().__await__() + + +def _make_mocks(pool, inner_ctx): + """Create mocks for PoolAcquireContext dependencies.""" + pool_state = MagicMock() + pool_state.host.return_value = "test-host:5432" + pool_state.acquire_from_pool.return_value = inner_ctx + + balancer = MagicMock() + balancer.get_pool = AsyncMock(return_value=pool) + + register = MagicMock() + unregister = MagicMock() + + return pool_state, balancer, register, unregister + + +def _make_ctx(pool_state, balancer, register, unregister, metrics, **kwargs): + """Create a PoolAcquireContext with the given mocks.""" + defaults = dict( + read_only=False, + fallback_master=False, + master_as_replica_weight=None, + timeout=1.0, + ) + defaults.update(kwargs) + return PoolAcquireContext( + pool_state=pool_state, + balancer=balancer, + register_connection=register, + unregister_connection=unregister, + metrics=metrics, + **defaults, + ) + + +async def test_timeout_acquire_context_aenter(): + conn = object() + ctx = TimeoutAcquireContext(FakeAcquireContext(conn), timeout=1.0) + result = await ctx.__aenter__() + assert result is conn + + +async def test_timeout_acquire_context_aexit_delegates(): + inner = FakeAcquireContext() + ctx = TimeoutAcquireContext(inner, timeout=1.0) + await ctx.__aenter__() + await ctx.__aexit__(None, None, None) + assert inner.exited + + +async def test_timeout_acquire_context_await(): + conn = object() + ctx = TimeoutAcquireContext(FakeAcquireContext(conn), timeout=1.0) + result = await ctx + assert result is conn + + +async def test_timeout_acquire_context_timeout_fires(): + class SlowContext: + async def __aenter__(self): + await asyncio.sleep(10) + return object() + + async def __aexit__(self, *exc): + pass + + def __await__(self): + return self.__aenter__().__await__() + + ctx = TimeoutAcquireContext(SlowContext(), timeout=0.01) + with pytest.raises(asyncio.TimeoutError): + await ctx + + +async def test_pool_acquire_context_aexit_removes_connection(): + metrics = CalculateMetrics() + pool = object() + inner_ctx = FakeAcquireContext() + pool_state, balancer, register, unregister = _make_mocks(pool, inner_ctx) + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + await ctx.__aenter__() + assert metrics._add_connections.get("test-host:5432") == 1 + + await ctx.__aexit__(None, None, None) + assert metrics._remove_connections.get("test-host:5432") == 1 + assert inner_ctx.exited + + +async def test_pool_acquire_context_await_registers_connection(): + metrics = CalculateMetrics() + pool = object() + conn = object() + inner_ctx = FakeAcquireContext(conn) + pool_state, balancer, register, unregister = _make_mocks(pool, inner_ctx) + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + result = await ctx + assert result is conn + register.assert_called_once_with(conn, pool) + assert metrics._add_connections.get("test-host:5432") == 1 + + +async def test_pool_acquire_context_remaining_timeout_raises(): + metrics = CalculateMetrics() + pool_state, balancer, register, unregister = ( + MagicMock(), MagicMock(), MagicMock(), MagicMock() + ) + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + past_deadline = asyncio.get_running_loop().time() - 1.0 + with pytest.raises(asyncio.TimeoutError): + ctx._remaining_timeout(past_deadline) + + +async def test_pool_acquire_context_deadline(): + metrics = CalculateMetrics() + pool_state, balancer, register, unregister = ( + MagicMock(), MagicMock(), MagicMock(), MagicMock() + ) + + ctx = _make_ctx( + pool_state, balancer, register, unregister, metrics, timeout=5.0, + ) + + now = asyncio.get_running_loop().time() + deadline = ctx._deadline() + assert deadline > now + assert deadline <= now + 5.1 + + +async def test_pool_acquire_context_aenter_cleans_up_on_register_failure(): + """If register_connection raises after driver_ctx.__aenter__ succeeds, + the driver context must still be cleaned up via __aexit__.""" + metrics = CalculateMetrics() + pool = object() + inner_ctx = FakeAcquireContext() + pool_state, balancer, register, unregister = _make_mocks(pool, inner_ctx) + register.side_effect = RuntimeError("register failed") + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + with pytest.raises(RuntimeError, match="register failed"): + await ctx.__aenter__() + + assert inner_ctx.entered + assert inner_ctx.exited, ( + "driver context __aexit__ must be called on failure" + ) + + +async def test_pool_acquire_context_await_cleans_up_on_register_failure(): + """If register_connection raises after await driver_ctx succeeds, + the connection must be released back to the pool.""" + metrics = CalculateMetrics() + pool = object() + inner_ctx = FakeAcquireContext() + pool_state, balancer, register, unregister = _make_mocks(pool, inner_ctx) + register.side_effect = RuntimeError("register failed") + pool_state.release_to_pool = AsyncMock() + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + with pytest.raises(RuntimeError, match="register failed"): + await ctx + + pool_state.release_to_pool.assert_awaited_once_with( + inner_ctx.conn, pool, + ) + + +async def test_timeout_acquire_context_aexit_propagates_return(): + """If inner __aexit__ returns True (suppress), it must propagate.""" + class SuppressingContext: + async def __aenter__(self): + return object() + + async def __aexit__(self, *exc): + return True # suppress exception + + def __await__(self): + return self.__aenter__().__await__() + + ctx = TimeoutAcquireContext(SuppressingContext(), timeout=1.0) + await ctx.__aenter__() + result = await ctx.__aexit__(ValueError, ValueError("x"), None) + assert result is True + + +async def test_pool_acquire_context_kwargs_passed_through(): + metrics = CalculateMetrics() + pool = object() + inner_ctx = FakeAcquireContext() + pool_state, balancer, register, unregister = _make_mocks(pool, inner_ctx) + + ctx = _make_ctx( + pool_state, balancer, register, unregister, metrics, + custom_kwarg="value", + ) + + await ctx + call_kwargs = pool_state.acquire_from_pool.call_args + assert call_kwargs.kwargs.get("custom_kwarg") == "value" + + +async def test_get_pool_raises_no_available_pool_error_when_none(): + metrics = CalculateMetrics() + pool_state, balancer, register, unregister = _make_mocks(None, None) + balancer.get_pool = AsyncMock(return_value=None) + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + with pytest.raises(NoAvailablePoolError): + await ctx + + +async def test_no_available_pool_error_message_contains_no_available_pool(): + metrics = CalculateMetrics() + pool_state, balancer, register, unregister = _make_mocks(None, None) + balancer.get_pool = AsyncMock(return_value=None) + + ctx = _make_ctx(pool_state, balancer, register, unregister, metrics) + + with pytest.raises(NoAvailablePoolError, match="No available pool"): + await ctx diff --git a/tests/test_policy.py b/tests/test_policy.py new file mode 100644 index 0000000..d5a2e94 --- /dev/null +++ b/tests/test_policy.py @@ -0,0 +1,10 @@ +import pytest + +from hasql.balancer_policy import AbstractBalancerPolicy + + +def test_abstract_balancer_policy_cannot_instantiate(): + with pytest.raises(TypeError, match="abstract"): + AbstractBalancerPolicy(None) + +