sqlspec 0.27.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/extensions/flask/extension.py

@@ -0,0 +1,389 @@
+"""Flask extension for SQLSpec database integration."""
+
+import atexit
+from typing import TYPE_CHECKING, Any, Literal
+
+from sqlspec.base import SQLSpec
+from sqlspec.config import AsyncDatabaseConfig, NoPoolAsyncConfig
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.flask._state import FlaskConfigState
+from sqlspec.extensions.flask._utils import get_or_create_session
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.portal import PortalProvider
+
+if TYPE_CHECKING:
+    from flask import Flask, Response
+
+__all__ = ("SQLSpecPlugin",)
+
+logger = get_logger("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(
+            pool_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
+
+        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 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")
+
+        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,
+        )
+
+    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()
+            logger.debug("Portal provider started for async adapters")
+
+        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
+
+        app.extensions["sqlspec"] = {"plugin": self, "pools": pools}
+
+        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()
+
+        logger.debug("SQLSpec Flask extension initialized")
+
+    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.
+        """
+        from flask import current_app, g
+
+        for config_state in self._config_states:
+            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]
+
+                setattr(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()
+
+            setattr(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 (unchanged).
+        """
+        from flask import g
+
+        for config_state in self._config_states:
+            if config_state.commit_mode == "manual":
+                continue
+
+            cache_key = f"sqlspec_session_cache_{config_state.session_key}"
+            session = getattr(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 and cleans up g object.
+
+        Args:
+            _exc: Exception that occurred (if any).
+        """
+        from flask import g
+
+        for config_state in self._config_states:
+            connection = getattr(g, config_state.connection_key, None)
+            ctx_key = f"{config_state.connection_key}_ctx"
+            conn_ctx = getattr(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:
+                    logger.exception("Error closing connection")
+
+                if hasattr(g, config_state.connection_key):
+                    delattr(g, config_state.connection_key)
+                if hasattr(g, ctx_key):
+                    delattr(g, ctx_key)
+
+            cache_key = f"sqlspec_session_cache_{config_state.session_key}"
+            if hasattr(g, cache_key):
+                delattr(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 getattr(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:
+                logger.exception("Error stopping portal during shutdown")
+            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:
+                    logger.debug(
+                        "Portal not initialized - skipping async pool shutdown for %s", config_state.session_key
+                    )
+                    return
+                _ = self._portal.portal.call(config_state.config.close_pool)  # type: ignore[arg-type]
+            else:
+                config_state.config.close_pool()
+        except Exception:
+            logger.exception("Error closing pool during shutdown for key %s", config_state.session_key)
+
+    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:
+            logger.exception("Error committing transaction")
+
+    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:
+            logger.debug("Rollback failed (may be no active transaction): %s", exc)

sqlspec/extensions/litestar/config.py

@@ -40,8 +40,7 @@ class LitestarConfig(TypedDict):
     in_memory: NotRequired[bool]
     """Enable in-memory table storage (Oracle-specific). Default: False.
 
-    When enabled, tables are created with the INMEMORY clause for Oracle Database,
-    which stores table data in columnar format in memory for faster query performance.
+    When enabled, tables are created with the in-memory attribute for databases that support it.
 
     This is an Oracle-specific feature that requires:
         - Oracle Database 12.1.0.2 or higher
@@ -62,8 +61,6 @@ class LitestarConfig(TypedDict):
             )
 
     Notes:
-        - Improves query performance for session lookups (10-100x faster)
-        - Tables created with INMEMORY clause
-        - Requires Oracle Database In-Memory option license
-        - Ignored by non-Oracle adapters
+        - Tables created with INMEMORY PRIORITY HIGH clause
+        - Ignored by unsupported adapters
     """

sqlspec/extensions/litestar/plugin.py

@@ -25,8 +25,9 @@ from sqlspec.extensions.litestar.handlers import (
     pool_provider_maker,
     session_provider_maker,
 )
-from sqlspec.typing import ConnectionT, PoolT
+from sqlspec.typing import NUMPY_INSTALLED, ConnectionT, PoolT, SchemaT
 from sqlspec.utils.logging import get_logger
+from sqlspec.utils.serializers import numpy_array_dec_hook, numpy_array_enc_hook, numpy_array_predicate
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable
@@ -82,6 +83,10 @@ class _PluginConfigState:
 class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
     """Litestar plugin for SQLSpec database integration.
 
+    Automatically configures NumPy array serialization when NumPy is installed,
+    enabling seamless bidirectional conversion between NumPy arrays and JSON
+    for vector embedding workflows.
+
     Session Table Migrations:
         The Litestar extension includes migrations for creating session storage tables.
         To include these migrations in your database migration workflow, add 'litestar'
@@ -225,6 +230,8 @@ class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
     def on_app_init(self, app_config: "AppConfig") -> "AppConfig":
         """Configure Litestar application with SQLSpec database integration.
 
+        Automatically registers NumPy array serialization when NumPy is installed.
+
         Args:
             app_config: The Litestar application configuration instance.
 
@@ -239,7 +246,7 @@ class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
         app_config.on_startup.append(store_sqlspec_in_state)
         app_config.signature_types.extend([SQLSpec, DatabaseConfigProtocol, SyncConfigT, AsyncConfigT])
 
-        signature_namespace = {"ConnectionT": ConnectionT, "PoolT": PoolT, "DriverT": DriverT}
+        signature_namespace = {"ConnectionT": ConnectionT, "PoolT": PoolT, "DriverT": DriverT, "SchemaT": SchemaT}
 
         for state in self._plugin_configs:
             state.annotation = type(state.config)
@@ -262,6 +269,23 @@ class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
         if signature_namespace:
             app_config.signature_namespace.update(signature_namespace)
 
+        if NUMPY_INSTALLED:
+            import numpy as np
+
+            if app_config.type_encoders is None:
+                app_config.type_encoders = {np.ndarray: numpy_array_enc_hook}
+            else:
+                encoders_dict = dict(app_config.type_encoders)
+                encoders_dict[np.ndarray] = numpy_array_enc_hook
+                app_config.type_encoders = encoders_dict
+
+            if app_config.type_decoders is None:
+                app_config.type_decoders = [(numpy_array_predicate, numpy_array_dec_hook)]  # type: ignore[list-item]
+            else:
+                decoders_list = list(app_config.type_decoders)
+                decoders_list.append((numpy_array_predicate, numpy_array_dec_hook))  # type: ignore[arg-type]
+                app_config.type_decoders = decoders_list
+
         return app_config
 
     def get_annotations(

sqlspec/extensions/starlette/__init__.py

@@ -0,0 +1,10 @@
+"""Starlette extension for SQLSpec.
+
+Provides middleware-based session management, automatic transaction handling,
+and connection pooling lifecycle management for Starlette applications.
+"""
+
+from sqlspec.extensions.starlette.extension import SQLSpecPlugin
+from sqlspec.extensions.starlette.middleware import SQLSpecAutocommitMiddleware, SQLSpecManualMiddleware
+
+__all__ = ("SQLSpecAutocommitMiddleware", "SQLSpecManualMiddleware", "SQLSpecPlugin")

sqlspec/extensions/starlette/_state.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from sqlspec.config import DatabaseConfigProtocol
+
+__all__ = ("CommitMode", "SQLSpecConfigState")
+
+CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
+
+
+@dataclass
+class SQLSpecConfigState:
+    """Internal state for each database configuration.
+
+    Tracks all configuration parameters needed for middleware and session management.
+    """
+
+    config: "DatabaseConfigProtocol[Any, Any, Any]"
+    connection_key: str
+    pool_key: str
+    session_key: str
+    commit_mode: CommitMode
+    extra_commit_statuses: "set[int] | None"
+    extra_rollback_statuses: "set[int] | None"

sqlspec/extensions/starlette/_utils.py

@@ -0,0 +1,52 @@
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+
+    from sqlspec.extensions.starlette._state import SQLSpecConfigState
+
+__all__ = ("get_connection_from_request", "get_or_create_session")
+
+
+def get_connection_from_request(request: "Request", config_state: "SQLSpecConfigState") -> Any:
+    """Get database connection from request state.
+
+    Args:
+        request: Starlette request instance.
+        config_state: Configuration state for the database.
+
+    Returns:
+        Database connection object.
+    """
+    return getattr(request.state, config_state.connection_key)
+
+
+def get_or_create_session(request: "Request", config_state: "SQLSpecConfigState") -> Any:
+    """Get or create database session for request.
+
+    Sessions are cached per request to ensure the same session instance
+    is returned for multiple calls within the same request.
+
+    Args:
+        request: Starlette request instance.
+        config_state: Configuration state for the database.
+
+    Returns:
+        Database session (driver instance).
+    """
+    session_instance_key = f"{config_state.session_key}_instance"
+
+    existing_session = getattr(request.state, session_instance_key, None)
+    if existing_session is not None:
+        return existing_session
+
+    connection = get_connection_from_request(request, config_state)
+
+    session = config_state.config.driver_type(
+        connection=connection,
+        statement_config=config_state.config.statement_config,
+        driver_features=config_state.config.driver_features,
+    )
+
+    setattr(request.state, session_instance_key, session)
+    return session