"""ADBC database configuration."""
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast
from typing_extensions import NotRequired
from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcCursor, AdbcSessionContext
from sqlspec.adapters.adbc.core import (
apply_driver_features,
build_connection_config,
detect_postgres_extensions,
get_statement_config,
is_postgres_dialect,
resolve_dialect_from_config,
resolve_dialect_name,
resolve_driver_connect_func,
)
from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler
from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig
from sqlspec.core import StatementConfig
from sqlspec.core.capabilities import TypeCoercionCapabilities
from sqlspec.core.config_runtime import (
build_postgres_extension_probe_names,
is_postgres_extension_active,
resolve_postgres_extension_state,
resolve_runtime_statement_config,
)
from sqlspec.driver import SyncPoolConnectionContext, SyncPoolSessionFactory
from sqlspec.exceptions import ImproperConfigurationError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.utils.config_tools import normalize_connection_config
if TYPE_CHECKING:
from collections.abc import Callable
from types import TracebackType
from sqlspec.observability import ObservabilityConfig
__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures")
[docs]
class AdbcConnectionParams(TypedDict):
"""ADBC connection parameters."""
uri: NotRequired[str]
dsn: NotRequired[str]
url: NotRequired[str]
connection_string: NotRequired[str]
db: NotRequired[str]
path: NotRequired[str]
driver_name: NotRequired[str]
db_kwargs: NotRequired[dict[str, Any]]
conn_kwargs: NotRequired[dict[str, Any]]
entrypoint: NotRequired[str]
profile: NotRequired[str]
adbc_driver_manager_entrypoint: NotRequired[str]
autocommit: NotRequired[bool]
isolation_level: NotRequired[str]
batch_size: NotRequired[int]
query_timeout: NotRequired[float]
connection_timeout: NotRequired[float]
ssl_mode: NotRequired[str]
ssl_cert: NotRequired[str]
ssl_key: NotRequired[str]
ssl_ca: NotRequired[str]
username: NotRequired[str]
password: NotRequired[str]
token: NotRequired[str]
project_id: NotRequired[str]
dataset_id: NotRequired[str]
account: NotRequired[str]
warehouse: NotRequired[str]
database: NotRequired[str]
schema: NotRequired[str]
role: NotRequired[str]
authorization_header: NotRequired[str]
grpc_options: NotRequired[dict[str, Any]]
grpc_max_message_size: NotRequired[int]
fetch_timeout: NotRequired[float]
gizmosql_backend: NotRequired[str]
tls_skip_verify: NotRequired[bool]
tls_root_certs: NotRequired[str]
mtls_cert_chain: NotRequired[str]
mtls_private_key: NotRequired[str]
with_cookie_middleware: NotRequired[bool]
headers: NotRequired[dict[str, str]]
extra: NotRequired[dict[str, Any]]
[docs]
class AdbcDriverFeatures(TypedDict):
"""ADBC driver feature configuration.
Controls optional type handling and serialization behavior for the ADBC adapter.
These features configure how data is converted between Python and Arrow types.
Attributes:
json_serializer: JSON serialization function to use.
Callable that takes Any and returns str (JSON string).
Default: sqlspec.utils.serializers.to_json
enable_cast_detection: Enable cast-aware parameter processing.
When True, detects SQL casts and applies appropriate
serialization. Currently used for PostgreSQL JSONB handling.
Default: True
enable_strict_type_coercion: Enforce strict type coercion rules.
When True, raises errors for unsupported type conversions.
When False, attempts best-effort conversion.
Default: False
strict_type_coercion: Alias for enable_strict_type_coercion.
enable_arrow_extension_types: Enable PyArrow extension type support.
When True, preserves Arrow extension type metadata when reading data.
When False, falls back to storage types.
Default: True
arrow_extension_types: Alias for enable_arrow_extension_types.
enable_pgvector: Enable automatic pgvector extension detection.
When True and the resolved dialect is PostgreSQL, queries ``pg_extension``
on the first connection to check for the ``vector`` extension.
Defaults to True when the ``pgvector`` Python package is installed.
enable_paradedb: Enable ParadeDB (pg_search) extension detection.
When True and the resolved dialect is PostgreSQL, queries ``pg_extension``
on the first connection to check for the ``pg_search`` extension.
Defaults to True. Independent of enable_pgvector.
enable_pg_textsearch: Enable pg_textsearch extension detection for BM25 search.
When True and the resolved dialect is PostgreSQL, queries ``pg_extension``
on the first connection to check for the ``pg_textsearch`` extension.
Defaults to True.
enable_events: Enable database event channel support.
Defaults to True when extension_config["events"] is configured.
Provides pub/sub capabilities via table-backed queue (ADBC has no native pub/sub).
Requires extension_config["events"] for migration setup.
on_connection_create: Callback executed when a connection is created.
Receives the raw ADBC connection for low-level driver configuration.
events_backend: Event channel backend selection.
Only option: "poll_queue" (durable table-backed queue with lease-based retries and acknowledgements).
ADBC does not have native pub/sub, so poll_queue is the only backend.
Defaults to "poll_queue".
"""
json_serializer: "NotRequired[Callable[[Any], str]]"
enable_cast_detection: NotRequired[bool]
enable_strict_type_coercion: NotRequired[bool]
strict_type_coercion: NotRequired[bool]
enable_arrow_extension_types: NotRequired[bool]
arrow_extension_types: NotRequired[bool]
enable_pgvector: NotRequired[bool]
enable_paradedb: NotRequired[bool]
enable_pg_textsearch: NotRequired[bool]
enable_events: NotRequired[bool]
on_connection_create: "NotRequired[Callable[[AdbcConnection], None]]"
events_backend: NotRequired[Literal["poll_queue"]]
class AdbcConnectionContext(SyncPoolConnectionContext):
"""Context manager for ADBC connections."""
__slots__ = ("_connection",)
def __init__(self, config: "AdbcConfig") -> None:
super().__init__(config)
self._connection: AdbcConnection | None = None
def __enter__(self) -> "AdbcConnection":
self._connection = self._config.create_connection()
return cast("AdbcConnection", self._connection)
def __exit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> bool | None:
if self._connection:
self._connection.close()
self._connection = None
return None
class _AdbcSessionConnectionHandler(SyncPoolSessionFactory):
__slots__ = ("_connection",)
def __init__(self, config: "AdbcConfig") -> None:
super().__init__(config)
self._connection: AdbcConnection | None = None
def acquire_connection(self) -> "AdbcConnection":
self._connection = self._config.create_connection()
return cast("AdbcConnection", self._connection)
def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None:
if self._connection is None:
return
self._connection.close()
self._connection = None
[docs]
class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
"""ADBC configuration for Arrow Database Connectivity.
ADBC provides an interface for connecting to multiple database systems
with Arrow-native data transfer.
Supports multiple database backends including PostgreSQL, SQLite, DuckDB,
BigQuery, and Snowflake with automatic driver detection and loading.
"""
driver_type: ClassVar[type[AdbcDriver]] = AdbcDriver
connection_type: "ClassVar[type[AdbcConnection]]" = AdbcConnection
supports_transactional_ddl: ClassVar[bool] = False
supports_native_arrow_export: "ClassVar[bool]" = True
supports_native_arrow_import: "ClassVar[bool]" = True
supports_arrow_streaming: "ClassVar[bool]" = True
supports_native_row_streaming: "ClassVar[bool]" = True
supports_native_parquet_export: "ClassVar[bool]" = True
supports_native_parquet_import: "ClassVar[bool]" = True
supports_reliable_rowcount: ClassVar[bool] = False
storage_partition_strategies: "ClassVar[tuple[str, ...]]" = ("fixed", "rows_per_chunk")
type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities(
datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=False, uuid_binding="native"
)
_connection_context_class: "ClassVar[type[AdbcConnectionContext]]" = AdbcConnectionContext
_session_factory_class: "ClassVar[type[_AdbcSessionConnectionHandler]]" = _AdbcSessionConnectionHandler
_session_context_class: "ClassVar[type[AdbcSessionContext]]" = AdbcSessionContext
_default_statement_config = StatementConfig()
__slots__ = (
"_default_session_config",
"_paradedb_available",
"_pg_textsearch_available",
"_pgvector_available",
"_resolved_dialect",
"_user_connection_hook",
)
[docs]
def __init__(
self,
*,
connection_config: "AdbcConnectionParams | dict[str, Any] | None" = None,
connection_instance: "AdbcConnection | None" = None,
migration_config: "dict[str, Any] | None" = None,
statement_config: StatementConfig | None = None,
driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None,
bind_key: str | None = None,
extension_config: "ExtensionConfigs | None" = None,
observability_config: "ObservabilityConfig | None" = None,
**kwargs: Any,
) -> None:
"""Initialize configuration.
Args:
connection_config: Connection configuration parameters
connection_instance: Pre-created connection instance to use instead of creating new one
migration_config: Migration configuration
statement_config: Default SQL statement configuration
driver_features: Driver feature configuration (AdbcDriverFeatures)
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.
"""
config_dict = normalize_connection_config(connection_config)
if "uri" not in config_dict:
uri = (
config_dict.pop("url", None)
or config_dict.pop("dsn", None)
or config_dict.pop("connection_string", None)
)
if uri is not None:
config_dict["uri"] = uri
else:
config_dict.pop("url", None)
config_dict.pop("dsn", None)
config_dict.pop("connection_string", None)
driver_name = config_dict.get("driver_name")
if (driver_name in {"sqlite", "adbc_driver_sqlite"} or driver_name is None) and "uri" not in config_dict:
db_path = (
config_dict.pop("database", None)
or config_dict.pop("db", None)
or config_dict.pop("path", None)
or config_dict.pop("file", None)
)
if db_path is not None:
config_dict["uri"] = db_path
elif driver_name in {"duckdb", "adbc_driver_duckdb"} and "path" not in config_dict:
duck_path = (
config_dict.pop("database", None)
or config_dict.pop("db", None)
or config_dict.pop("file", None)
or config_dict.pop("path", None)
)
if duck_path is not None:
config_dict["path"] = duck_path
self.connection_config = config_dict
self._pgvector_available: bool | None = None
self._paradedb_available: bool | None = None
self._pg_textsearch_available: bool | None = None
self._resolved_dialect = resolve_dialect_from_config(self.connection_config)
if statement_config is None:
statement_config = get_statement_config(self._resolved_dialect)
statement_config, driver_features = apply_driver_features(statement_config, driver_features)
features_dict = dict(driver_features) if driver_features else {}
self._user_connection_hook = cast(
"Callable[[AdbcConnection], None] | None", features_dict.pop("on_connection_create", None)
)
self._default_session_config = get_statement_config(self._resolved_dialect)
super().__init__(
connection_config=self.connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=features_dict,
bind_key=bind_key,
extension_config=extension_config,
observability_config=observability_config,
**kwargs,
)
@property
def supports_migration_schemas(self) -> bool: # type: ignore[override]
"""Migration schema support is only available for PostgreSQL-backed ADBC drivers."""
return is_postgres_dialect(self._resolved_dialect)
[docs]
def create_connection(self) -> AdbcConnection:
"""Create and return a new connection using the specified driver.
Returns:
A new connection instance.
Raises:
ImproperConfigurationError: If the connection could not be established.
"""
try:
driver_name = cast("str | None", self.connection_config.get("driver_name"))
uri = cast("str | None", self.connection_config.get("uri"))
connection = resolve_driver_connect_func(driver_name, uri)(
**build_connection_config(self.connection_config)
)
if self._user_connection_hook is not None:
self._user_connection_hook(cast("AdbcConnection", connection))
return cast("AdbcConnection", connection)
except Exception as e:
err_driver_name = self.connection_config.get("driver_name", "Unknown")
msg = f"Could not configure connection using driver '{err_driver_name}'. Error: {e}"
raise ImproperConfigurationError(msg) from e
@property
def pg_textsearch_available(self) -> bool:
"""Return True if the pg_textsearch extension is available."""
return bool(self._pg_textsearch_available)
def _detect_extensions_if_needed(self) -> None:
"""Detect postgres extensions on first call, caching results.
Only queries ``pg_extension`` when the resolved dialect is PostgreSQL
and detection has not yet run (``_pgvector_available is None``).
"""
if self._pgvector_available is not None:
return
dialect = resolve_dialect_name(self.statement_config.dialect)
if not is_postgres_dialect(dialect):
self._pgvector_available = False
self._paradedb_available = False
self._pg_textsearch_available = False
return
connection = self.create_connection()
try:
probe_names = build_postgres_extension_probe_names(self.driver_features)
pgvector_available, paradedb_available, pg_textsearch_available = detect_postgres_extensions(
connection,
enable_pgvector="vector" in probe_names,
enable_paradedb="pg_search" in probe_names,
enable_pg_textsearch="pg_textsearch" in probe_names,
)
finally:
connection.close()
detected_extensions: set[str] = set()
if pgvector_available:
detected_extensions.add("vector")
if paradedb_available:
detected_extensions.add("pg_search")
if pg_textsearch_available:
detected_extensions.add("pg_textsearch")
self.statement_config, self._pgvector_available, self._paradedb_available = resolve_postgres_extension_state(
self.statement_config, self.driver_features, detected_extensions
)
self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch")
[docs]
def provide_connection(self, *args: Any, **kwargs: Any) -> "AdbcConnectionContext":
"""Provide a connection context manager.
Args:
*args: Additional arguments.
**kwargs: Additional keyword arguments.
Returns:
A connection context manager.
"""
return AdbcConnectionContext(self)
[docs]
def provide_session(
self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any
) -> "AdbcSessionContext":
"""Provide a driver session context manager.
On first call with a PostgreSQL backend, detects pgvector/paradedb
extensions and updates the dialect accordingly.
Args:
*_args: Additional arguments.
statement_config: Optional statement configuration override.
**_kwargs: Additional keyword arguments.
Returns:
A context manager that yields an AdbcDriver instance.
"""
self._detect_extensions_if_needed()
statement_config = resolve_runtime_statement_config(
statement_config, self.statement_config, self._default_session_config
)
handler = _AdbcSessionConnectionHandler(self)
return AdbcSessionContext(
acquire_connection=handler.acquire_connection,
release_connection=handler.release_connection,
statement_config=statement_config,
driver_features=self.driver_features,
prepare_driver=self._prepare_driver,
)
[docs]
def get_signature_namespace(self) -> "dict[str, Any]":
"""Get the signature namespace for AdbcConfig types.
Returns:
Dictionary mapping type names to types.
"""
namespace = super().get_signature_namespace()
namespace.update({
"AdbcConnectionContext": AdbcConnectionContext,
"AdbcConnection": AdbcConnection,
"AdbcConnectionParams": AdbcConnectionParams,
"AdbcCursor": AdbcCursor,
"AdbcDriver": AdbcDriver,
"AdbcDriverFeatures": AdbcDriverFeatures,
"AdbcExceptionHandler": AdbcExceptionHandler,
"AdbcSessionContext": AdbcSessionContext,
})
return namespace
[docs]
def get_event_runtime_hints(self) -> "EventRuntimeHints":
"""Return polling defaults suitable for ADBC warehouses."""
return EventRuntimeHints(poll_interval=2.0, lease_seconds=60, retention_seconds=172_800)
