sqlspec 0.9.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

sqlspec/base.py

@@ -1,7 +1,9 @@
 # ruff: noqa: PLR6301
+import atexit
+import contextlib
 import re
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass, field
 from typing import (
     TYPE_CHECKING,
@@ -18,15 +20,14 @@ from typing import (
 
 from sqlspec.exceptions import NotFoundError
 from sqlspec.statement import SQLStatement
-from sqlspec.typing import ModelDTOT, StatementParameterType
+from sqlspec.typing import ConnectionT, ModelDTOT, PoolT, StatementParameterType, T
+from sqlspec.utils.sync_tools import maybe_async_
 
 if TYPE_CHECKING:
     from contextlib import AbstractAsyncContextManager, AbstractContextManager
 
-    from pyarrow import Table as ArrowTable
 
 __all__ = (
-    "AsyncArrowBulkOperationsMixin",
     "AsyncDatabaseConfig",
     "AsyncDriverAdapterProtocol",
     "CommonDriverAttributes",
@@ -36,15 +37,10 @@ __all__ = (
     "NoPoolSyncConfig",
     "SQLSpec",
     "SQLStatement",
-    "SyncArrowBulkOperationsMixin",
     "SyncDatabaseConfig",
     "SyncDriverAdapterProtocol",
 )
 
-T = TypeVar("T")
-ConnectionT = TypeVar("ConnectionT")
-PoolT = TypeVar("PoolT")
-PoolT_co = TypeVar("PoolT_co", covariant=True)
 AsyncConfigT = TypeVar("AsyncConfigT", bound="Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]]")
 SyncConfigT = TypeVar("SyncConfigT", bound="Union[SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]")
 ConfigT = TypeVar(
@@ -202,6 +198,15 @@ class SQLSpec:
 
     def __init__(self) -> None:
         self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
+        # Register the cleanup handler to run at program exit
+        atexit.register(self._cleanup_pools)
+
+    def _cleanup_pools(self) -> None:
+        """Clean up all open database pools at program exit."""
+        for config in self._configs.values():
+            if config.support_connection_pooling and config.pool_instance is not None:
+                with contextlib.suppress(Exception):
+                    maybe_async_(config.close_pool)()
 
     @overload
     def add_config(self, config: "SyncConfigT") -> "type[SyncConfigT]": ...
@@ -284,6 +289,24 @@ class SQLSpec:
         config = self.get_config(name)
         return config.create_connection()
 
+    @overload
+    def get_session(
+        self,
+        name: Union[
+            "type[NoPoolSyncConfig[ConnectionT, DriverT]]",
+            "type[SyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+    ) -> "DriverT": ...
+
+    @overload
+    def get_session(
+        self,
+        name: Union[
+            "type[NoPoolAsyncConfig[ConnectionT, DriverT]]",
+            "type[AsyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+    ) -> "Awaitable[DriverT]": ...
+
     def get_session(
         self,
         name: Union[
@@ -311,6 +334,28 @@ class SQLSpec:
             return _create_session()
         return cast("DriverT", config.driver_type(connection))  # pyright: ignore
 
+    @overload
+    def provide_connection(
+        self,
+        name: Union[
+            "type[NoPoolSyncConfig[ConnectionT, DriverT]]",
+            "type[SyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+        *args: Any,
+        **kwargs: Any,
+    ) -> "AbstractContextManager[ConnectionT]": ...
+
+    @overload
+    def provide_connection(
+        self,
+        name: Union[
+            "type[NoPoolAsyncConfig[ConnectionT, DriverT]]",
+            "type[AsyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+        *args: Any,
+        **kwargs: Any,
+    ) -> "AbstractAsyncContextManager[ConnectionT]": ...
+
     def provide_connection(
         self,
         name: Union[
@@ -335,6 +380,28 @@ class SQLSpec:
         config = self.get_config(name)
         return config.provide_connection(*args, **kwargs)
 
+    @overload
+    def provide_session(
+        self,
+        name: Union[
+            "type[NoPoolSyncConfig[ConnectionT, DriverT]]",
+            "type[SyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+        *args: Any,
+        **kwargs: Any,
+    ) -> "AbstractContextManager[DriverT]": ...
+
+    @overload
+    def provide_session(
+        self,
+        name: Union[
+            "type[NoPoolAsyncConfig[ConnectionT, DriverT]]",
+            "type[AsyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+        *args: Any,
+        **kwargs: Any,
+    ) -> "AbstractAsyncContextManager[DriverT]": ...
+
     def provide_session(
         self,
         name: Union[
@@ -393,6 +460,24 @@ class SQLSpec:
             return cast("Union[type[PoolT], Awaitable[type[PoolT]]]", config.create_pool())
         return None
 
+    @overload
+    def close_pool(
+        self,
+        name: Union[
+            "type[NoPoolSyncConfig[ConnectionT, DriverT]]",
+            "type[SyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+    ) -> "None": ...
+
+    @overload
+    def close_pool(
+        self,
+        name: Union[
+            "type[NoPoolAsyncConfig[ConnectionT, DriverT]]",
+            "type[AsyncDatabaseConfig[ConnectionT, PoolT, DriverT]]",
+        ],
+    ) -> "Awaitable[None]": ...
+
     def close_pool(
         self,
         name: Union[
@@ -466,40 +551,37 @@ class CommonDriverAttributes(Generic[ConnectionT]):
         return stmt.process()
 
 
-class SyncArrowBulkOperationsMixin(Generic[ConnectionT]):
-    """Mixin for sync drivers supporting bulk Apache Arrow operations."""
+class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
+    connection: "ConnectionT"
 
-    __supports_arrow__: "ClassVar[bool]" = True
+    def __init__(self, connection: "ConnectionT", **kwargs: Any) -> None:
+        self.connection = connection
 
+    @overload
     @abstractmethod
-    def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+    def select(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
         connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
         **kwargs: Any,
-    ) -> "ArrowTable":  # pyright: ignore[reportUnknownReturnType]
-        """Execute a SQL query and return results as an Apache Arrow Table.
+    ) -> "Sequence[dict[str, Any]]": ...
 
-        Args:
-            sql: The SQL query string.
-            parameters: Parameters for the query.
-            connection: Optional connection override.
-            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
-
-        Returns:
-            An Apache Arrow Table containing the query results.
-        """
-        raise NotImplementedError
-
-
-class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
-    connection: "ConnectionT"
-
-    def __init__(self, connection: "ConnectionT", **kwargs: Any) -> None:
-        self.connection = connection
+    @overload
+    @abstractmethod
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
 
     @abstractmethod
     def select(
@@ -511,7 +593,33 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         connection: "Optional[ConnectionT]" = None,
         schema_type: Optional[type[ModelDTOT]] = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]": ...
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]": ...
+
+    @overload
+    @abstractmethod
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+
+    @overload
+    @abstractmethod
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
 
     @abstractmethod
     def select_one(
@@ -525,6 +633,32 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]": ...
 
+    @overload
+    @abstractmethod
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+
+    @overload
+    @abstractmethod
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
+
     @abstractmethod
     def select_one_or_none(
         self,
@@ -537,6 +671,32 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]": ...
 
+    @overload
+    @abstractmethod
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+
+    @overload
+    @abstractmethod
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+
     @abstractmethod
     def select_value(
         self,
@@ -547,7 +707,33 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         connection: Optional[ConnectionT] = None,
         schema_type: Optional[type[T]] = None,
         **kwargs: Any,
-    ) -> "Union[Any, T]": ...
+    ) -> "Union[T, Any]": ...
+
+    @overload
+    @abstractmethod
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+
+    @overload
+    @abstractmethod
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
 
     @abstractmethod
     def select_value_or_none(
@@ -559,7 +745,7 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         connection: Optional[ConnectionT] = None,
         schema_type: Optional[type[T]] = None,
         **kwargs: Any,
-    ) -> "Optional[Union[Any, T]]": ...
+    ) -> "Optional[Union[T, Any]]": ...
 
     @abstractmethod
     def insert_update_delete(
@@ -572,6 +758,32 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         **kwargs: Any,
     ) -> int: ...
 
+    @overload
+    @abstractmethod
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+
+    @overload
+    @abstractmethod
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+
     @abstractmethod
     def insert_update_delete_returning(
         self,
@@ -582,7 +794,7 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
         connection: Optional[ConnectionT] = None,
         schema_type: Optional[type[ModelDTOT]] = None,
         **kwargs: Any,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]": ...
+    ) -> "Union[ModelDTOT, dict[str, Any]]": ...
 
     @abstractmethod
     def execute_script(
@@ -596,40 +808,37 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
     ) -> str: ...
 
 
-class AsyncArrowBulkOperationsMixin(Generic[ConnectionT]):
-    """Mixin for async drivers supporting bulk Apache Arrow operations."""
+class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
+    connection: "ConnectionT"
 
-    __supports_arrow__: "ClassVar[bool]" = True
+    def __init__(self, connection: "ConnectionT") -> None:
+        self.connection = connection
 
+    @overload
     @abstractmethod
-    async def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+    async def select(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
         connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
         **kwargs: Any,
-    ) -> "ArrowTable":  # pyright: ignore[reportUnknownReturnType]
-        """Execute a SQL query and return results as an Apache Arrow Table.
+    ) -> "Sequence[dict[str, Any]]": ...
 
-        Args:
-            sql: The SQL query string.
-            parameters: Parameters for the query.
-            connection: Optional connection override.
-            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
-
-        Returns:
-            An Apache Arrow Table containing the query results.
-        """
-        raise NotImplementedError
-
-
-class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
-    connection: "ConnectionT"
-
-    def __init__(self, connection: "ConnectionT") -> None:
-        self.connection = connection
+    @overload
+    @abstractmethod
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
 
     @abstractmethod
     async def select(
@@ -641,7 +850,33 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         connection: "Optional[ConnectionT]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]": ...
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]": ...
+
+    @overload
+    @abstractmethod
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+
+    @overload
+    @abstractmethod
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
 
     @abstractmethod
     async def select_one(
@@ -655,6 +890,32 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]": ...
 
+    @overload
+    @abstractmethod
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+
+    @overload
+    @abstractmethod
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
+
     @abstractmethod
     async def select_one_or_none(
         self,
@@ -667,6 +928,32 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]": ...
 
+    @overload
+    @abstractmethod
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+
+    @overload
+    @abstractmethod
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+
     @abstractmethod
     async def select_value(
         self,
@@ -677,7 +964,33 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         connection: "Optional[ConnectionT]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
-    ) -> "Union[Any, T]": ...
+    ) -> "Union[T, Any]": ...
+
+    @overload
+    @abstractmethod
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+
+    @overload
+    @abstractmethod
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
 
     @abstractmethod
     async def select_value_or_none(
@@ -689,7 +1002,7 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         connection: "Optional[ConnectionT]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[Any, T]]": ...
+    ) -> "Optional[Union[T, Any]]": ...
 
     @abstractmethod
     async def insert_update_delete(
@@ -702,6 +1015,32 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         **kwargs: Any,
     ) -> int: ...
 
+    @overload
+    @abstractmethod
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+
+    @overload
+    @abstractmethod
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[ConnectionT]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+
     @abstractmethod
     async def insert_update_delete_returning(
         self,
@@ -712,7 +1051,7 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
         connection: "Optional[ConnectionT]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]": ...
+    ) -> "Union[ModelDTOT, dict[str, Any]]": ...
 
     @abstractmethod
     async def execute_script(

sqlspec/exceptions.py

@@ -78,6 +78,15 @@ class SQLParsingError(SQLSpecError):
         super().__init__(message)
 
 
+class SQLConversionError(SQLSpecError):
+    """Issues converting SQL statements."""
+
+    def __init__(self, message: Optional[str] = None) -> None:
+        if message is None:
+            message = "Issues converting SQL statement."
+        super().__init__(message)
+
+
 class ParameterStyleMismatchError(SQLSpecError):
     """Error when parameter style doesn't match SQL placeholder style.
 

sqlspec/extensions/litestar/config.py

@@ -1,10 +1,6 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
 
-from sqlspec.base import (
-    ConnectionT,
-    PoolT,
-)
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.extensions.litestar.handlers import (
     autocommit_handler_maker,
@@ -23,13 +19,9 @@ if TYPE_CHECKING:
     from litestar.datastructures.state import State
     from litestar.types import BeforeMessageSendHookHandler, Scope
 
-    from sqlspec.base import (
-        AsyncConfigT,
-        ConnectionT,
-        DriverT,
-        PoolT,
-        SyncConfigT,
-    )
+    from sqlspec.base import AsyncConfigT, DriverT, SyncConfigT
+    from sqlspec.typing import ConnectionT, PoolT
+
 
 CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
 DEFAULT_COMMIT_MODE: CommitMode = "manual"