sqlspec 0.10.1__py3-none-any.whl → 0.11.1__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.filters import StatementFilter
+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 collections.abc import Generator, Mapping, Sequence
 
     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,14 +54,58 @@ 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.
+        """
+        data_params_for_statement: Optional[Union[Mapping[str, Any], Sequence[Any]]] = None
+        combined_filters_list: list[StatementFilter] = list(filters)
+
+        if parameters is not None:
+            if isinstance(parameters, StatementFilter):
+                combined_filters_list.insert(0, parameters)
+            else:
+                data_params_for_statement = parameters
+        if data_params_for_statement is not None and not isinstance(data_params_for_statement, (list, tuple, dict)):
+            data_params_for_statement = (data_params_for_statement,)
+        statement = SQLStatement(sql, data_params_for_statement, kwargs=kwargs, dialect=self.dialect)
+        for filter_obj in combined_filters_list:
+            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(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -66,8 +115,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -76,33 +124,32 @@ class DuckDBDriver(
         self,
         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 = [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]
+            column_names = [column[0] for column in cursor.description or []]
+            return self.to_schema([dict(zip(column_names, row)) for row in results], schema_type=schema_type)
 
     @overload
     def select_one(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -112,8 +159,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -121,33 +167,32 @@ class DuckDBDriver(
     def select_one(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
-        /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        parameters: "Optional[StatementParameterType]" = 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
-
-            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]
+            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 []]
+            return self.to_schema(dict(zip(column_names, result)), schema_type=schema_type)
 
     @overload
     def select_one_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -157,8 +202,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -166,33 +210,33 @@ class DuckDBDriver(
     def select_one_or_none(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
-        /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        parameters: "Optional[StatementParameterType]" = 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 = [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]
+            column_names = [column[0] for column in cursor.description or []]
+            return self.to_schema(dict(zip(column_names, result)), schema_type=schema_type)
 
     @overload
     def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -202,8 +246,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -212,29 +255,33 @@ class DuckDBDriver(
         self,
         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(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -244,8 +291,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
@@ -254,45 +300,43 @@ class DuckDBDriver(
         self,
         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,
-        /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        parameters: "Optional[StatementParameterType]" = 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(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
@@ -302,8 +346,7 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[DuckDBConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
@@ -312,38 +355,34 @@ class DuckDBDriver(
         self,
         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 []]
+            return self.to_schema(dict(zip(column_names, result[0])), schema_type=schema_type)
 
     def execute_script(
         self,
         sql: str,
-        parameters: Optional["StatementParameterType"] = None,
-        /,
-        *,
-        connection: Optional["DuckDBConnection"] = None,
+        parameters: "Optional[StatementParameterType]" = 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 ---
 
@@ -351,13 +390,36 @@ class DuckDBDriver(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         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.
+            *filters: Optional filters to apply to the SQL statement.
+            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)
+        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