From c495b4d097030e64e563a46bdb05cca786fadf43 Mon Sep 17 00:00:00 2001 From: Interitio Date: Mon, 1 Sep 2025 21:05:40 +1000 Subject: [PATCH] feat: Make pool size configurable. --- connector.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/connector.py b/connector.py index 6baf3b5..287b6b8 100644 --- a/connector.py +++ b/connector.py @@ -1,4 +1,4 @@ -from typing import Protocol, runtime_checkable, Callable, Awaitable, Optional +from typing import AsyncIterator, Protocol, runtime_checkable, Callable, Awaitable, Optional import logging from contextvars import ContextVar @@ -19,10 +19,13 @@ ctx_connection: Optional[ContextVar[psq.AsyncConnection]] = ContextVar('connecti class Connector: cursor_factory = AsyncLoggingCursor - def __init__(self, conn_args): + def __init__(self, conn_args, pool_min_size=1, pool_max_size=4): self._conn_args = conn_args self._conn_kwargs = dict(autocommit=True, row_factory=row_factory, cursor_factory=self.cursor_factory) + self.pool_min_size = pool_min_size + self.pool_max_size = pool_max_size + self.pool = self.make_pool() self.conn_hooks = [] @@ -47,8 +50,8 @@ class Connector: return AsyncConnectionPool( self._conn_args, open=False, - min_size=1, - max_size=4, + min_size=self.pool_min_size, + max_size=self.pool_max_size, configure=self._setup_connection, kwargs=self._conn_kwargs ) @@ -64,7 +67,7 @@ class Connector: old_pool = self.pool self.pool = self.make_pool() await self.pool.open() - logger.info(f"Old pool statistics: {self.pool.get_stats()}") + logger.info(f"Old pool statistics: {old_pool.get_stats()}") await old_pool.close() logger.info("Pool refresh complete.") @@ -95,7 +98,7 @@ class Connector: await self.pool.close() @asynccontextmanager - async def connection(self) -> psq.AsyncConnection: + async def connection(self) -> AsyncIterator[psq.AsyncConnection]: """ Asynchronous context manager to get and manage a connection.