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

sqlspec/_typing.py

@@ -377,18 +377,101 @@ class ArrowRecordBatchResult(Protocol):
         return None
 
 
+@runtime_checkable
+class ArrowSchemaProtocol(Protocol):
+    """Typed shim for pyarrow.Schema."""
+
+    def field(self, i: int) -> Any:
+        """Get field by index."""
+        ...
+
+    @property
+    def names(self) -> "list[str]":
+        """Get list of field names."""
+        ...
+
+    def __len__(self) -> int:
+        """Get number of fields."""
+        return 0
+
+
+@runtime_checkable
+class ArrowRecordBatchReaderProtocol(Protocol):
+    """Typed shim for pyarrow.RecordBatchReader."""
+
+    def read_all(self) -> Any:
+        """Read all batches into a table."""
+        ...
+
+    def read_next_batch(self) -> Any:
+        """Read next batch."""
+        ...
+
+    def __iter__(self) -> "Iterable[Any]":
+        """Iterate over batches."""
+        ...
+
+
 try:
     from pyarrow import RecordBatch as ArrowRecordBatch
+    from pyarrow import RecordBatchReader as ArrowRecordBatchReader
+    from pyarrow import Schema as ArrowSchema
     from pyarrow import Table as ArrowTable
 
     PYARROW_INSTALLED = True
 except ImportError:
     ArrowTable = ArrowTableResult  # type: ignore[assignment,misc]
     ArrowRecordBatch = ArrowRecordBatchResult  # type: ignore[assignment,misc]
+    ArrowSchema = ArrowSchemaProtocol  # type: ignore[assignment,misc]
+    ArrowRecordBatchReader = ArrowRecordBatchReaderProtocol  # type: ignore[assignment,misc]
 
     PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
 
 
+@runtime_checkable
+class PandasDataFrameProtocol(Protocol):
+    """Typed shim for pandas.DataFrame."""
+
+    def __len__(self) -> int:
+        """Get number of rows."""
+        ...
+
+    def __getitem__(self, key: Any) -> Any:
+        """Get column or row."""
+        ...
+
+
+@runtime_checkable
+class PolarsDataFrameProtocol(Protocol):
+    """Typed shim for polars.DataFrame."""
+
+    def __len__(self) -> int:
+        """Get number of rows."""
+        ...
+
+    def __getitem__(self, key: Any) -> Any:
+        """Get column or row."""
+        ...
+
+
+try:
+    from pandas import DataFrame as PandasDataFrame
+
+    PANDAS_INSTALLED = True
+except ImportError:
+    PandasDataFrame = PandasDataFrameProtocol  # type: ignore[assignment,misc]
+    PANDAS_INSTALLED = False
+
+
+try:
+    from polars import DataFrame as PolarsDataFrame
+
+    POLARS_INSTALLED = True
+except ImportError:
+    PolarsDataFrame = PolarsDataFrameProtocol  # type: ignore[assignment,misc]
+    POLARS_INSTALLED = False
+
+
 @runtime_checkable
 class NumpyArrayStub(Protocol):
     """Protocol stub for numpy.ndarray when numpy is not installed.
@@ -639,7 +722,9 @@ __all__ = (
     "OBSTORE_INSTALLED",
     "OPENTELEMETRY_INSTALLED",
     "ORJSON_INSTALLED",
+    "PANDAS_INSTALLED",
     "PGVECTOR_INSTALLED",
+    "POLARS_INSTALLED",
     "PROMETHEUS_INSTALLED",
     "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
@@ -650,7 +735,11 @@ __all__ = (
     "AiosqlSQLOperationType",
     "AiosqlSyncProtocol",
     "ArrowRecordBatch",
+    "ArrowRecordBatchReader",
+    "ArrowRecordBatchReaderProtocol",
     "ArrowRecordBatchResult",
+    "ArrowSchema",
+    "ArrowSchemaProtocol",
     "ArrowTable",
     "ArrowTableResult",
     "AttrsInstance",
@@ -670,6 +759,10 @@ __all__ = (
     "Histogram",
     "NumpyArray",
     "NumpyArrayStub",
+    "PandasDataFrame",
+    "PandasDataFrameProtocol",
+    "PolarsDataFrame",
+    "PolarsDataFrameProtocol",
     "Span",
     "Status",
     "StatusCode",

sqlspec/adapters/adbc/adk/store.py

@@ -639,31 +639,41 @@ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
             finally:
                 cursor.close()  # type: ignore[no-untyped-call]
 
-    def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
-        """List all sessions for a user in an app.
+    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).
-        """
-        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
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
+        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)
 
         try:
             with self._config.provide_connection() as conn:
                 cursor = conn.cursor()
                 try:
-                    cursor.execute(sql, (app_name, user_id))
+                    cursor.execute(sql, params)
                     rows = cursor.fetchall()
 
                     return [

sqlspec/adapters/adbc/data_dictionary.py

@@ -300,22 +300,44 @@ class AdbcDataDictionary(SyncDataDictionaryBase):
                 for row in result.data or []
             ]
 
+        if dialect == "postgres":
+            schema_name = schema or "public"
+            sql = """
+                SELECT
+                    a.attname::text AS column_name,
+                    pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                    CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
+                    pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
+                FROM pg_catalog.pg_attribute a
+                JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+                JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+                LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+                WHERE c.relname = ?
+                    AND n.nspname = ?
+                    AND a.attnum > 0
+                    AND NOT a.attisdropped
+                ORDER BY a.attnum
+            """
+            result = adbc_driver.execute(sql, (table, schema_name))
+            return result.data or []
+
         if schema:
-            sql = f"""
+            sql = """
                 SELECT column_name, data_type, is_nullable, column_default
                 FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                WHERE table_name = ? AND table_schema = ?
                 ORDER BY ordinal_position
             """
+            result = adbc_driver.execute(sql, (table, schema))
         else:
-            sql = f"""
+            sql = """
                 SELECT column_name, data_type, is_nullable, column_default
                 FROM information_schema.columns
-                WHERE table_name = '{table}'
+                WHERE table_name = ?
                 ORDER BY ordinal_position
             """
+            result = adbc_driver.execute(sql, (table,))
 
-        result = adbc_driver.execute(sql)
         return result.data or []
 
     def list_available_features(self) -> "list[str]":

sqlspec/adapters/adbc/driver.py

@@ -15,6 +15,7 @@ from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary
 from sqlspec.adapters.adbc.type_converter import ADBCTypeConverter
 from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
+from sqlspec.core.result import create_arrow_result
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import SyncDriverAdapterBase
 from sqlspec.exceptions import (
@@ -23,7 +24,6 @@ from sqlspec.exceptions import (
     DataError,
     ForeignKeyViolationError,
     IntegrityError,
-    MissingDependencyError,
     NotNullViolationError,
     SQLParsingError,
     SQLSpecError,
@@ -32,6 +32,7 @@ from sqlspec.exceptions import (
 )
 from sqlspec.typing import Empty
 from sqlspec.utils.logging import get_logger
+from sqlspec.utils.module_loader import ensure_pyarrow
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
@@ -39,9 +40,12 @@ if TYPE_CHECKING:
     from adbc_driver_manager.dbapi import Cursor
 
     from sqlspec.adapters.adbc._types import AdbcConnection
-    from sqlspec.core.result import SQLResult
+    from sqlspec.builder import QueryBuilder
+    from sqlspec.core import Statement, StatementFilter
+    from sqlspec.core.result import ArrowResult, SQLResult
     from sqlspec.driver import ExecutionResult
     from sqlspec.driver._sync import SyncDataDictionaryBase
+    from sqlspec.typing import StatementParameters
 
 __all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "get_adbc_statement_config")
 
@@ -507,18 +511,6 @@ class AdbcDriver(SyncDriverAdapterBase):
         self.dialect = statement_config.dialect
         self._data_dictionary: SyncDataDictionaryBase | None = None
 
-    @staticmethod
-    def _ensure_pyarrow_installed() -> None:
-        """Ensure PyArrow is installed.
-
-        Raises:
-            MissingDependencyError: If PyArrow is not installed
-        """
-        from sqlspec.typing import PYARROW_INSTALLED
-
-        if not PYARROW_INSTALLED:
-            raise MissingDependencyError(package="pyarrow", install_package="arrow")
-
     @staticmethod
     def _get_dialect(connection: "AdbcConnection") -> str:
         """Detect database dialect from connection information.
@@ -863,3 +855,80 @@ class AdbcDriver(SyncDriverAdapterBase):
         if self._data_dictionary is None:
             self._data_dictionary = AdbcDataDictionary()
         return self._data_dictionary
+
+    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,
+    ) -> "ArrowResult":
+        """Execute query and return results as Apache Arrow (ADBC native path).
+
+        ADBC provides zero-copy Arrow support via cursor.fetch_arrow_table().
+        This is 5-10x faster than the conversion path for large datasets.
+
+        Args:
+            statement: SQL statement, string, or QueryBuilder
+            *parameters: Query parameters or filters
+            statement_config: Optional statement configuration override
+            return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch
+            native_only: Ignored for ADBC (always uses native path)
+            batch_size: Batch size hint (for future streaming implementation)
+            arrow_schema: Optional pyarrow.Schema for type casting
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ArrowResult with native Arrow data
+
+        Raises:
+            MissingDependencyError: If pyarrow not installed
+            SQLExecutionError: If query execution fails
+
+        Example:
+            >>> result = driver.select_to_arrow(
+            ...     "SELECT * FROM users WHERE age > $1", 18
+            ... )
+            >>> df = result.to_pandas()  # Fast zero-copy conversion
+        """
+        ensure_pyarrow()
+
+        import pyarrow as pa
+
+        # Prepare statement
+        config = statement_config or self.statement_config
+        prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
+
+        # Use ADBC cursor for native Arrow
+        with self.with_cursor(self.connection) as cursor, self.handle_database_exceptions():
+            if cursor is None:
+                msg = "Failed to create cursor"
+                raise DatabaseConnectionError(msg)
+
+            # Get compiled SQL and parameters
+            sql, driver_params = self._get_compiled_sql(prepared_statement, config)
+
+            # Execute query
+            cursor.execute(sql, driver_params or ())
+
+            # Fetch as Arrow table (zero-copy!)
+            arrow_table = cursor.fetch_arrow_table()
+
+            # Apply schema casting if requested
+            if arrow_schema is not None:
+                arrow_table = arrow_table.cast(arrow_schema)
+
+            # Convert to batch if requested
+            if return_format == "batch":
+                batches = arrow_table.to_batches()
+                arrow_data: Any = batches[0] if batches else pa.RecordBatch.from_pydict({})
+            else:
+                arrow_data = arrow_table
+
+        # Create ArrowResult
+        return create_arrow_result(statement=prepared_statement, data=arrow_data, rows_affected=arrow_data.num_rows)

sqlspec/adapters/aiosqlite/adk/store.py

@@ -136,7 +136,7 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]):
         """
         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:
@@ -163,7 +163,7 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]):
             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:
@@ -228,11 +228,10 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]):
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self._config.provide_connection() as conn:
-            await self._enable_foreign_keys(conn)
-            await conn.executescript(self._get_create_sessions_table_sql())
-            await conn.executescript(self._get_create_events_table_sql())
-            await conn.commit()
+        async with self._config.provide_session() as driver:
+            await self._enable_foreign_keys(driver.connection)
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(
@@ -343,29 +342,39 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]):
             await conn.execute(sql, (state_json, now_julian, session_id))
             await conn.commit()
 
-    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).
-        """
-        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
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
+        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)
 
         async with self._config.provide_connection() as conn:
             await self._enable_foreign_keys(conn)
-            cursor = await conn.execute(sql, (app_name, user_id))
+            cursor = await conn.execute(sql, params)
             rows = await cursor.fetchall()
 
             return [

sqlspec/adapters/asyncmy/adk/store.py

@@ -106,7 +106,7 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
 
         return (col_def, fk_constraint)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get MySQL CREATE TABLE SQL for sessions.
 
         Returns:
@@ -145,7 +145,7 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
         ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get MySQL CREATE TABLE SQL for events.
 
         Returns:
@@ -199,9 +199,9 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
-            await cursor.execute(self._get_create_sessions_table_sql())
-            await cursor.execute(self._get_create_events_table_sql())
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(
@@ -326,29 +326,39 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
             await cursor.execute(sql, (session_id,))
             await conn.commit()
 
-    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).
-        """
-        sql = f"""
-        SELECT id, app_name, user_id, state, create_time, update_time
-        FROM {self._session_table}
-        WHERE app_name = %s AND user_id = %s
-        ORDER BY update_time DESC
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
+        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 = %s
+            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 = %s AND user_id = %s
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
 
         try:
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
-                await cursor.execute(sql, (app_name, user_id))
+                await cursor.execute(sql, params)
                 rows = await cursor.fetchall()
 
                 return [

sqlspec/adapters/asyncpg/adk/store.py

@@ -84,7 +84,7 @@ class AsyncpgADKStore(BaseAsyncADKStore[AsyncConfigT]):
         """
         super().__init__(config)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get PostgreSQL CREATE TABLE SQL for sessions.
 
         Returns:
@@ -125,7 +125,7 @@ class AsyncpgADKStore(BaseAsyncADKStore[AsyncConfigT]):
             WHERE state != '{{}}'::jsonb;
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get PostgreSQL CREATE TABLE SQL for events.
 
         Returns:
@@ -181,9 +181,9 @@ class AsyncpgADKStore(BaseAsyncADKStore[AsyncConfigT]):
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self.config.provide_connection() as conn:
-            await conn.execute(self._get_create_sessions_table_sql())
-            await conn.execute(self._get_create_events_table_sql())
+        async with self.config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(
@@ -294,29 +294,39 @@ class AsyncpgADKStore(BaseAsyncADKStore[AsyncConfigT]):
         async with self.config.provide_connection() as conn:
             await conn.execute(sql, session_id)
 
-    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).
-        """
-        sql = f"""
-        SELECT id, app_name, user_id, state, create_time, update_time
-        FROM {self._session_table}
-        WHERE app_name = $1 AND user_id = $2
-        ORDER BY update_time DESC
+            Uses composite index on (app_name, user_id) when user_id is provided.
         """
+        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 = $1
+            ORDER BY update_time DESC
+            """
+            params = [app_name]
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = $1 AND user_id = $2
+            ORDER BY update_time DESC
+            """
+            params = [app_name, user_id]
 
         try:
             async with self.config.provide_connection() as conn:
-                rows = await conn.fetch(sql, app_name, user_id)
+                rows = await conn.fetch(sql, *params)
 
                 return [
                     SessionRecord(