"""Generic data dictionary for arrow-odbc connections."""
from typing import TYPE_CHECKING, Any, ClassVar, Final
from mypy_extensions import mypyc_attr
from sqlspec.data_dictionary import (
ColumnMetadata,
DDLResult,
ForeignKeyMetadata,
IndexMetadata,
MetadataCapability,
MetadataCapabilityProfile,
MetadataFidelity,
MetadataSource,
MetadataSupport,
ObjectIdentity,
TableMetadata,
VersionInfo,
get_data_dictionary_loader,
get_dialect_config,
)
from sqlspec.driver import SyncDataDictionaryBase
from sqlspec.exceptions import SQLFileNotFoundError
from sqlspec.utils.text import normalize_identifier, quote_identifier
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlspec.adapters.arrow_odbc.driver import ArrowOdbcDriver
from sqlspec.core import SQL
from sqlspec.data_dictionary._types import DialectConfig
__all__ = ("ArrowOdbcDataDictionary",)
_ARROW_DECIMAL_FORMAT: Final = "DECIMAL({precision},{scale})"
_ODBC_METADATA_DOMAINS: Final = (
"schemas",
"objects",
"tables",
"columns",
"constraints",
"indexes",
"views",
"routines",
"privileges",
"dependencies",
"ddl",
"system",
"odbc_catalog",
)
_ODBC_CATALOG_UNAVAILABLE_WARNING: Final = (
"arrow-odbc does not expose SQLGetInfo, SQLGetFunctions, or raw ODBC catalog functions through its Python "
"Connection API."
)
_ODBC_DDL_UNSUPPORTED_WARNING: Final = "Arrow ODBC transport metadata is not a lossless DDL source."
_ODBC_COLUMN_PARTIAL_WARNING: Final = (
"Arrow ODBC column metadata comes from dialect SQL packs or zero-row Arrow schema probes, not SQLColumns."
)
def _arrow_type_to_sql(data_type: Any) -> str:
import pyarrow as pa
types = pa.types
if types.is_boolean(data_type):
return "BOOLEAN"
if types.is_int8(data_type) or types.is_int16(data_type) or types.is_uint8(data_type) or types.is_uint16(data_type):
return "SMALLINT"
if types.is_int32(data_type) or types.is_uint32(data_type):
return "INTEGER"
if types.is_int64(data_type) or types.is_uint64(data_type):
return "BIGINT"
if types.is_float16(data_type) or types.is_float32(data_type):
return "REAL"
if types.is_float64(data_type):
return "DOUBLE"
if types.is_decimal(data_type):
return _ARROW_DECIMAL_FORMAT.format(precision=data_type.precision, scale=data_type.scale)
if types.is_string(data_type) or types.is_large_string(data_type):
return "VARCHAR"
if types.is_binary(data_type) or types.is_large_binary(data_type) or types.is_fixed_size_binary(data_type):
return "VARBINARY"
if types.is_date(data_type):
return "DATE"
if types.is_time(data_type):
return "TIME"
if types.is_timestamp(data_type):
return "TIMESTAMP"
return str(data_type).upper()
@mypyc_attr(allow_interpreted_subclasses=True, native_class=False)
class ArrowOdbcDataDictionary(SyncDataDictionaryBase):
"""Runtime-dialect data dictionary for generic ODBC connections."""
dialect: ClassVar[str] = "sqlite"
def __init__(self, dialect: str = "sqlite") -> None:
super().__init__()
self._dialect = dialect