"""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
from typing_extensions import final
from sqlspec.adapters.duckdb._typing import DuckDBConnection
from sqlspec.adapters.duckdb._typing import duckdb_module as duckdb
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
SECRET_CONFLICT_ATTEMPTS: Final[int] = 5
SECRET_CONFLICT_BACKOFF: Final[float] = 0.005
_SECRET_LOCK: Final = threading.Lock()
_SECRET_VISIBLE_SETTINGS: Final[frozenset[str]] = frozenset({
"endpoint",
"http_proxy",
"http_proxy_username",
"key_id",
"region",
"url_style",
"use_ssl",
})
_SECRET_BOOLEAN_SETTINGS: Final[frozenset[str]] = frozenset({"use_ssl"})
_SECRET_BOOLEAN_TEXT: Final[dict[str, str]] = {
"true": "true",
"t": "true",
"yes": "true",
"y": "true",
"1": "true",
"false": "false",
"f": "false",
"no": "false",
"n": "false",
"0": "false",
}
[docs]
@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",
"_connection_registry",
"_extensions",
"_generation",
"_health_check_interval",
"_installed_signatures",
"_is_memory_db",
"_lock",
"_on_connection_create",
"_pool_id",
"_recycle",
"_secrets",
"_thread_local",
)
[docs]
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,
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
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._secrets = secrets or []
self._on_connection_create = on_connection_create
self._installed_signatures: set[tuple[Any, ...]] = set()
self._thread_local = threading.local()
self._connection_registry: set[DuckDBConnection] = set()
self._generation = 0
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 pool-owned connection and record it for shutdown."""
connection = self.new_connection(record_thread_state=True)
with self._lock:
self._connection_registry.add(connection)
return connection
[docs]
def new_connection(self, record_thread_state: bool = False) -> DuckDBConnection:
"""Open a standalone connection with extensions and secrets applied.
Unless the pool asks it to, the result is owned by the caller: it is not
thread-local, is not tracked for pool shutdown, and leaves the storage
setup already published for this thread untouched.
Args:
record_thread_state: Whether to publish storage setup for this thread.
Returns:
DuckDBConnection: A newly opened, fully configured connection.
"""
loaded_extensions: set[str] = set()
storage_protocols: set[str] = set()
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)
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:
loaded_extensions.add(ext_name)
storage_protocols.update(ext_config.get("storage_protocols", ()))
if ext_name == "gcs":
storage_protocols.add("gcss")
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,
)
created_secrets: list[dict[str, Any]] = [
{**secret_config, "value": dict(secret_config.get("value") or {})}
for secret_config in self._secrets
if _create_secret(connection, secret_config)
]
if self._on_connection_create:
self._on_connection_create(connection)
storage_protocols.update(connection.list_filesystems())
if record_thread_state:
self._thread_local.storage_extensions = frozenset(loaded_extensions)
self._thread_local.storage_secrets = tuple(created_secrets)
self._thread_local.storage_protocols = frozenset(storage_protocols)
return connection
def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]":
"""Return storage setup completed for this thread's current connection."""
state = self._thread_local.__dict__
if state.get("connection") is not connection:
return {}
return {
"_duckdb_storage_extensions": state.get("storage_extensions", frozenset()),
"_duckdb_storage_secrets": state.get("storage_secrets", ()),
"_duckdb_storage_protocols": state.get("storage_protocols", frozenset()),
}
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 _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 thread_state.get("generation") != self._generation:
stale = thread_state.pop("connection", None)
if stale is not None:
self._retire_connection(cast("DuckDBConnection", stale))
thread_state.pop("created_at", None)
thread_state.pop("last_used", None)
self._thread_local.generation = self._generation
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:
self._retire_connection(cast("DuckDBConnection", self._thread_local.connection))
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",
)
self._retire_connection(cast("DuckDBConnection", self._thread_local.connection))
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 _retire_connection(self, connection: DuckDBConnection) -> None:
"""Close a pool-owned connection and drop it from the shutdown registry."""
with self._lock:
self._connection_registry.discard(connection)
with suppress(Exception):
connection.close()
def _close_thread_connection(self) -> None:
"""Close the connection for the current thread."""
thread_state = self._thread_local.__dict__
thread_state.pop("storage_extensions", None)
thread_state.pop("storage_secrets", None)
thread_state.pop("storage_protocols", None)
if "connection" in thread_state:
self._retire_connection(cast("DuckDBConnection", self._thread_local.connection))
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
[docs]
@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:
with suppress(Exception):
connection.rollback()
if not self._is_memory_db:
self._close_thread_connection()
raise
else:
try:
connection.commit()
finally:
if not self._is_memory_db:
self._close_thread_connection()
[docs]
def close(self) -> None:
"""Close every connection this pool opened, on any thread."""
self._close_thread_connection()
with self._lock:
orphaned = list(self._connection_registry)
self._connection_registry.clear()
self._generation += 1
for connection in orphaned:
with suppress(Exception):
connection.close()
[docs]
def size(self) -> int:
"""Get current pool size (always 1 for thread-local)."""
return 1 if "connection" in self._thread_local.__dict__ else 0
[docs]
def checked_out(self) -> int:
"""Get number of checked out connections (always 0 for thread-local)."""
return 0
[docs]
def acquire(self) -> DuckDBConnection:
"""Acquire a thread-local connection.
Each thread gets its own dedicated DuckDB connection to prevent
thread-safety issues with concurrent cursor operations.
Returns:
DuckDBConnection: A thread-local connection
"""
return self._get_thread_connection()
[docs]
def release(self, connection: DuckDBConnection) -> None:
"""Release a connection (no-op for thread-local connections).
Args:
connection: The connection to release (ignored)
"""
def _validate_sql_identifier(value: str, field_name: str) -> None:
"""Raise ValueError if value is not safe to interpolate as a SQL identifier."""
if not _SQL_IDENTIFIER_RE.fullmatch(value):
msg = (
f"Invalid SQL identifier for {field_name!r}: {value!r}. "
"Must start with a letter and contain only letters, digits, and underscores."
)
raise ValueError(msg)
def _create_secret(connection: DuckDBConnection, secret_config: dict[str, Any]) -> bool:
secret_name = secret_config.get("name")
secret_type = secret_config.get("secret_type")
if not (secret_name and secret_type):
return False
required = bool(secret_config.get("required", False))
try:
_validate_sql_identifier(secret_name, "secret_name")
_validate_sql_identifier(secret_type, "secret_type")
sql = _secret_sql(secret_config, secret_name, secret_type)
mismatch = _apply_secret(connection, secret_config, secret_name, secret_type, sql)
except Exception:
if required:
raise
logger.warning("DuckDB secret %r creation failed (best-effort)", secret_name)
return False
if mismatch is None:
return True
if required:
raise RuntimeError(mismatch)
logger.warning("%s (best-effort)", mismatch)
return False
def _apply_secret(
connection: DuckDBConnection, secret_config: dict[str, Any], secret_name: str, secret_type: str, sql: str
) -> "str | None":
"""Create or reuse a declared secret under the process-wide secret lock.
Catalog write-write conflicts with connections outside that lock are retried with a
short backoff that runs while the lock is released.
Returns:
A description of how the visible secret differs from its declaration, or ``None`` when it matches.
"""
for attempt in range(1, SECRET_CONFLICT_ATTEMPTS):
applied, mismatch = _try_apply_secret(connection, secret_config, secret_name, secret_type, sql)
if applied:
return mismatch
time.sleep(SECRET_CONFLICT_BACKOFF * attempt)
with _SECRET_LOCK:
return _apply_secret_locked(connection, secret_config, secret_name, secret_type, sql)
def _try_apply_secret(
connection: DuckDBConnection, secret_config: dict[str, Any], secret_name: str, secret_type: str, sql: str
) -> "tuple[bool, str | None]":
"""Apply a declared secret once, returning ``(False, None)`` on a catalog write-write conflict."""
try:
with _SECRET_LOCK:
mismatch = _apply_secret_locked(connection, secret_config, secret_name, secret_type, sql)
except duckdb.TransactionException as exc:
if "write-write conflict" not in str(exc):
raise
return False, None
return True, mismatch
def _apply_secret_locked(
connection: DuckDBConnection, secret_config: dict[str, Any], secret_name: str, secret_type: str, sql: str
) -> "str | None":
"""Reuse a visible secret unless replacement is requested, otherwise create it, then compare it."""
expected_provider = secret_config.get("provider") or _default_secret_provider(connection, secret_type)
if not secret_config.get("replace", False):
existing = _visible_secret(connection, secret_config, secret_name, secret_type)
if existing is not None:
mismatch = _secret_mismatch(existing, secret_config, secret_name, secret_type, expected_provider, True)
if mismatch is None:
logger.debug("DuckDB secret %r already exists and is reused", secret_name)
return mismatch
connection.execute(sql)
created = _visible_secret(connection, secret_config, secret_name, secret_type)
if created is None:
storage = "persistent" if secret_config.get("persistent", False) else "temporary"
return f"DuckDB secret {secret_name!r} was not visible as a {storage} secret after creation"
return _secret_mismatch(created, secret_config, secret_name, secret_type, expected_provider, False)
def _secret_sql(secret_config: dict[str, Any], secret_name: str, secret_type: str) -> str:
parts = [f"TYPE {secret_type}"]
provider = secret_config.get("provider")
if provider:
_validate_sql_identifier(str(provider), "secret_provider")
parts.append(f"PROVIDER {provider}")
secret_value = secret_config.get("value") or {}
if not isinstance(secret_value, dict):
msg = "DuckDB secret value must be a dictionary"
raise TypeError(msg)
for key, value in secret_value.items():
parts.append(f"{_format_secret_key(key)} {_format_secret_literal(value)}")
scope = secret_config.get("scope")
if scope is not None:
parts.append(f"SCOPE {_format_secret_literal(scope)}")
kind = "PERSISTENT SECRET" if secret_config.get("persistent", False) else "SECRET"
if secret_config.get("replace", False):
create = f"CREATE OR REPLACE {kind} {secret_name}"
else:
create = f"CREATE {kind} IF NOT EXISTS {secret_name}"
body = ",\n ".join(parts)
return f"{create} (\n {body}\n)"
def _format_secret_key(key: Any) -> str:
key_text = str(key)
_validate_sql_identifier(key_text, "secret_value_key")
return key_text.upper()
def _format_secret_literal(value: Any) -> str:
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 _visible_secret(
connection: DuckDBConnection, secret_config: dict[str, Any], secret_name: str, secret_type: str
) -> "tuple[Any, ...] | None":
"""Return the type, provider, scope and redacted settings of a visible secret."""
return connection.execute(
"SELECT type, provider, scope, secret_string FROM duckdb_secrets() WHERE lower(name) = lower(?) AND persistent = ?",
(secret_name, bool(secret_config.get("persistent", False))),
).fetchone()
def _default_secret_provider(connection: DuckDBConnection, secret_type: str) -> "str | None":
"""Return DuckDB's default provider for a secret type, or ``None`` when DuckDB does not report it."""
try:
row = connection.execute(
"SELECT default_provider FROM duckdb_secret_types() WHERE type = lower(?)", (secret_type,)
).fetchone()
except duckdb.CatalogException:
return None
return str(row[0]) if row and row[0] else None
def _secret_mismatch(
row: "tuple[Any, ...]",
secret_config: dict[str, Any],
secret_name: str,
secret_type: str,
expected_provider: "str | None",
reused: bool,
) -> "str | None":
"""Describe how a visible secret differs from its declaration, or return ``None`` when it matches.
DuckDB redacts credential values, so the comparison covers the type, provider, declared
scope and the declared settings DuckDB reports unredacted. Settings the declaration omits
are not compared.
"""
storage = "persistent" if secret_config.get("persistent", False) else "temporary"
hint = "; set replace=True to apply the declaration" if reused else ""
actual_type, actual_provider, actual_scope, secret_string = row
if str(actual_type).lower() != secret_type.lower():
return (
f"DuckDB secret {secret_name!r} exists as a {storage} secret of type {actual_type!r}, "
f"but the declaration uses type {secret_type!r}{hint}"
)
differing: list[str] = []
if expected_provider and str(actual_provider).lower() != expected_provider.lower():
differing.append("provider")
scope = secret_config.get("scope")
if scope is not None and scope not in list(actual_scope or []):
differing.append("scope")
settings = f"{secret_string or ''};"
for key, value in (secret_config.get("value") or {}).items():
setting = str(key).lower()
if (
setting in _SECRET_VISIBLE_SETTINGS
and f";{setting}={_secret_setting_text(setting, value)};" not in settings
):
differing.append(setting)
if not differing:
return None
return (
f"DuckDB secret {secret_name!r} exists as a {storage} secret whose settings differ from the "
f"declaration: {', '.join(differing)}{hint}"
)
def _secret_setting_text(setting: str, value: Any) -> str:
"""Render a declared setting the way ``duckdb_secrets()`` reports it."""
if isinstance(value, bool):
return "true" if value else "false"
if setting in _SECRET_BOOLEAN_SETTINGS:
if isinstance(value, (int, float)):
return "true" if value else "false"
text = str(value).lower()
return _SECRET_BOOLEAN_TEXT.get(text, text)
return str(value)
