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

sqlspec/adapters/duckdb/config.py

@@ -14,7 +14,7 @@ if TYPE_CHECKING:
     from collections.abc import Generator, Sequence
 
 
-__all__ = ("DuckDB", "ExtensionConfig")
+__all__ = ("DuckDBConfig", "ExtensionConfig")
 
 
 class ExtensionConfig(TypedDict):
@@ -69,7 +69,7 @@ class SecretConfig(TypedDict):
 
 
 @dataclass
-class DuckDB(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
+class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
     """Configuration for DuckDB database connections.
 
     This class provides configuration options for DuckDB database connections, wrapping all parameters

sqlspec/adapters/duckdb/driver.py

@@ -1,24 +1,24 @@
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import SyncDriverAdapterProtocol, T
+from sqlspec.base import SyncArrowBulkOperationsMixin, SyncDriverAdapterProtocol, T
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Generator, Sequence
 
     from duckdb import DuckDBPyConnection
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType
 
 __all__ = ("DuckDBDriver",)
 
 
-class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
+class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDriverAdapterProtocol["DuckDBPyConnection"]):
     """DuckDB Sync Driver Adapter."""
 
     connection: "DuckDBPyConnection"
     use_cursor: bool = True
-    # param_style is inherited from CommonDriverAttributes
+    dialect: str = "duckdb"
 
     def __init__(self, connection: "DuckDBPyConnection", use_cursor: bool = True) -> None:
         self.connection = connection
@@ -27,7 +27,6 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
     # --- Helper Methods --- #
     def _cursor(self, connection: "DuckDBPyConnection") -> "DuckDBPyConnection":
         if self.use_cursor:
-            # Ignore lack of type hint on cursor()
             return connection.cursor()
         return connection
 
@@ -40,20 +39,43 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
             finally:
                 cursor.close()
         else:
-            yield connection  # Yield the connection directly
-
-    # --- Public API Methods (Original Implementation + _process_sql_params) --- #
+            yield connection
 
+    # --- Public API Methods --- #
+    @overload
     def select(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        connection: Optional["DuckDBPyConnection"] = None,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        **kwargs: Any,
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             results = cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
@@ -66,21 +88,44 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
                 return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore[reportUnknownArgumentType]
             return [dict(zip(column_names, row)) for row in results]  # pyright: ignore[reportUnknownArgumentType]
 
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     def select_one(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["DuckDBPyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            result = self.check_not_found(result)  #  pyright: ignore
+            result = self.check_not_found(result)  # pyright: ignore
 
             column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
             if schema_type is not None:
@@ -88,17 +133,40 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
             # Always return dictionaries
             return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
 
+    @overload
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = 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[DuckDBPyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     def select_one_or_none(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["DuckDBPyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
@@ -110,35 +178,82 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
                 return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
             return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
 
+    @overload
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
     def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[DuckDBPyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Union[T, Any]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
-
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
-            result = self.check_not_found(result)  #  pyright: ignore
+            result = self.check_not_found(result)  # pyright: ignore
             if schema_type is None:
                 return result[0]  # pyright: ignore
             return schema_type(result[0])  # type: ignore[call-arg]
 
+    @overload
     def select_value_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
         connection: "Optional[DuckDBPyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
@@ -153,73 +268,88 @@ class DuckDBDriver(SyncDriverAdapterProtocol["DuckDBPyConnection"]):
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["DuckDBPyConnection"] = None,
+        **kwargs: Any,
     ) -> int:
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             return getattr(cursor, "rowcount", -1)  # pyright: ignore[reportUnknownMemberType]
 
+    @overload
     def insert_update_delete_returning(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        connection: Optional["DuckDBPyConnection"] = None,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = 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[DuckDBPyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        **kwargs: Any,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             result = cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            if not result:
-                return None  # pyright: ignore[reportUnknownArgumentType]
+            result = self.check_not_found(result)  # pyright: ignore
             column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
             if schema_type is not None:
                 return cast("ModelDTOT", schema_type(**dict(zip(column_names, result[0]))))  # pyright: ignore[reportUnknownArgumentType]
             # Always return dictionaries
             return dict(zip(column_names, result[0]))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
 
-    def _process_sql_params(
-        self, sql: str, parameters: "Optional[StatementParameterType]" = None
-    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
-        """Process SQL query and parameters for DB-API execution.
-
-        Converts named parameters (:name) to positional parameters (?) for DuckDB.
-
-        Args:
-            sql: The SQL query string.
-            parameters: The parameters for the query (dict, tuple, list, or None).
-
-        Returns:
-            A tuple containing the processed SQL string and the processed parameters.
-        """
-        if not isinstance(parameters, dict) or not parameters:
-            # If parameters are not a dict, or empty dict, assume positional/no params
-            # Let the underlying driver handle tuples/lists directly
-            return sql, parameters
-
-        # Convert named parameters to positional parameters
-        processed_sql = sql
-        processed_params: list[Any] = []
-        for key, value in parameters.items():
-            # Replace :key with ? in the SQL
-            processed_sql = processed_sql.replace(f":{key}", "?")
-            processed_params.append(value)
-
-        return processed_sql, tuple(processed_params)
-
     def execute_script(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
+        *,
         connection: Optional["DuckDBPyConnection"] = None,
+        **kwargs: Any,
     ) -> str:
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         with self._with_cursor(connection) as cursor:
             cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
             return cast("str", getattr(cursor, "statusmessage", "DONE"))  # pyright: ignore[reportUnknownMemberType]
+
+    # --- Arrow Bulk Operations ---
+
+    def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[DuckDBPyConnection]" = None,
+        **kwargs: Any,
+    ) -> "ArrowTable":
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            return cast("ArrowTable", cursor.fetch_arrow_table())  # pyright: ignore[reportUnknownMemberType]

sqlspec/adapters/oracledb/__init__.py

@@ -1,16 +1,16 @@
 from sqlspec.adapters.oracledb.config import (
-    OracleAsync,
-    OracleAsyncPool,
-    OracleSync,
-    OracleSyncPool,
+    OracleAsyncConfig,
+    OracleAsyncPoolConfig,
+    OracleSyncConfig,
+    OracleSyncPoolConfig,
 )
 from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver
 
 __all__ = (
-    "OracleAsync",
+    "OracleAsyncConfig",
     "OracleAsyncDriver",
-    "OracleAsyncPool",
-    "OracleSync",
+    "OracleAsyncPoolConfig",
+    "OracleSyncConfig",
     "OracleSyncDriver",
-    "OracleSyncPool",
+    "OracleSyncPoolConfig",
 )

sqlspec/adapters/oracledb/config/__init__.py

@@ -1,9 +1,9 @@
-from sqlspec.adapters.oracledb.config._asyncio import OracleAsync, OracleAsyncPool
-from sqlspec.adapters.oracledb.config._sync import OracleSync, OracleSyncPool
+from sqlspec.adapters.oracledb.config._asyncio import OracleAsyncConfig, OracleAsyncPoolConfig
+from sqlspec.adapters.oracledb.config._sync import OracleSyncConfig, OracleSyncPoolConfig
 
 __all__ = (
-    "OracleAsync",
-    "OracleAsyncPool",
-    "OracleSync",
-    "OracleSyncPool",
+    "OracleAsyncConfig",
+    "OracleAsyncPoolConfig",
+    "OracleSyncConfig",
+    "OracleSyncPoolConfig",
 )

sqlspec/adapters/oracledb/config/_asyncio.py

@@ -1,6 +1,6 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from oracledb import create_pool_async as oracledb_create_pool  # pyright: ignore[reportUnknownVariableType]
 from oracledb.connection import AsyncConnection
@@ -18,18 +18,18 @@ if TYPE_CHECKING:
 
 
 __all__ = (
-    "OracleAsync",
-    "OracleAsyncPool",
+    "OracleAsyncConfig",
+    "OracleAsyncPoolConfig",
 )
 
 
 @dataclass
-class OracleAsyncPool(OracleGenericPoolConfig["AsyncConnection", "AsyncConnectionPool"]):
+class OracleAsyncPoolConfig(OracleGenericPoolConfig["AsyncConnection", "AsyncConnectionPool"]):
     """Async Oracle Pool Config"""
 
 
 @dataclass
-class OracleAsync(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool", "OracleAsyncDriver"]):
+class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool", "OracleAsyncDriver"]):
     """Oracle Async database Configuration.
 
     This class provides the base configuration for Oracle database connections, extending
@@ -42,7 +42,7 @@ class OracleAsync(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool",
     options.([2](https://python-oracledb.readthedocs.io/en/latest/user_guide/tuning.html))
     """
 
-    pool_config: "Optional[OracleAsyncPool]" = None
+    pool_config: "Optional[OracleAsyncPoolConfig]" = None
     """Oracle Pool configuration"""
     pool_instance: "Optional[AsyncConnectionPool]" = None
     """Optional pool to use.
@@ -112,7 +112,7 @@ class OracleAsync(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool",
         raise ImproperConfigurationError(msg)
 
     async def create_connection(self) -> "AsyncConnection":
-        """Create and return a new oracledb async connection.
+        """Create and return a new oracledb async connection from the pool.
 
         Returns:
             An AsyncConnection instance.
@@ -121,9 +121,8 @@ class OracleAsync(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool",
             ImproperConfigurationError: If the connection could not be created.
         """
         try:
-            import oracledb
-
-            return await oracledb.connect_async(**self.connection_config_dict)  # type: ignore[no-any-return]
+            pool = await self.provide_pool()
+            return cast("AsyncConnection", await pool.acquire())  # type: ignore[no-any-return,unused-ignore]
         except Exception as e:
             msg = f"Could not configure the Oracle async connection. Error: {e!s}"
             raise ImproperConfigurationError(msg) from e

sqlspec/adapters/oracledb/config/_sync.py

@@ -18,18 +18,18 @@ if TYPE_CHECKING:
 
 
 __all__ = (
-    "OracleSync",
-    "OracleSyncPool",
+    "OracleSyncConfig",
+    "OracleSyncPoolConfig",
 )
 
 
 @dataclass
-class OracleSyncPool(OracleGenericPoolConfig["Connection", "ConnectionPool"]):
+class OracleSyncPoolConfig(OracleGenericPoolConfig["Connection", "ConnectionPool"]):
     """Sync Oracle Pool Config"""
 
 
 @dataclass
-class OracleSync(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncDriver"]):
+class OracleSyncConfig(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncDriver"]):
     """Oracle Sync database Configuration.
 
     This class provides the base configuration for Oracle database connections, extending
@@ -42,7 +42,7 @@ class OracleSync(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncD
     options.([2](https://python-oracledb.readthedocs.io/en/latest/user_guide/tuning.html))
     """
 
-    pool_config: "Optional[OracleSyncPool]" = None
+    pool_config: "Optional[OracleSyncPoolConfig]" = None
     """Oracle Pool configuration"""
     pool_instance: "Optional[ConnectionPool]" = None
     """Optional pool to use.
@@ -112,7 +112,7 @@ class OracleSync(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncD
         raise ImproperConfigurationError(msg)
 
     def create_connection(self) -> "Connection":
-        """Create and return a new oracledb connection.
+        """Create and return a new oracledb connection from the pool.
 
         Returns:
             A Connection instance.
@@ -121,9 +121,8 @@ class OracleSync(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncD
             ImproperConfigurationError: If the connection could not be created.
         """
         try:
-            import oracledb
-
-            return oracledb.connect(**self.connection_config_dict)
+            pool = self.provide_pool()
+            return pool.acquire()
         except Exception as e:
             msg = f"Could not configure the Oracle connection. Error: {e!s}"
             raise ImproperConfigurationError(msg) from e