Clean Code Python: Connection Pooling and Database Resilience Under Load

from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine

from src.core.config import settings


def create_engine() -> AsyncEngine:
    """Create the primary async engine with production pool settings."""
    return create_async_engine(
        settings.database.url.get_secret_value(),

        # ── Pool sizing ──────────────────────────────────────────
        pool_size=20,        # Persistent connections kept open
        max_overflow=30,     # Temporary connections under burst
        # Total max: 20 + 30 = 50 connections
        # Rule of thumb: 1 connection handles ~1 concurrent query
        # 50 connections ≈ 50 concurrent database operations per instance

        # ── Timeouts ─────────────────────────────────────────────
        pool_timeout=10,     # Seconds to wait for a connection before raising
        # Default is 30 — too long. A user waiting 30s has already left.
        # 10s is enough for burst absorption without hiding pool exhaustion.

        # ── Connection health ────────────────────────────────────
        pool_pre_ping=True,  # SELECT 1 before every checkout
        # Adds ~1ms latency per checkout. Worth it.
        # Without it, a connection dropped by the database (idle timeout,
        # failover, network blip) causes a hard error on the NEXT query.
        # With it, the dead connection is silently replaced.

        # ── Connection recycling ─────────────────────────────────
        pool_recycle=1800,   # Replace connections older than 30 minutes
        # PostgreSQL and PgBouncer have idle connection timeouts.
        # Recycling before those timeouts prevents "connection reset" errors.
        # Also prevents long-lived connections from accumulating memory.

        # ── Diagnostics ──────────────────────────────────────────
        echo=settings.database.echo,  # SQL logging (off in prod)
        echo_pool="debug" if settings.debug else False,
    )

from collections.abc import AsyncIterator

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
)

from src.db.engine import create_engine

engine = create_engine()
AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,  # Prevent lazy loads after commit
)


async def get_session() -> AsyncIterator[AsyncSession]:
    """FastAPI dependency that provides a request-scoped session.

    The session is committed on success, rolled back on exception,
    and always closed — guaranteeing the connection returns to the pool.
    """
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise

from typing import Annotated

from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession

from src.db.session import get_session

SessionDep = Annotated[AsyncSession, Depends(get_session)]

from src.api.deps import SessionDep


@router.get("/books/{book_id}")
async def get_book(book_id: int, session: SessionDep):
    # session is created for this request, closed after response
    result = await session.execute(
        select(Book).where(Book.id == book_id)
    )
    return result.scalar_one_or_none()

import asyncio
import time
import logging
from enum import StrEnum
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger("shelfwise.db")


class CircuitState(StrEnum):
    CLOSED = "closed"      # Normal operation — requests flow through
    OPEN = "open"          # Failing — reject immediately
    HALF_OPEN = "half_open"  # Testing — allow one request through


class DatabaseCircuitBreaker:
    """Circuit breaker for database connections.

    Transitions:
    - CLOSED → OPEN: after `failure_threshold` consecutive failures
    - OPEN → HALF_OPEN: after `recovery_timeout` seconds
    - HALF_OPEN → CLOSED: on first success
    - HALF_OPEN → OPEN: on first failure
    """

    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: float = 30.0,
    ):
        self._failure_threshold = failure_threshold
        self._recovery_timeout = recovery_timeout
        self._failure_count = 0
        self._last_failure_time = 0.0
        self._state = CircuitState.CLOSED

    @property
    def state(self) -> CircuitState:
        if self._state == CircuitState.OPEN:
            # Check if recovery timeout has elapsed
            if time.monotonic() - self._last_failure_time >= self._recovery_timeout:
                self._state = CircuitState.HALF_OPEN
                logger.info("Circuit breaker → HALF_OPEN (testing recovery)")
        return self._state

    def record_success(self) -> None:
        self._failure_count = 0
        if self._state == CircuitState.HALF_OPEN:
            self._state = CircuitState.CLOSED
            logger.info("Circuit breaker → CLOSED (recovered)")

    def record_failure(self) -> None:
        self._failure_count += 1
        self._last_failure_time = time.monotonic()
        if self._failure_count >= self._failure_threshold:
            self._state = CircuitState.OPEN
            logger.warning(
                "Circuit breaker → OPEN after %d failures (blocking requests for %ds)",
                self._failure_count,
                self._recovery_timeout,
            )

    @asynccontextmanager
    async def protect(
        self, session_factory
    ) -> AsyncIterator[AsyncSession]:
        """Wrap session creation with circuit breaker logic."""
        if self.state == CircuitState.OPEN:
            raise ConnectionError(
                "Database circuit breaker is OPEN. "
                "Failing fast to prevent cascade."
            )

        try:
            async with session_factory() as session:
                yield session
                self.record_success()
        except Exception as exc:
            self.record_failure()
            raise


# Global instance
db_circuit_breaker = DatabaseCircuitBreaker(
    failure_threshold=5,
    recovery_timeout=30.0,
)

[databases]
shelfwise = host=postgres port=5432 dbname=shelfwise

[pgbouncer]
listen_port = 6432
pool_mode = transaction          ; connection released after each transaction
max_client_conn = 1000           ; total client connections PgBouncer accepts
default_pool_size = 50           ; PostgreSQL connections per database
reserve_pool_size = 10           ; extra connections for burst
reserve_pool_timeout = 3         ; seconds before using reserve pool
server_idle_timeout = 300        ; close idle server connections after 5 min

from sqlalchemy import event, text
from sqlalchemy.pool import Pool

from src.core.context import get_current_tenant, is_system_context


@event.listens_for(Pool, "checkout")
def set_tenant_search_path(dbapi_conn, connection_record, connection_proxy):
    """Set PostgreSQL search_path to the tenant's schema on checkout.

    This runs every time a connection is checked out from the pool,
    ensuring the correct schema regardless of which connection we get.
    """
    if is_system_context():
        # System context uses the public schema
        dbapi_conn.execute("SET search_path TO public")
        return

    try:
        tenant = get_current_tenant()
    except RuntimeError:
        # No tenant context (health check, startup) — use public
        dbapi_conn.execute("SET search_path TO public")
        return

    schema_name = f"tenant_{tenant.slug}"
    # Parameterized to prevent SQL injection via tenant slug
    dbapi_conn.execute(f"SET search_path TO {schema_name}, public")

import logging
from sqlalchemy import event
from sqlalchemy.pool import Pool

logger = logging.getLogger("shelfwise.pool")


@event.listens_for(Pool, "checkout")
def on_checkout(dbapi_conn, connection_record, connection_proxy):
    pool = connection_proxy._pool
    logger.debug(
        "Pool checkout",
        extra={
            "pool_size": pool.size(),
            "checked_in": pool.checkedin(),
            "checked_out": pool.checkedout(),
            "overflow": pool.overflow(),
        },
    )


@event.listens_for(Pool, "checkin")
def on_checkin(dbapi_conn, connection_record):
    logger.debug("Pool checkin")


def get_pool_stats(engine) -> dict:
    """Expose pool stats for health check endpoints and metrics."""
    pool = engine.pool
    return {
        "pool_size": pool.size(),
        "checked_in": pool.checkedin(),
        "checked_out": pool.checkedout(),
        "overflow": pool.overflow(),
        "utilization_pct": round(
            pool.checkedout() / (pool.size() + pool.overflow()) * 100
            if (pool.size() + pool.overflow()) > 0 else 0,
            1,
        ),
    }