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

sqlspec/driver/_sql_helpers.py

@@ -0,0 +1,148 @@
+"""Pure helper functions for SQL dialect translation.
+
+These functions are extracted from SQLTranslatorMixin to eliminate
+cross-trait attribute access that causes mypyc segmentation faults.
+"""
+
+from typing import TYPE_CHECKING, Final, NoReturn
+
+from sqlglot import exp, parse_one
+
+from sqlspec.core import SQL, Statement
+from sqlspec.exceptions import SQLConversionError
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+
+__all__ = (
+    "DEFAULT_PRETTY",
+    "convert_to_dialect",
+    "generate_sql_safely",
+    "parse_statement_safely",
+    "raise_conversion_error",
+    "raise_parse_error",
+    "raise_statement_parse_error",
+)
+
+
+DEFAULT_PRETTY: Final[bool] = True
+
+
+def parse_statement_safely(statement: "Statement", dialect: "DialectType | None") -> "exp.Expression":
+    """Parse statement with error handling.
+
+    Args:
+        statement: SQL statement to parse.
+        dialect: Source dialect for parsing.
+
+    Returns:
+        Parsed expression.
+
+    Raises:
+        SQLConversionError: If parsing fails.
+
+    """
+    try:
+        sql_string = str(statement)
+        return parse_one(sql_string, dialect=dialect)
+    except Exception as e:
+        raise_parse_error(e)
+
+
+def generate_sql_safely(expression: "exp.Expression", dialect: "DialectType | None", pretty: bool) -> str:
+    """Generate SQL with error handling.
+
+    Args:
+        expression: Parsed expression to convert.
+        dialect: Target SQL dialect.
+        pretty: Whether to format the output SQL.
+
+    Returns:
+        Generated SQL string.
+
+    Raises:
+        SQLConversionError: If generation fails.
+
+    """
+    try:
+        return expression.sql(dialect=dialect, pretty=pretty)
+    except Exception as e:
+        raise_conversion_error(dialect, e)
+
+
+def convert_to_dialect(
+    statement: "Statement",
+    source_dialect: "DialectType | None",
+    to_dialect: "DialectType | None" = None,
+    pretty: bool = DEFAULT_PRETTY,
+) -> str:
+    """Convert a statement to a target SQL dialect.
+
+    Args:
+        statement: SQL statement to convert.
+        source_dialect: Source dialect for parsing.
+        to_dialect: Target dialect (defaults to source_dialect).
+        pretty: Whether to format the output SQL.
+
+    Returns:
+        SQL string in target dialect.
+
+    Raises:
+        SQLConversionError: If conversion fails.
+
+    """
+    parsed_expression: exp.Expression | None = None
+
+    if statement is not None and isinstance(statement, SQL):
+        if statement.expression is None:
+            raise_statement_parse_error()
+        parsed_expression = statement.expression
+    elif isinstance(statement, exp.Expression):
+        parsed_expression = statement
+    else:
+        parsed_expression = parse_statement_safely(statement, source_dialect)
+
+    target_dialect = to_dialect or source_dialect
+
+    return generate_sql_safely(parsed_expression, target_dialect, pretty)
+
+
+def raise_statement_parse_error() -> "NoReturn":
+    """Raise error for unparsable statements.
+
+    Raises:
+        SQLConversionError: Always raised.
+
+    """
+    msg = "Statement could not be parsed"
+    raise SQLConversionError(msg)
+
+
+def raise_parse_error(e: Exception) -> "NoReturn":
+    """Raise error for parsing failures.
+
+    Args:
+        e: Original exception that caused the failure.
+
+    Raises:
+        SQLConversionError: Always raised.
+
+    """
+    error_msg = f"Failed to parse SQL statement: {e!s}"
+    raise SQLConversionError(error_msg) from e
+
+
+def raise_conversion_error(dialect: "DialectType | None", e: Exception) -> "NoReturn":
+    """Raise error for conversion failures.
+
+    Args:
+        dialect: Target dialect that caused the failure.
+        e: Original exception that caused the failure.
+
+    Raises:
+        SQLConversionError: Always raised.
+
+    """
+    error_msg = f"Failed to convert SQL expression to {dialect}: {e!s}"
+    raise SQLConversionError(error_msg) from e

sqlspec/driver/_storage_helpers.py

@@ -0,0 +1,144 @@
+"""Pure helper functions for storage operations.
+
+These functions are extracted from StorageDriverMixin to eliminate
+cross-trait attribute access that causes mypyc segmentation faults.
+"""
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Final, cast
+
+from sqlspec.storage import StorageBridgeJob, StorageTelemetry, create_storage_bridge_job
+from sqlspec.utils.arrow_helpers import arrow_table_to_rows as _arrow_table_to_rows_impl
+from sqlspec.utils.arrow_helpers import build_ingest_telemetry as _build_ingest_telemetry_impl
+from sqlspec.utils.arrow_helpers import coerce_arrow_table as _coerce_arrow_table_impl
+
+if TYPE_CHECKING:
+    from sqlspec.core.result import ArrowResult
+    from sqlspec.storage import StorageDestination
+    from sqlspec.typing import ArrowTable
+
+
+__all__ = (
+    "CAPABILITY_HINTS",
+    "arrow_table_to_rows",
+    "attach_partition_telemetry",
+    "build_ingest_telemetry",
+    "coerce_arrow_table",
+    "create_storage_job",
+    "stringify_storage_target",
+)
+
+
+CAPABILITY_HINTS: Final[dict[str, str]] = {
+    "arrow_export_enabled": "native Arrow export",
+    "arrow_import_enabled": "native Arrow import",
+    "parquet_export_enabled": "native Parquet export",
+    "parquet_import_enabled": "native Parquet import",
+}
+
+
+def stringify_storage_target(target: "StorageDestination | None") -> str | None:
+    """Convert storage target to string representation.
+
+    Args:
+        target: Storage destination path or None.
+
+    Returns:
+        String representation of the path or None.
+
+    """
+    if target is None:
+        return None
+    if isinstance(target, Path):
+        return target.as_posix()
+    return str(target)
+
+
+def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable":
+    """Coerce various sources to a PyArrow Table.
+
+    Args:
+        source: ArrowResult, PyArrow Table, RecordBatch, or iterable of dicts.
+
+    Returns:
+        PyArrow Table.
+
+    Raises:
+        TypeError: If source type is not supported.
+
+    """
+    return _coerce_arrow_table_impl(source)
+
+
+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.
+
+    Args:
+        table: Arrow table to convert.
+        columns: Optional list of columns to extract. Defaults to all columns.
+
+    Returns:
+        Tuple of (column_names, list of row tuples).
+
+    Raises:
+        ValueError: If table has no columns to import.
+
+    """
+    return _arrow_table_to_rows_impl(table, columns)
+
+
+def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "StorageTelemetry":
+    """Build telemetry dict from Arrow table statistics.
+
+    Args:
+        table: Arrow table to extract statistics from.
+        format_label: Format label for telemetry.
+
+    Returns:
+        StorageTelemetry dict with row/byte counts.
+
+    """
+    telemetry = _build_ingest_telemetry_impl(table, format_label=format_label)
+    return cast("StorageTelemetry", telemetry)
+
+
+def attach_partition_telemetry(telemetry: "StorageTelemetry", partitioner: "dict[str, object] | None") -> None:
+    """Attach partitioner info to telemetry dict (mutates in place).
+
+    Args:
+        telemetry: Telemetry dict to update.
+        partitioner: Partitioner configuration or None.
+
+    """
+    if not partitioner:
+        return
+    extra = dict(telemetry.get("extra", {}))
+    extra["partitioner"] = partitioner
+    telemetry["extra"] = extra
+
+
+def create_storage_job(
+    produced: "StorageTelemetry", provided: "StorageTelemetry | None" = None, *, status: str = "completed"
+) -> "StorageBridgeJob":
+    """Create a StorageBridgeJob from telemetry data.
+
+    Args:
+        produced: Telemetry from the production side of the operation.
+        provided: Optional telemetry from the source side.
+        status: Job status string.
+
+    Returns:
+        StorageBridgeJob instance.
+
+    """
+    merged = cast("StorageTelemetry", dict(produced))
+    if provided:
+        source_bytes = provided.get("bytes_processed")
+        if source_bytes is not None:
+            merged["bytes_processed"] = int(merged.get("bytes_processed", 0)) + int(source_bytes)
+        extra = dict(merged.get("extra", {}))
+        extra["source"] = provided
+        merged["extra"] = extra
+    return create_storage_bridge_job(status, merged)