"""ADBC driver implementation for Arrow Database Connectivity.
Provides database connectivity through ADBC with support for multiple
database dialects, parameter style conversion, and transaction management.
"""
import contextlib
from typing import TYPE_CHECKING, Any, Literal, cast
from typing_extensions import final
from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcSessionContext
from sqlspec.adapters.adbc.core import (
_prepare_batch_with_casts,
collect_rows,
create_mapped_exception,
detect_dialect,
driver_profile,
get_statement_config,
handle_postgres_rollback,
is_postgres_dialect,
normalize_postgres_empty_parameters,
normalize_script_rowcount,
prepare_postgres_parameters,
prepare_postgres_uuid_bindings,
resolve_column_names,
resolve_dialect_name,
resolve_many_rowcount,
resolve_parameter_casts,
resolve_rowcount,
)
from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary
from sqlspec.core import (
SQL,
StatementConfig,
build_arrow_result_from_reader,
build_arrow_result_from_table,
get_cache_config,
register_driver_profile,
)
from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream
from sqlspec.exceptions import DatabaseConnectionError, SQLSpecError
from sqlspec.utils.arrow_helpers import arrow_reader_with_deferred_close, arrow_table_to_pylist
from sqlspec.utils.logging import get_logger
from sqlspec.utils.module_loader import ensure_pyarrow
from sqlspec.utils.serializers import to_json
from sqlspec.utils.text import normalize_identifier, quote_identifier
if TYPE_CHECKING:
from collections.abc import Callable
from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcRawCursor
from sqlspec.builder import QueryBuilder
from sqlspec.core import ArrowResult, Statement, StatementFilter
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
from sqlspec.typing import ArrowReturnFormat, StatementParameters
__all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "AdbcSessionContext")
logger = get_logger("sqlspec.adapters.adbc")
_MULTI_ROW_BIND_UNSUPPORTED = "Binding multiple rows at once is not supported"
@final
class AdbcExceptionHandler(BaseSyncExceptionHandler):
"""Context manager for handling ADBC database exceptions.
ADBC propagates underlying database errors. Exception mapping
depends on the specific ADBC driver being used.
Uses deferred exception pattern for mypyc compatibility: exceptions
are stored in pending_exception rather than raised from __exit__
to avoid ABI boundary violations with compiled code.
"""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if issubclass(exc_type, AdbcNativeError):
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
class AdbcSelectStreamSource:
"""Native ADBC chunk source backed by a ``RecordBatchReader``."""
__slots__ = ("_cursor_manager", "_driver", "_parameters", "_reader", "_sql")
def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any) -> None:
self._driver = driver
self._sql = sql
self._parameters = parameters
self._cursor_manager: AdbcCursor | None = None
self._reader: Any = None
def start(self) -> None:
cursor_manager = self._driver.with_cursor(self._driver.connection)
reader: Any = None
try:
cursor = cursor_manager.__enter__()
handler = self._driver.handle_database_exceptions()
with handler:
execute_parameters = normalize_postgres_empty_parameters(self._driver._dialect_name, self._parameters)
cursor.execute(self._sql, parameters=execute_parameters)
reader = _fetch_record_batch(cursor)
self._driver._check_pending_exception(handler)
except BaseException:
with contextlib.suppress(Exception):
cursor_manager.__exit__(None, None, None)
raise
if reader is None:
msg = "ADBC did not return a record batch reader."
raise SQLSpecError(msg)
self._cursor_manager = cursor_manager
self._reader = iter(reader)
def fetch_chunk(self) -> "list[dict[str, Any]]":
reader = self._reader
if reader is None:
return []
while True:
try:
batch = next(reader)
except StopIteration:
return []
rows = arrow_table_to_pylist(
batch,
decode_arrow_extension_types=bool(
self._driver.driver_features.get("enable_arrow_extension_types", True)
),
)
if rows:
return rows
def close(self, error: bool = False) -> None:
self._reader = None
cursor_manager = self._cursor_manager
self._cursor_manager = None
if cursor_manager is not None:
with contextlib.suppress(Exception):
cursor_manager.__exit__(None, None, None)
@final
class AdbcDriver(SyncDriverAdapterBase):
"""ADBC driver for Arrow Database Connectivity.
Provides database connectivity through ADBC with support for multiple
database dialects, parameter style conversion, and transaction management.
"""
__slots__ = (
"_column_name_cache",
"_data_dictionary",
"_detected_dialect",
"_dialect_name",
"_is_flightsql",
"_is_postgres",
"_json_serializer",
"_transaction_active",
"dialect",
)
def __init__(
self,
connection: "AdbcConnection",
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
*,
dialect: "str | None" = None,
) -> None:
self._detected_dialect = detect_dialect(connection, logger, fallback_dialect=dialect)
self._is_flightsql = self._detect_flightsql_connection(connection)
if statement_config is None:
base_config = get_statement_config(self._detected_dialect)
statement_config = base_config.replace(enable_caching=get_cache_config().compiled_cache_enabled)
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self.dialect = statement_config.dialect
self._dialect_name = dialect or resolve_dialect_name(self.dialect)
self._is_postgres = is_postgres_dialect(self._dialect_name)
self._json_serializer = cast("Callable[[Any], str]", self.driver_features.get("json_serializer", to_json))
self._data_dictionary: AdbcDataDictionary | None = None
self._column_name_cache: dict[int, tuple[Any, list[str]]] = {}
self._transaction_active = False