284 lines
8.4 KiB
Python
284 lines
8.4 KiB
Python
"""
|
|
Database connection budget management.
|
|
|
|
Limits concurrent database connections per operation to prevent
|
|
a single operation (e.g., recall with parallel queries) from
|
|
exhausting the connection pool.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, AsyncIterator
|
|
|
|
if TYPE_CHECKING:
|
|
import asyncpg
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class OperationBudget:
|
|
"""
|
|
Tracks connection budget for a single operation.
|
|
|
|
Each operation gets a semaphore limiting its concurrent connections.
|
|
"""
|
|
|
|
operation_id: str
|
|
max_connections: int
|
|
semaphore: asyncio.Semaphore = field(init=False)
|
|
active_count: int = field(default=0, init=False)
|
|
|
|
def __post_init__(self):
|
|
self.semaphore = asyncio.Semaphore(self.max_connections)
|
|
|
|
|
|
class ConnectionBudgetManager:
|
|
"""
|
|
Manages per-operation connection budgets.
|
|
|
|
Usage:
|
|
manager = ConnectionBudgetManager(default_budget=4)
|
|
|
|
# Start an operation
|
|
async with manager.operation(max_connections=2) as op:
|
|
# Acquire connections within the budget
|
|
async with op.acquire(pool) as conn:
|
|
await conn.fetch(...)
|
|
|
|
# Multiple connections respect the budget
|
|
async with op.acquire(pool) as conn1, op.acquire(pool) as conn2:
|
|
# At most 2 concurrent connections for this operation
|
|
...
|
|
"""
|
|
|
|
def __init__(self, default_budget: int = 4):
|
|
"""
|
|
Initialize the budget manager.
|
|
|
|
Args:
|
|
default_budget: Default max connections per operation
|
|
"""
|
|
self.default_budget = default_budget
|
|
self._operations: dict[str, OperationBudget] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
@asynccontextmanager
|
|
async def operation(
|
|
self,
|
|
max_connections: int | None = None,
|
|
operation_id: str | None = None,
|
|
) -> AsyncIterator["BudgetedOperation"]:
|
|
"""
|
|
Create a budgeted operation context.
|
|
|
|
Args:
|
|
max_connections: Max concurrent connections for this operation.
|
|
Defaults to manager's default_budget.
|
|
operation_id: Optional custom operation ID. Auto-generated if not provided.
|
|
|
|
Yields:
|
|
BudgetedOperation context for acquiring connections
|
|
"""
|
|
op_id = operation_id or f"op-{uuid.uuid4().hex[:12]}"
|
|
budget = max_connections or self.default_budget
|
|
|
|
async with self._lock:
|
|
if op_id in self._operations:
|
|
raise ValueError(f"Operation {op_id} already exists")
|
|
self._operations[op_id] = OperationBudget(op_id, budget)
|
|
|
|
try:
|
|
yield BudgetedOperation(self, op_id)
|
|
finally:
|
|
async with self._lock:
|
|
self._operations.pop(op_id, None)
|
|
|
|
def _get_budget(self, operation_id: str) -> OperationBudget:
|
|
"""Get budget for an operation (internal use)."""
|
|
budget = self._operations.get(operation_id)
|
|
if not budget:
|
|
raise ValueError(f"Operation {operation_id} not found")
|
|
return budget
|
|
|
|
|
|
class BudgetedOperation:
|
|
"""
|
|
A single operation with connection budget.
|
|
|
|
Provides methods to acquire connections within the budget.
|
|
"""
|
|
|
|
def __init__(self, manager: ConnectionBudgetManager, operation_id: str):
|
|
self._manager = manager
|
|
self.operation_id = operation_id
|
|
|
|
@property
|
|
def budget(self) -> OperationBudget:
|
|
"""Get the budget for this operation."""
|
|
return self._manager._get_budget(self.operation_id)
|
|
|
|
@asynccontextmanager
|
|
async def acquire(self, pool: "asyncpg.Pool") -> AsyncIterator["asyncpg.Connection"]:
|
|
"""
|
|
Acquire a connection within the operation's budget.
|
|
|
|
Blocks if the operation has reached its connection limit.
|
|
|
|
Args:
|
|
pool: asyncpg connection pool
|
|
|
|
Yields:
|
|
Database connection
|
|
"""
|
|
budget = self.budget
|
|
async with budget.semaphore:
|
|
budget.active_count += 1
|
|
conn = await pool.acquire()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
budget.active_count -= 1
|
|
await pool.release(conn)
|
|
|
|
def wrap_pool(self, pool: "asyncpg.Pool") -> "BudgetedPool":
|
|
"""
|
|
Wrap a pool with this operation's budget.
|
|
|
|
The returned BudgetedPool can be passed to functions expecting a pool,
|
|
and all acquire() calls will be limited by this operation's budget.
|
|
|
|
Args:
|
|
pool: asyncpg connection pool to wrap
|
|
|
|
Returns:
|
|
BudgetedPool that limits connections to this operation's budget
|
|
"""
|
|
return BudgetedPool(pool, self)
|
|
|
|
async def acquire_many(
|
|
self,
|
|
pool: "asyncpg.Pool",
|
|
count: int,
|
|
) -> AsyncIterator[list["asyncpg.Connection"]]:
|
|
"""
|
|
Acquire multiple connections within the budget.
|
|
|
|
Note: This acquires connections sequentially to respect the budget.
|
|
For parallel acquisition, use multiple acquire() calls with asyncio.gather().
|
|
|
|
Args:
|
|
pool: asyncpg connection pool
|
|
count: Number of connections to acquire
|
|
|
|
Yields:
|
|
List of database connections
|
|
"""
|
|
connections = []
|
|
try:
|
|
for _ in range(count):
|
|
conn = await pool.acquire()
|
|
connections.append(conn)
|
|
yield connections
|
|
finally:
|
|
for conn in connections:
|
|
await pool.release(conn)
|
|
|
|
|
|
# Global default manager instance
|
|
_default_manager: ConnectionBudgetManager | None = None
|
|
|
|
|
|
def get_budget_manager(default_budget: int = 4) -> ConnectionBudgetManager:
|
|
"""
|
|
Get or create the global budget manager.
|
|
|
|
Args:
|
|
default_budget: Default max connections per operation
|
|
|
|
Returns:
|
|
Global ConnectionBudgetManager instance
|
|
"""
|
|
global _default_manager
|
|
if _default_manager is None:
|
|
_default_manager = ConnectionBudgetManager(default_budget=default_budget)
|
|
return _default_manager
|
|
|
|
|
|
@asynccontextmanager
|
|
async def budgeted_operation(
|
|
max_connections: int | None = None,
|
|
operation_id: str | None = None,
|
|
default_budget: int = 4,
|
|
) -> AsyncIterator[BudgetedOperation]:
|
|
"""
|
|
Convenience function to create a budgeted operation.
|
|
|
|
Args:
|
|
max_connections: Max concurrent connections for this operation
|
|
operation_id: Optional custom operation ID
|
|
default_budget: Default budget if manager not yet created
|
|
|
|
Yields:
|
|
BudgetedOperation context
|
|
|
|
Example:
|
|
async with budgeted_operation(max_connections=2) as op:
|
|
async with op.acquire(pool) as conn:
|
|
await conn.fetch(...)
|
|
"""
|
|
manager = get_budget_manager(default_budget)
|
|
async with manager.operation(max_connections, operation_id) as op:
|
|
yield op
|
|
|
|
|
|
class BudgetedPool:
|
|
"""
|
|
A pool wrapper that limits concurrent connection acquisitions.
|
|
|
|
This can be passed to functions expecting a pool, and acquire()
|
|
calls will be limited by the budget semaphore.
|
|
|
|
Usage:
|
|
async with budgeted_operation(max_connections=4) as op:
|
|
budgeted_pool = op.wrap_pool(pool)
|
|
# Pass budgeted_pool to functions that expect a pool
|
|
await some_function(budgeted_pool, ...)
|
|
"""
|
|
|
|
def __init__(self, pool: "asyncpg.Pool", operation: BudgetedOperation):
|
|
self._pool = pool
|
|
self._operation = operation
|
|
|
|
async def acquire(self) -> "asyncpg.Connection":
|
|
"""
|
|
Acquire a connection within the budget.
|
|
|
|
Note: Caller must release the connection when done.
|
|
Prefer using as context manager via acquire_with_retry or op.acquire().
|
|
"""
|
|
budget = self._operation.budget
|
|
await budget.semaphore.acquire()
|
|
budget.active_count += 1
|
|
try:
|
|
return await self._pool.acquire()
|
|
except Exception:
|
|
budget.active_count -= 1
|
|
budget.semaphore.release()
|
|
raise
|
|
|
|
async def release(self, conn: "asyncpg.Connection") -> None:
|
|
"""Release a connection back to the pool."""
|
|
budget = self._operation.budget
|
|
try:
|
|
await self._pool.release(conn)
|
|
finally:
|
|
budget.active_count -= 1
|
|
budget.semaphore.release()
|
|
|
|
def __getattr__(self, name):
|
|
"""Proxy other attributes to the underlying pool."""
|
|
return getattr(self._pool, name)
|