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

sqlspec/__init__.py

@@ -1,4 +1,4 @@
-from sqlspec import adapters, base, exceptions, extensions, filters, typing, utils
+from sqlspec import adapters, base, exceptions, extensions, filters, mixins, typing, utils
 from sqlspec.__metadata__ import __version__
 from sqlspec.base import SQLSpec
 
@@ -10,6 +10,7 @@ __all__ = (
     "exceptions",
     "extensions",
     "filters",
+    "mixins",
     "typing",
     "utils",
 )

sqlspec/adapters/adbc/__init__.py

@@ -1,7 +1,8 @@
 from sqlspec.adapters.adbc.config import AdbcConfig
-from sqlspec.adapters.adbc.driver import AdbcDriver
+from sqlspec.adapters.adbc.driver import AdbcConnection, AdbcDriver
 
 __all__ = (
     "AdbcConfig",
+    "AdbcConnection",
     "AdbcDriver",
 )

sqlspec/adapters/adbc/config.py

@@ -2,9 +2,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
 
-from adbc_driver_manager.dbapi import Connection
-
-from sqlspec.adapters.adbc.driver import AdbcDriver
+from sqlspec.adapters.adbc.driver import AdbcConnection, AdbcDriver
 from sqlspec.base import NoPoolSyncConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType
@@ -18,7 +16,7 @@ __all__ = ("AdbcConfig",)
 
 
 @dataclass
-class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
+class AdbcConfig(NoPoolSyncConfig["AdbcConnection", "AdbcDriver"]):
     """Configuration for ADBC connections.
 
     This class provides configuration options for ADBC database connections using the
@@ -33,20 +31,16 @@ class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
     """Additional database-specific connection parameters"""
     conn_kwargs: "Optional[dict[str, Any]]" = None
     """Additional database-specific connection parameters"""
-    connection_type: "type[Connection]" = field(init=False, default_factory=lambda: Connection)
+    connection_type: "type[AdbcConnection]" = field(init=False, default_factory=lambda: AdbcConnection)
     """Type of the connection object"""
     driver_type: "type[AdbcDriver]" = field(init=False, default_factory=lambda: AdbcDriver)  # type: ignore[type-abstract,unused-ignore]
     """Type of the driver object"""
-    pool_instance: None = field(init=False, default=None)
+    pool_instance: None = field(init=False, default=None, hash=False)
     """No connection pool is used for ADBC connections"""
-    _is_in_memory: bool = field(init=False, default=False)
-    """Flag indicating if the connection is for an in-memory database"""
 
     def _set_adbc(self) -> str:  # noqa: PLR0912
         """Identify the driver type based on the URI (if provided) or preset driver name.
 
-        Also sets the `_is_in_memory` flag for specific in-memory URIs.
-
         Raises:
             ImproperConfigurationError: If the driver name is not recognized or supported.
 
@@ -143,7 +137,7 @@ class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
             config["conn_kwargs"] = conn_kwargs
         return config
 
-    def _get_connect_func(self) -> "Callable[..., Connection]":
+    def _get_connect_func(self) -> "Callable[..., AdbcConnection]":
         self._set_adbc()
         driver_path = cast("str", self.driver_name)
         try:
@@ -166,7 +160,7 @@ class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
             raise ImproperConfigurationError(msg)
         return connect_func  # type: ignore[no-any-return]
 
-    def create_connection(self) -> "Connection":
+    def create_connection(self) -> "AdbcConnection":
         """Create and return a new database connection using the specific driver.
 
         Returns:
@@ -189,7 +183,7 @@ class AdbcConfig(NoPoolSyncConfig["Connection", "AdbcDriver"]):
             raise ImproperConfigurationError(msg) from e
 
     @contextmanager
-    def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[Connection, None, None]":
+    def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[AdbcConnection, None, None]":
         """Create and provide a database connection using the specific driver.
 
         Yields:

sqlspec/adapters/adbc/driver.py

@@ -1,21 +1,22 @@
 import contextlib
 import logging
 import re
-from collections.abc import Generator
+from collections.abc import Generator, Sequence
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union, cast, overload
 
 from adbc_driver_manager.dbapi import Connection, Cursor
 
-from sqlspec.base import SyncArrowBulkOperationsMixin, SyncDriverAdapterProtocol, T
+from sqlspec.base import SyncDriverAdapterProtocol
 from sqlspec.exceptions import ParameterStyleMismatchError, SQLParsingError
+from sqlspec.mixins import SQLTranslatorMixin, SyncArrowBulkOperationsMixin
 from sqlspec.statement import SQLStatement
 from sqlspec.typing import ArrowTable, StatementParameterType
 
 if TYPE_CHECKING:
-    from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType
+    from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType, T
 
-__all__ = ("AdbcDriver",)
+__all__ = ("AdbcConnection", "AdbcDriver")
 
 logger = logging.getLogger("sqlspec")
 
@@ -32,20 +33,26 @@ PARAM_REGEX = re.compile(
     re.VERBOSE | re.DOTALL,
 )
 
+AdbcConnection = Connection
 
-class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterProtocol["Connection"]):
+
+class AdbcDriver(
+    SyncArrowBulkOperationsMixin["AdbcConnection"],
+    SQLTranslatorMixin["AdbcConnection"],
+    SyncDriverAdapterProtocol["AdbcConnection"],
+):
     """ADBC Sync Driver Adapter."""
 
-    connection: Connection
+    connection: AdbcConnection
     __supports_arrow__: ClassVar[bool] = True
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "AdbcConnection") -> None:
         """Initialize the ADBC driver adapter."""
         self.connection = connection
         self.dialect = self._get_dialect(connection)
 
     @staticmethod
-    def _get_dialect(connection: "Connection") -> str:  # noqa: PLR0911
+    def _get_dialect(connection: "AdbcConnection") -> str:  # noqa: PLR0911
         """Get the database dialect based on the driver name.
 
         Args:
@@ -70,11 +77,11 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
         return "postgres"  # default to postgresql dialect
 
     @staticmethod
-    def _cursor(connection: "Connection", *args: Any, **kwargs: Any) -> "Cursor":
+    def _cursor(connection: "AdbcConnection", *args: Any, **kwargs: Any) -> "Cursor":
         return connection.cursor(*args, **kwargs)
 
     @contextmanager
-    def _with_cursor(self, connection: "Connection") -> Generator["Cursor", None, None]:
+    def _with_cursor(self, connection: "AdbcConnection") -> Generator["Cursor", None, None]:
         cursor = self._cursor(connection)
         try:
             yield cursor
@@ -160,16 +167,38 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
         stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
         return stmt.process()
 
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     def select(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Returns:
@@ -189,13 +218,35 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
                 return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
             return [dict(zip(column_names, row)) for row in results]  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
 
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     def select_one(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -215,13 +266,35 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
                 return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
             return schema_type(**dict(zip(column_names, result)))  # type: ignore[return-value]
 
+    @overload
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     def select_one_or_none(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -242,13 +315,35 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
                 return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
             return schema_type(**dict(zip(column_names, result)))  # type: ignore[return-value]
 
+    @overload
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     def select_value(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -267,13 +362,35 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
                 return result[0]  # pyright: ignore[reportUnknownVariableType]
             return schema_type(result[0])  # type: ignore[call-arg]
 
+    @overload
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
     def select_value_or_none(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -299,7 +416,7 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -314,13 +431,35 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             return cursor.rowcount if hasattr(cursor, "rowcount") else -1
 
+    @overload
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AdbcConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     def insert_update_delete_returning(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -353,7 +492,7 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AdbcConnection"] = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -376,7 +515,7 @@ class AdbcDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterPr
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AdbcConnection]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":
         """Execute a SQL query and return results as an Apache Arrow Table.

sqlspec/adapters/aiosqlite/__init__.py

@@ -1,7 +1,8 @@
 from sqlspec.adapters.aiosqlite.config import AiosqliteConfig
-from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver
+from sqlspec.adapters.aiosqlite.driver import AiosqliteConnection, AiosqliteDriver
 
 __all__ = (
     "AiosqliteConfig",
+    "AiosqliteConnection",
     "AiosqliteDriver",
 )

sqlspec/adapters/aiosqlite/config.py

@@ -2,16 +2,15 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from aiosqlite import Connection
+import aiosqlite
 
-from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver
+from sqlspec.adapters.aiosqlite.driver import AiosqliteConnection, AiosqliteDriver
 from sqlspec.base import NoPoolAsyncConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
-    from sqlite3 import Connection as SQLite3Connection
     from typing import Literal
 
 
@@ -19,7 +18,7 @@ __all__ = ("AiosqliteConfig",)
 
 
 @dataclass
-class AiosqliteConfig(NoPoolAsyncConfig["Connection", "AiosqliteDriver"]):
+class AiosqliteConfig(NoPoolAsyncConfig["AiosqliteConnection", "AiosqliteDriver"]):
     """Configuration for Aiosqlite database connections.
 
     This class provides configuration options for Aiosqlite database connections, wrapping all parameters
@@ -38,13 +37,11 @@ class AiosqliteConfig(NoPoolAsyncConfig["Connection", "AiosqliteDriver"]):
     """The isolation_level of the connection. This can be None for autocommit mode or one of "DEFERRED", "IMMEDIATE" or "EXCLUSIVE"."""
     check_same_thread: "Union[bool, EmptyType]" = field(default=Empty)
     """If True (default), ProgrammingError is raised if the database connection is used by a thread other than the one that created it. If False, the connection may be shared across multiple threads."""
-    factory: "Union[type[SQLite3Connection], EmptyType]" = field(default=Empty)
-    """A custom Connection class factory. If given, must be a callable that returns a Connection instance."""
     cached_statements: "Union[int, EmptyType]" = field(default=Empty)
     """The number of statements that SQLite will cache for this connection. The default is 128."""
     uri: "Union[bool, EmptyType]" = field(default=Empty)
     """If set to True, database is interpreted as a URI with supported options."""
-    connection_type: "type[Connection]" = field(init=False, default_factory=lambda: Connection)
+    connection_type: "type[AiosqliteConnection]" = field(init=False, default_factory=lambda: AiosqliteConnection)
     """Type of the connection object"""
     driver_type: "type[AiosqliteDriver]" = field(init=False, default_factory=lambda: AiosqliteDriver)  # type: ignore[type-abstract,unused-ignore]
     """Type of the driver object"""
@@ -57,10 +54,13 @@ class AiosqliteConfig(NoPoolAsyncConfig["Connection", "AiosqliteDriver"]):
             A string keyed dict of config kwargs for the aiosqlite.connect() function.
         """
         return dataclass_to_dict(
-            self, exclude_empty=True, convert_nested=False, exclude={"pool_instance", "connection_type", "driver_type"}
+            self,
+            exclude_empty=True,
+            convert_nested=False,
+            exclude={"pool_instance", "connection_type", "driver_type"},
         )
 
-    async def create_connection(self) -> "Connection":
+    async def create_connection(self) -> "AiosqliteConnection":
         """Create and return a new database connection.
 
         Returns:
@@ -69,8 +69,6 @@ class AiosqliteConfig(NoPoolAsyncConfig["Connection", "AiosqliteDriver"]):
         Raises:
             ImproperConfigurationError: If the connection could not be established.
         """
-        import aiosqlite
-
         try:
             return await aiosqlite.connect(**self.connection_config_dict)
         except Exception as e:
@@ -78,7 +76,7 @@ class AiosqliteConfig(NoPoolAsyncConfig["Connection", "AiosqliteDriver"]):
             raise ImproperConfigurationError(msg) from e
 
     @asynccontextmanager
-    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[Connection, None]":
+    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AiosqliteConnection, None]":
         """Create and provide a database connection.
 
         Yields: