sqlspec 0.10.1__py3-none-any.whl → 0.11.0__py3-none-any.whl

sqlspec/adapters/duckdb/driver.py

@@ -1,19 +1,24 @@
+import logging
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
 from duckdb import DuckDBPyConnection
 
 from sqlspec.base import SyncDriverAdapterProtocol
-from sqlspec.mixins import SQLTranslatorMixin, SyncArrowBulkOperationsMixin
+from sqlspec.mixins import ResultConverter, SQLTranslatorMixin, SyncArrowBulkOperationsMixin
+from sqlspec.statement import SQLStatement
 from sqlspec.typing import ArrowTable, StatementParameterType
 
 if TYPE_CHECKING:
     from collections.abc import Generator, Sequence
 
+    from sqlspec.filters import StatementFilter
     from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType, T
 
 __all__ = ("DuckDBConnection", "DuckDBDriver")
 
+logger = logging.getLogger("sqlspec")
+
 DuckDBConnection = DuckDBPyConnection
 
 
@@ -21,6 +26,7 @@ class DuckDBDriver(
     SyncArrowBulkOperationsMixin["DuckDBConnection"],
     SQLTranslatorMixin["DuckDBConnection"],
     SyncDriverAdapterProtocol["DuckDBConnection"],
+    ResultConverter,
 ):
     """DuckDB Sync Driver Adapter."""
 
@@ -32,7 +38,6 @@ class DuckDBDriver(
         self.connection = connection
         self.use_cursor = use_cursor
 
-    # --- Helper Methods --- #
     def _cursor(self, connection: "DuckDBConnection") -> "DuckDBConnection":
         if self.use_cursor:
             return connection.cursor()
@@ -49,6 +54,44 @@ class DuckDBDriver(
         else:
             yield connection
 
+    def _process_sql_params(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *filters: "StatementFilter",
+        **kwargs: Any,
+    ) -> "tuple[str, Optional[Union[tuple[Any, ...], list[Any], dict[str, Any]]]]":
+        """Process SQL and parameters for DuckDB using SQLStatement.
+
+        DuckDB supports both named (:name, $name) and positional (?) parameters.
+        This method processes the SQL with dialect-aware parsing and handles
+        parameters appropriately for DuckDB.
+
+        Args:
+            sql: SQL statement.
+            parameters: Query parameters.
+            *filters: Statement filters to apply.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Tuple of processed SQL and parameters.
+        """
+        statement = SQLStatement(sql, parameters, kwargs=kwargs, dialect=self.dialect)
+
+        # Apply any filters
+        for filter_obj in filters:
+            statement = statement.apply_filter(filter_obj)
+
+        processed_sql, processed_params, _ = statement.process()
+        if processed_params is None:
+            return processed_sql, None
+        if isinstance(processed_params, dict):
+            return processed_sql, processed_params
+        if isinstance(processed_params, (list, tuple)):
+            return processed_sql, tuple(processed_params)
+        return processed_sql, (processed_params,)  # type: ignore[unreachable]
+
     # --- Public API Methods --- #
     @overload
     def select(
@@ -56,7 +99,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -67,7 +110,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -77,24 +120,28 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[dict[str, Any], ModelDTOT]]":
+        """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, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
         with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            results = cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            cursor.execute(sql, [] if parameters is None else parameters)
+            results = cursor.fetchall()
             if not results:
                 return []
+            column_names = [column[0] for column in cursor.description or []]
 
-            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, row)))) for row in results]  # pyright: ignore[reportUnknownArgumentType]
-            return [dict(zip(column_names, row)) for row in results]  # pyright: ignore[reportUnknownArgumentType]
+            # Convert to dicts first
+            dict_results = [dict(zip(column_names, row)) for row in results]
+            return self.to_schema(dict_results, schema_type=schema_type)
 
     @overload
     def select_one(
@@ -102,7 +149,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -113,7 +160,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -121,25 +168,29 @@ class DuckDBDriver(
     def select_one(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[DuckDBConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Union[ModelDTOT, dict[str, Any]]":
+    ) -> "Union[dict[str, Any], ModelDTOT]":
+        """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, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **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
+            cursor.execute(sql, [] if parameters is None else parameters)
+            result = cursor.fetchone()
+            result = self.check_not_found(result)
+            column_names = [column[0] for column in cursor.description or []]
 
-            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))))  # pyright: ignore[reportUnknownArgumentType]
-            # Always return dictionaries
-            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+            # Convert to dict and use ResultConverter
+            dict_result = dict(zip(column_names, result))
+            return self.to_schema(dict_result, schema_type=schema_type)
 
     @overload
     def select_one_or_none(
@@ -147,7 +198,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -158,7 +209,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -166,25 +217,30 @@ class DuckDBDriver(
     def select_one_or_none(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[DuckDBConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query results, or None if no results.
+        """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
         with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            cursor.execute(sql, [] if parameters is None else parameters)
+            result = cursor.fetchone()
             if result is None:
                 return None
+            column_names = [column[0] for column in cursor.description or []]
 
-            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))))  # pyright: ignore[reportUnknownArgumentType]
-            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+            # Convert to dict and use ResultConverter
+            dict_result = dict(zip(column_names, result))
+            return self.to_schema(dict_result, schema_type=schema_type)
 
     @overload
     def select_value(
@@ -192,7 +248,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -203,7 +259,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -213,20 +269,26 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results.
+        """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **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
+            cursor.execute(sql, [] if parameters is None else parameters)
+            result = cursor.fetchone()
+            result = self.check_not_found(result)
+            result_value = result[0]
             if schema_type is None:
-                return result[0]  # pyright: ignore
-            return schema_type(result[0])  # type: ignore[call-arg]
+                return result_value
+            return schema_type(result_value)  # type: ignore[call-arg]
 
     @overload
     def select_value_or_none(
@@ -234,7 +296,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -245,7 +307,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -255,36 +317,37 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = 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, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
         with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            cursor.execute(sql, [] if parameters is None else parameters)
+            result = cursor.fetchone()
             if result is None:
                 return None
             if schema_type is None:
-                return result[0]  # pyright: ignore
+                return result[0]
             return schema_type(result[0])  # type: ignore[call-arg]
 
     def insert_update_delete(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        *filters: "StatementFilter",
+        connection: "Optional[DuckDBConnection]" = None,
         **kwargs: Any,
     ) -> int:
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
         with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            return getattr(cursor, "rowcount", -1)  # pyright: ignore[reportUnknownMemberType]
+            params = [] if parameters is None else parameters
+            cursor.execute(sql, params)
+            return getattr(cursor, "rowcount", -1)
 
     @overload
     def insert_update_delete_returning(
@@ -292,7 +355,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -303,7 +366,7 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -313,37 +376,38 @@ class DuckDBDriver(
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = 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, **kwargs)
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **kwargs)
         with self._with_cursor(connection) as cursor:
-            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
-            result = cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-            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]
+            params = [] if parameters is None else parameters
+            cursor.execute(sql, params)
+            result = cursor.fetchall()
+            result = self.check_not_found(result)
+            column_names = [col[0] for col in cursor.description or []]
+
+            # Convert to dict and use ResultConverter
+            dict_result = dict(zip(column_names, result[0]))
+            return self.to_schema(dict_result, schema_type=schema_type)
 
     def execute_script(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
+        parameters: "Optional[StatementParameterType]" = None,
         /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        connection: "Optional[DuckDBConnection]" = None,
         **kwargs: Any,
     ) -> str:
         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("str", getattr(cursor, "statusmessage", "DONE"))  # pyright: ignore[reportUnknownMemberType]
+            params = [] if parameters is None else parameters
+            cursor.execute(sql, params)
+            return cast("str", getattr(cursor, "statusmessage", "DONE"))
 
     # --- Arrow Bulk Operations ---
 
@@ -356,8 +420,35 @@ class DuckDBDriver(
         connection: "Optional[DuckDBConnection]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":
+        """Execute a SQL query and return results as an Apache Arrow Table.
+
+        Args:
+            sql: The SQL query string.
+            parameters: Parameters for the query.
+            connection: Optional connection override.
+            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
+
+        Returns:
+            An Apache Arrow Table containing the query results.
+        """
         connection = self._connection(connection)
-        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        # Extract filters from kwargs if present
+        filters = kwargs.pop("filters", [])
+
+        sql, parameters = self._process_sql_params(sql, parameters, *filters, **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]
+            params = [] if parameters is None else parameters
+            cursor.execute(sql, params)
+            return cast("ArrowTable", cursor.fetch_arrow_table())
+
+    def _connection(self, connection: "Optional[DuckDBConnection]" = None) -> "DuckDBConnection":
+        """Get the connection to use for the operation.
+
+        Args:
+            connection: Optional connection to use.
+
+        Returns:
+            The connection to use.
+        """
+        return connection or self.connection