Source code for sqlspec.adapters.arrow_odbc.config

"""arrow-odbc database configuration."""

from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast

from typing_extensions import NotRequired

from sqlspec.adapters.arrow_odbc._typing import ArrowOdbcConnection, ArrowOdbcSessionContext, arrow_odbc_connect
from sqlspec.adapters.arrow_odbc.core import (
    apply_driver_features,
    build_connection_config,
    build_statement_config,
    default_statement_config,
    resolve_dialect_from_dbms_name,
)
from sqlspec.adapters.arrow_odbc.driver import ArrowOdbcDriver
from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig
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
from sqlspec.utils.serializers import from_json, to_json
from sqlspec.utils.type_guards import supports_close

if TYPE_CHECKING:
    from collections.abc import Callable
    from types import TracebackType

    from sqlspec.core import StatementConfig
    from sqlspec.observability import ObservabilityConfig

__all__ = ("ArrowOdbcConfig", "ArrowOdbcConnectionParams", "ArrowOdbcDriverFeatures")



[docs]
class ArrowOdbcConnectionParams(TypedDict):
    """arrow-odbc connection parameters."""

    connection_string: NotRequired[str]
    dsn: NotRequired[str]
    driver: NotRequired[str]
    server: NotRequired[str]
    host: NotRequired[str]
    database: NotRequired[str]
    uid: NotRequired[str]
    pwd: NotRequired[str]
    user: NotRequired[str]
    password: NotRequired[str]
    trusted_connection: NotRequired[bool | str]
    trust_server_certificate: NotRequired[bool | str]
    encrypt: NotRequired[bool | str]
    login_timeout: NotRequired[int]
    login_timeout_sec: NotRequired[int]
    packet_size: NotRequired[int]
    autocommit: NotRequired[bool]
    extra: NotRequired[dict[str, Any]]

[docs]

class ArrowOdbcDriverFeatures(TypedDict): """arrow-odbc driver feature flags.""" chunk_size: NotRequired[int] max_bytes_per_batch: NotRequired[int] max_text_size: NotRequired[int] max_binary_size: NotRequired[int] fetch_concurrently: NotRequired[bool] query_timeout_sec: NotRequired[int] connection_string: NotRequired[str] dbms_name: NotRequired[str] json_serializer: "NotRequired[Callable[[Any], str]]" json_deserializer: "NotRequired[Callable[[str], Any]]" enable_events: NotRequired[bool] on_connection_create: "NotRequired[Callable[[ArrowOdbcConnection], None]]"

def _apply_json_serializer_override(statement_config: "StatementConfig", features: dict[str, Any]) -> "StatementConfig": serializer = cast("Callable[[Any], str] | None", features.get("json_serializer")) deserializer = cast("Callable[[str], Any] | None", features.get("json_deserializer")) if serializer is to_json and deserializer is from_json: return statement_config return statement_config.replace( parameter_config=statement_config.parameter_config.with_json_serializers( serializer or to_json, deserializer=deserializer ) ) class ArrowOdbcConnectionContext(SyncPoolConnectionContext): """Context manager for arrow-odbc connections.""" __slots__ = ("_connection",) def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._connection: ArrowOdbcConnection | None = None def __enter__(self) -> "ArrowOdbcConnection": self._connection = self._config.create_connection() return cast("ArrowOdbcConnection", self._connection) def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": if self._connection is not None: _close_arrow_odbc_connection(self._connection) self._connection = None return None class _ArrowOdbcSessionConnectionHandler(SyncPoolSessionFactory): """Session connection handler for no-pool arrow-odbc sessions.""" __slots__ = ("_connection",) def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._connection: ArrowOdbcConnection | None = None def acquire_connection(self) -> "ArrowOdbcConnection": self._connection = self._config.create_connection() return cast("ArrowOdbcConnection", self._connection) def release_connection(self, _conn: "ArrowOdbcConnection", **kwargs: Any) -> None: if self._connection is None: return _close_arrow_odbc_connection(self._connection) self._connection = None

[docs]

class ArrowOdbcConfig(NoPoolSyncConfig[ArrowOdbcConnection, ArrowOdbcDriver]): """Configuration for synchronous arrow-odbc connections.""" driver_type: "ClassVar[type[ArrowOdbcDriver]]" = ArrowOdbcDriver connection_type: "ClassVar[type[ArrowOdbcConnection]]" = ArrowOdbcConnection 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]" = False supports_native_parquet_import: "ClassVar[bool]" = False _connection_context_class: "ClassVar[type[ArrowOdbcConnectionContext]]" = ArrowOdbcConnectionContext _session_factory_class: "ClassVar[type[_ArrowOdbcSessionConnectionHandler]]" = _ArrowOdbcSessionConnectionHandler _session_context_class: "ClassVar[type[ArrowOdbcSessionContext]]" = ArrowOdbcSessionContext _default_statement_config = default_statement_config __slots__ = ("_user_connection_hook",)

[docs]

def __init__( self, *, connection_config: "ArrowOdbcConnectionParams | dict[str, Any] | None" = None, connection_instance: "ArrowOdbcConnection | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: "StatementConfig | None" = None, driver_features: "ArrowOdbcDriverFeatures | dict[str, Any] | None" = None, bind_key: str | None = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: """Initialize arrow-odbc configuration.""" normalized = normalize_connection_config(connection_config) provided_statement_config = statement_config statement_config, features = apply_driver_features( statement_config or default_statement_config, driver_features ) connection_string = normalized.get("connection_string") if connection_string is not None: features.setdefault("connection_string", str(connection_string)) elif normalized.get("driver") is not None: features.setdefault("dbms_name", str(normalized["driver"])) if provided_statement_config is None: statement_config = _resolve_statement_config(features) self._user_connection_hook = cast( "Callable[[ArrowOdbcConnection], None] | None", features.pop("on_connection_create", None) ) super().__init__( connection_config=normalized, connection_instance=connection_instance, migration_config=migration_config, statement_config=statement_config, driver_features=features, bind_key=bind_key, extension_config=extension_config, observability_config=observability_config, **kwargs, )

[docs]

def create_connection(self) -> "ArrowOdbcConnection": """Create and return a new arrow-odbc connection.""" if self.connection_instance is not None: return cast("ArrowOdbcConnection", self.connection_instance) connection_string, connect_kwargs = build_connection_config(self.connection_config) try: connection = cast("ArrowOdbcConnection", arrow_odbc_connect(connection_string, **connect_kwargs)) if self._user_connection_hook is not None: self._user_connection_hook(connection) return cast("ArrowOdbcConnection", connection) except Exception as exc: msg = f"Could not configure arrow-odbc connection. Error: {exc}" raise ImproperConfigurationError(msg) from exc