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

sqlspec/driver/_async.py

@@ -1,14 +1,22 @@
 """Asynchronous driver protocol implementation."""
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Final, NoReturn, Optional, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Final, TypeVar, overload
 
 from sqlspec.core import SQL, Statement
-from sqlspec.driver._common import CommonDriverAttributesMixin, DataDictionaryMixin, ExecutionResult, VersionInfo
-from sqlspec.driver.mixins import SQLTranslatorMixin, ToSchemaMixin
-from sqlspec.exceptions import NotFoundError
+from sqlspec.core.result import create_arrow_result
+from sqlspec.driver._common import (
+    CommonDriverAttributesMixin,
+    DataDictionaryMixin,
+    ExecutionResult,
+    VersionInfo,
+    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.type_guards import is_dict_row, is_indexable_row
+from sqlspec.utils.module_loader import ensure_pyarrow
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -16,7 +24,7 @@ if TYPE_CHECKING:
 
     from sqlspec.builder import QueryBuilder
     from sqlspec.core import SQLResult, StatementConfig, StatementFilter
-    from sqlspec.typing import ModelDTOT, StatementParameters
+    from sqlspec.typing import SchemaT, StatementParameters
 
 _LOGGER_NAME: Final[str] = "sqlspec"
 logger = get_logger(_LOGGER_NAME)
@@ -29,7 +37,7 @@ EMPTY_FILTERS: Final["list[StatementFilter]"] = []
 AsyncDriverT = TypeVar("AsyncDriverT", bound="AsyncDriverAdapterBase")
 
 
-class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, ToSchemaMixin):
+class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin):
     """Base class for asynchronous database drivers."""
 
     __slots__ = ()
@@ -96,7 +104,7 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
         """Commit the current transaction on the current connection."""
 
     @abstractmethod
-    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
+    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None":
         """Hook for database-specific special operations (e.g., PostgreSQL COPY, bulk operations).
 
         This method is called first in dispatch_statement_execution() to allow drivers to handle
@@ -169,10 +177,10 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
 
     async def execute(
         self,
-        statement: "Union[SQL, Statement, QueryBuilder]",
+        statement: "SQL | Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "SQLResult":
         """Execute a statement with parameter handling."""
@@ -183,11 +191,11 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
 
     async def execute_many(
         self,
-        statement: "Union[SQL, Statement, QueryBuilder]",
+        statement: "SQL | Statement | QueryBuilder",
         /,
         parameters: "Sequence[StatementParameters]",
-        *filters: "Union[StatementParameters, StatementFilter]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *filters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "SQLResult":
         """Execute statement multiple times with different parameters.
@@ -206,10 +214,10 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
 
     async def execute_script(
         self,
-        statement: "Union[str, SQL]",
+        statement: "str | SQL",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "SQLResult":
         """Execute a multi-statement script.
@@ -225,140 +233,209 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
     @overload
     async def select_one(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "type[ModelDTOT]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT]",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "ModelDTOT": ...
+    ) -> "SchemaT": ...
 
     @overload
     async def select_one(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
+        *parameters: "StatementParameters | StatementFilter",
         schema_type: None = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
 
     async def select_one(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT] | None" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "Union[dict[str, Any], ModelDTOT]":
+    ) -> "SchemaT | dict[str, Any]":
         """Execute a select statement and return exactly one row.
 
         Raises an exception if no rows or more than one row is returned.
         """
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
-        data = result.get_data()
-        data_len: int = len(data)
-        if data_len == 0:
-            self._raise_no_rows_found()
-        if data_len > 1:
-            self._raise_expected_one_row(data_len)
-        first_row = data[0]
-        return self.to_schema(first_row, schema_type=schema_type) if schema_type else first_row
+        try:
+            return result.one(schema_type=schema_type)
+        except ValueError as error:
+            handle_single_row_error(error)
 
     @overload
     async def select_one_or_none(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "type[ModelDTOT]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT]",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "Optional[ModelDTOT]": ...
+    ) -> "SchemaT | None": ...
 
     @overload
     async def select_one_or_none(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
+        *parameters: "StatementParameters | StatementFilter",
         schema_type: None = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "Optional[dict[str, Any]]": ...
+    ) -> "dict[str, Any] | None": ...
 
     async def select_one_or_none(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT] | None" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+    ) -> "SchemaT | dict[str, Any] | None":
         """Execute a select statement and return at most one row.
 
         Returns None if no rows are found.
         Raises an exception if more than one row is returned.
         """
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
-        data = result.get_data()
-        data_len: int = len(data)
-        if data_len == 0:
-            return None
-        if data_len > 1:
-            self._raise_expected_at_most_one_row(data_len)
-        first_row = data[0]
-        return cast(
-            "Optional[Union[dict[str, Any], ModelDTOT]]",
-            self.to_schema(first_row, schema_type=schema_type) if schema_type else first_row,
-        )
+        return result.one_or_none(schema_type=schema_type)
 
     @overload
     async def select(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "type[ModelDTOT]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT]",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "list[ModelDTOT]": ...
+    ) -> "list[SchemaT]": ...
 
     @overload
     async def select(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
+        *parameters: "StatementParameters | StatementFilter",
         schema_type: None = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "list[dict[str, Any]]": ...
 
     async def select(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT] | None" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "Union[list[dict[str, Any]], list[ModelDTOT]]":
+    ) -> "list[SchemaT] | list[dict[str, Any]]":
         """Execute a select statement and return all rows."""
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
-        return cast(
-            "Union[list[dict[str, Any]], list[ModelDTOT]]", self.to_schema(result.get_data(), schema_type=schema_type)
+        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: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> Any:
         """Execute a select statement and return a single scalar value.
@@ -368,28 +445,16 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
         """
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
         try:
-            row = result.one()
-        except ValueError as e:
-            self._raise_no_rows_found_from_exception(e)
-        if not row:
-            self._raise_no_rows_found()
-        if is_dict_row(row):
-            if not row:
-                self._raise_row_no_columns()
-            return next(iter(row.values()))
-        if is_indexable_row(row):
-            if not row:
-                self._raise_row_no_columns()
-            return row[0]
-        self._raise_unexpected_row_type(type(row))
-        return None
+            return result.scalar()
+        except ValueError as error:
+            handle_single_row_error(error)
 
     async def select_value_or_none(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> Any:
         """Execute a select statement and return a single scalar value or None.
@@ -399,53 +464,39 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
         Raises an exception if more than one row is returned.
         """
         result = await self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
-        data = result.get_data()
-        data_len: int = len(data)
-        if data_len == 0:
-            return None
-        if data_len > 1:
-            self._raise_expected_at_most_one_row(data_len)
-        row = data[0]
-        if is_dict_row(row):
-            if not row:
-                return None
-            return next(iter(row.values()))
-        if is_indexable_row(row):
-            return row[0]
-        self._raise_cannot_extract_value_from_row_type(type(row).__name__)
-        return None
+        return result.scalar_or_none()
 
     @overload
     async def select_with_total(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "type[ModelDTOT]",
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT]",
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "tuple[list[ModelDTOT], int]": ...
+    ) -> "tuple[list[SchemaT], int]": ...
 
     @overload
     async def select_with_total(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
+        *parameters: "StatementParameters | StatementFilter",
         schema_type: None = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
     ) -> "tuple[list[dict[str, Any]], int]": ...
 
     async def select_with_total(
         self,
-        statement: "Union[Statement, QueryBuilder]",
+        statement: "Statement | QueryBuilder",
         /,
-        *parameters: "Union[StatementParameters, StatementFilter]",
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        *parameters: "StatementParameters | StatementFilter",
+        schema_type: "type[SchemaT] | None" = None,
+        statement_config: "StatementConfig | None" = None,
         **kwargs: Any,
-    ) -> "tuple[Union[list[dict[str, Any]], list[ModelDTOT]], int]":
+    ) -> "tuple[list[SchemaT] | list[dict[str, Any]], int]":
         """Execute a select statement and return both the data and total count.
 
         This method is designed for pagination scenarios where you need both
@@ -469,42 +520,14 @@ class AsyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, To
         count_result = await self.dispatch_statement_execution(self._create_count_query(sql_statement), self.connection)
         select_result = await self.execute(sql_statement)
 
-        return (self.to_schema(select_result.get_data(), schema_type=schema_type), count_result.scalar())
-
-    def _raise_no_rows_found(self) -> NoReturn:
-        msg = "No rows found"
-        raise NotFoundError(msg)
-
-    def _raise_no_rows_found_from_exception(self, e: ValueError) -> NoReturn:
-        msg = "No rows found"
-        raise NotFoundError(msg) from e
-
-    def _raise_expected_one_row(self, data_len: int) -> NoReturn:
-        msg = f"Expected exactly one row, found {data_len}"
-        raise ValueError(msg)
-
-    def _raise_expected_at_most_one_row(self, data_len: int) -> NoReturn:
-        msg = f"Expected at most one row, found {data_len}"
-        raise ValueError(msg)
-
-    def _raise_row_no_columns(self) -> NoReturn:
-        msg = "Row has no columns"
-        raise ValueError(msg)
-
-    def _raise_unexpected_row_type(self, row_type: type) -> NoReturn:
-        msg = f"Unexpected row type: {row_type}"
-        raise ValueError(msg)
-
-    def _raise_cannot_extract_value_from_row_type(self, type_name: str) -> NoReturn:
-        msg = f"Cannot extract value from row type {type_name}"
-        raise TypeError(msg)
+        return (select_result.get_data(schema_type=schema_type), count_result.scalar())
 
 
 class AsyncDataDictionaryBase(DataDictionaryMixin):
     """Base class for asynchronous data dictionary implementations."""
 
     @abstractmethod
-    async def get_version(self, driver: "AsyncDriverAdapterBase") -> "Optional[VersionInfo]":
+    async def get_version(self, driver: "AsyncDriverAdapterBase") -> "VersionInfo | None":
         """Get database version information.
 
         Args:
@@ -538,7 +561,7 @@ class AsyncDataDictionaryBase(DataDictionaryMixin):
             Database-specific type name
         """
 
-    async def get_tables(self, driver: "AsyncDriverAdapterBase", schema: "Optional[str]" = None) -> "list[str]":
+    async def get_tables(self, driver: "AsyncDriverAdapterBase", schema: "str | None" = None) -> "list[str]":
         """Get list of tables in schema.
 
         Args:
@@ -552,7 +575,7 @@ class AsyncDataDictionaryBase(DataDictionaryMixin):
         return []
 
     async def get_columns(
-        self, driver: "AsyncDriverAdapterBase", table: str, schema: "Optional[str]" = None
+        self, driver: "AsyncDriverAdapterBase", table: str, schema: "str | None" = None
     ) -> "list[dict[str, Any]]":
         """Get column information for a table.
 
@@ -568,7 +591,7 @@ class AsyncDataDictionaryBase(DataDictionaryMixin):
         return []
 
     async def get_indexes(
-        self, driver: "AsyncDriverAdapterBase", table: str, schema: "Optional[str]" = None
+        self, driver: "AsyncDriverAdapterBase", table: str, schema: "str | None" = None
     ) -> "list[dict[str, Any]]":
         """Get index information for a table.