Source code for sqlspec.extensions.starlette.middleware

from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any

from starlette.middleware.base import BaseHTTPMiddleware

from sqlspec.core import CorrelationExtractor
from sqlspec.core.sqlcommenter import SQLCommenterContext
from sqlspec.extensions.starlette._utils import get_state_value, pop_state_value, set_state_value
from sqlspec.utils.correlation import CorrelationContext
from sqlspec.utils.sync_tools import ensure_async_, with_ensure_async_
from sqlspec.utils.type_guards import has_name

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from starlette.requests import Request
    from starlette.responses import Response

    from sqlspec.extensions.starlette._state import SQLSpecConfigState

__all__ = ("CorrelationMiddleware", "SQLCommenterMiddleware", "SQLSpecAutocommitMiddleware", "SQLSpecManualMiddleware")



[docs]
class SQLSpecManualMiddleware(BaseHTTPMiddleware):
    """Middleware for manual transaction mode.

    Acquires connection from pool, stores in request.state, releases after request.
    No automatic commit or rollback - user code must handle transactions.
    """


[docs]
    def __init__(self, app: Any, config_state: "SQLSpecConfigState") -> None:
        """Initialize middleware.

        Args:
            app: Starlette application instance.
            config_state: Configuration state for this database.
        """
        super().__init__(app)
        self.config_state = config_state

async def dispatch(self, request: "Request", call_next: Any) -> Any: """Process request with manual transaction mode. Args: request: Incoming HTTP request. call_next: Next middleware or route handler. Returns: HTTP response. """ async with self._connection_cm(request): return await call_next(request)

@asynccontextmanager async def _connection_cm(self, request: "Request") -> "AsyncIterator[Any]": """Yield a request-scoped connection for manual transaction mode.""" config = self.config_state.config connection_key = self.config_state.connection_key if config.supports_connection_pooling: pool = get_state_value(request.app.state, self.config_state.pool_key) async with with_ensure_async_(config.provide_connection(pool)) as connection: set_state_value(request.state, connection_key, connection) try: yield connection finally: pop_state_value(request.state, connection_key) else: connection = await ensure_async_(config.create_connection)() set_state_value(request.state, connection_key, connection) try: yield connection finally: pop_state_value(request.state, connection_key) await ensure_async_(connection.close)()

class SQLSpecAutocommitMiddleware(BaseHTTPMiddleware): """Middleware for autocommit transaction mode. Acquires connection, commits on success status codes, rollbacks on error status codes. """

def __init__(self, app: Any, config_state: "SQLSpecConfigState") -> None: """Initialize middleware. Args: app: Starlette application instance. config_state: Configuration state for this database. """ super().__init__(app) self.config_state = config_state

async def dispatch(self, request: "Request", call_next: Any) -> Any: """Process request with autocommit transaction mode. Args: request: Incoming HTTP request. call_next: Next middleware or route handler. Returns: HTTP response. """ async with self._connection_cm(request) as connection: try: response = await call_next(request) if self._should_commit(response.status_code): await ensure_async_(connection.commit)() else: await ensure_async_(connection.rollback)() except Exception: await ensure_async_(connection.rollback)() raise else: return response @asynccontextmanager async def _connection_cm(self, request: "Request") -> "AsyncIterator[Any]": """Yield a request-scoped connection for autocommit transaction mode.""" config = self.config_state.config connection_key = self.config_state.connection_key if config.supports_connection_pooling: pool = get_state_value(request.app.state, self.config_state.pool_key) async with with_ensure_async_(config.provide_connection(pool)) as connection: set_state_value(request.state, connection_key, connection) try: yield connection finally: pop_state_value(request.state, connection_key) else: connection = await ensure_async_(config.create_connection)() set_state_value(request.state, connection_key, connection) try: yield connection finally: pop_state_value(request.state, connection_key) await ensure_async_(connection.close)() def _should_commit(self, status_code: int) -> bool: """Determine if response status code should trigger commit. Args: status_code: HTTP status code. Returns: True if should commit, False if should rollback. """ return self.config_state.should_commit(status_code)

class CorrelationMiddleware(BaseHTTPMiddleware): """Middleware for correlation ID extraction and propagation. Extracts correlation IDs from request headers (or generates new ones) and propagates them through the request lifecycle via CorrelationContext. The middleware: 1. Extracts correlation ID from configurable headers 2. Sets it in the CorrelationContext for async/sync access 3. Stores it in request.state.correlation_id 4. Adds X-Correlation-ID header to the response 5. Cleans up the context on request completion """

def __init__( self, app: Any, *, primary_header: str = "x-request-id", additional_headers: tuple[str, ...] | None = None, auto_trace_headers: bool = True, max_length: int = 128, ) -> None: """Initialize correlation middleware. Args: app: Starlette application instance. primary_header: The primary header to check first. Defaults to "x-request-id". additional_headers: Additional headers to check after the primary header. auto_trace_headers: If True, include standard trace context headers as fallbacks. max_length: Maximum length for correlation IDs. Defaults to 128. """ super().__init__(app) self._extractor = CorrelationExtractor( primary_header=primary_header, additional_headers=additional_headers, auto_trace_headers=auto_trace_headers, max_length=max_length, )

async def dispatch(self, request: "Request", call_next: Any) -> "Response": """Extract correlation ID and propagate through request lifecycle. Args: request: Incoming HTTP request. call_next: Next middleware or route handler. Returns: HTTP response with X-Correlation-ID header. """ correlation_id = self._extractor.extract(lambda h: request.headers.get(h)) previous_id = CorrelationContext.get() CorrelationContext.set(correlation_id) set_state_value(request.state, "correlation_id", correlation_id) try: response: Response = await call_next(request) response.headers["X-Correlation-ID"] = correlation_id return response finally: CorrelationContext.set(previous_id) pop_state_value(request.state, "correlation_id") class SQLCommenterMiddleware(BaseHTTPMiddleware): """Middleware that populates SQLCommenterContext with request attributes. Extracts route and endpoint information from the Starlette/FastAPI request and sets them in :class:`~sqlspec.extensions.sqlcommenter.SQLCommenterContext` for the duration of the request. """ def __init__(self, app: Any, *, framework: str = "starlette") -> None: """Initialize SQLCommenter middleware. Args: app: Starlette application instance. framework: Framework name to include in attributes. """ super().__init__(app) self._framework = framework async def dispatch(self, request: "Request", call_next: Any) -> "Response": """Extract request context and set SQLCommenter attributes. Args: request: Incoming HTTP request. call_next: Next middleware or route handler. Returns: HTTP response. """ attrs: dict[str, str] = {"route": request.url.path, "framework": self._framework} endpoint = request.scope.get("endpoint") if has_name(endpoint): attrs["action"] = endpoint.__name__ previous = SQLCommenterContext.get() SQLCommenterContext.set(attrs) try: response: Response = await call_next(request) return response finally: SQLCommenterContext.set(previous)