sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/utils/arrow_helpers.py

@@ -0,0 +1,318 @@
+"""Arrow conversion helpers for dict-to-Arrow transformations.
+
+This module provides utilities for converting Python dictionaries to Apache Arrow
+format, handling empty results, NULL values, and type inference.
+
+NOTE: This module is excluded from mypyc compilation to avoid segmentation faults
+when compiled drivers touch Arrow objects.
+"""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow
+from sqlspec.utils.type_guards import has_arrow_table_stats, has_get_data
+
+if TYPE_CHECKING:
+    from sqlspec.core.result import ArrowResult
+    from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable, PandasDataFrame, PolarsDataFrame
+
+__all__ = (
+    "arrow_table_column_names",
+    "arrow_table_num_columns",
+    "arrow_table_num_rows",
+    "arrow_table_to_pandas",
+    "arrow_table_to_polars",
+    "arrow_table_to_pylist",
+    "arrow_table_to_return_format",
+    "arrow_table_to_rows",
+    "build_ingest_telemetry",
+    "cast_arrow_table_schema",
+    "coerce_arrow_table",
+    "convert_dict_to_arrow",
+    "convert_dict_to_arrow_with_schema",
+    "ensure_arrow_table",
+)
+
+
+@overload
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]", return_format: Literal["table"] = "table", batch_size: int | None = None
+) -> "ArrowTable": ...
+
+
+@overload
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]", return_format: Literal["reader"], batch_size: int | None = None
+) -> "ArrowRecordBatchReader": ...
+
+
+@overload
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]", return_format: Literal["batch"], batch_size: int | None = None
+) -> "ArrowRecordBatch": ...
+
+
+@overload
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]", return_format: Literal["batches"], batch_size: int | None = None
+) -> "list[ArrowRecordBatch]": ...
+
+
+def convert_dict_to_arrow(
+    data: "list[dict[str, Any]]",
+    return_format: Literal["table", "reader", "batch", "batches"] = "table",
+    batch_size: int | None = None,
+) -> "ArrowTable | ArrowRecordBatch | ArrowRecordBatchReader | list[ArrowRecordBatch]":
+    """Convert list of dictionaries to Arrow Table or RecordBatch.
+
+    Handles empty results, NULL values, and automatic type inference.
+    Used by adapters that don't have native Arrow support to convert
+    dict-based results to Arrow format.
+
+    Args:
+        data: List of dictionaries (one per row).
+        return_format: Output format - "table" for Table, "batch"/"batches" for RecordBatch.
+            "reader" returns a RecordBatchReader.
+        batch_size: Chunk size for batching (used when return_format="batch"/"batches").
+
+    Returns:
+        ArrowTable or ArrowRecordBatch depending on return_format.
+
+    Examples:
+        >>> data = [
+        ...     {"id": 1, "name": "Alice"},
+        ...     {"id": 2, "name": "Bob"},
+        ... ]
+        >>> table = convert_dict_to_arrow(data, return_format="table")
+        >>> print(table.num_rows)
+        2
+
+        >>> batch = convert_dict_to_arrow(data, return_format="batch")
+        >>> print(batch.num_rows)
+        2
+    """
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if not data:
+        empty_schema = pa.schema([])
+        empty_table = pa.Table.from_pydict({}, schema=empty_schema)
+
+        if return_format == "reader":
+            return pa.RecordBatchReader.from_batches(empty_table.schema, empty_table.to_batches())
+
+        if return_format in {"batch", "batches"}:
+            batches = empty_table.to_batches(max_chunksize=batch_size)
+            return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+        return empty_table
+
+    columns: dict[str, list[Any]] = {key: [row.get(key) for row in data] for key in data[0]}
+    arrow_table = pa.Table.from_pydict(columns)
+
+    if return_format == "reader":
+        batches = arrow_table.to_batches(max_chunksize=batch_size)
+        return pa.RecordBatchReader.from_batches(arrow_table.schema, batches)
+
+    if return_format == "batches":
+        return arrow_table.to_batches(max_chunksize=batch_size)
+
+    if return_format == "batch":
+        batches = arrow_table.to_batches(max_chunksize=batch_size)
+        return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+    return arrow_table
+
+
+def convert_dict_to_arrow_with_schema(
+    data: "list[dict[str, Any]]",
+    return_format: Literal["table", "reader", "batch", "batches"] = "table",
+    batch_size: int | None = None,
+    arrow_schema: Any = None,
+) -> "ArrowTable | ArrowRecordBatch | ArrowRecordBatchReader | list[ArrowRecordBatch]":
+    """Convert dict rows to Arrow and optionally cast to a schema."""
+    table = convert_dict_to_arrow(data, return_format="table", batch_size=batch_size)
+    if arrow_schema is not None:
+        ensure_pyarrow()
+        import pyarrow as pa
+
+        if not isinstance(arrow_schema, pa.Schema):
+            msg = f"arrow_schema must be a pyarrow.Schema, got {type(arrow_schema).__name__}"
+            raise TypeError(msg)
+        if not data:
+            table = pa.Table.from_pydict({name: [] for name in arrow_schema.names}, schema=arrow_schema)
+        else:
+            table = table.cast(arrow_schema)
+
+    if return_format == "table":
+        return table
+
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if return_format == "reader":
+        batches = table.to_batches(max_chunksize=batch_size)
+        return pa.RecordBatchReader.from_batches(table.schema, batches)
+    if return_format == "batches":
+        return table.to_batches(max_chunksize=batch_size)
+    batches = table.to_batches(max_chunksize=batch_size)
+    return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+
+def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable":
+    """Coerce various sources to a PyArrow Table."""
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if has_get_data(source):
+        table = source.get_data()
+        if isinstance(table, pa.Table):
+            return table
+        msg = "ArrowResult did not return a pyarrow.Table instance"
+        raise TypeError(msg)
+    if isinstance(source, pa.Table):
+        return source
+    if isinstance(source, pa.RecordBatch):
+        return pa.Table.from_batches([source])
+    if isinstance(source, Iterable):
+        return pa.Table.from_pylist(list(source))
+    msg = f"Unsupported Arrow source type: {type(source).__name__}"
+    raise TypeError(msg)
+
+
+def ensure_arrow_table(data: Any) -> "ArrowTable":
+    """Ensure data is a PyArrow Table."""
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if isinstance(data, pa.Table):
+        return data
+    msg = f"Expected an Arrow Table, but got {type(data).__name__}"
+    raise TypeError(msg)
+
+
+def cast_arrow_table_schema(table: "ArrowTable", arrow_schema: Any) -> "ArrowTable":
+    """Cast an Arrow table to a provided schema.
+
+    Args:
+        table: Arrow table to cast.
+        arrow_schema: Optional pyarrow.Schema for casting.
+
+    Returns:
+        Arrow table with updated schema.
+
+    Raises:
+        TypeError: If arrow_schema is not a pyarrow.Schema instance.
+    """
+    if arrow_schema is None:
+        return table
+
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if not isinstance(arrow_schema, pa.Schema):
+        msg = f"arrow_schema must be a pyarrow.Schema, got {type(arrow_schema).__name__}"
+        raise TypeError(msg)
+    return table.cast(arrow_schema)
+
+
+def arrow_table_to_return_format(
+    table: "ArrowTable",
+    return_format: Literal["table", "reader", "batch", "batches"] = "table",
+    batch_size: int | None = None,
+) -> "ArrowTable | ArrowRecordBatch | ArrowRecordBatchReader | list[ArrowRecordBatch]":
+    """Convert an Arrow table into the requested return format.
+
+    Args:
+        table: Arrow table to convert.
+        return_format: Output format (table, reader, batch, batches).
+        batch_size: Batch size for reader/batch outputs.
+
+    Returns:
+        Converted Arrow data in the requested format.
+    """
+    ensure_pyarrow()
+    import pyarrow as pa
+
+    if return_format == "table":
+        return table
+
+    batches = table.to_batches(max_chunksize=batch_size)
+    if return_format == "reader":
+        return pa.RecordBatchReader.from_batches(table.schema, batches)
+    if return_format == "batches":
+        return batches
+    return batches[0] if batches else pa.RecordBatch.from_pydict({})
+
+
+def arrow_table_to_rows(
+    table: "ArrowTable", columns: "list[str] | None" = None
+) -> "tuple[list[str], list[tuple[Any, ...]]]":
+    """Convert Arrow table to column names and row tuples."""
+    ensure_pyarrow()
+    resolved_columns = columns or list(table.column_names)
+    if not resolved_columns:
+        msg = "Arrow table has no columns to import"
+        raise ValueError(msg)
+    batches = table.to_pylist()
+    records: list[tuple[Any, ...]] = []
+    for row in batches:
+        record = tuple(row.get(col) for col in resolved_columns)
+        records.append(record)
+    return resolved_columns, records
+
+
+def arrow_table_to_pylist(table: "ArrowTable") -> "list[dict[str, Any]]":
+    """Convert Arrow table to list of dictionaries."""
+    return table.to_pylist()
+
+
+def arrow_table_column_names(table: "ArrowTable") -> "list[str]":
+    """Return Arrow table column names."""
+    return list(table.column_names)
+
+
+def arrow_table_num_rows(table: "ArrowTable") -> int:
+    """Return Arrow table row count."""
+    return int(table.num_rows)
+
+
+def arrow_table_num_columns(table: "ArrowTable") -> int:
+    """Return Arrow table column count."""
+    return int(table.num_columns)
+
+
+def arrow_table_to_pandas(table: "ArrowTable") -> "PandasDataFrame":
+    """Convert Arrow table to pandas DataFrame."""
+    ensure_pandas()
+    import pandas as pd
+
+    result = table.to_pandas()
+    if not isinstance(result, pd.DataFrame):
+        msg = f"Expected a pandas DataFrame, but got {type(result).__name__}"
+        raise TypeError(msg)
+    return result
+
+
+def arrow_table_to_polars(table: "ArrowTable") -> "PolarsDataFrame":
+    """Convert Arrow table to Polars DataFrame."""
+    ensure_polars()
+    import polars as pl
+
+    result = pl.from_arrow(table)
+    if not isinstance(result, pl.DataFrame):
+        msg = f"Expected a Polars DataFrame, but got {type(result).__name__}"
+        raise TypeError(msg)
+    return result
+
+
+def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "dict[str, int | str]":
+    """Build telemetry dict from Arrow table statistics."""
+    if has_arrow_table_stats(table):
+        rows = int(table.num_rows)
+        bytes_processed = int(table.nbytes)
+    else:
+        rows = 0
+        bytes_processed = 0
+    return {"rows_processed": rows, "bytes_processed": bytes_processed, "format": format_label}

sqlspec/utils/config_tools.py

@@ -0,0 +1,332 @@
+"""Configuration utilities for SQLSpec.
+
+This module consolidates configuration-related helpers:
+- pyproject.toml discovery for CLI convenience
+- dotted path resolution for config objects
+- connection config normalization for adapters
+"""
+
+import inspect
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.exceptions import ConfigResolverError, ImproperConfigurationError
+from sqlspec.utils.module_loader import import_string
+from sqlspec.utils.sync_tools import async_, await_
+from sqlspec.utils.type_guards import (
+    has_config_attribute,
+    has_connection_config,
+    has_database_url_and_bind_key,
+    has_migration_config,
+)
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
+
+__all__ = (
+    "discover_config_from_pyproject",
+    "find_pyproject_toml",
+    "normalize_connection_config",
+    "parse_pyproject_config",
+    "reject_pool_aliases",
+    "resolve_config_async",
+    "resolve_config_sync",
+)
+
+
+# =============================================================================
+# pyproject.toml Discovery
+# =============================================================================
+
+
+def discover_config_from_pyproject() -> str | None:
+    """Find and parse pyproject.toml for SQLSpec config.
+
+    Walks filesystem upward from current directory to find pyproject.toml.
+    Parses [tool.sqlspec] section for 'config' key.
+
+    Returns:
+        Config path(s) as string (comma-separated if list), or None if not found.
+    """
+    pyproject_path = find_pyproject_toml()
+    if pyproject_path is None:
+        return None
+
+    return parse_pyproject_config(pyproject_path)
+
+
+def find_pyproject_toml() -> "Path | None":
+    """Walk filesystem upward to find pyproject.toml.
+
+    Starts from current working directory and walks up to filesystem root.
+    Stops at .git directory boundary (repository root) if found.
+
+    Returns:
+        Path to pyproject.toml, or None if not found.
+    """
+    current = Path.cwd()
+
+    while True:
+        pyproject = current / "pyproject.toml"
+        if pyproject.exists():
+            return pyproject
+
+        # Stop at .git boundary (repository root)
+        if (current / ".git").exists():
+            return None
+
+        # Stop at filesystem root
+        if current == current.parent:
+            return None
+
+        current = current.parent
+
+
+def parse_pyproject_config(pyproject_path: "Path") -> str | None:
+    """Parse pyproject.toml for [tool.sqlspec] config.
+
+    Args:
+        pyproject_path: Path to pyproject.toml file.
+
+    Returns:
+        Config path(s) as string (converts list to comma-separated), or None if not found.
+
+    Raises:
+        ValueError: If [tool.sqlspec].config has invalid type (not str or list[str]).
+    """
+    try:
+        with pyproject_path.open("rb") as f:
+            data = tomllib.load(f)
+    except Exception as e:
+        msg = f"Failed to parse {pyproject_path}: {e}"
+        raise ValueError(msg) from e
+
+    # Navigate to [tool.sqlspec] section
+    tool_section = data.get("tool", {})
+    if not isinstance(tool_section, dict):
+        return None
+
+    sqlspec_section = tool_section.get("sqlspec", {})
+    if not isinstance(sqlspec_section, dict):
+        return None
+
+    # Extract config value
+    config = sqlspec_section.get("config")
+    if config is None:
+        return None
+
+    # Handle string config
+    if isinstance(config, str):
+        return config
+
+    # Handle list config (convert to comma-separated)
+    if isinstance(config, list):
+        if not all(isinstance(item, str) for item in config):
+            msg = f"Invalid [tool.sqlspec].config in {pyproject_path}: list items must be strings"
+            raise ValueError(msg)
+        return ",".join(config)
+
+    # Invalid type
+    msg = f"Invalid [tool.sqlspec].config in {pyproject_path}: must be string or list of strings, got {type(config).__name__}"
+    raise ValueError(msg)
+
+
+# =============================================================================
+# Config Resolution
+# =============================================================================
+
+
+async def resolve_config_async(
+    config_path: str,
+) -> "list[AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]] | AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]":
+    """Resolve config from dotted path, handling callables and direct instances.
+
+    This is the async-first version that handles both sync and async callables efficiently.
+
+    Args:
+        config_path: Dotted path to config object or callable function.
+
+    Returns:
+        Resolved config instance or list of config instances.
+
+    Raises:
+        ConfigResolverError: If config resolution fails.
+    """
+    try:
+        config_obj = import_string(config_path)
+    except ImportError as e:
+        msg = f"Failed to import config from path '{config_path}': {e}"
+        raise ConfigResolverError(msg) from e
+
+    if not callable(config_obj):
+        return _validate_config_result(config_obj, config_path)
+
+    try:
+        if inspect.iscoroutinefunction(config_obj):
+            result = await config_obj()
+        else:
+            result = await async_(config_obj)()
+    except Exception as e:
+        msg = f"Failed to execute callable config '{config_path}': {e}"
+        raise ConfigResolverError(msg) from e
+
+    return _validate_config_result(result, config_path)
+
+
+def resolve_config_sync(
+    config_path: str,
+) -> "list[AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]] | AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]":
+    """Synchronous wrapper for resolve_config.
+
+    Args:
+        config_path: Dotted path to config object or callable function.
+
+    Returns:
+        Resolved config instance or list of config instances.
+    """
+    try:
+        config_obj = import_string(config_path)
+    except ImportError as e:
+        msg = f"Failed to import config from path '{config_path}': {e}"
+        raise ConfigResolverError(msg) from e
+
+    if not callable(config_obj):
+        return _validate_config_result(config_obj, config_path)
+
+    try:
+        if inspect.iscoroutinefunction(config_obj):
+            result = await_(config_obj, raise_sync_error=False)()
+        else:
+            result = config_obj()
+    except Exception as e:
+        msg = f"Failed to execute callable config '{config_path}': {e}"
+        raise ConfigResolverError(msg) from e
+
+    return _validate_config_result(result, config_path)
+
+
+def _validate_config_result(
+    config_result: Any, config_path: str
+) -> "list[AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]] | AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]":
+    """Validate that the config result is a valid config or list of configs.
+
+    Args:
+        config_result: The result from config resolution.
+        config_path: Original config path for error messages.
+
+    Returns:
+        Validated config result.
+
+    Raises:
+        ConfigResolverError: If config result is invalid.
+    """
+    if config_result is None:
+        msg = f"Config '{config_path}' resolved to None. Expected config instance or list of configs."
+        raise ConfigResolverError(msg)
+
+    if isinstance(config_result, Sequence) and not isinstance(config_result, str):
+        if not config_result:
+            msg = f"Config '{config_path}' resolved to empty list. Expected at least one config."
+            raise ConfigResolverError(msg)
+
+        for i, config in enumerate(config_result):  # pyright: ignore
+            if not _is_valid_config(config):
+                msg = f"Config '{config_path}' returned invalid config at index {i}. Expected database config instance."
+                raise ConfigResolverError(msg)
+
+        return cast("list[AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]]", list(config_result))  # pyright: ignore
+
+    if not _is_valid_config(config_result):
+        msg = f"Config '{config_path}' returned invalid type '{type(config_result).__name__}'. Expected database config instance or list."
+        raise ConfigResolverError(msg)
+
+    return cast("AsyncDatabaseConfig[Any, Any, Any] | SyncDatabaseConfig[Any, Any, Any]", config_result)
+
+
+def _is_valid_config(config: Any) -> bool:
+    """Check if an object is a valid SQLSpec database config.
+
+    Args:
+        config: Object to validate.
+
+    Returns:
+        True if object is a valid config instance (not a class).
+    """
+    # Reject config classes - must be instances
+    if isinstance(config, type):
+        return False
+
+    if has_config_attribute(config):
+        nested_config = config.config
+        if has_migration_config(nested_config):
+            return True
+
+    if has_migration_config(config) and config.migration_config is not None:
+        if has_connection_config(config):
+            return True
+        if has_database_url_and_bind_key(config):
+            return True
+
+    return False
+
+
+# =============================================================================
+# Connection Config Normalization
+# =============================================================================
+
+
+def reject_pool_aliases(kwargs: "dict[str, Any]") -> None:
+    """Reject legacy pool_config/pool_instance aliases.
+
+    Args:
+        kwargs: Keyword arguments passed to the adapter config constructor.
+
+    Raises:
+        ImproperConfigurationError: If deprecated pool aliases are supplied.
+    """
+    if "pool_config" in kwargs or "pool_instance" in kwargs:
+        msg = (
+            "pool_config and pool_instance are no longer supported. "
+            "Use connection_config and connection_instance instead."
+        )
+        raise ImproperConfigurationError(msg)
+
+
+def normalize_connection_config(
+    connection_config: "Mapping[str, Any] | None", *, extra_key: str = "extra"
+) -> "dict[str, Any]":
+    """Normalize an adapter connection_config dictionary.
+
+    This function:
+    - Copies the provided mapping into a new dict.
+    - Merges any nested dict stored under ``extra_key`` into the top-level config.
+    - Ensures the extra mapping is a dictionary (or None).
+
+    Args:
+        connection_config: Raw connection configuration mapping.
+        extra_key: Key holding additional keyword arguments to merge.
+
+    Returns:
+        Normalized connection configuration.
+
+    Raises:
+        ImproperConfigurationError: If ``extra_key`` exists but is not a dictionary.
+    """
+    normalized: dict[str, Any] = dict(connection_config) if connection_config else {}
+    extras = normalized.pop(extra_key, {})
+    if extras is None:
+        return normalized
+    if not isinstance(extras, dict):
+        msg = f"The '{extra_key}' field in connection_config must be a dictionary."
+        raise ImproperConfigurationError(msg)
+    normalized.update(extras)
+    return normalized