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

sqlspec/adapters/asyncmy/driver.py

@@ -0,0 +1,255 @@
+# type: ignore
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+from sqlspec.base import AsyncDriverAdapterProtocol, T
+
+if TYPE_CHECKING:
+    from asyncmy import Connection
+    from asyncmy.cursors import Cursor
+
+    from sqlspec.typing import ModelDTOT, StatementParameterType
+
+__all__ = ("AsyncmyDriver",)
+
+
+class AsyncmyDriver(AsyncDriverAdapterProtocol["Connection"]):
+    """Asyncmy MySQL/MariaDB Driver Adapter."""
+
+    connection: "Connection"
+
+    def __init__(self, connection: "Connection") -> 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]:
+        cursor = connection.cursor()
+        try:
+            yield cursor
+        finally:
+            await cursor.close()
+
+    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]]]":
+        """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)
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)
+            results = await cursor.fetchall()
+            if not results:
+                return []
+            column_names = [c[0] for c in cursor.description or []]
+            if schema_type is None:
+                return [dict(zip(column_names, row)) for row in results]
+            return [schema_type(**dict(zip(column_names, row))) for row in results]
+
+    async def select_one(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["Connection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query 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()
+            result = self.check_not_found(result)
+            column_names = [c[0] for c in cursor.description or []]
+            if schema_type is None:
+                return dict(zip(column_names, result))
+            return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))
+
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["Connection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query 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 None:
+                return dict(zip(column_names, result))
+            return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))
+
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+    ) -> "Union[T, Any]":
+        """Fetch a single value from the database.
+
+        Returns:
+            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)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)
+            result = await cursor.fetchone()
+            result = self.check_not_found(result)
+
+            value = result[0]
+            if schema_type is not None:
+                return schema_type(value)  # type: ignore[call-arg]
+            return value
+
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+    ) -> "Optional[Union[T, Any]]":
+        """Fetch a single value from the database.
+
+        Returns:
+            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)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)
+            result = await cursor.fetchone()
+
+            if result is None:
+                return None
+
+            value = result[0]
+            if schema_type is not None:
+                return schema_type(value)  # type: ignore[call-arg]
+            return value
+
+    async def insert_update_delete(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["Connection"] = None,
+    ) -> int:
+        """Insert, update, or delete data from the database.
+
+        Returns:
+            Row count affected by the operation.
+        """
+        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)
+            return cursor.rowcount
+
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["Connection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Insert, update, or delete data from the database and return result.
+
+        Returns:
+            The first row of results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters)
+        column_names: list[str] = []
+
+        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))
+
+    async def execute_script(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["Connection"] = None,
+    ) -> str:
+        """Execute a script.
+
+        Returns:
+            Status message for the operation.
+        """
+        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)
+            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

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

sqlspec/adapters/asyncpg/config.py

@@ -1,13 +1,14 @@
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 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 Pool, PoolConnectionProxy
+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.base import AsyncDatabaseConfig, GenericPoolConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -17,21 +18,22 @@ if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 
     from asyncpg.connection import Connection
+    from asyncpg.pool import Pool
 
 
 __all__ = (
-    "AsyncPgConfig",
-    "AsyncPgPoolConfig",
+    "Asyncpg",
+    "AsyncpgPool",
 )
 
 
 T = TypeVar("T")
 
-PgConnection: TypeAlias = "Union[Connection, PoolConnectionProxy]"  # pyright: ignore[reportMissingTypeArgument]
+PgConnection: TypeAlias = "Union[Connection[Any], PoolConnectionProxy[Any]]"
 
 
 @dataclass
-class AsyncPgPoolConfig(GenericPoolConfig):
+class AsyncpgPool(GenericPoolConfig):
     """Configuration for Asyncpg's :class:`Pool <asyncpg.pool.Pool>`.
 
     For details see: https://magicstack.github.io/asyncpg/current/api/index.html#connection-pools
@@ -52,7 +54,7 @@ class AsyncPgPoolConfig(GenericPoolConfig):
     min_size: "Union[int, EmptyType]" = Empty
     """The number of connections to keep open inside the connection pool."""
     max_size: "Union[int, EmptyType]" = Empty
-    """The number of connections to allow in connection pool “overflow”, that is connections that can be opened above
+    """The number of connections to allow in connection pool "overflow", that is connections that can be opened above
     and beyond the pool_size setting, which defaults to 10."""
 
     max_queries: "Union[int, EmptyType]" = Empty
@@ -71,10 +73,10 @@ class AsyncPgPoolConfig(GenericPoolConfig):
 
 
 @dataclass
-class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore[reportMissingTypeArgument]
+class Asyncpg(AsyncDatabaseConfig["PgConnection", "Pool", "AsyncpgDriver"]):  # pyright: ignore[reportMissingTypeArgument]
     """Asyncpg Configuration."""
 
-    pool_config: "Optional[AsyncPgPoolConfig]" = None
+    pool_config: "Optional[AsyncpgPool]" = None
     """Asyncpg Pool configuration"""
     json_deserializer: "Callable[[str], Any]" = decode_json
     """For dialects that support the :class:`JSON <sqlalchemy.types.JSON>` datatype, this is a Python callable that will
@@ -83,11 +85,41 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
     json_serializer: "Callable[[Any], str]" = 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)
+    """Type of the connection object"""
+    driver_type: "type[AsyncpgDriver]" = field(init=False, default_factory=lambda: AsyncpgDriver)  # type: ignore[type-abstract,unused-ignore]
+    """Type of the driver object"""
     pool_instance: "Optional[Pool[Any]]" = None
-    """Optional pool to use.
+    """The connection pool instance. If set, this will be used instead of creating a new pool."""
 
-    If set, the plugin will use the provided pool rather than instantiate one.
-    """
+    @property
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Returns:
+            A string keyed dict of config kwargs for the asyncpg.connect function.
+
+        Raises:
+            ImproperConfigurationError: If the connection configuration is not provided.
+        """
+        if self.pool_config:
+            connect_dict: dict[str, Any] = {}
+
+            # Add dsn if available
+            if hasattr(self.pool_config, "dsn"):
+                connect_dict["dsn"] = self.pool_config.dsn
+
+            # Add any connect_kwargs if available
+            if (
+                hasattr(self.pool_config, "connect_kwargs")
+                and self.pool_config.connect_kwargs is not Empty
+                and isinstance(self.pool_config.connect_kwargs, dict)
+            ):
+                connect_dict.update(dict(self.pool_config.connect_kwargs.items()))
+
+            return connect_dict
+        msg = "You must provide a 'pool_config' for this adapter."
+        raise ImproperConfigurationError(msg)
 
     @property
     def pool_config_dict(self) -> "dict[str, Any]":
@@ -96,9 +128,17 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
         Returns:
             A string keyed dict of config kwargs for the Asyncpg :func:`create_pool <asyncpg.pool.create_pool>`
             function.
+
+        Raises:
+            ImproperConfigurationError: If no pool_config is provided but a pool_instance is set.
         """
         if self.pool_config:
-            return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                exclude={"pool_instance", "driver_type", "connection_type"},
+                convert_nested=False,
+            )
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
@@ -107,6 +147,10 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
 
         Returns:
             Getter that returns the pool instance used by the plugin.
+
+        Raises:
+            ImproperConfigurationError: If neither pool_config nor pool_instance are provided,
+                or if the pool could not be configured.
         """
         if self.pool_instance is not None:
             return self.pool_instance
@@ -117,11 +161,9 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
 
         pool_config = self.pool_config_dict
         self.pool_instance = await asyncpg_create_pool(**pool_config)
-        if self.pool_instance is None:
-            msg = "Could not configure the 'pool_instance'. Please check your configuration."
-            raise ImproperConfigurationError(
-                msg,
-            )
+        if self.pool_instance is None:  # pyright: ignore[reportUnnecessaryComparison]
+            msg = "Could not configure the 'pool_instance'. Please check your configuration."  # type: ignore[unreachable]
+            raise ImproperConfigurationError(msg)
         return self.pool_instance
 
     def provide_pool(self, *args: "Any", **kwargs: "Any") -> "Awaitable[Pool]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
@@ -132,13 +174,50 @@ class AsyncPgConfig(AsyncDatabaseConfig[PgConnection, Pool]):  # pyright: ignore
         """
         return self.create_pool()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
 
+    async def create_connection(self) -> "PgConnection":
+        """Create and return a new asyncpg connection.
+
+        Returns:
+            A Connection instance.
+
+        Raises:
+            ImproperConfigurationError: If the connection could not be created.
+        """
+        try:
+            import asyncpg
+
+            return await asyncpg.connect(**self.connection_config_dict)  # type: ignore[no-any-return]
+        except Exception as e:
+            msg = f"Could not configure the asyncpg connection. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
+
     @asynccontextmanager
-    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[PoolConnectionProxy, None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
+    async def provide_connection(
+        self, *args: "Any", **kwargs: "Any"
+    ) -> "AsyncGenerator[PoolConnectionProxy[Any], None]":  # pyright: ignore[reportMissingTypeArgument,reportUnknownParameterType]
         """Create a connection instance.
 
-        Returns:
+        Yields:
             A connection instance.
         """
         db_pool = await self.provide_pool(*args, **kwargs)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
         async with db_pool.acquire() as connection:  # pyright: ignore[reportUnknownVariableType]
             yield connection
+
+    async def close_pool(self) -> None:
+        """Close the pool."""
+        if self.pool_instance is not None:
+            await self.pool_instance.close()
+            self.pool_instance = None
+
+    @asynccontextmanager
+    async def provide_session(self, *args: Any, **kwargs: Any) -> "AsyncGenerator[AsyncpgDriver, None]":
+        """Create and provide a database session.
+
+        Yields:
+            A Aiosqlite driver instance.
+
+
+        """
+        async with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection)