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

sqlspec/extensions/starlette/extension.py

@@ -0,0 +1,254 @@
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.base import SQLSpec
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.starlette._state import SQLSpecConfigState
+from sqlspec.extensions.starlette._utils import get_or_create_session
+from sqlspec.extensions.starlette.middleware import SQLSpecAutocommitMiddleware, SQLSpecManualMiddleware
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from starlette.applications import Starlette
+    from starlette.requests import Request
+
+__all__ = ("SQLSpecPlugin",)
+
+logger = get_logger("extensions.starlette")
+
+DEFAULT_COMMIT_MODE = "manual"
+DEFAULT_CONNECTION_KEY = "db_connection"
+DEFAULT_POOL_KEY = "db_pool"
+DEFAULT_SESSION_KEY = "db_session"
+
+
+class SQLSpecPlugin:
+    """SQLSpec integration for Starlette applications.
+
+    Provides middleware-based session management, automatic transaction handling,
+    and connection pooling lifecycle management.
+
+    Example:
+        from starlette.applications import Starlette
+        from sqlspec import SQLSpec
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.extensions.starlette import SQLSpecPlugin
+
+        sqlspec = SQLSpec()
+        sqlspec.add_config(AsyncpgConfig(
+            bind_key="default",
+            pool_config={"dsn": "postgresql://localhost/mydb"},
+            extension_config={
+                "starlette": {
+                    "commit_mode": "autocommit",
+                    "session_key": "db"
+                }
+            }
+        ))
+
+        app = Starlette()
+        db_ext = SQLSpecPlugin(sqlspec, app)
+
+        @app.route("/users")
+        async def list_users(request):
+            db = db_ext.get_session(request)
+            result = await db.execute("SELECT * FROM users")
+            return JSONResponse({"users": result.all()})
+    """
+
+    __slots__ = ("_config_states", "_sqlspec")
+
+    def __init__(self, sqlspec: SQLSpec, app: "Starlette | None" = None) -> None:
+        """Initialize SQLSpec Starlette extension.
+
+        Args:
+            sqlspec: Pre-configured SQLSpec instance with registered configs.
+            app: Optional Starlette application to initialize immediately.
+        """
+        self._sqlspec = sqlspec
+        self._config_states: list[SQLSpecConfigState] = []
+
+        for cfg in self._sqlspec.configs.values():
+            settings = self._extract_starlette_settings(cfg)
+            state = self._create_config_state(cfg, settings)
+            self._config_states.append(state)
+
+        if app is not None:
+            self.init_app(app)
+
+    def _extract_starlette_settings(self, config: Any) -> "dict[str, Any]":
+        """Extract Starlette settings from config.extension_config.
+
+        Args:
+            config: Database configuration instance.
+
+        Returns:
+            Dictionary of Starlette-specific settings.
+        """
+        starlette_config = config.extension_config.get("starlette", {})
+
+        connection_key = starlette_config.get("connection_key", DEFAULT_CONNECTION_KEY)
+        pool_key = starlette_config.get("pool_key", DEFAULT_POOL_KEY)
+        session_key = starlette_config.get("session_key", DEFAULT_SESSION_KEY)
+        commit_mode = starlette_config.get("commit_mode", DEFAULT_COMMIT_MODE)
+
+        if not config.supports_connection_pooling and pool_key == DEFAULT_POOL_KEY:
+            pool_key = f"_{DEFAULT_POOL_KEY}_{id(config)}"
+
+        return {
+            "connection_key": connection_key,
+            "pool_key": pool_key,
+            "session_key": session_key,
+            "commit_mode": commit_mode,
+            "extra_commit_statuses": starlette_config.get("extra_commit_statuses"),
+            "extra_rollback_statuses": starlette_config.get("extra_rollback_statuses"),
+        }
+
+    def _create_config_state(self, config: Any, settings: "dict[str, Any]") -> SQLSpecConfigState:
+        """Create configuration state object.
+
+        Args:
+            config: Database configuration instance.
+            settings: Extracted Starlette settings.
+
+        Returns:
+            Configuration state instance.
+        """
+        return SQLSpecConfigState(
+            config=config,
+            connection_key=settings["connection_key"],
+            pool_key=settings["pool_key"],
+            session_key=settings["session_key"],
+            commit_mode=settings["commit_mode"],
+            extra_commit_statuses=settings["extra_commit_statuses"],
+            extra_rollback_statuses=settings["extra_rollback_statuses"],
+        )
+
+    def init_app(self, app: "Starlette") -> None:
+        """Initialize Starlette application with SQLSpec.
+
+        Validates configuration, wraps lifespan, and adds middleware.
+
+        Args:
+            app: Starlette application instance.
+        """
+        self._validate_unique_keys()
+
+        original_lifespan = app.router.lifespan_context
+
+        @asynccontextmanager
+        async def combined_lifespan(app: "Starlette") -> "AsyncGenerator[None, None]":
+            async with self.lifespan(app), original_lifespan(app):
+                yield
+
+        app.router.lifespan_context = combined_lifespan
+
+        for config_state in self._config_states:
+            self._add_middleware(app, config_state)
+
+    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.pool_key, state.session_key}
+            duplicates = all_keys & keys
+
+            if duplicates:
+                msg = f"Duplicate state keys found: {duplicates}"
+                raise ImproperConfigurationError(msg)
+
+            all_keys.update(keys)
+
+    def _add_middleware(self, app: "Starlette", config_state: SQLSpecConfigState) -> None:
+        """Add transaction middleware for configuration.
+
+        Args:
+            app: Starlette application instance.
+            config_state: Configuration state.
+        """
+        if config_state.commit_mode == "manual":
+            app.add_middleware(SQLSpecManualMiddleware, config_state=config_state)
+        elif config_state.commit_mode == "autocommit":
+            app.add_middleware(SQLSpecAutocommitMiddleware, config_state=config_state, include_redirect=False)
+        elif config_state.commit_mode == "autocommit_include_redirect":
+            app.add_middleware(SQLSpecAutocommitMiddleware, config_state=config_state, include_redirect=True)
+
+    @asynccontextmanager
+    async def lifespan(self, app: "Starlette") -> "AsyncGenerator[None, None]":
+        """Manage connection pool lifecycle.
+
+        Args:
+            app: Starlette application instance.
+
+        Yields:
+            None
+        """
+        for config_state in self._config_states:
+            if config_state.config.supports_connection_pooling:
+                pool = await config_state.config.create_pool()
+                setattr(app.state, config_state.pool_key, pool)
+
+        try:
+            yield
+        finally:
+            for config_state in self._config_states:
+                if config_state.config.supports_connection_pooling:
+                    close_result = config_state.config.close_pool()
+                    if close_result is not None:
+                        await close_result
+
+    def get_session(self, request: "Request", key: "str | None" = None) -> Any:
+        """Get or create database session for request.
+
+        Sessions are cached per request to ensure consistency.
+
+        Args:
+            request: Starlette request instance.
+            key: Optional session key to retrieve specific database session.
+
+        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(request, config_state)
+
+    def get_connection(self, request: "Request", key: "str | None" = None) -> Any:
+        """Get database connection from request state.
+
+        Args:
+            request: Starlette request instance.
+            key: Optional session key to retrieve specific database connection.
+
+        Returns:
+            Database connection object.
+        """
+        config_state = self._config_states[0] if key is None else self._get_config_state_by_key(key)
+
+        return getattr(request.state, config_state.connection_key)
+
+    def _get_config_state_by_key(self, key: str) -> SQLSpecConfigState:
+        """Get configuration state by session key.
+
+        Args:
+            key: Session key to search for.
+
+        Returns:
+            Configuration state matching the key.
+
+        Raises:
+            ValueError: If no configuration found with the specified key.
+        """
+        for state in self._config_states:
+            if state.session_key == key:
+                return state
+
+        msg = f"No configuration found with session_key: {key}"
+        raise ValueError(msg)

sqlspec/extensions/starlette/middleware.py

@@ -0,0 +1,154 @@
+from typing import TYPE_CHECKING, Any
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+
+    from sqlspec.extensions.starlette._state import SQLSpecConfigState
+
+__all__ = ("SQLSpecAutocommitMiddleware", "SQLSpecManualMiddleware")
+
+logger = get_logger("extensions.starlette.middleware")
+
+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 = getattr(request.app.state, self.config_state.pool_key)
+            async with config.provide_connection(pool) as connection:  # type: ignore[union-attr]
+                setattr(request.state, connection_key, connection)
+                try:
+                    return await call_next(request)
+                finally:
+                    delattr(request.state, connection_key)
+        else:
+            connection = await config.create_connection()
+            setattr(request.state, connection_key, connection)
+            try:
+                return await call_next(request)
+            finally:
+                await 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 = getattr(request.app.state, self.config_state.pool_key)
+            async with config.provide_connection(pool) as connection:  # type: ignore[union-attr]
+                setattr(request.state, connection_key, connection)
+                try:
+                    response = await call_next(request)
+
+                    if self._should_commit(response.status_code):
+                        await connection.commit()
+                    else:
+                        await connection.rollback()
+                except Exception:
+                    await connection.rollback()
+                    raise
+                else:
+                    return response
+                finally:
+                    delattr(request.state, connection_key)
+        else:
+            connection = await config.create_connection()
+            setattr(request.state, connection_key, connection)
+            try:
+                response = await call_next(request)
+
+                if self._should_commit(response.status_code):
+                    await connection.commit()
+                else:
+                    await connection.rollback()
+            except Exception:
+                await connection.rollback()
+                raise
+            else:
+                return response
+            finally:
+                await 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)

sqlspec/protocols.py

@@ -39,6 +39,7 @@ __all__ = (
     "ParameterValueProtocol",
     "SQLBuilderProtocol",
     "SelectBuilderProtocol",
+    "SupportsArrowResults",
     "WithMethodProtocol",
 )
 
@@ -440,3 +441,42 @@ class SelectBuilderProtocol(SQLBuilderProtocol, Protocol):
     def select(self, *columns: "str | exp.Expression") -> Self:
         """Add SELECT columns to the query."""
         ...
+
+
+@runtime_checkable
+class SupportsArrowResults(Protocol):
+    """Protocol for adapters that support Arrow result format.
+
+    Adapters implementing this protocol can return query results in Apache Arrow
+    format via the select_to_arrow() method, enabling zero-copy data transfer and
+    efficient integration with data science tools.
+    """
+
+    def select_to_arrow(
+        self,
+        statement: Any,
+        /,
+        *parameters: Any,
+        statement_config: Any | None = None,
+        return_format: str = "table",
+        native_only: bool = False,
+        batch_size: int | None = None,
+        arrow_schema: Any | None = None,
+        **kwargs: Any,
+    ) -> "ArrowTable | ArrowRecordBatch":
+        """Execute query and return results as Apache Arrow Table or RecordBatch.
+
+        Args:
+            statement: SQL statement to execute.
+            *parameters: Query parameters and filters.
+            statement_config: Optional statement configuration override.
+            return_format: Output format - "table", "reader", or "batches".
+            native_only: If True, raise error when native Arrow path unavailable.
+            batch_size: Chunk size for streaming modes.
+            arrow_schema: Optional target Arrow schema for type casting.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            ArrowResult containing Arrow data.
+        """
+        ...

sqlspec/storage/_utils.py

@@ -2,23 +2,10 @@
 
 from typing import TYPE_CHECKING
 
-from sqlspec.exceptions import MissingDependencyError
-from sqlspec.typing import PYARROW_INSTALLED
-
 if TYPE_CHECKING:
     from pathlib import Path
 
-__all__ = ("ensure_pyarrow", "resolve_storage_path")
-
-
-def ensure_pyarrow() -> None:
-    """Ensure PyArrow is available for Arrow operations.
-
-    Raises:
-        MissingDependencyError: If pyarrow is not installed.
-    """
-    if not PYARROW_INSTALLED:
-        raise MissingDependencyError(package="pyarrow", install_package="pyarrow")
+__all__ = ("resolve_storage_path",)
 
 
 def resolve_storage_path(

sqlspec/storage/backends/fsspec.py

@@ -5,9 +5,8 @@ from typing import TYPE_CHECKING, Any
 
 from mypy_extensions import mypyc_attr
 
-from sqlspec.exceptions import MissingDependencyError
-from sqlspec.storage._utils import ensure_pyarrow, resolve_storage_path
-from sqlspec.typing import FSSPEC_INSTALLED
+from sqlspec.storage._utils import resolve_storage_path
+from sqlspec.utils.module_loader import ensure_fsspec, ensure_pyarrow
 from sqlspec.utils.sync_tools import async_
 
 if TYPE_CHECKING:
@@ -105,8 +104,7 @@ class FSSpecBackend:
     __slots__ = ("_fs_uri", "backend_type", "base_path", "fs", "protocol")
 
     def __init__(self, uri: str, **kwargs: Any) -> None:
-        if not FSSPEC_INSTALLED:
-            raise MissingDependencyError(package="fsspec", install_package="fsspec")
+        ensure_fsspec()
 
         base_path = kwargs.pop("base_path", "")
 

sqlspec/storage/backends/local.py

@@ -12,7 +12,7 @@ from urllib.parse import unquote, urlparse
 
 from mypy_extensions import mypyc_attr
 
-from sqlspec.storage._utils import ensure_pyarrow
+from sqlspec.utils.module_loader import ensure_pyarrow
 from sqlspec.utils.sync_tools import async_
 
 if TYPE_CHECKING:

sqlspec/storage/backends/obstore.py

@@ -8,19 +8,17 @@ import fnmatch
 import io
 import logging
 from collections.abc import AsyncIterator, Iterator
-from typing import TYPE_CHECKING, Any, Final, cast
+from pathlib import Path
+from typing import Any, Final, cast
 from urllib.parse import urlparse
 
-from sqlspec.utils.sync_tools import async_
-
-if TYPE_CHECKING:
-    from pathlib import Path
-
 from mypy_extensions import mypyc_attr
 
-from sqlspec.exceptions import MissingDependencyError, StorageOperationFailedError
-from sqlspec.storage._utils import ensure_pyarrow, resolve_storage_path
-from sqlspec.typing import OBSTORE_INSTALLED, ArrowRecordBatch, ArrowTable
+from sqlspec.exceptions import StorageOperationFailedError
+from sqlspec.storage._utils import resolve_storage_path
+from sqlspec.typing import ArrowRecordBatch, ArrowTable
+from sqlspec.utils.module_loader import ensure_obstore, ensure_pyarrow
+from sqlspec.utils.sync_tools import async_
 
 __all__ = ("ObStoreBackend",)
 
@@ -119,11 +117,8 @@ class ObStoreBackend:
             uri: Storage URI (e.g., 's3://bucket', 'file:///path', 'gs://bucket')
             **kwargs: Additional options including base_path and obstore configuration
 
-        Raises:
-            MissingDependencyError: If obstore is not installed.
         """
-        if not OBSTORE_INSTALLED:
-            raise MissingDependencyError(package="obstore", install_package="obstore")
+        ensure_obstore()
 
         try:
             # Extract base_path from kwargs
@@ -144,8 +139,6 @@ class ObStoreBackend:
 
                 self.store = MemoryStore()
             elif uri.startswith("file://"):
-                from pathlib import Path as PathlibPath
-
                 from obstore.store import LocalStore
 
                 # Parse URI to extract path
@@ -155,7 +148,7 @@ class ObStoreBackend:
                 # Append fragment if present (handles paths with '#' character)
                 if parsed.fragment:
                     path_str = f"{path_str}#{parsed.fragment}"
-                path_obj = PathlibPath(path_str)
+                path_obj = Path(path_str)
 
                 # If path points to a file, use its parent as the base directory
                 if path_obj.is_file():
@@ -194,9 +187,8 @@ class ObStoreBackend:
 
     def _resolve_path_for_local_store(self, path: "str | Path") -> str:
         """Resolve path for LocalStore which expects relative paths from its root."""
-        from pathlib import Path as PathlibPath
 
-        path_obj = PathlibPath(str(path))
+        path_obj = Path(str(path))
 
         # If absolute path, try to make it relative to LocalStore root
         if path_obj.is_absolute() and self._local_store_root:

sqlspec/typing.py

@@ -16,7 +16,9 @@ from sqlspec._typing import (
     OBSTORE_INSTALLED,
     OPENTELEMETRY_INSTALLED,
     ORJSON_INSTALLED,
+    PANDAS_INSTALLED,
     PGVECTOR_INSTALLED,
+    POLARS_INSTALLED,
     PROMETHEUS_INSTALLED,
     PYARROW_INSTALLED,
     PYDANTIC_INSTALLED,
@@ -26,6 +28,10 @@ from sqlspec._typing import (
     AiosqlSQLOperationType,
     AiosqlSyncProtocol,
     ArrowRecordBatch,
+    ArrowRecordBatchReader,
+    ArrowRecordBatchReaderProtocol,
+    ArrowSchema,
+    ArrowSchemaProtocol,
     ArrowTable,
     AttrsInstance,
     AttrsInstanceStub,
@@ -41,6 +47,8 @@ from sqlspec._typing import (
     Gauge,
     Histogram,
     NumpyArray,
+    PandasDataFrame,
+    PolarsDataFrame,
     Span,
     Status,
     StatusCode,
@@ -136,7 +144,9 @@ __all__ = (
     "OBSTORE_INSTALLED",
     "OPENTELEMETRY_INSTALLED",
     "ORJSON_INSTALLED",
+    "PANDAS_INSTALLED",
     "PGVECTOR_INSTALLED",
+    "POLARS_INSTALLED",
     "PROMETHEUS_INSTALLED",
     "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
@@ -147,6 +157,10 @@ __all__ = (
     "AiosqlSQLOperationType",
     "AiosqlSyncProtocol",
     "ArrowRecordBatch",
+    "ArrowRecordBatchReader",
+    "ArrowRecordBatchReaderProtocol",
+    "ArrowSchema",
+    "ArrowSchemaProtocol",
     "ArrowTable",
     "AttrsInstance",
     "BaseModel",
@@ -162,6 +176,8 @@ __all__ = (
     "Gauge",
     "Histogram",
     "NumpyArray",
+    "PandasDataFrame",
+    "PolarsDataFrame",
     "PoolT",
     "SchemaT",
     "Span",

sqlspec/utils/__init__.py

@@ -1,10 +1,31 @@
 """Utility functions and classes for SQLSpec.
 
 This package provides various utility modules for deprecation handling,
-fixture loading, logging, module loading, singleton patterns, sync/async
-conversion, text processing, and type guards.
+fixture loading, logging, module loading (including dependency checking),
+portal pattern for async bridging, singleton patterns, sync/async conversion,
+text processing, and type guards.
 """
 
-from sqlspec.utils import deprecation, fixtures, logging, module_loader, singleton, sync_tools, text, type_guards
+from sqlspec.utils import (
+    deprecation,
+    fixtures,
+    logging,
+    module_loader,
+    portal,
+    singleton,
+    sync_tools,
+    text,
+    type_guards,
+)
 
-__all__ = ("deprecation", "fixtures", "logging", "module_loader", "singleton", "sync_tools", "text", "type_guards")
+__all__ = (
+    "deprecation",
+    "fixtures",
+    "logging",
+    "module_loader",
+    "portal",
+    "singleton",
+    "sync_tools",
+    "text",
+    "type_guards",
+)