sqlspec 0.9.0__py3-none-any.whl → 0.9.1__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,
@@ -19,6 +21,7 @@ from typing import (
 from sqlspec.exceptions import NotFoundError
 from sqlspec.statement import SQLStatement
 from sqlspec.typing import ModelDTOT, StatementParameterType
+from sqlspec.utils.sync_tools import maybe_async_
 
 if TYPE_CHECKING:
     from contextlib import AbstractAsyncContextManager, AbstractContextManager
@@ -202,6 +205,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 +296,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 +341,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 +387,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 +467,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[
@@ -501,6 +593,32 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
     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: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+
+    @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(
         self,
@@ -511,7 +629,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 +669,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 +707,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 +743,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 +781,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 +794,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 +830,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(
@@ -631,6 +879,32 @@ class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Gener
     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: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+
+    @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(
         self,
@@ -641,7 +915,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 +955,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 +993,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 +1029,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 +1067,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 +1080,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 +1116,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/extensions/litestar/plugin.py

@@ -7,6 +7,7 @@ from sqlspec.base import (
     AsyncConfigT,
     ConnectionT,
     DatabaseConfigProtocol,
+    DriverT,
     PoolT,
     SyncConfigT,
 )
@@ -75,6 +76,7 @@ class SQLSpec(InitPluginProtocol, SQLSpecBase):
                 SQLSpec,
                 ConnectionT,
                 PoolT,
+                DriverT,
                 DatabaseConfig,
                 DatabaseConfigProtocol,
                 SyncConfigT,