sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/extensions/flask/_state.py

@@ -0,0 +1,76 @@
+"""Flask configuration state management."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from sqlspec.config import DatabaseConfigProtocol
+
+__all__ = ("FlaskConfigState",)
+
+HTTP_SUCCESS_MIN = 200
+HTTP_SUCCESS_MAX = 300
+HTTP_REDIRECT_MAX = 400
+
+
+@dataclass
+class FlaskConfigState:
+    """Internal state for each database configuration in Flask extension.
+
+    Holds configuration, state keys, commit settings, and transaction logic.
+    """
+
+    config: "DatabaseConfigProtocol[Any, Any, Any]"
+    connection_key: str
+    session_key: str
+    commit_mode: Literal["manual", "autocommit", "autocommit_include_redirect"]
+    extra_commit_statuses: "set[int] | None"
+    extra_rollback_statuses: "set[int] | None"
+    is_async: bool
+    disable_di: bool
+    enable_correlation_middleware: bool = False
+    correlation_header: str = "x-request-id"
+    correlation_headers: "tuple[str, ...] | None" = None
+    auto_trace_headers: bool = True
+
+    def should_commit(self, status_code: int) -> bool:
+        """Determine if HTTP status code should trigger commit.
+
+        Args:
+            status_code: HTTP response status code.
+
+        Returns:
+            True if status should trigger commit, False otherwise.
+        """
+        if self.extra_commit_statuses and status_code in self.extra_commit_statuses:
+            return True
+
+        if self.extra_rollback_statuses and status_code in self.extra_rollback_statuses:
+            return False
+
+        if self.commit_mode == "manual":
+            return False
+
+        if self.commit_mode == "autocommit":
+            return HTTP_SUCCESS_MIN <= status_code < HTTP_SUCCESS_MAX
+
+        if self.commit_mode == "autocommit_include_redirect":
+            return HTTP_SUCCESS_MIN <= status_code < HTTP_REDIRECT_MAX
+
+        return False
+
+    def should_rollback(self, status_code: int) -> bool:
+        """Determine if HTTP status code should trigger rollback.
+
+        In autocommit modes, anything that doesn't commit should rollback.
+
+        Args:
+            status_code: HTTP response status code.
+
+        Returns:
+            True if status should trigger rollback, False otherwise.
+        """
+        if self.commit_mode == "manual":
+            return False
+
+        return not self.should_commit(status_code)

sqlspec/extensions/flask/_utils.py

@@ -0,0 +1,71 @@
+"""Helper utilities for Flask extension."""
+
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from sqlspec.extensions.flask._state import FlaskConfigState
+    from sqlspec.protocols import DictProtocol
+    from sqlspec.utils.portal import Portal
+
+__all__ = ("get_context_value", "get_or_create_session", "has_context_value", "pop_context_value", "set_context_value")
+
+_MISSING = object()
+
+
+def _get_context_dict(target: Any) -> dict[str, Any]:
+    """Return the underlying context dictionary."""
+    return cast("DictProtocol", target).__dict__
+
+
+def get_context_value(target: Any, key: str, default: Any = _MISSING) -> Any:
+    """Get a value from a Flask context object."""
+    data = _get_context_dict(target)
+    if default is _MISSING:
+        return data[key]
+    return data.get(key, default)
+
+
+def set_context_value(target: Any, key: str, value: Any) -> None:
+    """Set a value on a Flask context object."""
+    _get_context_dict(target)[key] = value
+
+
+def pop_context_value(target: Any, key: str) -> Any | None:
+    """Remove a value from a Flask context object."""
+    return _get_context_dict(target).pop(key, None)
+
+
+def has_context_value(target: Any, key: str) -> bool:
+    """Check if a Flask context object has a stored value."""
+    return key in _get_context_dict(target)
+
+
+def get_or_create_session(config_state: "FlaskConfigState", portal: "Portal | None") -> Any:
+    """Get or create database session for current request.
+
+    Sessions are cached per request in Flask g object to ensure
+    the same session is reused throughout the request lifecycle.
+
+    Args:
+        config_state: Configuration state for this database.
+        portal: Portal for async operations (None for sync).
+
+    Returns:
+        Database session (driver instance).
+    """
+    from flask import g
+
+    cache_key = f"sqlspec_session_cache_{config_state.session_key}"
+
+    cached_session = get_context_value(g, cache_key, None)
+    if cached_session is not None:
+        return cached_session
+
+    connection = get_context_value(g, config_state.connection_key)
+
+    session = config_state.config.driver_type(
+        connection=connection, statement_config=config_state.config.statement_config
+    )
+
+    set_context_value(g, cache_key, session)
+    return session

sqlspec/extensions/flask/extension.py

@@ -0,0 +1,519 @@
+"""Flask extension for SQLSpec database integration."""
+
+import atexit
+import logging
+from typing import TYPE_CHECKING, Any, Literal
+
+from sqlspec.base import SQLSpec
+from sqlspec.config import AsyncDatabaseConfig, NoPoolAsyncConfig
+from sqlspec.core import CorrelationExtractor
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.flask._state import FlaskConfigState
+from sqlspec.extensions.flask._utils import (
+    get_context_value,
+    get_or_create_session,
+    has_context_value,
+    pop_context_value,
+    set_context_value,
+)
+from sqlspec.utils.correlation import CorrelationContext
+from sqlspec.utils.logging import get_logger, log_with_context
+from sqlspec.utils.portal import PortalProvider
+
+if TYPE_CHECKING:
+    from flask import Flask, Response
+
+__all__ = ("SQLSpecPlugin",)
+
+logger = get_logger("sqlspec.extensions.flask")
+
+DEFAULT_COMMIT_MODE: Literal["manual"] = "manual"
+DEFAULT_SESSION_KEY = "db_session"
+
+
+class SQLSpecPlugin:
+    """Flask extension for SQLSpec database integration.
+
+    Provides request-scoped session management, automatic transaction handling,
+    and async adapter support via portal pattern.
+
+    Example:
+        from flask import Flask
+        from sqlspec import SQLSpec
+        from sqlspec.adapters.sqlite import SqliteConfig
+        from sqlspec.extensions.flask import SQLSpecPlugin
+
+        sqlspec = SQLSpec()
+        config = SqliteConfig(
+            connection_config={"database": "app.db"},
+            extension_config={
+                "flask": {
+                    "commit_mode": "autocommit",
+                    "session_key": "db"
+                }
+            }
+        )
+        sqlspec.add_config(config)
+
+        app = Flask(__name__)
+        plugin = SQLSpecPlugin(sqlspec, app)
+
+        @app.route("/users")
+        def list_users():
+            db = plugin.get_session()
+            result = db.execute("SELECT * FROM users")
+            return {"users": result.all()}
+    """
+
+    def __init__(self, sqlspec: SQLSpec, app: "Flask | None" = None) -> None:
+        """Initialize Flask extension with SQLSpec instance.
+
+        Args:
+            sqlspec: SQLSpec instance with registered configs.
+            app: Optional Flask application to initialize immediately.
+        """
+        self._sqlspec = sqlspec
+        self._config_states: list[FlaskConfigState] = []
+        self._portal: PortalProvider | None = None
+        self._has_async_configs = False
+        self._cleanup_registered = False
+        self._shutdown_complete = False
+        self._enable_correlation = False
+        self._extractor: CorrelationExtractor | None = None
+
+        for cfg in self._sqlspec.configs.values():
+            state = self._create_config_state(cfg)
+            self._config_states.append(state)
+
+            if state.is_async:
+                self._has_async_configs = True
+
+            if state.enable_correlation_middleware and not self._enable_correlation:
+                self._enable_correlation = True
+                self._extractor = CorrelationExtractor(
+                    primary_header=state.correlation_header,
+                    additional_headers=state.correlation_headers,
+                    auto_trace_headers=state.auto_trace_headers,
+                )
+
+        if app is not None:
+            self.init_app(app)
+
+    def _create_config_state(self, config: Any) -> FlaskConfigState:
+        """Create configuration state from database config.
+
+        Args:
+            config: Database configuration instance.
+
+        Returns:
+            FlaskConfigState instance.
+        """
+        flask_config = config.extension_config.get("flask", {})
+
+        session_key = flask_config.get("session_key", DEFAULT_SESSION_KEY)
+        connection_key = flask_config.get("connection_key", f"sqlspec_connection_{session_key}")
+        commit_mode = flask_config.get("commit_mode", DEFAULT_COMMIT_MODE)
+        extra_commit_statuses = flask_config.get("extra_commit_statuses")
+        extra_rollback_statuses = flask_config.get("extra_rollback_statuses")
+        disable_di = flask_config.get("disable_di", False)
+
+        enable_correlation = flask_config.get("enable_correlation_middleware", False)
+        correlation_header = flask_config.get("correlation_header", "x-request-id")
+        correlation_headers = flask_config.get("correlation_headers")
+        if correlation_headers is not None:
+            correlation_headers = tuple(correlation_headers)
+        auto_trace_headers = flask_config.get("auto_trace_headers", True)
+
+        is_async = isinstance(config, (AsyncDatabaseConfig, NoPoolAsyncConfig))
+
+        return FlaskConfigState(
+            config=config,
+            connection_key=connection_key,
+            session_key=session_key,
+            commit_mode=commit_mode,
+            extra_commit_statuses=extra_commit_statuses,
+            extra_rollback_statuses=extra_rollback_statuses,
+            is_async=is_async,
+            disable_di=disable_di,
+            enable_correlation_middleware=enable_correlation,
+            correlation_header=correlation_header,
+            correlation_headers=correlation_headers,
+            auto_trace_headers=auto_trace_headers,
+        )
+
+    def init_app(self, app: "Flask") -> None:
+        """Initialize Flask application with SQLSpec.
+
+        Validates configuration, creates portal if needed, creates pools,
+        and registers hooks.
+
+        Args:
+            app: Flask application to initialize.
+
+        Raises:
+            ImproperConfigurationError: If extension already registered or keys not unique.
+        """
+        if "sqlspec" in app.extensions:
+            msg = "SQLSpec extension already registered on this Flask application"
+            raise ImproperConfigurationError(msg)
+
+        self._validate_unique_keys()
+
+        if self._has_async_configs:
+            self._portal = PortalProvider()
+            self._portal.start()
+            log_with_context(logger, logging.DEBUG, "extension.init", framework="flask", stage="portal_started")
+
+        pools: dict[str, Any] = {}
+        for config_state in self._config_states:
+            if config_state.config.supports_connection_pooling:
+                if config_state.is_async:
+                    pool = self._portal.portal.call(config_state.config.create_pool)  # type: ignore[union-attr,arg-type]
+                else:
+                    pool = config_state.config.create_pool()
+                pools[config_state.session_key] = pool
+                log_with_context(
+                    logger, logging.DEBUG, "session.create", framework="flask", session_key=config_state.session_key
+                )
+
+        app.extensions["sqlspec"] = {"plugin": self, "pools": pools}
+
+        if any(not state.disable_di for state in self._config_states):
+            app.before_request(self._before_request_handler)
+            app.after_request(self._after_request_handler)
+            app.teardown_appcontext(self._teardown_appcontext_handler)
+
+        self._register_shutdown_hook()
+
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "extension.init",
+            framework="flask",
+            stage="configured",
+            config_count=len(self._config_states),
+            async_enabled=self._has_async_configs,
+        )
+
+    def _validate_unique_keys(self) -> None:
+        """Validate that all state keys are unique across configs.
+
+        Raises:
+            ImproperConfigurationError: If duplicate keys found.
+        """
+        all_keys: set[str] = set()
+
+        for state in self._config_states:
+            keys = {state.connection_key, state.session_key}
+            duplicates = all_keys & keys
+
+            if duplicates:
+                msg = f"Duplicate state keys found: {duplicates}. Use unique session_key values."
+                raise ImproperConfigurationError(msg)
+
+            all_keys.update(keys)
+
+    def _register_shutdown_hook(self) -> None:
+        """Register shutdown hook for pool and portal cleanup."""
+
+        if self._cleanup_registered:
+            return
+
+        atexit.register(self.shutdown)
+        self._cleanup_registered = True
+
+    def _before_request_handler(self) -> None:
+        """Acquire connection before request.
+
+        Stores connection in Flask g object for each configured database.
+        Also stores context managers for proper cleanup.
+        Extracts correlation ID if correlation middleware is enabled.
+        """
+        from flask import current_app, g, request
+
+        if self._enable_correlation and self._extractor is not None:
+            correlation_id = self._extractor.extract(lambda h: request.headers.get(h))
+            set_context_value(g, "correlation_id", correlation_id)
+            CorrelationContext.set(correlation_id)
+
+        for config_state in self._config_states:
+            if config_state.disable_di:
+                continue
+
+            if config_state.config.supports_connection_pooling:
+                pool = current_app.extensions["sqlspec"]["pools"][config_state.session_key]
+                conn_ctx = config_state.config.provide_connection(pool)
+
+                if config_state.is_async:
+                    connection = self._portal.portal.call(conn_ctx.__aenter__)  # type: ignore[union-attr]
+                else:
+                    connection = conn_ctx.__enter__()  # type: ignore[union-attr]
+
+                set_context_value(g, f"{config_state.connection_key}_ctx", conn_ctx)
+            elif config_state.is_async:
+                connection = self._portal.portal.call(config_state.config.create_connection)  # type: ignore[union-attr,arg-type]
+            else:
+                connection = config_state.config.create_connection()
+
+            set_context_value(g, config_state.connection_key, connection)
+
+    def _after_request_handler(self, response: "Response") -> "Response":
+        """Handle transaction after request based on response status.
+
+        Args:
+            response: Flask response object.
+
+        Returns:
+            Response object with correlation ID header if enabled.
+        """
+        from flask import g
+
+        if self._enable_correlation:
+            correlation_id = get_context_value(g, "correlation_id", None)
+            if correlation_id:
+                response.headers["X-Correlation-ID"] = correlation_id
+
+        for config_state in self._config_states:
+            if config_state.disable_di:
+                continue
+
+            if config_state.commit_mode == "manual":
+                continue
+
+            cache_key = f"sqlspec_session_cache_{config_state.session_key}"
+            session = get_context_value(g, cache_key, None)
+
+            if session is None:
+                continue
+
+            if config_state.should_commit(response.status_code):
+                self._execute_commit(session, config_state)
+            elif config_state.should_rollback(response.status_code):
+                self._execute_rollback(session, config_state)
+
+        return response
+
+    def _teardown_appcontext_handler(self, _exc: "BaseException | None" = None) -> None:
+        """Clean up connections when request context ends.
+
+        Closes all connections, cleans up g object, and clears correlation context.
+
+        Args:
+            _exc: Exception that occurred (if any).
+        """
+        from flask import g
+
+        if self._enable_correlation:
+            CorrelationContext.clear()
+            if has_context_value(g, "correlation_id"):
+                pop_context_value(g, "correlation_id")
+
+        for config_state in self._config_states:
+            if config_state.disable_di:
+                continue
+
+            connection = get_context_value(g, config_state.connection_key, None)
+            ctx_key = f"{config_state.connection_key}_ctx"
+            conn_ctx = get_context_value(g, ctx_key, None)
+
+            if connection is not None:
+                try:
+                    if conn_ctx is not None:
+                        if config_state.is_async:
+                            self._portal.portal.call(conn_ctx.__aexit__, None, None, None)  # type: ignore[union-attr]
+                        else:
+                            conn_ctx.__exit__(None, None, None)
+                    elif config_state.is_async:
+                        self._portal.portal.call(connection.close)  # type: ignore[union-attr]
+                    else:
+                        connection.close()
+                except Exception as exc:
+                    log_with_context(
+                        logger,
+                        logging.ERROR,
+                        "session.close",
+                        framework="flask",
+                        session_key=config_state.session_key,
+                        operation="connection",
+                        status="failed",
+                        error_type=type(exc).__name__,
+                    )
+
+                if has_context_value(g, config_state.connection_key):
+                    pop_context_value(g, config_state.connection_key)
+                if has_context_value(g, ctx_key):
+                    pop_context_value(g, ctx_key)
+
+            cache_key = f"sqlspec_session_cache_{config_state.session_key}"
+            if has_context_value(g, cache_key):
+                pop_context_value(g, cache_key)
+
+    def get_session(self, key: "str | None" = None) -> Any:
+        """Get or create database session for current request.
+
+        Sessions are cached per request for consistency.
+
+        Args:
+            key: Session key for multi-database configs. Defaults to first config if None.
+
+        Returns:
+            Database session (driver instance).
+        """
+        config_state = self._config_states[0] if key is None else self._get_config_state_by_key(key)
+
+        return get_or_create_session(config_state, self._portal.portal if self._portal else None)
+
+    def get_connection(self, key: "str | None" = None) -> Any:
+        """Get database connection for current request.
+
+        Args:
+            key: Session key for multi-database configs. Defaults to first config if None.
+
+        Returns:
+            Raw database connection.
+        """
+        from flask import g
+
+        config_state = self._config_states[0] if key is None else self._get_config_state_by_key(key)
+
+        return get_context_value(g, config_state.connection_key)
+
+    def _get_config_state_by_key(self, key: str) -> FlaskConfigState:
+        """Get config state by session key.
+
+        Args:
+            key: Session key to look up.
+
+        Returns:
+            FlaskConfigState for the key.
+
+        Raises:
+            ImproperConfigurationError: If key not found.
+        """
+        for state in self._config_states:
+            if state.session_key == key:
+                return state
+
+        msg = f"No configuration found for key: {key}"
+        raise ImproperConfigurationError(msg)
+
+    def shutdown(self) -> None:
+        """Dispose connection pools and stop async portal."""
+
+        if self._shutdown_complete:
+            return
+
+        self._shutdown_complete = True
+
+        for config_state in self._config_states:
+            if config_state.config.supports_connection_pooling:
+                self._close_pool_state(config_state)
+
+        if self._portal is not None:
+            try:
+                self._portal.stop()
+            except Exception as exc:
+                log_with_context(
+                    logger,
+                    logging.ERROR,
+                    "extension.init",
+                    framework="flask",
+                    stage="shutdown",
+                    status="failed",
+                    error_type=type(exc).__name__,
+                )
+            finally:
+                self._portal = None
+
+    def _close_pool_state(self, config_state: FlaskConfigState) -> None:
+        """Close pool associated with configuration state."""
+
+        try:
+            if config_state.is_async:
+                if self._portal is None:
+                    log_with_context(
+                        logger,
+                        logging.DEBUG,
+                        "session.close",
+                        framework="flask",
+                        session_key=config_state.session_key,
+                        operation="pool",
+                        status="skipped",
+                        reason="portal_not_initialized",
+                    )
+                    return
+                _ = self._portal.portal.call(config_state.config.close_pool)  # type: ignore[arg-type]
+            else:
+                config_state.config.close_pool()
+            log_with_context(
+                logger,
+                logging.DEBUG,
+                "session.close",
+                framework="flask",
+                session_key=config_state.session_key,
+                operation="pool",
+                status="complete",
+            )
+        except Exception as exc:
+            log_with_context(
+                logger,
+                logging.ERROR,
+                "session.close",
+                framework="flask",
+                session_key=config_state.session_key,
+                operation="pool",
+                status="failed",
+                error_type=type(exc).__name__,
+            )
+
+    def _execute_commit(self, session: Any, config_state: FlaskConfigState) -> None:
+        """Execute commit on session.
+
+        Args:
+            session: Database session.
+            config_state: Configuration state.
+        """
+        try:
+            if config_state.is_async:
+                connection = self.get_connection(config_state.session_key)
+                self._portal.portal.call(connection.commit)  # type: ignore[union-attr]
+            else:
+                connection = self.get_connection(config_state.session_key)
+                connection.commit()
+        except Exception as exc:
+            log_with_context(
+                logger,
+                logging.ERROR,
+                "session.close",
+                framework="flask",
+                session_key=config_state.session_key,
+                operation="commit",
+                status="failed",
+                error_type=type(exc).__name__,
+            )
+
+    def _execute_rollback(self, session: Any, config_state: FlaskConfigState) -> None:
+        """Execute rollback on session.
+
+        Args:
+            session: Database session.
+            config_state: Configuration state.
+        """
+        try:
+            if config_state.is_async:
+                connection = self.get_connection(config_state.session_key)
+                self._portal.portal.call(connection.rollback)  # type: ignore[union-attr]
+            else:
+                connection = self.get_connection(config_state.session_key)
+                connection.rollback()
+        except Exception as exc:
+            log_with_context(
+                logger,
+                logging.DEBUG,
+                "session.close",
+                framework="flask",
+                session_key=config_state.session_key,
+                operation="rollback",
+                status="failed",
+                error_type=type(exc).__name__,
+            )

sqlspec/extensions/litestar/__init__.py

@@ -0,0 +1,28 @@
+from sqlspec.extensions.litestar._utils import get_sqlspec_scope_state, set_sqlspec_scope_state
+from sqlspec.extensions.litestar.channels import SQLSpecChannelsBackend
+from sqlspec.extensions.litestar.cli import database_group
+from sqlspec.extensions.litestar.config import LitestarConfig
+from sqlspec.extensions.litestar.plugin import (
+    DEFAULT_COMMIT_MODE,
+    DEFAULT_CONNECTION_KEY,
+    DEFAULT_POOL_KEY,
+    DEFAULT_SESSION_KEY,
+    CommitMode,
+    SQLSpecPlugin,
+)
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+
+__all__ = (
+    "DEFAULT_COMMIT_MODE",
+    "DEFAULT_CONNECTION_KEY",
+    "DEFAULT_POOL_KEY",
+    "DEFAULT_SESSION_KEY",
+    "BaseSQLSpecStore",
+    "CommitMode",
+    "LitestarConfig",
+    "SQLSpecChannelsBackend",
+    "SQLSpecPlugin",
+    "database_group",
+    "get_sqlspec_scope_state",
+    "set_sqlspec_scope_state",
+)