Source code for sqlspec.adapters.duckdb.config

"""DuckDB database configuration with connection pooling."""

from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast

from typing_extensions import NotRequired

from sqlspec.adapters.duckdb._typing import DuckDBConnection, DuckDBCursor, DuckDBSessionContext
from sqlspec.adapters.duckdb.core import (
    apply_driver_features,
    build_connection_config,
    build_statement_config,
    default_statement_config,
)
from sqlspec.adapters.duckdb.driver import DuckDBDriver, DuckDBExceptionHandler
from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool
from sqlspec.config import ExtensionConfigs, SyncDatabaseConfig
from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.utils.config_tools import normalize_connection_config
from sqlspec.utils.serializers import to_json

if TYPE_CHECKING:
    from sqlspec.core import StatementConfig
    from sqlspec.observability import ObservabilityConfig
__all__ = (
    "DuckDBConfig",
    "DuckDBConnectionParams",
    "DuckDBDriverFeatures",
    "DuckDBExtensionConfig",
    "DuckDBPoolParams",
    "DuckDBSecretConfig",
)
EXTENSION_FLAG_KEYS: "tuple[str, ...]" = (
    "allow_community_extensions",
    "allow_unsigned_extensions",
    "enable_external_access",
)


class DuckDBConnectionParams(TypedDict):
    """DuckDB connection parameters.

    Mirrors the keyword arguments accepted by duckdb.connect so callers can drive every DuckDB
    configuration switch directly through SQLSpec. All keys are optional and forwarded verbatim
    to DuckDB, either as top-level parameters or via the nested ``config`` dictionary when DuckDB
    expects them there.
    """

    database: NotRequired[str]
    read_only: NotRequired[bool]
    config: NotRequired["dict[str, Any]"]
    memory_limit: NotRequired[str]
    threads: NotRequired[int]
    temp_directory: NotRequired[str]
    max_temp_directory_size: NotRequired[str]
    autoload_known_extensions: NotRequired[bool]
    autoinstall_known_extensions: NotRequired[bool]
    allow_community_extensions: NotRequired[bool]
    allow_unsigned_extensions: NotRequired[bool]
    extension_directory: NotRequired[str]
    custom_extension_repository: NotRequired[str]
    autoinstall_extension_repository: NotRequired[str]
    allow_persistent_secrets: NotRequired[bool]
    enable_external_access: NotRequired[bool]
    secret_directory: NotRequired[str]
    enable_object_cache: NotRequired[bool]
    parquet_metadata_cache: NotRequired[bool]
    enable_external_file_cache: NotRequired[bool]
    checkpoint_threshold: NotRequired[str]
    enable_progress_bar: NotRequired[bool]
    progress_bar_time: NotRequired[int]
    enable_logging: NotRequired[bool]
    log_query_path: NotRequired[str]
    logging_level: NotRequired[str]
    preserve_insertion_order: NotRequired[bool]
    default_null_order: NotRequired[str]
    default_order: NotRequired[str]
    ieee_floating_point_ops: NotRequired[bool]
    binary_as_string: NotRequired[bool]
    arrow_large_buffer_size: NotRequired[bool]
    errors_as_json: NotRequired[bool]
    extra: NotRequired["dict[str, Any]"]


class DuckDBPoolParams(DuckDBConnectionParams):
    """Complete pool configuration for DuckDB adapter.

    Extends DuckDBConnectionParams with lifecycle settings consumed by SQLSpec's
    thread-local DuckDB connection manager.
    """

    pool_recycle_seconds: NotRequired[int]
    health_check_interval: NotRequired[float]



[docs]
class DuckDBExtensionConfig(TypedDict):
    """DuckDB extension configuration for auto-management."""

    name: str
    """Name of the extension to install/load."""

    version: NotRequired[str]
    """Specific version of the extension."""

    repository: NotRequired[str]
    """Repository for the extension (core, community, or custom URL)."""

    repository_url: NotRequired[str]
    """Custom repository URL for the extension."""

    force_install: NotRequired[bool]
    """Force reinstallation of the extension."""

    install: NotRequired[bool]
    """Force an explicit install_extension() call even for a name-only config."""

    required: NotRequired[bool]
    """When True, install/load failure raises instead of best-effort WARNING."""

[docs]

class DuckDBSecretConfig(TypedDict): """DuckDB secret configuration for AI/API integrations.""" secret_type: str """Type of secret.""" name: str """Name of the secret.""" value: NotRequired["dict[str, Any]"] """Secret configuration values.""" provider: NotRequired[str] """Secret provider, such as config or credential_chain.""" scope: NotRequired[str] """Optional path or storage-prefix scope for the secret.""" persistent: NotRequired[bool] """Persist the secret to DuckDB's configured secret directory.""" required: NotRequired[bool] """When True, secret-creation failure raises (and is verified). Default best-effort."""

class DuckDBDriverFeatures(TypedDict): """TypedDict for DuckDB driver features configuration. Attributes: extensions: List of extensions to install/load on connection creation. secrets: List of secrets to create for AI/API integrations. on_connection_create: Callback executed when connection is created. json_serializer: Custom JSON serializer for dict/list parameter conversion. Defaults to sqlspec.utils.serializers.to_json if not provided. enable_uuid_conversion: Enable automatic UUID string conversion. When True (default), UUID strings are automatically converted to UUID objects. When False, UUID strings are treated as regular strings. extension_flags: Connection-level flags applied via SET statements immediately after connection creation. enable_events: Enable database event channel support. Defaults to True when extension_config["events"] is configured. Provides pub/sub capabilities via table-backed queue (DuckDB has no native pub/sub). Requires extension_config["events"] for migration setup. events_backend: Event channel backend selection. Only option: "poll_queue" (durable table-backed queue with lease-based retries and acknowledgements). DuckDB does not have native pub/sub, so poll_queue is the only backend. Defaults to "poll_queue". """ extensions: NotRequired[Sequence[DuckDBExtensionConfig]] secrets: NotRequired[Sequence[DuckDBSecretConfig]] on_connection_create: NotRequired["Callable[[DuckDBConnection], DuckDBConnection | None]"] json_serializer: NotRequired["Callable[[Any], str]"] enable_uuid_conversion: NotRequired[bool] extension_flags: NotRequired["dict[str, Any]"] enable_events: NotRequired[bool] events_backend: NotRequired[Literal["poll_queue"]] class DuckDBConnectionContext(SyncPoolConnectionContext): """Context manager for DuckDB connections.""" __slots__ = () class _DuckDBSessionConnectionHandler(SyncPoolSessionFactory): __slots__ = ()

[docs]

class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, DuckDBDriver]): """DuckDB configuration with connection pooling. This configuration supports DuckDB's features including: - Connection pooling - Extension management and installation - Secret management for API integrations - Auto configuration settings - Arrow integration - Direct file querying capabilities - Configurable type handlers for JSON serialization and UUID conversion DuckDB Connection Pool Configuration: - Default pool size is 1-4 connections (DuckDB uses single connection by default) - Connection recycling is set to 24 hours by default (set to 0 to disable) - Shared memory databases use `:memory:shared_db` for proper concurrency Type Handler Configuration via driver_features: - `json_serializer`: Custom JSON serializer for dict/list parameters. Defaults to `sqlspec.utils.serializers.to_json` if not provided. Accepts serializer callables that return text. - `enable_uuid_conversion`: Enable automatic UUID string conversion (default: True). When True, UUID strings in query results are automatically converted to UUID objects. When False, UUID strings are treated as regular strings. """ driver_type: "ClassVar[type[DuckDBDriver]]" = DuckDBDriver connection_type: "ClassVar[type[DuckDBConnection]]" = DuckDBConnection supports_transactional_ddl: "ClassVar[bool]" = True supports_migration_schemas: "ClassVar[bool]" = True supports_native_arrow_export: "ClassVar[bool]" = True supports_native_arrow_import: "ClassVar[bool]" = True supports_arrow_streaming: "ClassVar[bool]" = True supports_native_parquet_export: "ClassVar[bool]" = True supports_native_parquet_import: "ClassVar[bool]" = True storage_partition_strategies: "ClassVar[tuple[str, ...]]" = ("fixed", "rows_per_chunk", "manifest") _connection_context_class: "ClassVar[type[DuckDBConnectionContext]]" = DuckDBConnectionContext _session_factory_class: "ClassVar[type[_DuckDBSessionConnectionHandler]]" = _DuckDBSessionConnectionHandler _session_context_class: "ClassVar[type[DuckDBSessionContext]]" = DuckDBSessionContext _default_statement_config = default_statement_config

[docs]

def __init__( self, *, connection_config: "DuckDBPoolParams | dict[str, Any] | None" = None, connection_instance: "DuckDBConnectionPool | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: "StatementConfig | None" = None, driver_features: "DuckDBDriverFeatures | dict[str, Any] | None" = None, bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: """Initialize DuckDB configuration. Args: connection_config: Connection and pool configuration parameters connection_instance: Pre-created pool instance migration_config: Migration configuration statement_config: Statement configuration override driver_features: DuckDB-specific driver features including json_serializer and enable_uuid_conversion options bind_key: Optional unique identifier for this configuration extension_config: Extension-specific configuration observability_config: Adapter-level observability overrides for lifecycle hooks and observers **kwargs: Additional keyword arguments passed to the base configuration. """ connection_config = normalize_connection_config(connection_config) connection_config.setdefault("database", ":memory:shared_db") if connection_config.get("database") in {":memory:", ""}: connection_config["database"] = ":memory:shared_db" extension_flags: dict[str, Any] = {} for key in tuple(connection_config.keys()): if key in EXTENSION_FLAG_KEYS: extension_flags[key] = connection_config.pop(key) features: dict[str, Any] = dict(driver_features) if driver_features else {} self._user_connection_hook = cast( "Callable[[DuckDBConnection], DuckDBConnection | None] | None", features.pop("on_connection_create", None) ) features.setdefault("enable_uuid_conversion", True) serializer = features.setdefault("json_serializer", to_json) if extension_flags: existing_flags = cast("dict[str, Any]", features.get("extension_flags", {})) merged_flags = {**existing_flags, **extension_flags} features["extension_flags"] = merged_flags statement_config = statement_config or build_statement_config( json_serializer=cast("Callable[[Any], str]", serializer) ) statement_config, features = apply_driver_features(statement_config, features) super().__init__( bind_key=bind_key, connection_config=connection_config, connection_instance=connection_instance, migration_config=migration_config, statement_config=statement_config, driver_features=features, extension_config=extension_config, observability_config=observability_config, **kwargs, )

def _create_pool(self) -> DuckDBConnectionPool: """Create connection pool from configuration.""" connection_config = build_connection_config(self.connection_config) extensions = cast("list[dict[str, Any]] | None", self.driver_features.get("extensions", None)) secrets = cast("list[dict[str, Any]] | None", self.driver_features.get("secrets", None)) extension_flags = cast("dict[str, Any] | None", self.driver_features.get("extension_flags", None)) extensions_dicts = [dict(ext) for ext in extensions] if extensions else None secrets_dicts = [dict(secret) for secret in secrets] if secrets else None extension_flags_dict = dict(extension_flags) if extension_flags else None pool_recycle_seconds = self.connection_config.get("pool_recycle_seconds") health_check_interval = self.connection_config.get("health_check_interval") pool_kwargs: dict[str, Any] = {} if pool_recycle_seconds is not None: pool_kwargs["pool_recycle_seconds"] = pool_recycle_seconds if health_check_interval is not None: pool_kwargs["health_check_interval"] = health_check_interval return DuckDBConnectionPool( connection_config=connection_config, extensions=extensions_dicts, extension_flags=extension_flags_dict, secrets=secrets_dicts, on_connection_create=self._user_connection_hook, **pool_kwargs, ) def _close_pool(self) -> None: """Close the connection pool.""" if self.connection_instance: self.connection_instance.close()

[docs]

def create_connection(self) -> DuckDBConnection: """Get a DuckDB connection from the pool. This method ensures the pool is created and returns a connection from the pool. The connection is checked out from the pool and must be properly managed by the caller. Returns: DuckDBConnection: A connection from the pool """ pool = self.provide_pool() return pool.acquire()