"""DuckDB driver implementation."""
import contextlib
from time import perf_counter
from typing import TYPE_CHECKING, Any, Literal, cast
from sqlglot import exp
from sqlspec.adapters.duckdb._typing import DuckDBCursor, DuckDBSessionContext
from sqlspec.adapters.duckdb._typing import duckdb_module as duckdb
from sqlspec.adapters.duckdb.core import (
_build_storage_copy_sql,
_build_storage_read_sql,
_DuckDBStreamSource,
_resolve_native_storage_target,
_restore_uuid_columns,
collect_rows,
create_mapped_exception,
default_statement_config,
driver_profile,
format_identifier,
normalize_execute_parameters,
resolve_rowcount,
)
from sqlspec.adapters.duckdb.data_dictionary import DuckDBDataDictionary
from sqlspec.core import (
SQL,
ParameterStyle,
StatementConfig,
build_arrow_result_from_reader,
build_arrow_result_from_table,
get_cache_config,
register_driver_profile,
)
from sqlspec.core.result import DMLResult
from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream
from sqlspec.exceptions import SQLSpecError
from sqlspec.typing import import_optional
from sqlspec.utils.logging import get_logger
from sqlspec.utils.module_loader import ensure_pyarrow
from sqlspec.utils.text import quote_identifier
from sqlspec.utils.uuids import uuid4
if TYPE_CHECKING:
from collections.abc import Sequence
import pyarrow as pa
from sqlspec.adapters.duckdb._typing import DuckDBConnection
from sqlspec.builder import QueryBuilder
from sqlspec.core import ArrowResult, SQLResult, Statement, StatementFilter
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
from sqlspec.typing import ArrowReturnFormat, StatementParameters
else:
pa = import_optional("pyarrow")
__all__ = ("DuckDBCursor", "DuckDBDriver", "DuckDBExceptionHandler", "DuckDBSessionContext")
logger = get_logger("sqlspec.adapters.duckdb")
def _is_plain_values_insert(expression: exp.Insert, column_count: int) -> bool:
if any(expression.args.get(key) for key in ("conflict", "alternative", "with_")):
return False
values = expression.expression
if not isinstance(values, exp.Values) or len(values.expressions) != 1:
return False
row = values.expressions[0]
if not isinstance(row, exp.Tuple) or not column_count or len(row.expressions) != column_count:
return False
for index, item in enumerate(row.expressions, 1):
if isinstance(item, exp.Placeholder):
if item.name not in ("", "?", str(index)):
return False
elif not isinstance(item, exp.Parameter) or item.name != str(index):
return False
return True
class DuckDBExceptionHandler(BaseSyncExceptionHandler):
"""Context manager for handling DuckDB database exceptions.
Uses exception type and message-based detection to map DuckDB errors
to specific SQLSpec exceptions for better error handling.
Uses deferred exception pattern for mypyc compatibility: exceptions
are stored in pending_exception rather than raised from __exit__
to avoid ABI boundary violations with compiled code.
"""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if issubclass(exc_type, duckdb.Error):
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
[docs]
class DuckDBDriver(SyncDriverAdapterBase):
"""Synchronous DuckDB database driver.
Provides SQL statement execution, transaction management, and result handling
for DuckDB databases. Supports multiple parameter styles including QMARK,
NUMERIC, and NAMED_DOLLAR formats.
The driver handles script execution, batch operations, and integrates with
the sqlspec.core modules for statement processing and caching.
"""
__slots__ = ("_data_dictionary", "_transaction_active")
dialect = "duckdb"
[docs]
def __init__(
self,
connection: "DuckDBConnection",
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = default_statement_config.replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
driver_features = dict(driver_features) if driver_features else {}
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self._data_dictionary: DuckDBDataDictionary | None = None
self._transaction_active = False
[docs]
def dispatch_execute(self, cursor: "DuckDBConnection", statement: SQL) -> "ExecutionResult":
"""Execute single SQL statement with data handling.
Executes a SQL statement with parameter binding and processes the results.
Handles both data-returning queries and data modification operations.
Args:
cursor: DuckDB cursor object
statement: SQL statement to execute
Returns:
ExecutionResult with execution metadata
"""
sql, prepared_parameters = self._compiled_sql(statement, self.statement_config)
cursor.execute(sql, normalize_execute_parameters(prepared_parameters))
is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor)
if is_select_like:
arrow_table = cursor.to_arrow_table()
data = arrow_table.to_pylist()
_restore_uuid_columns(data, cursor.description)
column_names = list(arrow_table.column_names)
return self.create_execution_result(
cursor,
selected_data=data,
column_names=column_names,
data_row_count=len(data),
is_select_result=True,
row_format="dict",
)
row_count = resolve_rowcount(cursor)
return self.create_execution_result(cursor, rowcount_override=row_count)
[docs]
def dispatch_execute_many(self, cursor: "DuckDBConnection", statement: SQL) -> "ExecutionResult":
"""Execute SQL with multiple parameter sets using batch processing.
Uses DuckDB's executemany method for batch operations and calculates
row counts for both data modification and query operations.
Args:
cursor: DuckDB cursor object
statement: SQL statement with multiple parameter sets
Returns:
ExecutionResult with batch execution metadata
"""
sql, prepared_parameters = self._compiled_sql(statement, self.statement_config)
if prepared_parameters:
parameter_sets = cast("list[Any]", prepared_parameters)
cursor.executemany(sql, parameter_sets)
row_count = len(parameter_sets) if statement.is_modifying_operation() else resolve_rowcount(cursor)
else:
row_count = 0
return self.create_execution_result(cursor, rowcount_override=row_count, is_many_result=True)
[docs]
def execute_many(
self,
statement: "SQL | Statement | QueryBuilder",
/,
parameters: "Sequence[StatementParameters]",
*filters: "StatementParameters | StatementFilter",
statement_config: "StatementConfig | None" = None,
**kwargs: Any,
) -> "SQLResult":
"""Execute many with a DuckDB bulk insert fast path for simple INSERT batches."""
config = statement_config or self.statement_config
if isinstance(statement, str) and not filters and not kwargs and config is self.statement_config:
prepared_statement = SQL(statement, tuple(parameters), statement_config=config, is_many=True)
cached_statement, prepared_parameters = self._compiled_statement(prepared_statement, config)
parsed_expression = cached_statement.expression
if isinstance(parsed_expression, exp.Insert) and not parsed_expression.args.get("returning"):
bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters)
if bulk_result is not None:
return bulk_result
return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs)
[docs]
def dispatch_execute_script(self, cursor: "DuckDBConnection", statement: SQL) -> "ExecutionResult":
"""Execute SQL script with statement splitting and parameter handling.
Parses multi-statement scripts and executes each statement sequentially
with the provided parameters.
Args:
cursor: DuckDB cursor object
statement: SQL statement with script content
Returns:
ExecutionResult with script execution metadata
"""
sql, prepared_parameters = self._compiled_sql(statement, self.statement_config)
statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
successful_count = 0
last_result = None
for stmt in statements:
last_result = cursor.execute(stmt, normalize_execute_parameters(prepared_parameters))
successful_count += 1
return self.create_execution_result(
last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True
)
[docs]
def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None":
"""Return a native DuckDB row stream backed by Arrow record batches."""
if not statement.returns_rows():
return None
sql, prepared_parameters = self._compiled_sql(statement, self.statement_config)
return SyncRowStream(_DuckDBStreamSource(self._open_stream_reader, sql, prepared_parameters, chunk_size))
[docs]
def begin(self) -> None:
"""Begin a database transaction."""
try:
self.connection.execute("BEGIN TRANSACTION")
except duckdb.Error as e:
msg = f"Failed to begin DuckDB transaction: {e}"
raise SQLSpecError(msg) from e
self._transaction_active = True
[docs]
def commit(self) -> None:
"""Commit the current transaction."""
try:
self.connection.commit()
except duckdb.Error as e:
msg = f"Failed to commit DuckDB transaction: {e}"
raise SQLSpecError(msg) from e
self._transaction_active = False
[docs]
def rollback(self) -> None:
"""Rollback the current transaction."""
try:
self.connection.rollback()
except duckdb.Error as e:
msg = f"Failed to rollback DuckDB transaction: {e}"
raise SQLSpecError(msg) from e
finally:
self._transaction_active = False
[docs]
def create_savepoint(self, name: str) -> None:
"""Raise because DuckDB does not support savepoints.
Raises:
NotImplementedError: Always.
"""
msg = "DuckDB does not support savepoints."
raise NotImplementedError(msg)
[docs]
def release_savepoint(self, name: str) -> None:
"""Raise because DuckDB does not support savepoints.
Raises:
NotImplementedError: Always.
"""
msg = "DuckDB does not support savepoints."
raise NotImplementedError(msg)
[docs]
def rollback_to_savepoint(self, name: str) -> None:
"""Raise because DuckDB does not support savepoints.
Raises:
NotImplementedError: Always.
"""
msg = "DuckDB does not support savepoints."
raise NotImplementedError(msg)
[docs]
def set_migration_session_schema(self, schema: str) -> None:
"""Set DuckDB search_path for migration SQL."""
self.connection.execute(f"SET search_path = {quote_identifier(schema)}")
[docs]
def has_schema(self, schema: str) -> bool:
"""Return whether a DuckDB schema exists."""
result = self.connection.execute(
"SELECT 1 FROM information_schema.schemata WHERE schema_name = ?", [schema]
).fetchone()
return result is not None
[docs]
def with_cursor(self, connection: "DuckDBConnection") -> "DuckDBCursor":
"""Create context manager for DuckDB cursor.
Args:
connection: DuckDB connection instance
Returns:
DuckDBCursor context manager instance
"""
return DuckDBCursor(connection)
[docs]
def handle_database_exceptions(self) -> "DuckDBExceptionHandler":
"""Handle database-specific exceptions and wrap them appropriately.
Returns:
Exception handler with deferred exception pattern for mypyc compatibility.
"""
return DuckDBExceptionHandler()
[docs]
def select_to_arrow(
self,
statement: "Statement | QueryBuilder",
/,
*parameters: "StatementParameters | StatementFilter",
statement_config: "StatementConfig | None" = None,
return_format: "ArrowReturnFormat" = "table",
native_only: bool = False,
batch_size: int | None = None,
arrow_schema: Any = None,
**kwargs: Any,
) -> "ArrowResult":
"""Execute query and return results as Apache Arrow (DuckDB native path).
DuckDB provides native Arrow support via cursor.arrow().
This is the fastest path due to DuckDB's columnar architecture.
Args:
statement: SQL statement, string, or QueryBuilder
*parameters: Query parameters or filters
statement_config: Optional statement configuration override
return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch,
"batches" for list of RecordBatch, "reader" for RecordBatchReader
native_only: Ignored for DuckDB (always uses native path)
batch_size: Batch size hint (for future streaming implementation)
arrow_schema: Optional pyarrow.Schema for type casting
**kwargs: Additional keyword arguments
Returns:
ArrowResult with native Arrow data
"""
ensure_pyarrow()
config = statement_config or self.statement_config
prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
exc_handler = self.handle_database_exceptions()
arrow_result: ArrowResult | None = None
with self.with_cursor(self.connection) as cursor, exc_handler:
sql, driver_params = self._compiled_sql(prepared_statement, config)
cursor.execute(sql, driver_params or ())
if return_format in {"reader", "batches"}:
arrow_reader = (
cursor.to_arrow_reader(batch_size) if batch_size is not None else cursor.to_arrow_reader()
)
arrow_result = build_arrow_result_from_reader(
prepared_statement,
arrow_reader,
return_format=return_format,
batch_size=batch_size,
arrow_schema=arrow_schema,
)
else:
arrow_table = cursor.to_arrow_table()
arrow_result = build_arrow_result_from_table(
prepared_statement,
arrow_table,
return_format=return_format,
batch_size=batch_size,
arrow_schema=arrow_schema,
)
self._check_pending_exception(exc_handler)
if arrow_result is None:
msg = "Unreachable"
raise RuntimeError(msg) # pragma: no cover
return arrow_result
[docs]
def select_to_storage(
self,
statement: "Statement | QueryBuilder | SQL | str",
destination: "StorageDestination",
/,
*parameters: "StatementParameters | StatementFilter",
statement_config: "StatementConfig | None" = None,
partitioner: "dict[str, object] | None" = None,
format_hint: "StorageFormat | None" = None,
telemetry: "StorageTelemetry | None" = None,
**kwargs: Any,
) -> "StorageBridgeJob":
"""Persist a query with native object-store COPY when settings are compatible."""
self._require_capability("arrow_export_enabled")
sync_pipeline = self._storage_pipeline()
format_choice = format_hint or "parquet"
if (
self.storage_pipeline_factory is None
and format_choice in {"parquet", "csv"}
and not {"arrow_schema", "return_format", "batch_size", "native_only"}.intersection(kwargs)
):
target = _resolve_native_storage_target(sync_pipeline, destination, self.driver_features, write=True)
if target is not None:
config = statement_config or self.statement_config
config = config.replace(
parameter_config=config.parameter_config.replace(
default_execution_parameter_style=ParameterStyle.QMARK
)
)
prepared = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
sql, driver_params = self._compiled_sql(prepared, config)
copy_sql = _build_storage_copy_sql(sql, format_choice)
if copy_sql is not None:
return self._native_storage_job(
copy_sql,
[target.uri, *cast("Sequence[Any]", driver_params or ())],
destination=str(destination),
format_choice=format_choice,
operation="write",
partitioner=partitioner,
telemetry=telemetry,
)
arrow_result = self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs)
telemetry_payload = self._write_storage_result(
arrow_result, destination, format_hint=format_hint, pipeline=sync_pipeline
)
self._attach_partition_telemetry(telemetry_payload, partitioner)
return self._storage_job(telemetry_payload, telemetry)
[docs]
def load_from_arrow(
self,
table: str,
source: "ArrowResult | Any",
*,
partitioner: "dict[str, object] | None" = None,
overwrite: bool = False,
telemetry: "StorageTelemetry | None" = None,
) -> "StorageBridgeJob":
"""Load Arrow data into DuckDB using temporary table registration."""
self._require_capability("arrow_import_enabled")
ensure_pyarrow()
source_data = source.get_data() if hasattr(source, "get_data") else source
arrow_table = None
arrow_source: object
if isinstance(source_data, pa.RecordBatchReader):
arrow_source = source_data
else:
arrow_table = self._coerce_arrow_table(source_data)
arrow_source = arrow_table
table_name = format_identifier(table)
temp_view = f"_sqlspec_arrow_{uuid4().hex}"
if overwrite:
self.connection.execute(f"TRUNCATE TABLE {table_name}")
self.connection.register(temp_view, arrow_source)
inserted_rows = 0
try:
insert_result = self.connection.execute(f"INSERT INTO {table_name} SELECT * FROM {temp_view}")
inserted_rows = _resolve_duckdb_inserted_rows(insert_result)
finally:
with contextlib.suppress(Exception):
self.connection.unregister(temp_view)
if isinstance(source_data, pa.RecordBatchReader):
telemetry_payload: StorageTelemetry = {
"rows_processed": inserted_rows,
"bytes_processed": 0,
"format": "arrow",
}
else:
if arrow_table is None:
msg = "DuckDB Arrow load did not resolve an Arrow table."
raise SQLSpecError(msg)
telemetry_payload = self._ingest_telemetry(arrow_table)
telemetry_payload["destination"] = table
self._attach_partition_telemetry(telemetry_payload, partitioner)
return self._storage_job(telemetry_payload, telemetry)
[docs]
def load_from_storage(
self,
table: str,
source: "StorageDestination",
*,
file_format: "StorageFormat",
partitioner: "dict[str, object] | None" = None,
overwrite: bool = False,
) -> "StorageBridgeJob":
"""Read an artifact from storage and load it into DuckDB."""
self._require_capability("arrow_import_enabled")
if not overwrite and file_format == "parquet" and self.storage_pipeline_factory is None:
pipeline = self._storage_pipeline()
target = _resolve_native_storage_target(pipeline, source, self.driver_features, write=False)
if target is not None:
read_sql = _build_storage_read_sql(table, target.uri, file_format)
if read_sql is not None:
return self._native_storage_job(
read_sql,
[target.uri],
destination=table,
format_choice=file_format,
operation="read",
partitioner=partitioner,
)
arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format)
return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound)
@property
def data_dictionary(self) -> "DuckDBDataDictionary":
"""Get the data dictionary for this driver.
Returns:
Data dictionary instance for metadata queries
"""
if self._data_dictionary is None:
self._data_dictionary = DuckDBDataDictionary()
return self._data_dictionary
def _native_storage_job(
self,
sql: str,
parameters: "list[Any]",
*,
destination: str,
format_choice: "StorageFormat",
operation: "Literal['read', 'write']",
partitioner: "dict[str, object] | None",
telemetry: "StorageTelemetry | None" = None,
) -> "StorageBridgeJob":
runtime = self.observability
span = runtime.start_storage_span(operation, destination=destination, format_label=format_choice)
start = perf_counter()
try:
rows = self._execute_storage_sql(sql, parameters)
path = destination
if operation == "write":
path = destination.partition("://")[2]
if destination.startswith("alias://"):
path = path.partition("/")[2].strip().lstrip("/")
payload: StorageTelemetry = {
"destination": path,
"rows_processed": rows,
"duration_s": perf_counter() - start,
"format": format_choice,
"backend": "duckdb",
}
payload = runtime.annotate_storage_telemetry(payload)
except Exception as exc:
runtime.end_storage_span(span, error=exc)
raise
runtime.end_storage_span(span, telemetry=payload)
self._attach_partition_telemetry(payload, partitioner)
return self._storage_job(payload, telemetry)
def _execute_storage_sql(self, sql: str, parameters: "list[Any]") -> int:
handler = self.handle_database_exceptions()
result = None
with self.with_cursor(self.connection) as cursor, handler:
cursor.execute(sql, parameters)
result = cursor.fetchone()
self._check_pending_exception(handler)
if result is None:
msg = "DuckDB storage transfer did not report its row count."
raise SQLSpecError(msg)
return int(result[0])
[docs]
def collect_rows(self, cursor: "DuckDBConnection", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]":
"""Collect DuckDB rows for the direct execution path."""
data, column_names = collect_rows(cast("list[Any] | None", fetched), cursor.description)
return data, column_names, len(data)
[docs]
def resolve_rowcount(self, cursor: "DuckDBConnection") -> int:
"""Resolve rowcount from DuckDB cursor for the direct execution path."""
return resolve_rowcount(cursor)
def _execute_bulk_insert_many(self, expression: exp.Insert, prepared_parameters: Any) -> "DMLResult | None":
"""Execute a batch INSERT via Arrow registration when the payload is simple."""
if not isinstance(prepared_parameters, list) or not prepared_parameters:
return None
if not isinstance(expression.this, exp.Schema):
return None
if not _is_plain_values_insert(expression, len(expression.this.expressions)):
return None
if not isinstance(prepared_parameters[0], (list, tuple)):
return None
table_expr = expression.this.this
if not isinstance(table_expr, exp.Table):
return None
if table_expr.alias:
return None
rows = prepared_parameters
column_names = [column.name for column in expression.this.expressions]
arrow_table = self._build_arrow_table(rows, column_names)
if arrow_table is None:
return None
target_table = table_expr.sql(dialect="duckdb")
column_sql = ", ".join(column.sql(dialect="duckdb") for column in expression.this.expressions)
temp_view = f"_sqlspec_batch_{uuid4().hex}"
exc_handler = self.handle_database_exceptions()
with exc_handler:
self.connection.register(temp_view, arrow_table)
try:
self.connection.execute(f"INSERT INTO {target_table} ({column_sql}) SELECT * FROM {temp_view}")
finally:
with contextlib.suppress(Exception):
self.connection.unregister(temp_view)
self._check_pending_exception(exc_handler)
return DMLResult("INSERT", len(rows))
@staticmethod
def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | None:
"""Build a pyarrow table from batch rows when they share a stable shape."""
if not rows or pa is None:
return None
first_row = rows[0]
if isinstance(first_row, dict):
keys = column_names or list(first_row.keys())
if any(not isinstance(row, dict) for row in rows):
return None
try:
return pa.table({key: [row.get(key) for row in rows] for key in keys})
except Exception:
return None
if isinstance(first_row, (list, tuple)):
values = list(first_row)
if not column_names:
column_names = [f"col_{index}" for index in range(len(values))]
if any(not isinstance(row, (list, tuple)) or len(row) != len(column_names) for row in rows):
return None
try:
return pa.Table.from_arrays([pa.array(col) for col in zip(*rows, strict=True)], names=column_names)
except Exception:
return None
return None
def _connection_in_transaction(self) -> bool:
"""Check if connection is in transaction.
DuckDB does not expose native transaction state, so it is tracked via a
flag toggled in begin/commit/rollback.
Returns:
True when a transaction is active.
"""
return self._transaction_active
def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tuple[Any, list[Any] | None]":
"""Open a native Arrow reader through DuckDB's exception boundary."""
ensure_pyarrow()
handler = self.handle_database_exceptions()
description: list[Any] | None = None
reader: Any | None = None
with handler:
result = self.connection.execute(sql, normalize_execute_parameters(parameters))
description = result.description
reader = result.to_arrow_reader(chunk_size)
self._check_pending_exception(handler)
if reader is None:
msg = "DuckDB did not return an Arrow record batch reader."
raise SQLSpecError(msg)
return reader, description
def _resolve_duckdb_inserted_rows(result: object) -> int:
fetchall = getattr(result, "fetchall", None)
if not callable(fetchall):
return 0
try:
rows = fetchall()
except Exception:
return 0
if not isinstance(rows, list) or not rows:
return 0
first_row = rows[0]
if isinstance(first_row, (tuple, list)) and first_row and isinstance(first_row[0], int):
return max(first_row[0], 0)
if isinstance(first_row, int):
return max(first_row, 0)
return 0
register_driver_profile("duckdb", driver_profile)
