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

sqlspec/adapters/asyncpg/driver.py

@@ -0,0 +1,288 @@
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+from asyncpg import Connection
+from typing_extensions import TypeAlias
+
+from sqlspec.base import AsyncDriverAdapterProtocol, T
+
+if TYPE_CHECKING:
+    from asyncpg.connection import Connection
+    from asyncpg.pool import PoolConnectionProxy
+
+    from sqlspec.typing import ModelDTOT, StatementParameterType
+
+__all__ = ("AsyncpgDriver",)
+
+
+PgConnection: TypeAlias = "Union[Connection[Any], PoolConnectionProxy[Any]]"  # pyright: ignore[reportMissingTypeArgument]
+
+
+class AsyncpgDriver(AsyncDriverAdapterProtocol["PgConnection"]):
+    """AsyncPG Postgres Driver Adapter."""
+
+    connection: "PgConnection"
+
+    def __init__(self, connection: "PgConnection") -> None:
+        self.connection = connection
+
+    def _process_sql_params(
+        self, sql: str, parameters: "Optional[StatementParameterType]" = None
+    ) -> "tuple[str, Union[tuple[Any, ...], list[Any], dict[str, Any]]]":
+        sql, parameters = super()._process_sql_params(sql, parameters)
+        return sql, parameters if parameters is not None else ()
+
+    async def select(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch data from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+
+        Returns:
+            List of row data as either model instances or dictionaries.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters)
+
+        results = await connection.fetch(sql, *parameters)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        if not results:
+            return []
+        if schema_type is None:
+            return [dict(row.items()) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return [cast("ModelDTOT", schema_type(**dict(row.items()))) for row in results]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+    async def select_one(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
+        """Fetch one row from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        result = await connection.fetchrow(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        result = self.check_not_found(result)
+
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # type: ignore[attr-defined]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # type: ignore[attr-defined]
+
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch one row from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters)
+
+        result = await connection.fetchrow(sql, *parameters)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        result = self.check_not_found(result)
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # type: ignore[attr-defined]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # type: ignore[attr-defined]
+
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        connection: "Optional[PgConnection]" = 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, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        result = await connection.fetchval(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        result = self.check_not_found(result)
+        if schema_type is None:
+            return result[0]
+        return schema_type(result[0])  # type: ignore[call-arg]
+
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        connection: "Optional[PgConnection]" = 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, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        result = await connection.fetchval(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        if result is None:
+            return None
+        if schema_type is None:
+            return result[0]
+        return schema_type(result[0])  # type: ignore[call-arg]
+
+    async def insert_update_delete(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+    ) -> int:
+        """Insert, update, or delete data from the database.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+
+        Returns:
+            Row count affected by the operation.
+        """
+        connection = self._connection(connection)
+        sql, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        status = await connection.execute(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        # AsyncPG returns a string like "INSERT 0 1" where the last number is the affected rows
+        try:
+            return int(status.split()[-1])  # pyright: ignore[reportUnknownMemberType]
+        except (ValueError, IndexError, AttributeError):
+            return -1  # Fallback if we can't parse the status
+
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Insert, update, or delete data from the database and return result.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+
+        Returns:
+            The first row of results.
+        """
+        connection = self._connection(connection)
+        sql, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        result = await connection.fetchrow(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        if result is None:
+            return None
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType]
+
+    async def execute_script(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+    ) -> str:
+        """Execute a script.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+
+        Returns:
+            Status message for the operation.
+        """
+        connection = self._connection(connection)
+        sql, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        return await connection.execute(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+    async def execute_script_returning(
+        self,
+        sql: str,
+        parameters: Optional["StatementParameterType"] = None,
+        /,
+        connection: Optional["PgConnection"] = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Execute a script and return result.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            connection: Optional connection to use.
+            schema_type: Optional schema class for the result.
+
+        Returns:
+            The first row of results.
+        """
+        connection = self._connection(connection)
+        sql, params = self._process_sql_params(sql, parameters)
+        # Use empty tuple if params is None
+        params = params if params is not None else ()
+
+        result = await connection.fetchrow(sql, *params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        if result is None:
+            return None
+        if schema_type is None:
+            # Always return as dictionary
+            return dict(result.items())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        return cast("ModelDTOT", schema_type(**dict(result.items())))  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType]

sqlspec/adapters/duckdb/__init__.py

@@ -1,3 +1,7 @@
-from sqlspec.adapters.duckdb.config import DuckDBConfig
+from sqlspec.adapters.duckdb.config import DuckDB
+from sqlspec.adapters.duckdb.driver import DuckDBDriver
 
-__all__ = ("DuckDBConfig",)
+__all__ = (
+    "DuckDB",
+    "DuckDBDriver",
+)

sqlspec/adapters/duckdb/config.py

@@ -1,10 +1,11 @@
 from contextlib import contextmanager
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Union, cast
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
 
 from duckdb import DuckDBPyConnection
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import Literal, NotRequired, TypedDict
 
+from sqlspec.adapters.duckdb.driver import DuckDBDriver
 from sqlspec.base import NoPoolSyncConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -13,7 +14,7 @@ if TYPE_CHECKING:
     from collections.abc import Generator, Sequence
 
 
-__all__ = ("DuckDBConfig", "ExtensionConfig")
+__all__ = ("DuckDB", "ExtensionConfig")
 
 
 class ExtensionConfig(TypedDict):
@@ -29,6 +30,8 @@ class ExtensionConfig(TypedDict):
     """The name of the extension to install"""
     config: "NotRequired[dict[str, Any]]"
     """Optional configuration settings to apply after installation"""
+    install_if_missing: "NotRequired[bool]"
+    """Whether to install if missing"""
     force_install: "NotRequired[bool]"
     """Whether to force reinstall if already present"""
     repository: "NotRequired[str]"
@@ -39,8 +42,34 @@ class ExtensionConfig(TypedDict):
     """Optional version of the extension to install"""
 
 
+class SecretConfig(TypedDict):
+    """Configuration for a secret to store in a connection.
+
+    This class provides configuration options for storing a secret in a connection for later retrieval.
+
+    For details see: https://duckdb.org/docs/stable/configuration/secrets_manager
+    """
+
+    secret_type: Union[
+        Literal[
+            "azure", "gcs", "s3", "r2", "huggingface", "http", "mysql", "postgres", "bigquery", "openai", "open_prompt"  # noqa: PYI051
+        ],
+        str,
+    ]
+    provider: NotRequired[str]
+    """The provider of the secret"""
+    name: str
+    """The name of the secret to store"""
+    value: dict[str, Any]
+    """The secret value to store"""
+    persist: NotRequired[bool]
+    """Whether to persist the secret"""
+    replace_if_exists: NotRequired[bool]
+    """Whether to replace the secret if it already exists"""
+
+
 @dataclass
-class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
+class DuckDB(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
     """Configuration for DuckDB database connections.
 
     This class provides configuration options for DuckDB database connections, wrapping all parameters
@@ -49,7 +78,7 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
     For details see: https://duckdb.org/docs/api/python/overview#connection-options
     """
 
-    database: "Union[str, EmptyType]" = Empty
+    database: "Union[str, EmptyType]" = field(default=":memory:")
     """The path to the database file to be opened. Pass ":memory:" to open a connection to a database that resides in RAM instead of on disk. If not specified, an in-memory database will be created."""
 
     read_only: "Union[bool, EmptyType]" = Empty
@@ -63,6 +92,18 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
 
     extensions: "Union[Sequence[ExtensionConfig], ExtensionConfig, EmptyType]" = Empty
     """A sequence of extension configurations to install and configure upon connection creation."""
+    secrets: "Union[Sequence[SecretConfig], SecretConfig , EmptyType]" = Empty
+    """A dictionary of secrets to store in the connection for later retrieval."""
+    auto_update_extensions: "bool" = False
+    """Whether to automatically update on connection creation"""
+    on_connection_create: "Optional[Callable[[DuckDBPyConnection], Optional[DuckDBPyConnection]]]" = None
+    """A callable to be called after the connection is created."""
+    connection_type: "type[DuckDBPyConnection]" = field(init=False, default_factory=lambda: DuckDBPyConnection)
+    """The type of connection to create. Defaults to DuckDBPyConnection."""
+    driver_type: "type[DuckDBDriver]" = field(init=False, default_factory=lambda: DuckDBDriver)  # type: ignore[type-abstract,unused-ignore]
+    """The type of driver to use. Defaults to DuckDBDriver."""
+    pool_instance: "None" = field(init=False, default=None)
+    """The pool instance to use. Defaults to None."""
 
     def __post_init__(self) -> None:
         """Post-initialization validation and processing.
@@ -73,9 +114,10 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
         """
         if self.config is Empty:
             self.config = {}
-
         if self.extensions is Empty:
             self.extensions = []
+        if self.secrets is Empty:
+            self.secrets = []
         if isinstance(self.extensions, dict):
             self.extensions = [self.extensions]
         # this is purely for mypy
@@ -103,8 +145,8 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
         Args:
             connection: The DuckDB connection to configure.
         """
-        for config in cast("list[str]", self.config):
-            connection.execute(config)
+        for key, value in cast("dict[str,Any]", self.config).items():
+            connection.execute(f"SET {key}='{value}'")
 
     def _configure_extensions(self, connection: "DuckDBPyConnection") -> None:
         """Configure extensions for the connection.
@@ -119,9 +161,104 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
 
         for extension in cast("list[ExtensionConfig]", self.extensions):
             self._configure_extension(connection, extension)
+        if self.auto_update_extensions:
+            connection.execute("update extensions")
 
     @staticmethod
-    def _configure_extension(connection: "DuckDBPyConnection", extension: ExtensionConfig) -> None:
+    def _secret_exists(connection: "DuckDBPyConnection", name: "str") -> bool:
+        """Check if a secret exists in the connection.
+
+        Args:
+            connection: The DuckDB connection to check for the secret.
+            name: The name of the secret to check for.
+
+        Returns:
+            bool: True if the secret exists, False otherwise.
+        """
+        results = connection.execute("select 1 from duckdb_secrets() where name=?", [name]).fetchone()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        return results is not None
+
+    @classmethod
+    def _is_community_extension(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
+        """Check if an extension is a community extension.
+
+        Args:
+            connection: The DuckDB connection to check for the extension.
+            name: The name of the extension to check.
+
+        Returns:
+            bool: True if the extension is a community extension, False otherwise.
+        """
+        results = connection.execute(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            "select 1 from duckdb_extensions() where extension_name=?", [name]
+        ).fetchone()
+        return results is None
+
+    @classmethod
+    def _extension_installed(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
+        """Check if a extension exists in the connection.
+
+        Args:
+            connection: The DuckDB connection to check for the secret.
+            name: The name of the secret to check for.
+
+        Returns:
+            bool: True if the extension is installed, False otherwise.
+        """
+        results = connection.execute(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            "select 1 from duckdb_extensions() where extension_name=? and installed=true", [name]
+        ).fetchone()
+        return results is not None
+
+    @classmethod
+    def _extension_loaded(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
+        """Check if a extension is loaded in the connection.
+
+        Args:
+            connection: The DuckDB connection to check for the extension.
+            name: The name of the extension to check for.
+
+        Returns:
+            bool: True if the extension is loaded, False otherwise.
+        """
+        results = connection.execute(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            "select 1 from duckdb_extensions() where extension_name=? and loaded=true", [name]
+        ).fetchone()
+        return results is not None
+
+    @classmethod
+    def _configure_secrets(
+        cls,
+        connection: "DuckDBPyConnection",
+        secrets: "Sequence[SecretConfig]",
+    ) -> None:
+        """Configure persistent secrets for the connection.
+
+        Args:
+            connection: The DuckDB connection to configure secrets for.
+            secrets: The list of secrets to store in the connection.
+
+        Raises:
+            ImproperConfigurationError: If a secret could not be stored in the connection.
+        """
+        try:
+            for secret in secrets:
+                secret_exists = cls._secret_exists(connection, secret["name"])
+                if not secret_exists or secret.get("replace_if_exists", False):
+                    provider_type = "" if not secret.get("provider") else f"provider {secret.get('provider')},"
+                    connection.execute(
+                        f"""create or replace {"persistent" if secret.get("persist", False) else ""} secret {secret["name"]} (
+                        type {secret["secret_type"]},
+                        {provider_type}
+                        {" ,".join([f"{k} '{v}'" for k, v in secret["value"].items()])}
+                    ) """
+                    )
+        except Exception as e:
+            msg = f"Failed to store secret. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
+
+    @classmethod
+    def _configure_extension(cls, connection: "DuckDBPyConnection", extension: "ExtensionConfig") -> None:
         """Configure a single extension for the connection.
 
         Args:
@@ -132,16 +269,32 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
             ImproperConfigurationError: If extension installation or configuration fails.
         """
         try:
-            if extension.get("force_install"):
+            # Install extension if needed
+            if (
+                not cls._extension_installed(connection, extension["name"])
+                and extension.get("install_if_missing", True)
+            ) or extension.get("force_install", False):
+                repository = extension.get("repository", None)
+                repository_url = (
+                    "https://community-extensions.duckdb.org"
+                    if repository is None
+                    and cls._is_community_extension(connection, extension["name"])
+                    and extension.get("repository_url") is None
+                    else extension.get("repository_url", None)
+                )
                 connection.install_extension(
                     extension=extension["name"],
                     force_install=extension.get("force_install", False),
-                    repository=extension.get("repository"),
-                    repository_url=extension.get("repository_url"),
+                    repository=repository,
+                    repository_url=repository_url,
                     version=extension.get("version"),
                 )
-            connection.load_extension(extension["name"])
 
+            # Load extension if not already loaded
+            if not cls._extension_loaded(connection, extension["name"]):
+                connection.load_extension(extension["name"])
+
+            # Apply any configuration settings
             if extension.get("config"):
                 for key, value in extension.get("config", {}).items():
                     connection.execute(f"SET {key}={value}")
@@ -156,7 +309,20 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
         Returns:
             A string keyed dict of config kwargs for the duckdb.connect() function.
         """
-        config = dataclass_to_dict(self, exclude_empty=True, exclude={"extensions"}, convert_nested=False)
+        config = dataclass_to_dict(
+            self,
+            exclude_empty=True,
+            exclude={
+                "extensions",
+                "pool_instance",
+                "secrets",
+                "on_connection_create",
+                "auto_update_extensions",
+                "driver_type",
+                "connection_type",
+            },
+            convert_nested=False,
+        )
         if not config.get("database"):
             config["database"] = ":memory:"
         return config
@@ -175,7 +341,11 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
         try:
             connection = duckdb.connect(**self.connection_config_dict)  # pyright: ignore[reportUnknownMemberType]
             self._configure_extensions(connection)
+            self._configure_secrets(connection, cast("list[SecretConfig]", self.secrets))
             self._configure_connection(connection)
+            if self.on_connection_create:
+                self.on_connection_create(connection)
+
         except Exception as e:
             msg = f"Could not configure the DuckDB connection. Error: {e!s}"
             raise ImproperConfigurationError(msg) from e
@@ -196,3 +366,15 @@ class DuckDBConfig(NoPoolSyncConfig[DuckDBPyConnection]):
             yield connection
         finally:
             connection.close()
+
+    @contextmanager
+    def provide_session(self, *args: Any, **kwargs: Any) -> "Generator[DuckDBDriver, None, None]":
+        """Create and provide a database connection.
+
+        Yields:
+            A DuckDB connection instance.
+
+
+        """
+        with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection, use_cursor=True)