Source code for sqlspec.adapters.adbc.config

"""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,
    build_postgres_extension_probe_names,
    detect_postgres_extensions,
    get_statement_config,
    is_postgres_dialect,
    resolve_dialect_from_config,
    resolve_dialect_name,
    resolve_driver_connect_func,
    resolve_postgres_extension_state,
    resolve_runtime_statement_config,
)
from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler
from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig
from sqlspec.core import StatementConfig
from sqlspec.driver._sync 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")


class AdbcConnectionParams(TypedDict):
    """ADBC connection parameters."""

    uri: 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]]
    gizmosql_backend: NotRequired[str]
    tls_skip_verify: NotRequired[bool]
    extra: NotRequired[dict[str, Any]]




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_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_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

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 storage_partition_strategies: "ClassVar[tuple[str, ...]]" = ("fixed", "rows_per_chunk") _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", "_pgvector_available", "_resolved_dialect", "_user_connection_hook", )

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. """ self.connection_config = normalize_connection_config(connection_config) self._pgvector_available: bool | None = None self._paradedb_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, )