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

sqlspec/adapters/aiosqlite/driver.py

@@ -1,49 +1,77 @@
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+import aiosqlite
 
-if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
+from sqlspec.base import AsyncDriverAdapterProtocol
+from sqlspec.mixins import SQLTranslatorMixin
 
-    from aiosqlite import Connection, Cursor
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Sequence
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
-__all__ = ("AiosqliteDriver",)
+__all__ = ("AiosqliteConnection", "AiosqliteDriver")
+AiosqliteConnection = aiosqlite.Connection
 
 
-class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
+class AiosqliteDriver(
+    SQLTranslatorMixin["AiosqliteConnection"],
+    AsyncDriverAdapterProtocol["AiosqliteConnection"],
+):
     """SQLite Async Driver Adapter."""
 
-    connection: "Connection"
+    connection: "AiosqliteConnection"
     dialect: str = "sqlite"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "AiosqliteConnection") -> None:
         self.connection = connection
 
     @staticmethod
-    async def _cursor(connection: "Connection", *args: Any, **kwargs: Any) -> "Cursor":
+    async def _cursor(connection: "AiosqliteConnection", *args: Any, **kwargs: Any) -> "aiosqlite.Cursor":
         return await connection.cursor(*args, **kwargs)
 
     @asynccontextmanager
-    async def _with_cursor(self, connection: "Connection") -> "AsyncGenerator[Cursor, None]":
+    async def _with_cursor(self, connection: "AiosqliteConnection") -> "AsyncGenerator[aiosqlite.Cursor, None]":
         cursor = await self._cursor(connection)
         try:
             yield cursor
         finally:
             await cursor.close()
 
+    # --- Public API Methods --- #
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     async def select(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = 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:
@@ -61,13 +89,35 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return [dict(zip(column_names, row)) for row in results]  # pyright: ignore[reportUnknownArgumentType]
             return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def select_one(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -87,13 +137,35 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
             return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     async def select_one_or_none(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -114,13 +186,35 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
             return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     async def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -139,13 +233,35 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return result[0]
             return schema_type(result[0])  # type: ignore[call-arg]
 
+    @overload
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
     async def select_value_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -171,7 +287,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -186,13 +302,35 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
             await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             return cursor.rowcount if hasattr(cursor, "rowcount") else -1  # pyright: ignore[reportUnknownVariableType, reportGeneralTypeIssues]
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AiosqliteConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def insert_update_delete_returning(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -220,7 +358,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -241,7 +379,7 @@ class AiosqliteDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AiosqliteConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":

sqlspec/adapters/asyncmy/__init__.py

@@ -1,8 +1,9 @@
 from sqlspec.adapters.asyncmy.config import AsyncmyConfig, AsyncmyPoolConfig
-from sqlspec.adapters.asyncmy.driver import AsyncmyDriver  # type: ignore[attr-defined]
+from sqlspec.adapters.asyncmy.driver import AsyncmyConnection, AsyncmyDriver  # type: ignore[attr-defined]
 
 __all__ = (
     "AsyncmyConfig",
+    "AsyncmyConnection",
     "AsyncmyDriver",
     "AsyncmyPoolConfig",
 )

sqlspec/adapters/asyncmy/driver.py

@@ -1,51 +1,77 @@
 # type: ignore
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+from asyncmy import Connection
+
+from sqlspec.base import AsyncDriverAdapterProtocol
+from sqlspec.mixins import SQLTranslatorMixin
 
 if TYPE_CHECKING:
-    from asyncmy import Connection
     from asyncmy.cursors import Cursor
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
 __all__ = ("AsyncmyDriver",)
 
+AsyncmyConnection = Connection
+
 
-class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
+class AsyncmyDriver(
+    SQLTranslatorMixin["AsyncmyConnection"],
+    AsyncDriverAdapterProtocol["AsyncmyConnection"],
+):
     """Asyncmy MySQL/MariaDB Driver Adapter."""
 
-    connection: "Connection"
+    connection: "AsyncmyConnection"
     dialect: str = "mysql"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "AsyncmyConnection") -> None:
         self.connection = connection
 
-    @staticmethod
-    async def _cursor(connection: "Connection") -> "Cursor":
-        return await connection.cursor()
-
     @staticmethod
     @asynccontextmanager
-    async def _with_cursor(connection: "Connection") -> AsyncGenerator["Cursor", None]:
+    async def _with_cursor(connection: "AsyncmyConnection") -> AsyncGenerator["Cursor", None]:
         cursor = connection.cursor()
         try:
             yield cursor
         finally:
             await cursor.close()
 
+    # --- Public API Methods --- #
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     async def select(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = 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:
@@ -63,13 +89,35 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return [dict(zip(column_names, row)) for row in results]
             return [schema_type(**dict(zip(column_names, row))) for row in results]
 
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def select_one(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -89,13 +137,35 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return dict(zip(column_names, result))
             return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))
 
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     async def select_one_or_none(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -116,13 +186,35 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return dict(zip(column_names, result))
             return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))
 
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     async def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -144,13 +236,35 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
                 return schema_type(value)  # type: ignore[call-arg]
             return value
 
+    @overload
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
     async def select_value_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -180,7 +294,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -195,13 +309,35 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             await cursor.execute(sql, parameters)
             return cursor.rowcount
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncmyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def insert_update_delete_returning(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -230,7 +366,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["AsyncmyConnection"] = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.

sqlspec/adapters/asyncpg/config.py

@@ -190,9 +190,7 @@ class AsyncpgConfig(AsyncDatabaseConfig["AsyncpgConnection", "Pool", "AsyncpgDri
             raise ImproperConfigurationError(msg) from e
 
     @asynccontextmanager
-    async def provide_connection(
-        self, *args: "Any", **kwargs: "Any"
-    ) -> "AsyncGenerator[PoolConnectionProxy[Any], None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
+    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AsyncpgConnection, None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
         """Create a connection instance.
 
         Yields: