Clean Code Python: Performance — Profiling, N+1 Detection, and Load Testing

from contextvars import ContextVar
from typing import Any

from sqlalchemy import event
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.orm import Session

import structlog

logger = structlog.get_logger()

_query_count: ContextVar[int] = ContextVar("query_count", default=0)
_query_threshold: int = 10


def reset_query_count() -> None:
    """Reset at the start of each request."""
    _query_count.set(0)


def get_query_count() -> int:
    """Return current query count for the request."""
    return _query_count.get()


def register_query_counter(engine: Engine) -> None:
    """Register event listener that counts queries per request."""

    @event.listens_for(engine, "before_cursor_execute")
    def _count_query(
        conn: Connection,
        cursor: Any,
        statement: str,
        parameters: Any,
        context: Any,
        executemany: bool,
    ) -> None:
        current = _query_count.get()
        _query_count.set(current + 1)

        if current + 1 > _query_threshold:
            logger.warning(
                "high_query_count",
                query_number=current + 1,
                threshold=_query_threshold,
                statement=statement[:200],
            )

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response

from src.core.config import settings
from src.db.query_counter import get_query_count, reset_query_count


class QueryCountMiddleware(BaseHTTPMiddleware):
    """Track and expose query counts per request."""

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        reset_query_count()
        response = await call_next(request)
        count = get_query_count()

        if settings.debug:
            response.headers["X-Query-Count"] = str(count)

        if count > 20:
            # Emit metric for monitoring dashboards
            from src.metrics import high_query_count_total
            high_query_count_total.labels(
                endpoint=request.url.path,
                method=request.method,
            ).inc()

        return response

from sqlalchemy import select
from sqlalchemy.orm import selectinload, joinedload
from sqlalchemy.ext.asyncio import AsyncSession

from src.models.book import Book
from src.models.author import Author


class CatalogRepository:
    """Repository with explicit eager loading strategies."""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def list_books_with_authors(
        self,
        tenant_id: UUID,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Book]:
        """Fetch books with authors in exactly 2 queries.

        Uses selectinload to avoid the cartesian explosion
        that joinedload would cause with 50K books.
        """
        stmt = (
            select(Book)
            .options(selectinload(Book.author))
            .where(Book.tenant_id == tenant_id)
            .order_by(Book.title)
            .limit(limit)
            .offset(offset)
        )
        result = await self._session.execute(stmt)
        return list(result.scalars().all())

    async def get_book_with_full_details(
        self, book_id: UUID
    ) -> Book | None:
        """Single book with all relations — joinedload is fine here.

        One book = one author = one publisher. No cartesian risk.
        """
        stmt = (
            select(Book)
            .options(
                joinedload(Book.author),
                joinedload(Book.publisher),
            )
            .where(Book.id == book_id)
        )
        result = await self._session.execute(stmt)
        return result.scalar_one_or_none()

import orjson
from fastapi.responses import ORJSONResponse


class ShelfWiseResponse(ORJSONResponse):
    """Default response class using orjson for serialization.

    orjson natively handles:
    - UUID (no str() conversion needed)
    - datetime (ISO 8601, no isoformat() call)
    - Decimal (as JSON number, not string)
    - dataclasses (no dict() conversion)
    """
    media_type = "application/json"

from fastapi import FastAPI
from src.core.serialization import ShelfWiseResponse

app = FastAPI(default_response_class=ShelfWiseResponse)

from contextlib import asynccontextmanager
from collections.abc import AsyncIterator

import httpx


@asynccontextmanager
async def managed_http_client() -> AsyncIterator[httpx.AsyncClient]:
    """Shared HTTP client with connection pooling.

    Reuses TCP connections across requests.
    Limits concurrent connections to prevent exhaustion.
    """
    async with httpx.AsyncClient(
        timeout=httpx.Timeout(10.0, connect=5.0),
        limits=httpx.Limits(
            max_connections=100,
            max_keepalive_connections=20,
            keepalive_expiry=30.0,
        ),
        http2=True,  # Multiplexing over single connection
    ) as client:
        yield client

import random
from uuid import uuid4

from locust import HttpUser, between, task

# Simulate 5 tenants with different catalog sizes
TENANTS = [
    {"id": str(uuid4()), "token": "jwt_tenant_a", "catalog_size": 500},
    {"id": str(uuid4()), "token": "jwt_tenant_b", "catalog_size": 5_000},
    {"id": str(uuid4()), "token": "jwt_tenant_c", "catalog_size": 50_000},
    {"id": str(uuid4()), "token": "jwt_tenant_d", "catalog_size": 100},
    {"id": str(uuid4()), "token": "jwt_tenant_e", "catalog_size": 10_000},
]


class ShelfWiseUser(HttpUser):
    """Simulates a tenant user browsing the catalog and placing orders."""

    wait_time = between(0.5, 2.0)

    def on_start(self) -> None:
        self.tenant = random.choice(TENANTS)
        self.client.headers.update({
            "Authorization": f"Bearer {self.tenant['token']}",
        })

    @task(10)
    def browse_catalog(self) -> None:
        """Most common operation: paginated catalog browse."""
        page = random.randint(1, 10)
        self.client.get(f"/api/v2/books?page={page}&limit=50")

    @task(5)
    def view_book(self) -> None:
        """View a single book's details."""
        book_id = str(uuid4())  # Use seeded IDs in real tests
        self.client.get(f"/api/v2/books/{book_id}")

    @task(2)
    def search_catalog(self) -> None:
        """Full-text search — expensive operation."""
        queries = ["python", "machine learning", "distributed", "design"]
        self.client.get(f"/api/v2/books/search?q={random.choice(queries)}")

    @task(1)
    def place_order(self) -> None:
        """Least common but most expensive operation."""
        self.client.post(
            "/api/v2/orders",
            json={
                "items": [
                    {"book_id": str(uuid4()), "quantity": random.randint(1, 3)}
                    for _ in range(random.randint(1, 5))
                ],
            },
        )

import json
import statistics
from pathlib import Path

import pytest


@pytest.fixture
def baseline() -> dict[str, float]:
    """Load baseline p99 latencies from previous successful run."""
    baseline_path = Path("tests/performance/baseline.json")
    return json.loads(baseline_path.read_text())


@pytest.fixture
def current_results() -> dict[str, float]:
    """Load current run results from locust stats."""
    results_path = Path("tests/performance/current_stats.json")
    return json.loads(results_path.read_text())


REGRESSION_THRESHOLD = 1.2  # 20% regression is actionable


def test_catalog_p99_no_regression(
    baseline: dict[str, float],
    current_results: dict[str, float],
) -> None:
    """Catalog endpoint p99 must not regress more than 20%."""
    baseline_p99 = baseline["GET /api/v2/books"]["p99"]
    current_p99 = current_results["GET /api/v2/books"]["p99"]

    assert current_p99 < baseline_p99 * REGRESSION_THRESHOLD, (
        f"p99 regression: {baseline_p99:.0f}ms → {current_p99:.0f}ms "
        f"({current_p99 / baseline_p99:.1%} of baseline)"
    )