sqlspec 0.8.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

sqlspec/adapters/asyncmy/driver.py

@@ -1,7 +1,7 @@
 # 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
 
@@ -18,6 +18,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
     """Asyncmy MySQL/MariaDB Driver Adapter."""
 
     connection: "Connection"
+    dialect: str = "mysql"
 
     def __init__(self, connection: "Connection") -> None:
         self.connection = connection
@@ -35,21 +36,46 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         finally:
             await cursor.close()
 
+    # --- Public API Methods --- #
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     async def select(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["Connection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        **kwargs: Any,
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch data from the database.
 
         Returns:
             List of row data as either model instances or dictionaries.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             results = await cursor.fetchall()
@@ -60,13 +86,37 @@ 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[Connection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def select_one(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["Connection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
         """Fetch one row from the database.
 
@@ -74,7 +124,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             result = await cursor.fetchone()
@@ -84,13 +134,37 @@ 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[Connection]" = 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[Connection]" = 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,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
         """Fetch one row from the database.
 
@@ -98,7 +172,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             The first row of the query results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             result = await cursor.fetchone()
@@ -109,13 +183,37 @@ 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[Connection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
     async def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
         connection: "Optional[Connection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Union[T, Any]":
         """Fetch a single value from the database.
 
@@ -123,7 +221,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             The first value from the first row of results, or None if no results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -135,13 +233,37 @@ 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[Connection]" = 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,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[Connection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
         """Fetch a single value from the database.
 
@@ -149,7 +271,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             The first value from the first row of results, or None if no results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
@@ -168,7 +290,9 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["Connection"] = None,
+        **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
 
@@ -176,19 +300,43 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             Row count affected by the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             return cursor.rowcount
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = 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[Connection]" = 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,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
         """Insert, update, or delete data from the database and return result.
 
@@ -196,7 +344,7 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             The first row of results.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         column_names: list[str] = []
 
         async with self._with_cursor(connection) as cursor:
@@ -214,7 +362,9 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["Connection"] = None,
+        **kwargs: Any,
     ) -> str:
         """Execute a script.
 
@@ -222,34 +372,8 @@ class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
             Status message for the operation.
         """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
 
         async with self._with_cursor(connection) as cursor:
             await cursor.execute(sql, parameters)
             return "DONE"
-
-    async def execute_script_returning(
-        self,
-        sql: str,
-        parameters: Optional["StatementParameterType"] = None,
-        /,
-        connection: Optional["Connection"] = None,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
-        """Execute a script and return result.
-
-        Returns:
-            The first row of results.
-        """
-        connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
-        async with self._with_cursor(connection) as cursor:
-            await cursor.execute(sql, parameters)
-            result = await cursor.fetchone()
-            if result is None:
-                return None
-            column_names = [c[0] for c in cursor.description or []]
-            if schema_type is not None:
-                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))
-            return dict(zip(column_names, result))

sqlspec/adapters/asyncpg/__init__.py

@@ -1,9 +1,9 @@
-from sqlspec.adapters.asyncpg.config import Asyncpg, AsyncpgPool, PgConnection
-from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
+from sqlspec.adapters.asyncpg.config import AsyncpgConfig, AsyncpgPoolConfig
+from sqlspec.adapters.asyncpg.driver import AsyncpgConnection, AsyncpgDriver
 
 __all__ = (
-    "Asyncpg",
+    "AsyncpgConfig",
+    "AsyncpgConnection",
     "AsyncpgDriver",
-    "AsyncpgPool",
-    "PgConnection",
+    "AsyncpgPoolConfig",
 )

sqlspec/adapters/asyncpg/config.py

@@ -5,10 +5,9 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 from asyncpg import Record
 from asyncpg import create_pool as asyncpg_create_pool
 from asyncpg.pool import PoolConnectionProxy
-from typing_extensions import TypeAlias
 
 from sqlspec._serialization import decode_json, encode_json
-from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
+from sqlspec.adapters.asyncpg.driver import AsyncpgConnection, AsyncpgDriver
 from sqlspec.base import AsyncDatabaseConfig, GenericPoolConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -22,18 +21,16 @@ if TYPE_CHECKING:
 
 
 __all__ = (
-    "Asyncpg",
-    "AsyncpgPool",
+    "AsyncpgConfig",
+    "AsyncpgPoolConfig",
 )
 
 
 T = TypeVar("T")
 
-PgConnection: TypeAlias = "Union[Connection[Any], PoolConnectionProxy[Any]]"
-
 
 @dataclass
-class AsyncpgPool(GenericPoolConfig):
+class AsyncpgPoolConfig(GenericPoolConfig):
     """Configuration for Asyncpg's :class:`Pool <asyncpg.pool.Pool>`.
 
     For details see: https://magicstack.github.io/asyncpg/current/api/index.html#connection-pools
@@ -73,23 +70,25 @@ class AsyncpgPool(GenericPoolConfig):
 
 
 @dataclass
-class Asyncpg(AsyncDatabaseConfig["PgConnection", "Pool", "AsyncpgDriver"]):  # pyright: ignore[reportMissingTypeArgument]
+class AsyncpgConfig(AsyncDatabaseConfig["AsyncpgConnection", "Pool", "AsyncpgDriver"]):  # pyright: ignore[reportMissingTypeArgument]
     """Asyncpg Configuration."""
 
-    pool_config: "Optional[AsyncpgPool]" = None
+    pool_config: "Optional[AsyncpgPoolConfig]" = field(default=None)
     """Asyncpg Pool configuration"""
-    json_deserializer: "Callable[[str], Any]" = decode_json
+    json_deserializer: "Callable[[str], Any]" = field(hash=False, default=decode_json)
     """For dialects that support the :class:`JSON <sqlalchemy.types.JSON>` datatype, this is a Python callable that will
     convert a JSON string to a Python object. By default, this is set to SQLSpec's
     :attr:`decode_json() <sqlspec._serialization.decode_json>` function."""
-    json_serializer: "Callable[[Any], str]" = encode_json
+    json_serializer: "Callable[[Any], str]" = field(hash=False, default=encode_json)
     """For dialects that support the JSON datatype, this is a Python callable that will render a given object as JSON.
     By default, SQLSpec's :attr:`encode_json() <sqlspec._serialization.encode_json>` is used."""
-    connection_type: "type[PgConnection]" = field(init=False, default_factory=lambda: PoolConnectionProxy)
+    connection_type: "type[AsyncpgConnection]" = field(
+        hash=False, init=False, default_factory=lambda: PoolConnectionProxy
+    )
     """Type of the connection object"""
-    driver_type: "type[AsyncpgDriver]" = field(init=False, default_factory=lambda: AsyncpgDriver)  # type: ignore[type-abstract,unused-ignore]
+    driver_type: "type[AsyncpgDriver]" = field(hash=False, init=False, default_factory=lambda: AsyncpgDriver)  # type: ignore[type-abstract,unused-ignore]
     """Type of the driver object"""
-    pool_instance: "Optional[Pool[Any]]" = None
+    pool_instance: "Optional[Pool[Any]]" = field(hash=False, default=None)
     """The connection pool instance. If set, this will be used instead of creating a new pool."""
 
     @property
@@ -174,8 +173,8 @@ class Asyncpg(AsyncDatabaseConfig["PgConnection", "Pool", "AsyncpgDriver"]):  #
         """
         return self.create_pool()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
 
-    async def create_connection(self) -> "PgConnection":
-        """Create and return a new asyncpg connection.
+    async def create_connection(self) -> "AsyncpgConnection":
+        """Create and return a new asyncpg connection from the pool.
 
         Returns:
             A Connection instance.
@@ -184,9 +183,8 @@ class Asyncpg(AsyncDatabaseConfig["PgConnection", "Pool", "AsyncpgDriver"]):  #
             ImproperConfigurationError: If the connection could not be created.
         """
         try:
-            import asyncpg
-
-            return await asyncpg.connect(**self.connection_config_dict)  # type: ignore[no-any-return]
+            pool = await self.provide_pool()
+            return await pool.acquire()
         except Exception as e:
             msg = f"Could not configure the asyncpg connection. Error: {e!s}"
             raise ImproperConfigurationError(msg) from e