Source code for sqlspec.adapters.duckdb.pool

"""DuckDB connection pool with thread-local connections."""

import logging
import re
import threading
import time
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING, Any, Final, cast

import duckdb
from typing_extensions import final

from sqlspec.adapters.duckdb._typing import DuckDBConnection
from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context
from sqlspec.utils.uuids import uuid4

if TYPE_CHECKING:
    from collections.abc import Callable, Generator

__all__ = ("DuckDBConnectionPool",)


_SQL_IDENTIFIER_RE: Final[re.Pattern[str]] = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")
_EXPLICIT_INSTALL_KEYS: Final[tuple[str, ...]] = ("version", "repository", "repository_url")


logger = get_logger(POOL_LOGGER_NAME)
_ADAPTER_NAME = "duckdb"

DEFAULT_MIN_POOL: Final[int] = 1
DEFAULT_MAX_POOL: Final[int] = 4
POOL_TIMEOUT: Final[float] = 30.0
POOL_RECYCLE: Final[int] = 86400
HEALTH_CHECK_INTERVAL: Final[float] = 30.0




@final
class DuckDBConnectionPool:
    """Thread-local connection manager for DuckDB.

    Uses thread-local storage to ensure each thread gets its own DuckDB connection,
    preventing the thread-safety issues that cause segmentation faults when
    multiple cursors share the same connection concurrently.

    This design trades traditional pooling for thread safety, which is essential
    for DuckDB since connections and cursors are not thread-safe.
    """

    __slots__ = (
        "_connection_config",
        "_extension_flags",
        "_extensions",
        "_health_check_interval",
        "_installed_signatures",
        "_is_memory_db",
        "_lock",
        "_on_connection_create",
        "_pool_id",
        "_recycle",
        "_secrets",
        "_thread_local",
    )



    def __init__(
        self,
        connection_config: "dict[str, Any]",
        pool_recycle_seconds: int = POOL_RECYCLE,
        health_check_interval: float = HEALTH_CHECK_INTERVAL,
        extensions: "list[dict[str, Any]] | None" = None,
        extension_flags: "dict[str, Any] | None" = None,
        secrets: "list[dict[str, Any]] | None" = None,
        on_connection_create: "Callable[[DuckDBConnection], DuckDBConnection | None] | None" = None,
    ) -> None:
        """Initialize the thread-local connection manager.

        Args:
            connection_config: DuckDB connection configuration
            pool_recycle_seconds: Connection recycle time in seconds
            health_check_interval: Seconds of idle time before running health check
            extensions: List of extensions to install/load
            extension_flags: Connection-level SET statements applied after creation
            secrets: List of secrets to create
            on_connection_create: Callback executed when connection is created
        """
        self._connection_config = connection_config
        self._recycle = pool_recycle_seconds
        self._health_check_interval = health_check_interval
        self._extensions = extensions or []
        self._extension_flags = extension_flags or {}
        self._secrets = secrets or []
        self._on_connection_create = on_connection_create
        self._installed_signatures: set[tuple[Any, ...]] = set()
        self._thread_local = threading.local()
        self._lock = threading.RLock()
        self._pool_id = str(uuid4())[:8]
        database = connection_config.get("database", "")
        self._is_memory_db = database.startswith(":memory:") or database == ""

@property def _database_name(self) -> str: """Get sanitized database name for logging.""" db = self._connection_config.get("database", "") if db.startswith(":memory:") or db == "": return ":memory:" return str(db) def _create_connection(self) -> DuckDBConnection: """Create a new DuckDB connection with extensions and secrets.""" connect_parameters = {} config_dict = {} for key, value in self._connection_config.items(): if key in {"database", "read_only"}: connect_parameters[key] = value elif key == "config" and isinstance(value, dict): config_dict.update(value) else: config_dict[key] = value if config_dict: connect_parameters["config"] = config_dict connection = duckdb.connect(**connect_parameters) self._apply_extension_flags(connection) for ext_config in self._extensions: ext_name = ext_config.get("name") if not ext_name: continue required = bool(ext_config.get("required", False)) install_kwargs: dict[str, Any] = {k: ext_config[k] for k in _EXPLICIT_INSTALL_KEYS if k in ext_config} force_install = bool(ext_config.get("force_install", False)) explicit_install = bool(ext_config.get("install", False)) or force_install or bool(install_kwargs) install_error: str | None = None if explicit_install: install_error = self._install_extension_once( connection, ext_name, install_kwargs, force_install, required ) try: connection.load_extension(ext_name) except Exception as exc: if required: raise if explicit_install: log_with_context( logger, logging.WARNING, "pool.extension.setup.failed", adapter=_ADAPTER_NAME, pool_id=self._pool_id, database=self._database_name, extension=ext_name, install_error=install_error, load_error=str(exc), ) else: log_with_context( logger, logging.WARNING, "pool.extension.load.failed", adapter=_ADAPTER_NAME, pool_id=self._pool_id, database=self._database_name, extension=ext_name, error=str(exc), ) else: if install_error is not None: log_with_context( logger, logging.DEBUG, "pool.extension.install.failed", adapter=_ADAPTER_NAME, pool_id=self._pool_id, database=self._database_name, extension=ext_name, error=install_error, ) for secret_config in self._secrets: _create_secret(connection, secret_config) if self._on_connection_create: self._on_connection_create(connection) return connection def _install_extension_once( self, connection: "DuckDBConnection", ext_name: str, install_kwargs: "dict[str, Any]", force_install: bool, required: bool, ) -> "str | None": """Install an extension once per pool per signature, best-effort unless required. Returns the install error string on a best-effort failure so the caller can combine it with the load outcome into a single setup result; returns ``None`` when the install succeeds or is skipped. Raises when ``required`` is set. """ signature = ( ext_name, install_kwargs.get("version"), install_kwargs.get("repository"), install_kwargs.get("repository_url"), ) with self._lock: if not force_install and signature in self._installed_signatures: return None try: if force_install: connection.install_extension(ext_name, force_install=True, **install_kwargs) elif install_kwargs: connection.install_extension(ext_name, **install_kwargs) else: connection.install_extension(ext_name) except Exception as exc: if required: raise return str(exc) self._installed_signatures.add(signature) return None def _apply_extension_flags(self, connection: DuckDBConnection) -> None: """Apply connection-level extension flags via SET statements.""" if not self._extension_flags: return for key, value in self._extension_flags.items(): if not key or not key.replace("_", "").isalnum(): continue normalized = self._normalize_flag_value(value) try: connection.execute(f"SET {key} = {normalized}") except Exception as exc: # pragma: no cover log_with_context( logger, logging.DEBUG, "pool.flag.set.failed", adapter=_ADAPTER_NAME, pool_id=self._pool_id, database=self._database_name, flag=key, error=str(exc), ) @staticmethod def _normalize_flag_value(value: Any) -> str: """Convert Python value to DuckDB SET literal.""" if isinstance(value, bool): return "TRUE" if value else "FALSE" if isinstance(value, (int, float)): return str(value) escaped = str(value).replace("'", "''") return f"'{escaped}'" def _get_thread_connection(self) -> DuckDBConnection: """Get or create a connection for the current thread. Each thread gets its own dedicated DuckDB connection to prevent thread-safety issues with concurrent cursor operations. """ thread_state = self._thread_local.__dict__ if "connection" not in thread_state: self._thread_local.connection = self._create_connection() self._thread_local.created_at = time.time() self._thread_local.last_used = time.time() return cast("DuckDBConnection", self._thread_local.connection) if self._recycle > 0 and time.time() - self._thread_local.created_at > self._recycle: with suppress(Exception): self._thread_local.connection.close() self._thread_local.connection = self._create_connection() self._thread_local.created_at = time.time() self._thread_local.last_used = time.time() return cast("DuckDBConnection", self._thread_local.connection) idle_time = time.time() - thread_state.get("last_used", 0) if idle_time > self._health_check_interval and not self._is_connection_alive(self._thread_local.connection): log_with_context( logger, logging.DEBUG, "pool.connection.recycle", adapter=_ADAPTER_NAME, pool_id=self._pool_id, database=self._database_name, idle_seconds=round(idle_time, 1), reason="failed_health_check", ) with suppress(Exception): self._thread_local.connection.close() self._thread_local.connection = self._create_connection() self._thread_local.created_at = time.time() self._thread_local.last_used = time.time() return cast("DuckDBConnection", self._thread_local.connection) def _close_thread_connection(self) -> None: """Close the connection for the current thread.""" thread_state = self._thread_local.__dict__ if "connection" in thread_state: with suppress(Exception): self._thread_local.connection.close() del self._thread_local.connection if "created_at" in thread_state: del self._thread_local.created_at if "last_used" in thread_state: del self._thread_local.last_used def _is_connection_alive(self, connection: DuckDBConnection) -> bool: """Check if a connection is still alive and usable. Args: connection: Connection to check Returns: True if connection is alive, False otherwise """ try: cursor = connection.cursor() cursor.close() except Exception: return False return True

@contextmanager def get_connection(self) -> "Generator[DuckDBConnection, None, None]": """Get a thread-local connection. Each thread gets its own dedicated DuckDB connection to prevent thread-safety issues with concurrent cursor operations. For file-based databases, the connection is closed when the context manager exits to release DuckDB's file lock, allowing subsequent connections with different configurations. For in-memory databases, connections are kept alive to preserve data, as in-memory data is lost when the last connection closes. Yields: DuckDBConnection: A thread-local connection. """ connection = self._get_thread_connection() try: yield connection except Exception: self._close_thread_connection() raise else: if not self._is_memory_db: self._close_thread_connection()