fleet-memory/hindsight-api-slim/hindsight_api/engine/db_budget.py

"""
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)