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

sqlspec/driver/_async.py

@@ -1,14 +1,18 @@
 """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.driver._common import (
+    CommonDriverAttributesMixin,
+    DataDictionaryMixin,
+    ExecutionResult,
+    VersionInfo,
+    handle_single_row_error,
+)
+from sqlspec.driver.mixins import SQLTranslatorMixin
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.type_guards import is_dict_row, is_indexable_row
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -16,7 +20,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 +33,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 +100,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 +173,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 +187,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 +210,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 +229,124 @@ 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_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 +356,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 +375,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 +431,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 +472,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 +486,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 +502,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.