feat: Make pool size configurable.
This commit is contained in:
15
connector.py
15
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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user