sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/extensions/starlette/middleware.py

@@ -0,0 +1,281 @@
+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_
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+    from starlette.responses import Response
+
+    from sqlspec.extensions.starlette._state import SQLSpecConfigState
+
+__all__ = ("CorrelationMiddleware", "SQLCommenterMiddleware", "SQLSpecAutocommitMiddleware", "SQLSpecManualMiddleware")
+
+HTTP_200_OK = 200
+HTTP_300_MULTIPLE_CHOICES = 300
+HTTP_400_BAD_REQUEST = 400
+
+
+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.
+    """
+
+    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.
+        """
+        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:
+                    return await call_next(request)
+                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:
+                return await call_next(request)
+            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", include_redirect: bool = False) -> None:
+        """Initialize middleware.
+
+        Args:
+            app: Starlette application instance.
+            config_state: Configuration state for this database.
+            include_redirect: If True, commit on 3xx status codes as well.
+        """
+        super().__init__(app)
+        self.config_state = config_state
+        self.include_redirect = include_redirect
+
+    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.
+        """
+        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:
+                    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
+                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:
+                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
+            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.
+        """
+        extra_commit = self.config_state.extra_commit_statuses or set()
+        extra_rollback = self.config_state.extra_rollback_statuses or set()
+
+        if status_code in extra_commit:
+            return True
+        if status_code in extra_rollback:
+            return False
+
+        if HTTP_200_OK <= status_code < HTTP_300_MULTIPLE_CHOICES:
+            return True
+        return bool(self.include_redirect and HTTP_300_MULTIPLE_CHOICES <= status_code < HTTP_400_BAD_REQUEST)
+
+
+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
+
+    Example:
+        ```python
+        from starlette.applications import Starlette
+        from sqlspec.extensions.starlette.middleware import (
+            CorrelationMiddleware,
+        )
+
+        app = Starlette()
+        app.add_middleware(
+            CorrelationMiddleware,
+            primary_header="x-request-id",
+            auto_trace_headers=True,
+        )
+        ```
+    """
+
+    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 (e.g. "starlette", "fastapi").
+        """
+        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 endpoint is not None and hasattr(endpoint, "__name__"):
+            attrs["action"] = endpoint.__name__
+
+        previous = SQLCommenterContext.get()
+        SQLCommenterContext.set(attrs)
+        try:
+            response: Response = await call_next(request)
+            return response
+        finally:
+            SQLCommenterContext.set(previous)