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

sqlspec/adapters/sqlite/adk/store.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any
 from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.serializers import from_json, to_json
-from sqlspec.utils.sync_tools import async_
+from sqlspec.utils.sync_tools import async_, run_
 
 if TYPE_CHECKING:
     from sqlspec.adapters.sqlite.config import SqliteConfig
@@ -140,7 +140,7 @@ class SqliteADKStore(BaseAsyncADKStore["SqliteConfig"]):
         """
         super().__init__(config)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get SQLite CREATE TABLE SQL for sessions.
 
         Returns:
@@ -172,7 +172,7 @@ class SqliteADKStore(BaseAsyncADKStore["SqliteConfig"]):
             ON {self._session_table}(update_time DESC);
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get SQLite CREATE TABLE SQL for events.
 
         Returns:
@@ -237,10 +237,10 @@ class SqliteADKStore(BaseAsyncADKStore["SqliteConfig"]):
 
     def _create_tables(self) -> None:
         """Synchronous implementation of create_tables."""
-        with self._config.provide_connection() as conn:
-            self._enable_foreign_keys(conn)
-            conn.executescript(self._get_create_sessions_table_sql())
-            conn.executescript(self._get_create_events_table_sql())
+        with self._config.provide_session() as driver:
+            driver.connection.execute("PRAGMA foreign_keys = ON")
+            driver.execute_script(run_(self._get_create_sessions_table_sql)())
+            driver.execute_script(run_(self._get_create_events_table_sql)())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_tables(self) -> None:
@@ -370,18 +370,28 @@ class SqliteADKStore(BaseAsyncADKStore["SqliteConfig"]):
         """
         await async_(self._update_session_state)(session_id, state)
 
-    def _list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
+    def _list_sessions(self, app_name: str, user_id: "str | None") -> "list[SessionRecord]":
         """Synchronous implementation of list_sessions."""
-        sql = f"""
-        SELECT id, app_name, user_id, state, create_time, update_time
-        FROM {self._session_table}
-        WHERE app_name = ? AND user_id = ?
-        ORDER BY update_time DESC
-        """
+        if user_id is None:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ?
+            ORDER BY update_time DESC
+            """
+            params: tuple[str, ...] = (app_name,)
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ? AND user_id = ?
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
 
         with self._config.provide_connection() as conn:
             self._enable_foreign_keys(conn)
-            cursor = conn.execute(sql, (app_name, user_id))
+            cursor = conn.execute(sql, params)
             rows = cursor.fetchall()
 
             return [
@@ -396,18 +406,18 @@ class SqliteADKStore(BaseAsyncADKStore["SqliteConfig"]):
                 for row in rows
             ]
 
-    async def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
-        """List all sessions for a user in an app.
+    async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        """List sessions for an app, optionally filtered by user.
 
         Args:
             app_name: Application name.
-            user_id: User identifier.
+            user_id: User identifier. If None, lists all sessions for the app.
 
         Returns:
             List of session records ordered by update_time DESC.
 
         Notes:
-            Uses composite index on (app_name, user_id).
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
         return await async_(self._list_sessions)(app_name, user_id)
 

sqlspec/config.py

@@ -26,11 +26,13 @@ __all__ = (
     "ConfigT",
     "DatabaseConfigProtocol",
     "DriverT",
+    "FlaskConfig",
     "LifecycleConfig",
     "LitestarConfig",
     "MigrationConfig",
     "NoPoolAsyncConfig",
     "NoPoolSyncConfig",
+    "StarletteConfig",
     "SyncConfigT",
     "SyncDatabaseConfig",
 )
@@ -98,6 +100,49 @@ class MigrationConfig(TypedDict):
     """Wrap migrations in transactions when supported. When enabled (default for adapters that support it), each migration runs in a transaction that is committed on success or rolled back on failure. This prevents partial migrations from leaving the database in an inconsistent state. Requires adapter support for transactional DDL. Defaults to True for PostgreSQL, SQLite, and DuckDB; False for MySQL, Oracle, and BigQuery. Individual migrations can override this with a '-- transactional: false' comment."""
 
 
+class FlaskConfig(TypedDict):
+    """Configuration options for Flask SQLSpec extension.
+
+    All fields are optional with sensible defaults. Use in extension_config["flask"]:
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+
+        config = AsyncpgConfig(
+            pool_config={"dsn": "postgresql://localhost/mydb"},
+            extension_config={
+                "flask": {
+                    "commit_mode": "autocommit",
+                    "session_key": "db"
+                }
+            }
+        )
+
+    Notes:
+        This TypedDict provides type safety for extension config.
+        Flask extension uses g object for request-scoped storage.
+    """
+
+    connection_key: NotRequired[str]
+    """Key for storing connection in Flask g object. Default: auto-generated from session_key."""
+
+    session_key: NotRequired[str]
+    """Key for accessing session via plugin.get_session(). Default: 'db_session'."""
+
+    commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
+    """Transaction commit mode. Default: 'manual'.
+    - manual: No automatic commits, user handles explicitly
+    - autocommit: Commits on 2xx status, rollback otherwise
+    - autocommit_include_redirect: Commits on 2xx-3xx status, rollback otherwise
+    """
+
+    extra_commit_statuses: NotRequired[set[int]]
+    """Additional HTTP status codes that trigger commit. Default: None."""
+
+    extra_rollback_statuses: NotRequired[set[int]]
+    """Additional HTTP status codes that trigger rollback. Default: None."""
+
+
 class LitestarConfig(TypedDict):
     """Configuration options for Litestar SQLSpec plugin.
 
@@ -126,6 +171,61 @@ class LitestarConfig(TypedDict):
     """Additional HTTP status codes that trigger rollback. Default: set()"""
 
 
+class StarletteConfig(TypedDict):
+    """Configuration options for Starlette and FastAPI extensions.
+
+    All fields are optional with sensible defaults. Use in extension_config["starlette"]:
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+
+        config = AsyncpgConfig(
+            pool_config={"dsn": "postgresql://localhost/mydb"},
+            extension_config={
+                "starlette": {
+                    "commit_mode": "autocommit",
+                    "session_key": "db"
+                }
+            }
+        )
+
+    Notes:
+        Both Starlette and FastAPI extensions use the "starlette" key.
+        This TypedDict provides type safety for extension config.
+    """
+
+    connection_key: NotRequired[str]
+    """Key for storing connection in request.state. Default: 'db_connection'"""
+
+    pool_key: NotRequired[str]
+    """Key for storing connection pool in app.state. Default: 'db_pool'"""
+
+    session_key: NotRequired[str]
+    """Key for storing session in request.state. Default: 'db_session'"""
+
+    commit_mode: NotRequired[Literal["manual", "autocommit", "autocommit_include_redirect"]]
+    """Transaction commit mode. Default: 'manual'
+
+    - manual: No automatic commit/rollback
+    - autocommit: Commit on 2xx, rollback otherwise
+    - autocommit_include_redirect: Commit on 2xx-3xx, rollback otherwise
+    """
+
+    extra_commit_statuses: NotRequired[set[int]]
+    """Additional HTTP status codes that trigger commit. Default: set()
+
+    Example:
+        extra_commit_statuses={201, 202}
+    """
+
+    extra_rollback_statuses: NotRequired[set[int]]
+    """Additional HTTP status codes that trigger rollback. Default: set()
+
+    Example:
+        extra_rollback_statuses={409}
+    """
+
+
 class ADKConfig(TypedDict):
     """Configuration options for ADK session store extension.
 
@@ -356,8 +456,6 @@ class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
         migration_config = self.migration_config or {}
         script_location = migration_config.get("script_location", "migrations")
 
-        from pathlib import Path
-
         migration_path = Path(script_location)
         if migration_path.exists() and not self._migration_loader.list_files():
             self._migration_loader.load_sql(migration_path)

sqlspec/core/filters.py

@@ -576,10 +576,13 @@ class LimitOffsetFilter(PaginationFilter):
         limit_placeholder = exp.Placeholder(this=limit_param_name)
         offset_placeholder = exp.Placeholder(this=offset_param_name)
 
-        try:
-            current_statement = sqlglot.parse_one(statement.raw_sql, dialect=statement.dialect)
-        except Exception:
-            current_statement = exp.Select().from_(f"({statement.raw_sql})")
+        if statement.statement_expression is not None:
+            current_statement = statement.statement_expression
+        else:
+            try:
+                current_statement = sqlglot.parse_one(statement.raw_sql, dialect=statement.dialect)
+            except Exception:
+                current_statement = exp.Select().from_(f"({statement.raw_sql})")
 
         if isinstance(current_statement, exp.Select):
             new_statement = current_statement.limit(limit_placeholder).offset(offset_placeholder)
@@ -587,7 +590,6 @@ class LimitOffsetFilter(PaginationFilter):
             new_statement = exp.Select().from_(current_statement).limit(limit_placeholder).offset(offset_placeholder)
 
         result = statement.copy(statement=new_statement)
-
         result = result.add_named_parameter(limit_param_name, self.limit)
         return result.add_named_parameter(offset_param_name, self.offset)
 
@@ -628,12 +630,18 @@ class OrderByFilter(StatementFilter):
         col_expr = exp.column(self.field_name)
         order_expr = col_expr.desc() if converted_sort_order == "desc" else col_expr.asc()
 
-        if statement.statement_expression is None:
-            new_statement = exp.Select().order_by(order_expr)
-        elif isinstance(statement.statement_expression, exp.Select):
-            new_statement = statement.statement_expression.order_by(order_expr)
+        if statement.statement_expression is not None:
+            current_statement = statement.statement_expression
+        else:
+            try:
+                current_statement = sqlglot.parse_one(statement.raw_sql, dialect=statement.dialect)
+            except Exception:
+                current_statement = exp.Select().from_(f"({statement.raw_sql})")
+
+        if isinstance(current_statement, exp.Select):
+            new_statement = current_statement.order_by(order_expr)
         else:
-            new_statement = exp.Select().from_(statement.statement_expression).order_by(order_expr)
+            new_statement = exp.Select().from_(current_statement).order_by(order_expr)
 
         return statement.copy(statement=new_statement)
 

sqlspec/core/result.py

@@ -16,13 +16,14 @@ from mypy_extensions import mypyc_attr
 from typing_extensions import TypeVar
 
 from sqlspec.core.compiler import OperationType
+from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow
 from sqlspec.utils.schema import to_schema
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     from sqlspec.core.statement import SQL
-    from sqlspec.typing import SchemaT
+    from sqlspec.typing import ArrowTable, PandasDataFrame, PolarsDataFrame, SchemaT
 
 
 __all__ = ("ArrowResult", "SQLResult", "StatementResult")
@@ -618,7 +619,7 @@ class ArrowResult(StatementResult):
         """
         return self.data is not None
 
-    def get_data(self) -> Any:
+    def get_data(self) -> "ArrowTable":
         """Get the Apache Arrow Table from the result.
 
         Returns:
@@ -626,10 +627,19 @@ class ArrowResult(StatementResult):
 
         Raises:
             ValueError: If no Arrow table is available.
+            TypeError: If data is not an Arrow Table.
         """
         if self.data is None:
             msg = "No Arrow table available for this result"
             raise ValueError(msg)
+
+        ensure_pyarrow()
+
+        import pyarrow as pa
+
+        if not isinstance(self.data, pa.Table):
+            msg = f"Expected an Arrow Table, but got {type(self.data).__name__}"
+            raise TypeError(msg)
         return self.data
 
     @property
@@ -680,6 +690,127 @@ class ArrowResult(StatementResult):
 
         return cast("int", self.data.num_columns)
 
+    def to_pandas(self) -> "PandasDataFrame":
+        """Convert Arrow data to pandas DataFrame.
+
+        Returns:
+            pandas DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_pandas()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_pandas()
+
+        import pandas as pd
+
+        result = self.data.to_pandas()
+        if not isinstance(result, pd.DataFrame):
+            msg = f"Expected a pandas DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_polars(self) -> "PolarsDataFrame":
+        """Convert Arrow data to Polars DataFrame.
+
+        Returns:
+            Polars DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_polars()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_polars()
+
+        import polars as pl
+
+        result = pl.from_arrow(self.data)
+        if not isinstance(result, pl.DataFrame):
+            msg = f"Expected a Polars DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_dict(self) -> "list[dict[str, Any]]":
+        """Convert Arrow data to list of dictionaries.
+
+        Returns:
+            List of dictionaries, one per row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> rows = result.to_dict()
+            >>> print(rows[0])
+            {'id': 1, 'name': 'Alice'}
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("list[dict[str, Any]]", self.data.to_pylist())
+
+    def __len__(self) -> int:
+        """Return number of rows in the Arrow table.
+
+        Returns:
+            Number of rows.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> print(len(result))
+            100
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("int", self.data.num_rows)
+
+    def __iter__(self) -> "Iterator[dict[str, Any]]":
+        """Iterate over rows as dictionaries.
+
+        Yields:
+            Dictionary for each row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> for row in result:
+            ...     print(row["name"])
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        yield from self.data.to_pylist()
+
 
 def create_sql_result(
     statement: "SQL",

sqlspec/driver/_async.py

@@ -4,6 +4,7 @@ from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Final, TypeVar, overload
 
 from sqlspec.core import SQL, Statement
+from sqlspec.core.result import create_arrow_result
 from sqlspec.driver._common import (
     CommonDriverAttributesMixin,
     DataDictionaryMixin,
@@ -12,7 +13,10 @@ from sqlspec.driver._common import (
     handle_single_row_error,
 )
 from sqlspec.driver.mixins import SQLTranslatorMixin
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.utils.arrow_helpers import convert_dict_to_arrow
 from sqlspec.utils.logging import get_logger
+from sqlspec.utils.module_loader import ensure_pyarrow
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -341,6 +345,91 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin):
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
         return result.get_data(schema_type=schema_type)
 
+    async def select_to_arrow(
+        self,
+        statement: "Statement | QueryBuilder",
+        /,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
+        return_format: str = "table",
+        native_only: bool = False,
+        batch_size: int | None = None,
+        arrow_schema: Any = None,
+        **kwargs: Any,
+    ) -> "Any":
+        """Execute query and return results as Apache Arrow format (async).
+
+        This base implementation uses the conversion path: execute() → dict → Arrow.
+        Adapters with native Arrow support (ADBC, DuckDB, BigQuery) override this
+        method to use zero-copy native paths for 5-10x performance improvement.
+
+        Args:
+            statement: SQL query string, Statement, or QueryBuilder
+            *parameters: Query parameters (same format as execute()/select())
+            statement_config: Optional statement configuration override
+            return_format: "table" for pyarrow.Table (default), "reader" for RecordBatchReader,
+                         "batches" for iterator of RecordBatches
+            native_only: If True, raise error if native Arrow unavailable (default: False)
+            batch_size: Rows per batch for "batches" format (default: None = all rows)
+            arrow_schema: Optional pyarrow.Schema for type casting
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ArrowResult containing pyarrow.Table, RecordBatchReader, or RecordBatches
+
+        Raises:
+            ImproperConfigurationError: If native_only=True and adapter doesn't support native Arrow
+
+        Examples:
+            >>> result = await driver.select_to_arrow(
+            ...     "SELECT * FROM users WHERE age > ?", 18
+            ... )
+            >>> df = result.to_pandas()
+            >>> print(df.head())
+
+            >>> # Force native Arrow path (raises error if unavailable)
+            >>> result = await driver.select_to_arrow(
+            ...     "SELECT * FROM users", native_only=True
+            ... )
+        """
+        # Check pyarrow is available
+        ensure_pyarrow()
+
+        # Check if native_only requested but not supported
+        if native_only:
+            msg = (
+                f"Adapter '{self.__class__.__name__}' does not support native Arrow results. "
+                f"Use native_only=False to allow conversion path, or switch to an adapter "
+                f"with native Arrow support (ADBC, DuckDB, BigQuery)."
+            )
+            raise ImproperConfigurationError(msg)
+
+        # Execute query using standard path
+        result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+
+        # Convert dict results to Arrow
+        arrow_data = convert_dict_to_arrow(
+            result.data,
+            return_format=return_format,  # type: ignore[arg-type]
+            batch_size=batch_size,
+        )
+        if arrow_schema is not None:
+            import pyarrow as pa
+
+            if not isinstance(arrow_schema, pa.Schema):
+                msg = f"arrow_schema must be a pyarrow.Schema, got {type(arrow_schema).__name__}"
+                raise TypeError(msg)
+
+            arrow_data = arrow_data.cast(arrow_schema)
+        return create_arrow_result(
+            statement=result.statement,
+            data=arrow_data,
+            rows_affected=result.rows_affected,
+            last_inserted_id=result.last_inserted_id,
+            execution_time=result.execution_time,
+            metadata=result.metadata,
+        )
+
     async def select_value(
         self,
         statement: "Statement | QueryBuilder",