From 192db7f176c5aa770be3811484e78d880e915e3a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 17:27:27 +0000 Subject: [PATCH 01/14] feat(arrow): add streaming RecordBatch coercion and batch telemetry helpers --- sqlspec/driver/_common.py | 47 +++++++++++- sqlspec/driver/_storage_helpers.py | 93 ++++++++++++++++++++++- sqlspec/storage/pipeline.py | 14 ++++ sqlspec/utils/arrow_helpers.py | 70 ++++++++++++++++- tests/unit/driver/test_storage_helpers.py | 51 +++++++++++++ tests/unit/utils/test_arrow_helpers.py | 56 ++++++++++++++ 6 files changed, 325 insertions(+), 6 deletions(-) create mode 100644 tests/unit/driver/test_storage_helpers.py diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index f0f220288..c039afb8c 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -77,12 +77,15 @@ from sqlspec.driver._query_cache import CachedQuery, QueryCache, StatementCacheHolder, resolve_statement_cache_size from sqlspec.driver._storage_helpers import ( CAPABILITY_HINTS, + arrow_batch_to_rows, arrow_table_needs_parameter_preparation, arrow_table_to_rows, attach_partition_telemetry, build_ingest_telemetry, + coerce_arrow_batches, coerce_arrow_table, create_storage_job, + ingest_telemetry_from_batches, records_to_arrow_table, ) from sqlspec.exceptions import ( @@ -124,7 +127,7 @@ StorageTelemetry, SyncStoragePipeline, ) - from sqlspec.typing import ArrowTable, SchemaT, StatementParameters + from sqlspec.typing import ArrowRecordBatch, ArrowTable, SchemaT, StatementParameters __all__ = ( @@ -3133,6 +3136,18 @@ def _coerce_arrow_table(self, source: "ArrowResult | Any") -> "ArrowTable": """Coerce various sources to a PyArrow Table.""" return coerce_arrow_table(source) + @staticmethod + def _coerce_arrow_batches( + source: "ArrowResult | Any", *, batch_size: int = 65_536 + ) -> "tuple[list[str], abc.Iterator[ArrowRecordBatch]]": + """Coerce an Arrow source into column names and a lazy RecordBatch iterator.""" + return coerce_arrow_batches(source, batch_size=batch_size) + + @staticmethod + def _arrow_batch_to_rows(batch: "ArrowRecordBatch", columns: "list[str] | None" = None) -> "list[tuple[Any, ...]]": + """Convert a single Arrow RecordBatch into row tuples.""" + return arrow_batch_to_rows(batch, columns) + @staticmethod def _arrow_table_to_rows( table: "ArrowTable", columns: "list[str] | None" = None @@ -3156,6 +3171,36 @@ def _ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "S """Build telemetry dict from Arrow table statistics.""" return build_ingest_telemetry(table, format_label=format_label) + @staticmethod + def _ingest_telemetry_from_batches( + rows_processed: int, + bytes_processed: int, + *, + duration_s: float | None = None, + format_label: str = "arrow", + read_duration_s: float | None = None, + encode_duration_s: float | None = None, + write_duration_s: float | None = None, + commit_duration_s: float | None = None, + shards_processed: int | None = None, + shards_skipped: int | None = None, + wal_bytes: int | None = None, + ) -> "StorageTelemetry": + """Build telemetry dict from accumulated batch metrics.""" + return ingest_telemetry_from_batches( + rows_processed, + bytes_processed, + duration_s=duration_s, + format_label=format_label, + read_duration_s=read_duration_s, + encode_duration_s=encode_duration_s, + write_duration_s=write_duration_s, + commit_duration_s=commit_duration_s, + shards_processed=shards_processed, + shards_skipped=shards_skipped, + wal_bytes=wal_bytes, + ) + @staticmethod def _records_to_arrow_table( records: "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]", columns: "list[str] | None" diff --git a/sqlspec/driver/_storage_helpers.py b/sqlspec/driver/_storage_helpers.py index 12b0d397d..e9ae08a7d 100644 --- a/sqlspec/driver/_storage_helpers.py +++ b/sqlspec/driver/_storage_helpers.py @@ -8,28 +8,33 @@ 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_batch_to_rows as _arrow_batch_to_rows_impl from sqlspec.utils.arrow_helpers import arrow_table_needs_parameter_preparation as _arrow_rows_need_preparation 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 _ingest_telemetry_impl +from sqlspec.utils.arrow_helpers import coerce_arrow_batches as _coerce_arrow_batches_impl from sqlspec.utils.arrow_helpers import coerce_arrow_table as _coerce_arrow_table_impl from sqlspec.utils.arrow_helpers import records_to_arrow_table as _records_to_arrow_table_impl if TYPE_CHECKING: - from collections.abc import Iterable, Mapping + from collections.abc import Iterable, Iterator, Mapping from sqlspec.core.result import ArrowResult from sqlspec.storage import StorageDestination - from sqlspec.typing import ArrowTable + from sqlspec.typing import ArrowRecordBatch, ArrowTable __all__ = ( "CAPABILITY_HINTS", + "arrow_batch_to_rows", "arrow_table_needs_parameter_preparation", "arrow_table_to_rows", "attach_partition_telemetry", "build_ingest_telemetry", + "coerce_arrow_batches", "coerce_arrow_table", "create_storage_job", + "ingest_telemetry_from_batches", "records_to_arrow_table", "stringify_storage_target", ) @@ -71,6 +76,34 @@ def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable": return _coerce_arrow_table_impl(source) +def coerce_arrow_batches( + source: "ArrowResult | Any", *, batch_size: int = 65_536 +) -> "tuple[list[str], Iterator[ArrowRecordBatch]]": + """Coerce an Arrow source into column names and a lazy RecordBatch iterator. + + Args: + source: ArrowResult, RecordBatchReader, RecordBatch, Table, or row iterable. + batch_size: Maximum rows per batch when chunking an in-memory Table. + + Returns: + Tuple of ``(column_names, batch_iterator)``. + """ + return _coerce_arrow_batches_impl(source, batch_size=batch_size) + + +def arrow_batch_to_rows(batch: "ArrowRecordBatch", columns: "list[str] | None" = None) -> "list[tuple[Any, ...]]": + """Convert a single Arrow RecordBatch into row tuples. + + Args: + batch: Arrow RecordBatch to convert. + columns: Optional subset or order of columns to extract. + + Returns: + List of row tuples. + """ + return _arrow_batch_to_rows_impl(batch, columns) + + def arrow_table_to_rows( table: "ArrowTable", columns: "list[str] | None" = None ) -> "tuple[list[str], list[tuple[Any, ...]]]": @@ -105,6 +138,62 @@ def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") return cast("StorageTelemetry", telemetry) +def ingest_telemetry_from_batches( + rows_processed: int, + bytes_processed: int, + *, + duration_s: float | None = None, + format_label: str = "arrow", + read_duration_s: float | None = None, + encode_duration_s: float | None = None, + write_duration_s: float | None = None, + commit_duration_s: float | None = None, + shards_processed: int | None = None, + shards_skipped: int | None = None, + wal_bytes: int | None = None, +) -> "StorageTelemetry": + """Build a StorageTelemetry payload from accumulated batch metrics. + + Args: + rows_processed: Total number of rows processed across all batches. + bytes_processed: Total uncompressed or wire bytes processed across all batches. + duration_s: Optional total wall-clock duration in seconds. + format_label: Payload format label (default ``"arrow"``). + read_duration_s: Optional source read time in seconds. + encode_duration_s: Optional payload encoding time in seconds. + write_duration_s: Optional database/storage write time in seconds. + commit_duration_s: Optional commit/finalization time in seconds. + shards_processed: Optional count of processed storage shards. + shards_skipped: Optional count of checkpoint-skipped storage shards. + wal_bytes: Optional WAL byte delta recorded during ingestion. + + Returns: + Populated ``StorageTelemetry`` dictionary. + """ + telemetry: StorageTelemetry = { + "rows_processed": int(rows_processed), + "bytes_processed": int(bytes_processed), + "format": format_label, + } + if duration_s is not None: + telemetry["duration_s"] = float(duration_s) + if read_duration_s is not None: + telemetry["read_duration_s"] = float(read_duration_s) + if encode_duration_s is not None: + telemetry["encode_duration_s"] = float(encode_duration_s) + if write_duration_s is not None: + telemetry["write_duration_s"] = float(write_duration_s) + if commit_duration_s is not None: + telemetry["commit_duration_s"] = float(commit_duration_s) + if shards_processed is not None: + telemetry["shards_processed"] = int(shards_processed) + if shards_skipped is not None: + telemetry["shards_skipped"] = int(shards_skipped) + if wal_bytes is not None: + telemetry["wal_bytes"] = int(wal_bytes) + return telemetry + + def records_to_arrow_table( records: "Iterable[Mapping[str, Any]] | Iterable[Iterable[Any]]", columns: "list[str] | None" ) -> "ArrowTable": diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index 4548cd4aa..ea343881b 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -94,6 +94,20 @@ class StorageTelemetry(TypedDict, total=False): """Number of partitioned output objects written during export.""" duration_s: float """Elapsed wall-clock transfer time in seconds.""" + read_duration_s: float + """Elapsed time spent reading source storage or Arrow batches in seconds.""" + encode_duration_s: float + """Elapsed time spent encoding batches into wire payloads in seconds.""" + write_duration_s: float + """Elapsed time spent writing payloads to the database or storage sink in seconds.""" + commit_duration_s: float + """Elapsed time spent committing the transaction or finalizing the sink in seconds.""" + shards_processed: int + """Number of storage shards loaded during a multi-shard operation.""" + shards_skipped: int + """Number of already-checkpointed storage shards skipped during resume.""" + wal_bytes: int + """Estimated or observed Write-Ahead Log bytes generated during ingestion.""" format: str """Serialized wire/file format (such as ``parquet``, ``arrow``, ``csv``, or ``jsonl``).""" extra: "dict[str, object]" diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index db48de372..3c706c2b9 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -8,7 +8,7 @@ """ import contextlib -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping from functools import lru_cache from typing import TYPE_CHECKING, Any, Literal, cast, overload @@ -25,6 +25,7 @@ from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable, PandasDataFrame, PolarsDataFrame __all__ = ( + "arrow_batch_to_rows", "arrow_reader_to_return_format", "arrow_reader_with_deferred_close", "arrow_table_column_names", @@ -39,6 +40,7 @@ "arrow_type_from_token", "build_ingest_telemetry", "cast_arrow_table_schema", + "coerce_arrow_batches", "coerce_arrow_table", "convert_dict_to_arrow", "convert_dict_to_arrow_with_schema", @@ -392,13 +394,18 @@ def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable": if has_get_data(source): table = source.get_data() - if _get_arrow_table_coercer().get(table) is _coerce_arrow_table_identity: - return cast("ArrowTable", table) + coercer = _get_arrow_table_coercer().get(table) + if coercer is not None: + return cast("ArrowTable", coercer(table)) msg = "ArrowResult did not return a pyarrow.Table instance" raise TypeError(msg) coercer = _get_arrow_table_coercer().get(source) if coercer is not None: return cast("ArrowTable", coercer(source)) + if hasattr(source, "__arrow_c_stream__"): + import pyarrow as pa + + return pa.RecordBatchReader.from_stream(source).read_all() if isinstance(source, Mapping): import pyarrow as pa @@ -411,6 +418,58 @@ def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable": raise TypeError(msg) +def coerce_arrow_batches( + source: "ArrowResult | Any", *, batch_size: int = 65_536 +) -> "tuple[list[str], Iterator[ArrowRecordBatch]]": + """Coerce an Arrow source into column names and a lazy RecordBatch iterator. + + Streams ``pyarrow.RecordBatchReader`` and ``__arrow_c_stream__`` sources one + batch at a time without materializing the full table into memory. + + Args: + source: ArrowResult, RecordBatchReader, RecordBatch, Table, C stream, or dict/row iterable. + batch_size: Maximum rows per batch when splitting an in-memory Table. + + Returns: + Tuple of ``(column_names, batch_iterator)``. + """ + ensure_pyarrow() + import pyarrow as pa + + unwrapped = source.get_data() if has_get_data(source) else source + if isinstance(unwrapped, pa.RecordBatchReader): + return list(unwrapped.schema.names), iter(unwrapped) + if isinstance(unwrapped, pa.RecordBatch): + return list(unwrapped.schema.names), iter((unwrapped,)) + if isinstance(unwrapped, pa.Table): + return list(unwrapped.schema.names), iter(unwrapped.to_batches(max_chunksize=batch_size)) + if hasattr(unwrapped, "__arrow_c_stream__"): + reader = pa.RecordBatchReader.from_stream(unwrapped) + return list(reader.schema.names), iter(reader) + table = coerce_arrow_table(unwrapped) + return list(table.schema.names), iter(table.to_batches(max_chunksize=batch_size)) + + +def arrow_batch_to_rows(batch: "ArrowRecordBatch", columns: "list[str] | None" = None) -> "list[tuple[Any, ...]]": + """Convert a single Arrow RecordBatch to a list of row tuples. + + Args: + batch: Arrow RecordBatch to extract rows from. + columns: Optional subset or ordering of column names. + + Returns: + List of row tuples for the batch. + """ + ensure_pyarrow() + resolved_columns = columns or list(batch.schema.names) + if not resolved_columns or batch.num_rows == 0: + return [] + col_data = [batch.column(col).to_pylist() for col in resolved_columns] + if not col_data or not col_data[0]: + return [] + return [tuple(row) for row in zip(*col_data, strict=False)] + + def ensure_arrow_table(data: Any) -> "ArrowTable": """Ensure data is a PyArrow Table.""" ensure_pyarrow() @@ -759,6 +818,10 @@ def _coerce_arrow_record_batch(source: Any) -> Any: return pa.Table.from_batches([source]) +def _coerce_arrow_record_batch_reader(source: Any) -> Any: + return source.read_all() + + def _get_arrow_table_coercer() -> "TypeDispatcher[Any]": global _ARROW_TABLE_COERCER if _ARROW_TABLE_COERCER is None: @@ -768,6 +831,7 @@ def _get_arrow_table_coercer() -> "TypeDispatcher[Any]": dispatcher = TypeDispatcher[Any]() dispatcher.register(pa.Table, _coerce_arrow_table_identity) dispatcher.register(pa.RecordBatch, _coerce_arrow_record_batch) + dispatcher.register(pa.RecordBatchReader, _coerce_arrow_record_batch_reader) _ARROW_TABLE_COERCER = dispatcher return _ARROW_TABLE_COERCER diff --git a/tests/unit/driver/test_storage_helpers.py b/tests/unit/driver/test_storage_helpers.py new file mode 100644 index 000000000..32c638405 --- /dev/null +++ b/tests/unit/driver/test_storage_helpers.py @@ -0,0 +1,51 @@ +"""Unit tests for driver storage helper utilities.""" + +import pytest + +from sqlspec.driver._common import CommonDriverAttributesMixin +from sqlspec.driver._storage_helpers import arrow_batch_to_rows, coerce_arrow_batches, ingest_telemetry_from_batches +from sqlspec.typing import PYARROW_INSTALLED + +pytestmark = pytest.mark.skipif(not PYARROW_INSTALLED, reason="pyarrow not installed") + + +def test_ingest_telemetry_from_batches_populates_fields() -> None: + """ingest_telemetry_from_batches builds a StorageTelemetry dictionary with stage metrics.""" + telemetry = ingest_telemetry_from_batches( + 1024, + 65536, + duration_s=1.25, + format_label="parquet", + read_duration_s=0.25, + encode_duration_s=0.35, + write_duration_s=0.50, + commit_duration_s=0.15, + shards_processed=4, + shards_skipped=1, + wal_bytes=4096, + ) + assert telemetry["rows_processed"] == 1024 + assert telemetry["bytes_processed"] == 65536 + assert telemetry["duration_s"] == 1.25 + assert telemetry["format"] == "parquet" + assert telemetry.get("read_duration_s") == 0.25 + assert telemetry.get("encode_duration_s") == 0.35 + assert telemetry.get("write_duration_s") == 0.50 + assert telemetry.get("commit_duration_s") == 0.15 + assert telemetry.get("shards_processed") == 4 + assert telemetry.get("shards_skipped") == 1 + assert telemetry.get("wal_bytes") == 4096 + + +def test_driver_common_exposes_streaming_arrow_batch_helpers() -> None: + """CommonDriverAttributesMixin exposes _coerce_arrow_batches and _arrow_batch_to_rows.""" + import pyarrow as pa + + table = pa.table({"k": [10, 20, 30], "v": ["x", "y", "z"]}) + cols, batches = CommonDriverAttributesMixin._coerce_arrow_batches(table, batch_size=2) + batch_list = list(batches) + assert cols == ["k", "v"] + assert len(batch_list) == 2 + assert CommonDriverAttributesMixin._arrow_batch_to_rows(batch_list[0]) == [(10, "x"), (20, "y")] + assert coerce_arrow_batches is not None + assert arrow_batch_to_rows is not None diff --git a/tests/unit/utils/test_arrow_helpers.py b/tests/unit/utils/test_arrow_helpers.py index cc377beb9..13855c6b7 100644 --- a/tests/unit/utils/test_arrow_helpers.py +++ b/tests/unit/utils/test_arrow_helpers.py @@ -363,3 +363,59 @@ def test_sql_type_for_category_decimal_limits() -> None: sql_type_for_category("decimal", "spanner", precision=50, scale=4) with pytest.raises(SQLBuilderError, match="Unknown SQL type category"): sql_type_for_category("geometry", "postgres") + + +def test_coerce_arrow_batches_streams_record_batch_reader_lazily() -> None: + """coerce_arrow_batches yields RecordBatches lazily from a RecordBatchReader without calling read_all().""" + import pyarrow as pa + + from sqlspec.utils.arrow_helpers import coerce_arrow_batches + + yielded_indices: list[int] = [] + + def _batch_gen() -> Any: + for idx in range(3): + yielded_indices.append(idx) + yield pa.record_batch({"id": pa.array([idx * 2, idx * 2 + 1], type=pa.int64())}) + + schema = pa.schema([pa.field("id", pa.int64())]) + reader = pa.RecordBatchReader.from_batches(schema, _batch_gen()) + + columns, batch_iter = coerce_arrow_batches(reader) + assert columns == ["id"] + assert yielded_indices == [] + + first_batch = next(batch_iter) + assert first_batch.num_rows == 2 + assert yielded_indices == [0] + + remaining = list(batch_iter) + assert len(remaining) == 2 + assert yielded_indices == [0, 1, 2] + + +def test_coerce_arrow_batches_accepts_table_record_batch_and_arrow_result() -> None: + """coerce_arrow_batches handles Table, RecordBatch, ArrowResult, and Arrow C stream sources.""" + import pyarrow as pa + + from sqlspec.core.result import ArrowResult + from sqlspec.core.statement import SQL + from sqlspec.utils.arrow_helpers import arrow_batch_to_rows, coerce_arrow_batches + + table = pa.table({"id": [1, 2, 3, 4, 5], "name": ["a", "b", "c", "d", "e"]}) + cols, batches = coerce_arrow_batches(table, batch_size=2) + batch_list = list(batches) + assert cols == ["id", "name"] + assert [b.num_rows for b in batch_list] == [2, 2, 1] + assert arrow_batch_to_rows(batch_list[0]) == [(1, "a"), (2, "b")] + assert arrow_batch_to_rows(batch_list[0], ["name"]) == [("a",), ("b",)] + + single_batch = batch_list[0] + cols_single, iter_single = coerce_arrow_batches(single_batch) + assert cols_single == ["id", "name"] + assert len(list(iter_single)) == 1 + + arrow_result = ArrowResult(statement=SQL("SELECT 1"), data=table.to_reader(max_chunksize=2)) + cols_res, iter_res = coerce_arrow_batches(arrow_result) + assert cols_res == ["id", "name"] + assert [b.num_rows for b in iter_res] == [2, 2, 1] From 9b22b9d11023039c0b8881c4012763cb0d31ea51 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 17:47:37 +0000 Subject: [PATCH 02/14] feat(postgres): add zero-loop columnar PGCOPY binary encoding and streaming COPY --- sqlspec/adapters/asyncpg/driver.py | 35 +- sqlspec/adapters/cockroach_psycopg/core.py | 3 +- sqlspec/adapters/psqlpy/driver.py | 42 +- sqlspec/adapters/psycopg/core.py | 4 +- sqlspec/adapters/psycopg/driver.py | 63 ++- sqlspec/driver/_common.py | 67 +++ sqlspec/driver/_storage_helpers.py | 117 +--- sqlspec/utils/arrow_helpers.py | 509 +++++++++++++++++- .../adapters/test_postgres_pgcopy_binary.py | 450 ++++++++++++++++ tests/unit/storage/test_bridge.py | 15 + 10 files changed, 1163 insertions(+), 142 deletions(-) create mode 100644 tests/unit/adapters/test_postgres_pgcopy_binary.py diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 307e35e54..81f000fb4 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from contextlib import suppress from io import BytesIO +from itertools import chain from typing import TYPE_CHECKING, Any, cast from sqlglot import exp, parse_one @@ -355,7 +356,6 @@ async def load_from_arrow( ) -> "StorageBridgeJob": """Load Arrow data into a PostgreSQL table via COPY.""" self._require_capability("arrow_import_enabled") - arrow_table = self._coerce_arrow_table(source) table_name, schema_name, quoted_target = self._copy_target(table) if overwrite: try: @@ -363,12 +363,33 @@ async def load_from_arrow( except AsyncpgPostgresError as exc: msg = f"Failed to truncate table '{table}': {exc}" raise SQLSpecError(msg) from exc - columns, records = self._arrow_table_to_rows(arrow_table) - if records: - await self.connection.copy_records_to_table( - table_name, records=records, columns=columns, schema_name=schema_name - ) - telemetry_payload = self._ingest_telemetry(arrow_table) + columns, batch_iter = self._coerce_arrow_batches(source) + first_batch = next((b for b in batch_iter if b.num_rows > 0), None) + metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} + if first_batch is not None: + all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + if use_binary: + async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) + await self.connection.copy_to_table( + table_name, + source=async_chunks, + columns=columns, + schema_name=schema_name, + format="binary", + ) + else: + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if records: + await self.connection.copy_records_to_table( + table_name, records=records, columns=columns, schema_name=schema_name + ) + telemetry_payload = self._ingest_telemetry_from_batches( + metrics["rows_processed"], metrics["bytes_processed"] + ) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) diff --git a/sqlspec/adapters/cockroach_psycopg/core.py b/sqlspec/adapters/cockroach_psycopg/core.py index 89f7b4ec4..d1a7ea0f3 100644 --- a/sqlspec/adapters/cockroach_psycopg/core.py +++ b/sqlspec/adapters/cockroach_psycopg/core.py @@ -9,7 +9,7 @@ from sqlglot.tokenizer_core import TokenType from sqlspec.adapters.psycopg.core import apply_driver_features as psycopg_apply_driver_features -from sqlspec.adapters.psycopg.core import build_statement_config, driver_profile +from sqlspec.adapters.psycopg.core import build_copy_from_command, build_statement_config, driver_profile from sqlspec.exceptions import ImproperConfigurationError, SerializationConflictError, SQLSpecError from sqlspec.utils.config_tools import normalize_connection_config from sqlspec.utils.text import quote_identifier, split_qualified_identifier @@ -26,6 +26,7 @@ "CockroachPsycopgRetryConfig", "apply_driver_features", "build_connection_config", + "build_copy_from_command", "build_native_export", "build_native_import", "build_statement_config", diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 5df67ce27..1f70b7071 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -4,6 +4,7 @@ and transaction management. """ +from itertools import chain from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.psqlpy._typing import ( @@ -343,7 +344,6 @@ async def load_from_arrow( """Load Arrow-formatted data into PostgreSQL via psqlpy binary COPY.""" self._require_capability("arrow_import_enabled") - arrow_table = self._coerce_arrow_table(source) if overwrite: qualified = format_table_identifier(table) exc_handler = self.handle_database_exceptions() @@ -352,25 +352,45 @@ async def load_from_arrow( if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - columns, records = self._arrow_table_to_rows(arrow_table) - if records: + columns, batch_iter = self._coerce_arrow_batches(source) + first_batch = next((b for b in batch_iter if b.num_rows > 0), None) + metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} + if first_batch is not None: + all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) schema_name, table_name = split_schema_and_table(table) + use_binary_wire = bool(self.driver_features.get("enable_pgcopy_binary", False)) exc_handler = self.handle_database_exceptions() async with exc_handler, self.with_cursor(self.connection) as cursor: copy_kwargs: dict[str, Any] = {"columns": columns} if schema_name: copy_kwargs["schema_name"] = schema_name - json_columns = await self._resolve_json_columns(schema_name, table_name) - deserializer = cast("Callable[[str], Any]", self.driver_features.get("json_deserializer", from_json)) - await cursor.copy_records_to_table( - table_name, - coerce_json_columns(records, columns, json_columns, deserializer=deserializer), - **copy_kwargs, - ) + if use_binary_wire and hasattr(cursor, "binary_copy_to_table"): + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + payload = self._encode_arrow_batch_pg_binary( + batch, columns=columns, include_header=True, include_trailer=True + ) + await cursor.binary_copy_to_table(payload, table_name, **copy_kwargs) + else: + json_columns = await self._resolve_json_columns(schema_name, table_name) + deserializer = cast("Callable[[str], Any]", self.driver_features.get("json_deserializer", from_json)) + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if records: + await cursor.copy_records_to_table( + table_name, + coerce_json_columns(records, columns, json_columns, deserializer=deserializer), + **copy_kwargs, + ) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry(arrow_table) + telemetry_payload = self._ingest_telemetry_from_batches( + metrics["rows_processed"], metrics["bytes_processed"] + ) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 4459e54b0..2de7d9a12 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -122,9 +122,11 @@ def pipeline_supported() -> bool: return False -def build_copy_from_command(table: str, columns: "list[str]") -> "PsycopgComposed": +def build_copy_from_command(table: str, columns: "list[str]", *, binary: bool = False) -> "PsycopgComposed": table_identifier = _compose_table_identifier(table) column_sql = PsycopgSQL(", ").join([PsycopgIdentifier(column) for column in columns]) + if binary: + return PsycopgSQL("COPY {} ({}) FROM STDIN WITH (FORMAT BINARY)").format(table_identifier, column_sql) return PsycopgSQL("COPY {} ({}) FROM STDIN").format(table_identifier, column_sql) diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index a400ac40a..5e1f6eebf 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -2,6 +2,7 @@ from collections.abc import Sized from contextlib import AsyncExitStack, ExitStack +from itertools import chain from typing import TYPE_CHECKING, Any, LiteralString, cast from sqlspec.adapters.psycopg._typing import ( @@ -522,7 +523,6 @@ def load_from_arrow( """Load Arrow data into PostgreSQL using COPY.""" self._require_capability("arrow_import_enabled") - arrow_table = self._coerce_arrow_table(source) if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() @@ -530,23 +530,33 @@ def load_from_arrow( cursor.execute(truncate_sql) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - if arrow_table.num_rows > 0: - columns, records = self._arrow_table_to_rows(arrow_table) - if self._arrow_rows_need_preparation(arrow_table): - records = cast( - "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) - ) - copy_sql = build_copy_from_command(table, columns) + columns, batch_iter = self._coerce_arrow_batches(source) + first_batch = next((b for b in batch_iter if b.num_rows > 0), None) + metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} + if first_batch is not None: + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + copy_sql = build_copy_from_command(table, columns, binary=use_binary) + all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() with ExitStack() as stack: stack.enter_context(exc_handler) cursor = stack.enter_context(self.with_cursor(self.connection)) copy_ctx = stack.enter_context(cursor.copy(copy_sql)) - for record in records: - copy_ctx.write_row(record) + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): + copy_ctx.write(chunk) + else: + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + for record in records: + copy_ctx.write_row(record) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry(arrow_table) + telemetry_payload = self._ingest_telemetry_from_batches( + metrics["rows_processed"], metrics["bytes_processed"] + ) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) @@ -1054,7 +1064,6 @@ async def load_from_arrow( """Load Arrow data into PostgreSQL asynchronously via COPY.""" self._require_capability("arrow_import_enabled") - arrow_table = self._coerce_arrow_table(source) if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() @@ -1062,23 +1071,33 @@ async def load_from_arrow( await cursor.execute(truncate_sql) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - if arrow_table.num_rows > 0: - columns, records = self._arrow_table_to_rows(arrow_table) - if self._arrow_rows_need_preparation(arrow_table): - records = cast( - "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) - ) - copy_sql = build_copy_from_command(table, columns) + columns, batch_iter = self._coerce_arrow_batches(source) + first_batch = next((b for b in batch_iter if b.num_rows > 0), None) + metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} + if first_batch is not None: + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + copy_sql = build_copy_from_command(table, columns, binary=use_binary) + all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() async with AsyncExitStack() as stack: await stack.enter_async_context(exc_handler) cursor = await stack.enter_async_context(self.with_cursor(self.connection)) copy_ctx = await stack.enter_async_context(cursor.copy(copy_sql)) - for record in records: - await copy_ctx.write_row(record) + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): + await copy_ctx.write(chunk) + else: + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + for record in records: + await copy_ctx.write_row(record) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry(arrow_table) + telemetry_payload = self._ingest_telemetry_from_batches( + metrics["rows_processed"], metrics["bytes_processed"] + ) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index c039afb8c..1f534b770 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -85,7 +85,10 @@ coerce_arrow_batches, coerce_arrow_table, create_storage_job, + encode_arrow_batch_pg_binary, ingest_telemetry_from_batches, + iter_pgcopy_binary_chunks, + iter_pgcopy_binary_chunks_async, records_to_arrow_table, ) from sqlspec.exceptions import ( @@ -3148,6 +3151,70 @@ def _arrow_batch_to_rows(batch: "ArrowRecordBatch", columns: "list[str] | None" """Convert a single Arrow RecordBatch into row tuples.""" return arrow_batch_to_rows(batch, columns) + @staticmethod + def _encode_arrow_batch_pg_binary( + batch: "ArrowRecordBatch", + *, + columns: "abc.Sequence[str] | None" = None, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + include_header: bool = False, + include_trailer: bool = False, + ) -> bytes: + """Encode an Arrow RecordBatch into PostgreSQL COPY BINARY format.""" + return encode_arrow_batch_pg_binary( + batch, + columns=columns, + vector_columns=vector_columns, + jsonb_columns=jsonb_columns, + include_header=include_header, + include_trailer=include_trailer, + ) + + @staticmethod + def _iter_pgcopy_binary_chunks( + source: "ArrowResult | Any", + *, + columns: "abc.Sequence[str] | None" = None, + batch_size: int = 65_536, + chunk_bytes: int = 8_388_608, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + metrics: "dict[str, int] | None" = None, + ) -> "abc.Iterator[bytes]": + """Yield PostgreSQL COPY BINARY chunks from an Arrow source in bounded memory.""" + return iter_pgcopy_binary_chunks( + source, + columns=columns, + batch_size=batch_size, + chunk_bytes=chunk_bytes, + vector_columns=vector_columns, + jsonb_columns=jsonb_columns, + metrics=metrics, + ) + + @staticmethod + def _iter_pgcopy_binary_chunks_async( + source: "ArrowResult | Any", + *, + columns: "abc.Sequence[str] | None" = None, + batch_size: int = 65_536, + chunk_bytes: int = 8_388_608, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + metrics: "dict[str, int] | None" = None, + ) -> Any: + """Return an async iterator of PostgreSQL COPY BINARY chunks from an Arrow source.""" + return iter_pgcopy_binary_chunks_async( + source, + columns=columns, + batch_size=batch_size, + chunk_bytes=chunk_bytes, + vector_columns=vector_columns, + jsonb_columns=jsonb_columns, + metrics=metrics, + ) + @staticmethod def _arrow_table_to_rows( table: "ArrowTable", columns: "list[str] | None" = None diff --git a/sqlspec/driver/_storage_helpers.py b/sqlspec/driver/_storage_helpers.py index e9ae08a7d..179d86f0b 100644 --- a/sqlspec/driver/_storage_helpers.py +++ b/sqlspec/driver/_storage_helpers.py @@ -1,31 +1,32 @@ -"""Pure helper functions for storage operations. - -These functions are extracted from StorageDriverMixin to eliminate -cross-trait attribute access that causes mypyc segmentation faults. -""" +"""Pure helper functions for storage operations.""" 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_batch_to_rows as _arrow_batch_to_rows_impl -from sqlspec.utils.arrow_helpers import arrow_table_needs_parameter_preparation as _arrow_rows_need_preparation -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 _ingest_telemetry_impl -from sqlspec.utils.arrow_helpers import coerce_arrow_batches as _coerce_arrow_batches_impl -from sqlspec.utils.arrow_helpers import coerce_arrow_table as _coerce_arrow_table_impl -from sqlspec.utils.arrow_helpers import records_to_arrow_table as _records_to_arrow_table_impl +from sqlspec.utils.arrow_helpers import ( + PG_BINARY_COPY_HEADER, + PG_BINARY_COPY_TRAILER, + arrow_batch_to_rows, + arrow_table_needs_parameter_preparation, + arrow_table_to_rows, + build_ingest_telemetry, + coerce_arrow_batches, + coerce_arrow_table, + encode_arrow_batch_pg_binary, + iter_pgcopy_binary_chunks, + iter_pgcopy_binary_chunks_async, + records_to_arrow_table, +) if TYPE_CHECKING: - from collections.abc import Iterable, Iterator, Mapping - - from sqlspec.core.result import ArrowResult from sqlspec.storage import StorageDestination - from sqlspec.typing import ArrowRecordBatch, ArrowTable __all__ = ( "CAPABILITY_HINTS", + "PG_BINARY_COPY_HEADER", + "PG_BINARY_COPY_TRAILER", "arrow_batch_to_rows", "arrow_table_needs_parameter_preparation", "arrow_table_to_rows", @@ -34,7 +35,10 @@ "coerce_arrow_batches", "coerce_arrow_table", "create_storage_job", + "encode_arrow_batch_pg_binary", "ingest_telemetry_from_batches", + "iter_pgcopy_binary_chunks", + "iter_pgcopy_binary_chunks_async", "records_to_arrow_table", "stringify_storage_target", ) @@ -64,80 +68,6 @@ def stringify_storage_target(target: "StorageDestination | None") -> str | None: 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. - """ - return _coerce_arrow_table_impl(source) - - -def coerce_arrow_batches( - source: "ArrowResult | Any", *, batch_size: int = 65_536 -) -> "tuple[list[str], Iterator[ArrowRecordBatch]]": - """Coerce an Arrow source into column names and a lazy RecordBatch iterator. - - Args: - source: ArrowResult, RecordBatchReader, RecordBatch, Table, or row iterable. - batch_size: Maximum rows per batch when chunking an in-memory Table. - - Returns: - Tuple of ``(column_names, batch_iterator)``. - """ - return _coerce_arrow_batches_impl(source, batch_size=batch_size) - - -def arrow_batch_to_rows(batch: "ArrowRecordBatch", columns: "list[str] | None" = None) -> "list[tuple[Any, ...]]": - """Convert a single Arrow RecordBatch into row tuples. - - Args: - batch: Arrow RecordBatch to convert. - columns: Optional subset or order of columns to extract. - - Returns: - List of row tuples. - """ - return _arrow_batch_to_rows_impl(batch, columns) - - -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). - """ - return _arrow_table_to_rows_impl(table, columns) - - -def arrow_table_needs_parameter_preparation(table: "ArrowTable") -> bool: - """Return whether Arrow rows may contain nested values needing preparation.""" - return _arrow_rows_need_preparation(table) - - -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 = _ingest_telemetry_impl(table, format_label=format_label) - return cast("StorageTelemetry", telemetry) - - def ingest_telemetry_from_batches( rows_processed: int, bytes_processed: int, @@ -194,13 +124,6 @@ def ingest_telemetry_from_batches( return telemetry -def records_to_arrow_table( - records: "Iterable[Mapping[str, Any]] | Iterable[Iterable[Any]]", columns: "list[str] | None" -) -> "ArrowTable": - """Normalize mapping or positional records into a PyArrow table for ingest.""" - return _records_to_arrow_table_impl(records, columns) - - def attach_partition_telemetry(telemetry: "StorageTelemetry", partitioner: "dict[str, object] | None") -> None: """Attach partitioner info to telemetry dict (mutates in place). diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 3c706c2b9..ec23c9d5c 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -8,23 +8,28 @@ """ import contextlib -from collections.abc import Callable, Iterable, Iterator, Mapping +import struct +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import lru_cache -from typing import TYPE_CHECKING, Any, Literal, cast, overload +from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload from sqlspec.data_dictionary._registry import get_dialect_config from sqlspec.exceptions import ImproperConfigurationError, SQLBuilderError from sqlspec.utils.dispatch import TypeDispatcher from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow +from sqlspec.utils.serializers import to_json from sqlspec.utils.type_guards import has_arrow_table_stats, has_get_data from sqlspec.utils.uuids import uuid_from_bytes if TYPE_CHECKING: from sqlspec.core.result import ArrowResult from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.storage import StorageTelemetry from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable, PandasDataFrame, PolarsDataFrame __all__ = ( + "PG_BINARY_COPY_HEADER", + "PG_BINARY_COPY_TRAILER", "arrow_batch_to_rows", "arrow_reader_to_return_format", "arrow_reader_with_deferred_close", @@ -44,11 +49,19 @@ "coerce_arrow_table", "convert_dict_to_arrow", "convert_dict_to_arrow_with_schema", + "encode_arrow_batch_pg_binary", "ensure_arrow_table", + "iter_pgcopy_binary_chunks", + "iter_pgcopy_binary_chunks_async", "records_to_arrow_table", "sql_type_for_category", "sql_type_from_arrow_type", ) +PG_BINARY_COPY_HEADER: Final[bytes] = b"PGCOPY\n\xff\r\n\x00\x00\x00\x00\x00\x00\x00\x00\x00" +PG_BINARY_COPY_TRAILER: Final[bytes] = b"\xff\xff" +_PG_DATE_EPOCH_DAYS: Final[int] = 10957 +_PG_EPOCH_OFFSET_US: Final[int] = 946_684_800_000_000 +_PG_NULL_MARKER: Final[bytes] = b"\xff\xff\xff\xff" _ARROW_TABLE_COERCER: "TypeDispatcher[Any] | None" = None _ARROW_SCHEMA_DECISION_CACHE_SIZE = 512 _CONFIG_CATEGORY_DEFAULTS: "dict[str, str]" = { @@ -610,7 +623,7 @@ def arrow_table_to_polars(table: "ArrowTable") -> "PolarsDataFrame": return result -def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "dict[str, int | str]": +def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "StorageTelemetry": """Build telemetry dict from Arrow table statistics.""" if has_arrow_table_stats(table): rows = int(table.num_rows) @@ -660,6 +673,478 @@ def arrow_reader_to_return_format( return shaped, int(table.num_rows) +def encode_arrow_batch_pg_binary( + batch: "ArrowRecordBatch", + *, + columns: "Sequence[str] | None" = None, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + include_header: bool = False, + include_trailer: bool = False, +) -> bytes: + """Encode an Arrow RecordBatch into PostgreSQL COPY BINARY format. + + Uses zero-loop structured NumPy dtype packing when all columns have fixed + wire widths and no nulls, falling back to direct Arrow buffer slicing for + variable-width and nullable columns. + + Args: + batch: Arrow RecordBatch to encode. + columns: Optional subset or order of columns to encode. + vector_columns: Column names using pgvector binary wire format. + jsonb_columns: Column names requiring PostgreSQL jsonb version byte prefix. + include_header: Whether to prepend ``PG_BINARY_COPY_HEADER``. + include_trailer: Whether to append ``PG_BINARY_COPY_TRAILER``. + + Returns: + PostgreSQL ``COPY ... WITH (FORMAT BINARY)`` byte payload. + """ + ensure_pyarrow() + resolved_columns = list(columns) if columns is not None else list(batch.schema.names) + body = b"" + if resolved_columns and batch.num_rows > 0: + fast_body = _encode_pg_binary_fast_path_1(batch, resolved_columns, jsonb_columns) + if fast_body is not None: + body = fast_body + else: + body = _encode_pg_binary_fast_path_2(batch, resolved_columns, vector_columns, jsonb_columns) + parts: list[bytes] = [] + if include_header: + parts.append(PG_BINARY_COPY_HEADER) + if body: + parts.append(body) + if include_trailer: + parts.append(PG_BINARY_COPY_TRAILER) + return b"".join(parts) + + +def iter_pgcopy_binary_chunks( + source: "ArrowResult | Any", + *, + columns: "Sequence[str] | None" = None, + batch_size: int = 65_536, + chunk_bytes: int = 8_388_608, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + metrics: "dict[str, int] | None" = None, +) -> "Iterator[bytes]": + """Yield PostgreSQL COPY BINARY chunks from an Arrow source in bounded memory. + + Args: + source: Arrow Table, RecordBatch, RecordBatchReader, ArrowResult, or RecordBatch iterator. + columns: Optional explicit column order. + batch_size: Maximum rows per batch when splitting a Table. + chunk_bytes: Target byte size for yielded chunks (default 8 MiB). + vector_columns: Column names using pgvector binary wire format. + jsonb_columns: Column names requiring PostgreSQL jsonb version byte prefix. + metrics: Optional mutable dict updated in-place with ``rows_processed`` and ``bytes_processed``. + + Yields: + Byte chunks forming a complete ``PGCOPY`` binary stream. + """ + ensure_pyarrow() + import pyarrow as pa + + if columns is not None and isinstance(source, Iterator) and not isinstance(source, pa.RecordBatchReader): + resolved_columns = list(columns) + batches: Iterator[ArrowRecordBatch] = source + else: + inferred_columns, batches = coerce_arrow_batches(source, batch_size=batch_size) + resolved_columns = list(columns) if columns is not None else inferred_columns + + buffer = bytearray(PG_BINARY_COPY_HEADER) + threshold = max(1, int(chunk_bytes)) + total_rows = 0 + total_bytes = 0 + for batch in batches: + if batch.num_rows == 0: + continue + total_rows += int(batch.num_rows) + total_bytes += int(batch.nbytes) + encoded = encode_arrow_batch_pg_binary( + batch, + columns=resolved_columns, + vector_columns=vector_columns, + jsonb_columns=jsonb_columns, + include_header=False, + include_trailer=False, + ) + buffer.extend(encoded) + if len(buffer) >= threshold: + yield bytes(buffer) + buffer.clear() + if metrics is not None: + metrics["rows_processed"] = total_rows + metrics["bytes_processed"] = total_bytes + buffer.extend(PG_BINARY_COPY_TRAILER) + yield bytes(buffer) + + +def iter_pgcopy_binary_chunks_async( + source: "ArrowResult | Any", + *, + columns: "Sequence[str] | None" = None, + batch_size: int = 65_536, + chunk_bytes: int = 8_388_608, + vector_columns: frozenset[str] = frozenset(), + jsonb_columns: frozenset[str] = frozenset(), + metrics: "dict[str, int] | None" = None, +) -> Any: + """Return an async byte-chunk iterator of PostgreSQL COPY BINARY payloads.""" + return _AsyncBytesIterator( + iter_pgcopy_binary_chunks( + source, + columns=columns, + batch_size=batch_size, + chunk_bytes=chunk_bytes, + vector_columns=vector_columns, + jsonb_columns=jsonb_columns, + metrics=metrics, + ) + ) + + +def _extract_pg_fixed_column_array(col: Any, dtype: Any, num_rows: int) -> "tuple[int, str, Any] | None": + """Return ``(byte_width, numpy_dtype_str, numpy_array)`` for fixed-width scalar columns.""" + import numpy as np + import pyarrow as pa + + is_extension = getattr(pa.types, "is_extension", None) + if is_extension is not None and is_extension(dtype): + dtype = dtype.storage_type + col = col.storage if hasattr(col, "storage") else col + + if pa.types.is_boolean(dtype): + filled = col.fill_null(False) if col.null_count > 0 else col + arr = filled.to_numpy(zero_copy_only=False).astype("|u1", copy=False) + return 1, "|u1", arr + if pa.types.is_int8(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.int8, count=col.offset + num_rows)[col.offset :] + return 2, ">i2", raw.astype(">i2") + if pa.types.is_uint8(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.uint8, count=col.offset + num_rows)[col.offset :] + return 2, ">i2", raw.astype(">i2") + if pa.types.is_int16(dtype) or pa.types.is_uint16(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.int16, count=col.offset + num_rows)[col.offset :] + return 2, ">i2", raw.astype(">i2", copy=False) + if pa.types.is_int32(dtype) or pa.types.is_uint32(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.int32, count=col.offset + num_rows)[col.offset :] + return 4, ">i4", raw.astype(">i4", copy=False) + if pa.types.is_int64(dtype) or pa.types.is_uint64(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.int64, count=col.offset + num_rows)[col.offset :] + return 8, ">i8", raw.astype(">i8", copy=False) + if pa.types.is_float16(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.float16, count=col.offset + num_rows)[col.offset :] + return 4, ">f4", raw.astype(">f4") + if pa.types.is_float32(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.float32, count=col.offset + num_rows)[col.offset :] + return 4, ">f4", raw.astype(">f4", copy=False) + if pa.types.is_float64(dtype): + raw = np.frombuffer(col.buffers()[1], dtype=np.float64, count=col.offset + num_rows)[col.offset :] + return 8, ">f8", raw.astype(">f8", copy=False) + if pa.types.is_date32(dtype): + buf = col.buffers()[1] + raw = np.frombuffer(buf, dtype=np.int32, count=col.offset + num_rows)[col.offset :] + return 4, ">i4", (raw - _PG_DATE_EPOCH_DAYS).astype(">i4") + if pa.types.is_date64(dtype): + buf = col.buffers()[1] + raw = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] + return 4, ">i4", ((raw // 86_400_000) - _PG_DATE_EPOCH_DAYS).astype(">i4") + if pa.types.is_timestamp(dtype): + buf = col.buffers()[1] + raw = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] + unit = dtype.unit + if unit == "us": + shifted = raw - _PG_EPOCH_OFFSET_US + elif unit == "ms": + shifted = raw * 1_000 - _PG_EPOCH_OFFSET_US + elif unit == "s": + shifted = raw * 1_000_000 - _PG_EPOCH_OFFSET_US + else: + shifted = (raw // 1_000) - _PG_EPOCH_OFFSET_US + return 8, ">i8", shifted.astype(">i8") + if pa.types.is_fixed_size_binary(dtype): + width = int(dtype.byte_width) + buf = col.buffers()[1] + arr = np.frombuffer(buf, dtype=f"|V{width}", count=col.offset + num_rows)[col.offset :] + return width, f"|V{width}", arr + return None + + +def _encode_pg_binary_fast_path_1( + batch: "ArrowRecordBatch", columns: "list[str]", jsonb_columns: frozenset[str] +) -> bytes | None: + """Pack non-null fixed-width and FixedSizeList float vector columns via structured NumPy dtype.""" + import numpy as np + import pyarrow as pa + + num_rows = int(batch.num_rows) + dtype_fields: list[Any] = [("ncols", ">i2")] + assignments: list[tuple[str, Any]] = [("ncols", len(columns))] + + for idx, col_name in enumerate(columns): + if col_name in jsonb_columns: + return None + col = batch.column(col_name) + if col.null_count > 0: + return None + if isinstance(col, pa.ChunkedArray): + col = col.combine_chunks() + field_type = col.type + fixed_info = _extract_pg_fixed_column_array(col, field_type, num_rows) + if fixed_info is not None: + byte_width, np_dtype, arr = fixed_info + len_key = f"c{idx}_len" + val_key = f"c{idx}_val" + dtype_fields.append((len_key, ">i4")) + dtype_fields.append((val_key, np_dtype)) + assignments.append((len_key, byte_width)) + assignments.append((val_key, arr)) + continue + + if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): + dim = int(field_type.list_size) + flat_values = col.values[col.offset * dim : (col.offset + num_rows) * dim] + if flat_values.null_count > 0: + return None + is_f16 = pa.types.is_float16(field_type.value_type) + elem_dtype = ">f2" if is_f16 else ">f4" + elem_size = 2 if is_f16 else 4 + arr_2d = flat_values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).reshape(num_rows, dim) + len_key = f"c{idx}_len" + dim_key = f"c{idx}_dim" + unused_key = f"c{idx}_unused" + val_key = f"c{idx}_val" + dtype_fields.extend([ + (len_key, ">i4"), + (dim_key, ">u2"), + (unused_key, ">u2"), + (val_key, elem_dtype, (dim,)), + ]) + assignments.extend([ + (len_key, 4 + elem_size * dim), + (dim_key, dim), + (unused_key, 0), + (val_key, arr_2d), + ]) + continue + + return None + + row_dtype = np.dtype(dtype_fields) + out = np.empty(num_rows, dtype=row_dtype) + for field_key, value in assignments: + out[field_key] = value + return out.tobytes() + + +def _encode_pg_binary_fast_path_2( + batch: "ArrowRecordBatch", + columns: "list[str]", + vector_columns: frozenset[str], + jsonb_columns: frozenset[str], +) -> bytes: + """Encode variable-width and nullable columns by slicing Arrow buffers directly.""" + import numpy as np + import pyarrow as pa + + num_rows = int(batch.num_rows) + row_header = struct.pack(">h", len(columns)) + col_encoders: list[Callable[[int, bytearray], None]] = [] + + for col_name in columns: + col = batch.column(col_name) + if isinstance(col, pa.ChunkedArray): + col = col.combine_chunks() + field_type = col.type + has_nulls = col.null_count > 0 + null_mask = col.is_null().to_numpy(zero_copy_only=False) if has_nulls else None + + if col_name not in jsonb_columns: + fixed_info = _extract_pg_fixed_column_array(col, field_type, num_rows) + if fixed_info is not None: + byte_width, np_dtype, arr = fixed_info + packed = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("val", np_dtype)])) + packed["len"] = byte_width + packed["val"] = arr + packed_mv = memoryview(packed.tobytes()) + stride = 4 + byte_width + + def _emit_fixed( + row_idx: int, + out_buf: bytearray, + _mv: memoryview = packed_mv, + _stride: int = stride, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + else: + start = row_idx * _stride + out_buf.extend(_mv[start : start + _stride]) + + col_encoders.append(_emit_fixed) + continue + + if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): + dim = int(field_type.list_size) + flat_values = col.values[col.offset * dim : (col.offset + num_rows) * dim] + is_f16 = pa.types.is_float16(field_type.value_type) + elem_dtype = ">f2" if is_f16 else ">f4" + elem_size = 2 if is_f16 else 4 + arr_2d = ( + flat_values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).reshape(num_rows, dim) + ) + packed = np.empty( + num_rows, + dtype=np.dtype([ + ("len", ">i4"), + ("dim", ">u2"), + ("unused", ">u2"), + ("val", elem_dtype, (dim,)), + ]), + ) + packed["len"] = 4 + elem_size * dim + packed["dim"] = dim + packed["unused"] = 0 + packed["val"] = arr_2d + packed_mv = memoryview(packed.tobytes()) + stride = 8 + elem_size * dim + + def _emit_fixed_vec( + row_idx: int, + out_buf: bytearray, + _mv: memoryview = packed_mv, + _stride: int = stride, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + else: + start = row_idx * _stride + out_buf.extend(_mv[start : start + _stride]) + + col_encoders.append(_emit_fixed_vec) + continue + + if ( + pa.types.is_string(field_type) + or pa.types.is_large_string(field_type) + or pa.types.is_binary(field_type) + or pa.types.is_large_binary(field_type) + ): + offset_dtype = ( + np.int64 if (pa.types.is_large_string(field_type) or pa.types.is_large_binary(field_type)) else np.int32 + ) + offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[ + col.offset : + ] + data_buf = col.buffers()[2] + data_mv = memoryview(data_buf) if data_buf is not None else memoryview(b"") + is_jsonb = col_name in jsonb_columns + diffs = np.diff(offsets) + lengths_mv = memoryview((diffs + (1 if is_jsonb else 0)).astype(">i4").tobytes()) + + def _emit_varlen( + row_idx: int, + out_buf: bytearray, + _offsets: Any = offsets, + _data_mv: memoryview = data_mv, + _len_mv: memoryview = lengths_mv, + _nulls: Any = null_mask, + _is_jsonb: bool = is_jsonb, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + return + out_buf.extend(_len_mv[row_idx * 4 : (row_idx + 1) * 4]) + if _is_jsonb: + out_buf.append(1) + start = int(_offsets[row_idx]) + end = int(_offsets[row_idx + 1]) + if end > start: + out_buf.extend(_data_mv[start:end]) + + col_encoders.append(_emit_varlen) + continue + + if (pa.types.is_list(field_type) or pa.types.is_large_list(field_type)) and ( + col_name in vector_columns or pa.types.is_floating(field_type.value_type) + ): + offset_dtype = np.int64 if pa.types.is_large_list(field_type) else np.int32 + offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[ + col.offset : + ] + is_f16 = pa.types.is_float16(field_type.value_type) + elem_dtype = ">f2" if is_f16 else ">f4" + elem_size = 2 if is_f16 else 4 + flat_bytes = col.values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).tobytes() + flat_mv = memoryview(flat_bytes) + dims = np.diff(offsets) + headers = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("dim", ">u2"), ("unused", ">u2")])) + headers["len"] = 4 + elem_size * dims + headers["dim"] = dims + headers["unused"] = 0 + headers_mv = memoryview(headers.tobytes()) + + def _emit_var_vec( + row_idx: int, + out_buf: bytearray, + _offsets: Any = offsets, + _flat_mv: memoryview = flat_mv, + _hdr_mv: memoryview = headers_mv, + _elem_size: int = elem_size, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + return + out_buf.extend(_hdr_mv[row_idx * 8 : (row_idx + 1) * 8]) + start = int(_offsets[row_idx]) * _elem_size + end = int(_offsets[row_idx + 1]) * _elem_size + if end > start: + out_buf.extend(_flat_mv[start:end]) + + col_encoders.append(_emit_var_vec) + continue + + py_values = col.to_pylist() + is_jsonb = col_name in jsonb_columns + + def _emit_fallback( + row_idx: int, + out_buf: bytearray, + _vals: list[Any] = py_values, + _is_jsonb: bool = is_jsonb, + ) -> None: + val = _vals[row_idx] + if val is None: + out_buf.extend(_PG_NULL_MARKER) + return + if isinstance(val, (bytes, bytearray, memoryview)): + raw = bytes(val) + elif isinstance(val, str): + raw = val.encode("utf-8") + elif isinstance(val, (dict, list, tuple)): + raw = to_json(val).encode("utf-8") + else: + raw = str(val).encode("utf-8") + if _is_jsonb: + out_buf.extend(struct.pack(">iB", len(raw) + 1, 1)) + else: + out_buf.extend(struct.pack(">i", len(raw))) + out_buf.extend(raw) + + col_encoders.append(_emit_fallback) + + out_buf = bytearray() + for row_idx in range(num_rows): + out_buf.extend(row_header) + for encoder in col_encoders: + encoder(row_idx, out_buf) + return bytes(out_buf) + + @lru_cache(maxsize=_ARROW_SCHEMA_DECISION_CACHE_SIZE) def _arrow_schema_has_opaque_uuid(schema: Any) -> bool: return any( @@ -899,3 +1384,21 @@ def _finalize(self) -> None: self._closed = True with contextlib.suppress(Exception): self._close_callback() + + +class _AsyncBytesIterator: + """Wrap a synchronous byte-chunk iterator as an async iterator.""" + + __slots__ = ("_iterator",) + + def __init__(self, iterator: "Iterator[bytes]") -> None: + self._iterator = iterator + + def __aiter__(self) -> "_AsyncBytesIterator": + return self + + async def __anext__(self) -> bytes: + try: + return next(self._iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc diff --git a/tests/unit/adapters/test_postgres_pgcopy_binary.py b/tests/unit/adapters/test_postgres_pgcopy_binary.py new file mode 100644 index 000000000..54f04b1c5 --- /dev/null +++ b/tests/unit/adapters/test_postgres_pgcopy_binary.py @@ -0,0 +1,450 @@ +# pyright: reportPrivateUsage=false +"""Unit tests for vectorized PGCOPY binary encoding and streaming COPY across PostgreSQL & CockroachDB drivers.""" + +import datetime as dt +import struct +import uuid +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pyarrow as pa +import pytest + +from sqlspec.adapters.asyncpg import AsyncpgDriver, default_statement_config as asyncpg_statement_config +from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgDriver +from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver +from sqlspec.adapters.cockroach_psycopg.core import build_copy_from_command as crdb_build_copy_from_command +from sqlspec.adapters.psqlpy import PsqlpyDriver, default_statement_config as psqlpy_statement_config +from sqlspec.adapters.psycopg import PsycopgAsyncDriver, PsycopgSyncDriver, default_statement_config as psycopg_statement_config +from sqlspec.adapters.psycopg.core import build_copy_from_command as pg_build_copy_from_command +from sqlspec.utils.arrow_helpers import ( + PG_BINARY_COPY_HEADER, + PG_BINARY_COPY_TRAILER, + encode_arrow_batch_pg_binary, + iter_pgcopy_binary_chunks, +) + +_CAPS: dict[str, Any] = { + "arrow_export_enabled": True, + "arrow_import_enabled": True, + "parquet_export_enabled": True, + "parquet_import_enabled": True, + "partition_strategies": ["fixed"], +} + + +def test_pg_binary_copy_constants() -> None: + """PGCOPY binary header and trailer match the PostgreSQL wire specification.""" + assert len(PG_BINARY_COPY_HEADER) == 19 + assert PG_BINARY_COPY_HEADER[:11] == b"PGCOPY\n\xff\r\n\x00" + assert PG_BINARY_COPY_HEADER[11:15] == b"\x00\x00\x00\x00" + assert PG_BINARY_COPY_HEADER[15:19] == b"\x00\x00\x00\x00" + assert PG_BINARY_COPY_TRAILER == b"\xff\xff" + + +def test_encode_arrow_batch_pg_binary_fast_path_1_fixed_width_and_pgvector() -> None: + """Fast Path 1 packs non-null fixed-width and FixedSizeList vector columns via structured NumPy dtype.""" + uid = uuid.UUID("12345678-1234-5678-1234-567812345678") + batch = pa.RecordBatch.from_arrays( + [ + pa.array([True, False], type=pa.bool_()), + pa.array([10, -20], type=pa.int16()), + pa.array([1000, -2000], type=pa.int32()), + pa.array([100_000, -200_000], type=pa.int64()), + pa.array([1.5, -2.5], type=pa.float32()), + pa.array([3.125, -6.25], type=pa.float64()), + pa.array([dt.date(2000, 1, 1), dt.date(2000, 1, 3)], type=pa.date32()), + pa.array( + [ + dt.datetime(2000, 1, 1, 0, 0, 0, tzinfo=dt.timezone.utc), + dt.datetime(2000, 1, 1, 0, 0, 2, tzinfo=dt.timezone.utc), + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array([uid.bytes, uid.bytes], type=pa.binary(16)), + pa.FixedSizeListArray.from_arrays(pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], type=pa.float32()), 3), + pa.FixedSizeListArray.from_arrays( + pa.array([np.float16(0.5), np.float16(1.5), np.float16(2.5), np.float16(3.5)], type=pa.float16()), 2 + ), + ], + names=[ + "flag", + "i16", + "i32", + "i64", + "f32", + "f64", + "dt", + "ts", + "uid", + "vec_f32", + "vec_f16", + ], + ) + + payload = encode_arrow_batch_pg_binary(batch, include_header=True, include_trailer=True) + assert payload.startswith(PG_BINARY_COPY_HEADER) + assert payload.endswith(PG_BINARY_COPY_TRAILER) + + body = payload[len(PG_BINARY_COPY_HEADER) : -len(PG_BINARY_COPY_TRAILER)] + expected_row_0 = b"".join([ + struct.pack(">h", 11), + struct.pack(">iB", 1, 1), + struct.pack(">ih", 2, 10), + struct.pack(">ii", 4, 1000), + struct.pack(">iq", 8, 100_000), + struct.pack(">if", 4, 1.5), + struct.pack(">id", 8, 3.125), + struct.pack(">ii", 4, 0), + struct.pack(">iq", 8, 0), + struct.pack(">i16s", 16, uid.bytes), + struct.pack(">iHH3f", 16, 3, 0, 1.0, 2.0, 3.0), + struct.pack(">iHH2e", 8, 2, 0, 0.5, 1.5), + ]) + expected_row_1 = b"".join([ + struct.pack(">h", 11), + struct.pack(">iB", 1, 0), + struct.pack(">ih", 2, -20), + struct.pack(">ii", 4, -2000), + struct.pack(">iq", 8, -200_000), + struct.pack(">if", 4, -2.5), + struct.pack(">id", 8, -6.25), + struct.pack(">ii", 4, 2), + struct.pack(">iq", 8, 2_000_000), + struct.pack(">i16s", 16, uid.bytes), + struct.pack(">iHH3f", 16, 3, 0, 4.0, 5.0, 6.0), + struct.pack(">iHH2e", 8, 2, 0, 2.5, 3.5), + ]) + assert body == expected_row_0 + expected_row_1 + + +def test_encode_arrow_batch_pg_binary_fast_path_2_variable_width_jsonb_and_nulls() -> None: + """Fast Path 2 slices Arrow UTF-8/binary buffers, prepends jsonb version byte, and writes -1 for NULLs.""" + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, None, 3], type=pa.int64()), + pa.array(["alpha", None, "gamma"], type=pa.string()), + pa.array(['{"a":1}', None, '{"b":2}'], type=pa.string()), + pa.array([[1.0, 2.0], None, [3.0, 4.0]], type=pa.list_(pa.float32())), + ], + names=["id", "label", "payload", "embedding"], + ) + + payload = encode_arrow_batch_pg_binary( + batch, + jsonb_columns=frozenset({"payload"}), + vector_columns=frozenset({"embedding"}), + include_header=False, + include_trailer=False, + ) + + expected_row_0 = b"".join([ + struct.pack(">h", 4), + struct.pack(">iq", 8, 1), + struct.pack(">i5s", 5, b"alpha"), + struct.pack(">iB7s", 8, 1, b'{"a":1}'), + struct.pack(">iHH2f", 12, 2, 0, 1.0, 2.0), + ]) + expected_row_1 = b"".join([ + struct.pack(">h", 4), + struct.pack(">i", -1), + struct.pack(">i", -1), + struct.pack(">i", -1), + struct.pack(">i", -1), + ]) + expected_row_2 = b"".join([ + struct.pack(">h", 4), + struct.pack(">iq", 8, 3), + struct.pack(">i5s", 5, b"gamma"), + struct.pack(">iB7s", 8, 1, b'{"b":2}'), + struct.pack(">iHH2f", 12, 2, 0, 3.0, 4.0), + ]) + assert payload == expected_row_0 + expected_row_1 + expected_row_2 + + +def test_iter_pgcopy_binary_chunks_streams_reader_batches() -> None: + """iter_pgcopy_binary_chunks emits header, batch payloads, and trailer across RecordBatchReader.""" + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + batches = [ + pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["a", "b"]}, schema=schema), + pa.RecordBatch.from_pydict({"id": [3, 4], "name": ["c", "d"]}, schema=schema), + ] + reader = pa.RecordBatchReader.from_batches(schema, batches) + + chunks = list(iter_pgcopy_binary_chunks(reader, chunk_bytes=16)) + combined = b"".join(chunks) + + assert combined.startswith(PG_BINARY_COPY_HEADER) + assert combined.endswith(PG_BINARY_COPY_TRAILER) + assert len(chunks) >= 2 + + +def test_build_copy_from_command_supports_binary_and_text_modes() -> None: + """psycopg and cockroach_psycopg build_copy_from_command emit FORMAT BINARY when binary=True.""" + pg_bin = pg_build_copy_from_command("public.items", ["id", "name"], binary=True) + pg_txt = pg_build_copy_from_command("public.items", ["id", "name"], binary=False) + crdb_bin = crdb_build_copy_from_command("public.items", ["id", "name"], binary=True) + + assert "WITH (FORMAT BINARY)" in repr(pg_bin) + assert "WITH (FORMAT BINARY)" not in repr(pg_txt) + assert "WITH (FORMAT BINARY)" in repr(crdb_bin) + + +class _DummyPsycopgCopy: + def __init__(self, writes: list[bytes], rows: list[tuple[Any, ...]]) -> None: + self._writes = writes + self._rows = rows + + def __enter__(self) -> "_DummyPsycopgCopy": + return self + + def __exit__(self, *_args: Any) -> None: + return None + + async def __aenter__(self) -> "_DummyPsycopgCopy": + return self + + async def __aexit__(self, *_args: Any) -> None: + return None + + def write(self, data: bytes) -> None: + self._writes.append(bytes(data)) + + def write_row(self, row: tuple[Any, ...]) -> None: + self._rows.append(tuple(row)) + + +class _DummyPsycopgAsyncCopy: + def __init__(self, writes: list[bytes], rows: list[tuple[Any, ...]]) -> None: + self._writes = writes + self._rows = rows + + async def __aenter__(self) -> "_DummyPsycopgAsyncCopy": + return self + + async def __aexit__(self, *_args: Any) -> None: + return None + + async def write(self, data: bytes) -> None: + self._writes.append(bytes(data)) + + async def write_row(self, row: tuple[Any, ...]) -> None: + self._rows.append(tuple(row)) + + +class _DummyPsycopgSyncCursor: + def __init__(self, conn: "_DummyPsycopgSyncConn") -> None: + self._conn = conn + + def __enter__(self) -> "_DummyPsycopgSyncCursor": + return self + + def __exit__(self, *_args: Any) -> None: + return None + + def execute(self, sql: Any) -> None: + self._conn.executed.append(sql) + + def copy(self, sql: Any) -> _DummyPsycopgCopy: + self._conn.copy_sqls.append(sql) + return _DummyPsycopgCopy(self._conn.binary_writes, self._conn.row_writes) + + +class _DummyPsycopgSyncConn: + def __init__(self) -> None: + self.executed: list[Any] = [] + self.copy_sqls: list[Any] = [] + self.binary_writes: list[bytes] = [] + self.row_writes: list[tuple[Any, ...]] = [] + self.info = SimpleNamespace(transaction_status=0) + + def cursor(self) -> _DummyPsycopgSyncCursor: + return _DummyPsycopgSyncCursor(self) + + +class _DummyPsycopgAsyncCursor: + def __init__(self, conn: "_DummyPsycopgAsyncConn") -> None: + self._conn = conn + + async def __aenter__(self) -> "_DummyPsycopgAsyncCursor": + return self + + async def __aexit__(self, *_args: Any) -> None: + return None + + async def execute(self, sql: Any) -> None: + self._conn.executed.append(sql) + + def copy(self, sql: Any) -> _DummyPsycopgAsyncCopy: + self._conn.copy_sqls.append(sql) + return _DummyPsycopgAsyncCopy(self._conn.binary_writes, self._conn.row_writes) + + +class _DummyPsycopgAsyncConn: + def __init__(self) -> None: + self.executed: list[Any] = [] + self.copy_sqls: list[Any] = [] + self.binary_writes: list[bytes] = [] + self.row_writes: list[tuple[Any, ...]] = [] + self.info = SimpleNamespace(transaction_status=0) + + def cursor(self) -> _DummyPsycopgAsyncCursor: + return _DummyPsycopgAsyncCursor(self) + + +def test_psycopg_and_cockroach_sync_load_from_arrow_streams_binary_copy() -> None: + """PsycopgSyncDriver and CockroachPsycopgSyncDriver stream RecordBatchReader via binary COPY.""" + schema = pa.schema([("id", pa.int64()), ("label", pa.string())]) + for driver_cls in (PsycopgSyncDriver, CockroachPsycopgSyncDriver): + conn = _DummyPsycopgSyncConn() + driver = driver_cls( + connection=cast("Any", conn), + statement_config=psycopg_statement_config, + driver_features={"storage_capabilities": _CAPS}, + ) + reader = pa.RecordBatchReader.from_batches( + schema, + [ + pa.RecordBatch.from_pydict({"id": [1, 2], "label": ["a", "b"]}, schema=schema), + pa.RecordBatch.from_pydict({"id": [3], "label": ["c"]}, schema=schema), + ], + ) + job = driver.load_from_arrow("public.items", reader) + + assert job.telemetry["rows_processed"] == 3 + assert len(conn.copy_sqls) == 1 + assert "WITH (FORMAT BINARY)" in repr(conn.copy_sqls[0]) + combined = b"".join(conn.binary_writes) + assert combined.startswith(PG_BINARY_COPY_HEADER) + assert combined.endswith(PG_BINARY_COPY_TRAILER) + assert conn.row_writes == [] + + +async def test_psycopg_and_cockroach_async_load_from_arrow_streams_binary_copy() -> None: + """PsycopgAsyncDriver and CockroachPsycopgAsyncDriver stream RecordBatchReader via async binary COPY.""" + schema = pa.schema([("id", pa.int64()), ("label", pa.string())]) + for driver_cls in (PsycopgAsyncDriver, CockroachPsycopgAsyncDriver): + conn = _DummyPsycopgAsyncConn() + driver = driver_cls( + connection=cast("Any", conn), + statement_config=psycopg_statement_config, + driver_features={"storage_capabilities": _CAPS}, + ) + reader = pa.RecordBatchReader.from_batches( + schema, + [pa.RecordBatch.from_pydict({"id": [10, 20], "label": ["x", "y"]}, schema=schema)], + ) + job = await driver.load_from_arrow("public.items", reader) + + assert job.telemetry["rows_processed"] == 2 + assert len(conn.copy_sqls) == 1 + assert "WITH (FORMAT BINARY)" in repr(conn.copy_sqls[0]) + combined = b"".join(conn.binary_writes) + assert combined.startswith(PG_BINARY_COPY_HEADER) + assert combined.endswith(PG_BINARY_COPY_TRAILER) + + +class _DummyAsyncpgBinaryConn: + def __init__(self) -> None: + self.copy_to_table_calls: list[dict[str, Any]] = [] + self.chunks: list[bytes] = [] + + async def copy_to_table( + self, + table_name: str, + *, + source: Any, + columns: list[str] | None = None, + schema_name: str | None = None, + format: str | None = None, + ) -> None: + async for chunk in source: + self.chunks.append(bytes(chunk)) + self.copy_to_table_calls.append({ + "table_name": table_name, + "columns": columns, + "schema_name": schema_name, + "format": format, + }) + + +async def test_asyncpg_and_cockroach_asyncpg_load_from_arrow_streams_binary_copy() -> None: + """AsyncpgDriver and CockroachAsyncpgDriver stream RecordBatchReader into copy_to_table(format='binary').""" + schema = pa.schema([("id", pa.int64()), ("vec", pa.list_(pa.float32(), 2))]) + for driver_cls in (AsyncpgDriver, CockroachAsyncpgDriver): + conn = _DummyAsyncpgBinaryConn() + driver = driver_cls( + connection=cast("Any", conn), + statement_config=asyncpg_statement_config, + driver_features={"storage_capabilities": _CAPS}, + ) + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, 2], type=pa.int64()), + pa.FixedSizeListArray.from_arrays(pa.array([1.0, 2.0, 3.0, 4.0], type=pa.float32()), 2), + ], + schema=schema, + ) + reader = pa.RecordBatchReader.from_batches(schema, [batch]) + job = await driver.load_from_arrow("analytics.embeddings", reader) + + assert job.telemetry["rows_processed"] == 2 + assert len(conn.copy_to_table_calls) == 1 + assert conn.copy_to_table_calls[0] == { + "table_name": "embeddings", + "columns": ["id", "vec"], + "schema_name": "analytics", + "format": "binary", + } + combined = b"".join(conn.chunks) + assert combined.startswith(PG_BINARY_COPY_HEADER) + assert combined.endswith(PG_BINARY_COPY_TRAILER) + + +class _DummyPsqlpyBinaryConn: + def __init__(self) -> None: + self.binary_calls: list[dict[str, Any]] = [] + + async def binary_copy_to_table( + self, + source: bytes, + table_name: str, + *, + columns: list[str] | None = None, + schema_name: str | None = None, + ) -> int: + self.binary_calls.append({ + "source": bytes(source), + "table_name": table_name, + "columns": columns, + "schema_name": schema_name, + }) + return 2 + + async def fetch(self, _sql: str, _params: Any = None) -> Any: + return SimpleNamespace(result=list) + + +async def test_psqlpy_load_from_arrow_supports_binary_copy_to_table() -> None: + """PsqlpyDriver streams PGCOPY binary payloads to binary_copy_to_table when enable_pgcopy_binary=True.""" + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + reader = pa.RecordBatchReader.from_batches( + schema, + [pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["alpha", "beta"]}, schema=schema)], + ) + conn = _DummyPsqlpyBinaryConn() + driver = PsqlpyDriver( + connection=cast("Any", conn), + statement_config=psqlpy_statement_config, + driver_features={"storage_capabilities": _CAPS, "enable_pgcopy_binary": True}, + ) + + job = await driver.load_from_arrow("public.items", reader) + + assert job.telemetry["rows_processed"] == 2 + assert len(conn.binary_calls) == 1 + assert conn.binary_calls[0]["table_name"] == "items" + assert conn.binary_calls[0]["columns"] == ["id", "name"] + assert conn.binary_calls[0]["source"].startswith(PG_BINARY_COPY_HEADER) + assert conn.binary_calls[0]["source"].endswith(PG_BINARY_COPY_TRAILER) diff --git a/tests/unit/storage/test_bridge.py b/tests/unit/storage/test_bridge.py index ac64f5716..cbb2748e1 100644 --- a/tests/unit/storage/test_bridge.py +++ b/tests/unit/storage/test_bridge.py @@ -69,12 +69,27 @@ def test_encode_row_payload_jsonl_returns_line_delimited_bytes() -> None: class DummyAsyncpgConnection: def __init__(self) -> None: self.calls: list[tuple[str, str | None, list[tuple[object, ...]], list[str]]] = [] + self.binary_chunks: list[bytes] = [] async def copy_records_to_table( self, table: str, *, records: list[tuple[object, ...]], columns: list[str], schema_name: str | None = None ) -> None: self.calls.append((table, schema_name, records, columns)) + async def copy_to_table( + self, + table: str, + *, + source: Any, + columns: list[str] | None = None, + schema_name: str | None = None, + format: str | None = None, + ) -> None: + _ = format + async for chunk in source: + self.binary_chunks.append(bytes(chunk)) + self.calls.append((table, schema_name, [], list(columns or []))) + class DummyPsqlpyConnection: def __init__(self) -> None: From fa978d704002d9e81c0d150367ed8ee57626cbc6 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 18:01:24 +0000 Subject: [PATCH 03/14] feat(storage): add streaming arrow batch pipeline, vector parquet defaults, and schema-pinned jsonl --- sqlspec/core/result/_base.py | 11 +- sqlspec/storage/_arrow_payload.py | 77 ++++++++- sqlspec/storage/_arrow_stream.py | 96 ++++++++++- sqlspec/storage/backends/fsspec.py | 6 +- sqlspec/storage/backends/local.py | 9 +- sqlspec/storage/backends/obstore.py | 46 +++-- sqlspec/storage/pipeline.py | 163 +++++++++++++++++- tests/unit/storage/test_pipeline_streaming.py | 141 +++++++++++++++ 8 files changed, 518 insertions(+), 31 deletions(-) create mode 100644 tests/unit/storage/test_pipeline_streaming.py diff --git a/sqlspec/core/result/_base.py b/sqlspec/core/result/_base.py index 4a9528ddf..5d2601c08 100644 --- a/sqlspec/core/result/_base.py +++ b/sqlspec/core/result/_base.py @@ -34,6 +34,7 @@ arrow_table_to_pylist, arrow_table_to_return_format, cast_arrow_table_schema, + coerce_arrow_table, convert_dict_to_arrow, ensure_arrow_table, ) @@ -987,7 +988,10 @@ def write_to_storage_sync( compression: str | None = None, pipeline: "SyncStoragePipeline | None" = None, ) -> "StorageTelemetry": - table = self._as_table() + if self.data is None: + msg = "No Arrow table available for this result" + raise ValueError(msg) + table = coerce_arrow_table(self.data) active_pipeline = pipeline or SyncStoragePipeline() return active_pipeline.write_arrow( table, destination, format_hint=format_hint, storage_options=storage_options, compression=compression @@ -1002,7 +1006,10 @@ async def write_to_storage_async( compression: str | None = None, pipeline: "AsyncStoragePipeline | None" = None, ) -> "StorageTelemetry": - table = self._as_table() + if self.data is None: + msg = "No Arrow table available for this result" + raise ValueError(msg) + table = coerce_arrow_table(self.data) active_pipeline = pipeline or AsyncStoragePipeline() return await active_pipeline.write_arrow( table, destination, format_hint=format_hint, storage_options=storage_options, compression=compression diff --git a/sqlspec/storage/_arrow_payload.py b/sqlspec/storage/_arrow_payload.py index 72c90627e..4aaaa40ca 100644 --- a/sqlspec/storage/_arrow_payload.py +++ b/sqlspec/storage/_arrow_payload.py @@ -8,7 +8,13 @@ if TYPE_CHECKING: from sqlspec.typing import ArrowTable -__all__ = ("StorageFormat", "decode_arrow_payload", "encode_arrow_payload") +__all__ = ( + "StorageFormat", + "batches_to_arrow_table", + "decode_arrow_payload", + "encode_arrow_payload", + "resolve_parquet_use_dictionary", +) StorageFormat = Literal["jsonl", "json", "parquet", "arrow-ipc", "csv"] @@ -16,6 +22,43 @@ _PYARROW_JSON_BLOCK_SIZE = 1 << 20 +def batches_to_arrow_table(batches: Any) -> "ArrowTable": + """Coalesce an Arrow RecordBatchReader, RecordBatch, Table, or batch iterable into an ArrowTable.""" + pa = import_pyarrow() + if isinstance(batches, pa.Table): + return batches + if isinstance(batches, pa.RecordBatchReader): + return batches.read_all() + if isinstance(batches, pa.RecordBatch): + return pa.Table.from_batches([batches]) + batch_list = list(batches) + if not batch_list: + return pa.table({}) + return pa.Table.from_batches(batch_list, schema=batch_list[0].schema) + + +def resolve_parquet_use_dictionary( + schema: Any, explicit_use_dictionary: "bool | list[str] | None" = None +) -> "bool | list[str]": + """Resolve Parquet dictionary encoding, disabling dictionary construction on float vector columns. + + Args: + schema: PyArrow Schema of the table or batch being written. + explicit_use_dictionary: Caller-supplied ``use_dictionary`` override. + + Returns: + ``True`` when no float vector columns exist, ``explicit_use_dictionary`` when provided, + or the list of non-vector column names to dictionary-encode. + """ + if explicit_use_dictionary is not None: + return explicit_use_dictionary + pa = import_pyarrow() + has_float_vector = any(_is_float_vector_field(field, pa) for field in schema) + if not has_float_vector: + return True + return [field.name for field in schema if not _is_float_vector_field(field, pa)] + + def encode_arrow_payload( table: "ArrowTable", format_choice: StorageFormat, @@ -39,36 +82,54 @@ def encode_arrow_payload( pa_csv.write_csv(table, sink, write_options=csv_opts) else: pq = import_pyarrow_parquet() - pq.write_table(table, sink, compression=compression) + parquet_opts = dict(write_options) if write_options else {} + parquet_opts["use_dictionary"] = resolve_parquet_use_dictionary( + table.schema, parquet_opts.get("use_dictionary") + ) + pq.write_table(table, sink, compression=compression, **parquet_opts) buffer = sink.getvalue() result_bytes: bytes = buffer.to_pybytes() return result_bytes -def decode_arrow_payload(payload: bytes, format_choice: StorageFormat) -> "ArrowTable": +def decode_arrow_payload(payload: bytes, format_choice: StorageFormat, *, arrow_schema: Any = None) -> "ArrowTable": """Decode bytes into an Arrow table using optional PyArrow dependencies.""" pa = import_pyarrow() if format_choice == "parquet": pq = import_pyarrow_parquet() - return cast("ArrowTable", pq.read_table(pa.BufferReader(payload))) + table = cast("ArrowTable", pq.read_table(pa.BufferReader(payload))) + return table.cast(arrow_schema) if arrow_schema is not None else table if format_choice == "arrow-ipc": with pa.ipc.open_file(pa.BufferReader(payload)) as reader: - return cast("ArrowTable", reader.read_all()) + table = cast("ArrowTable", reader.read_all()) + return table.cast(arrow_schema) if arrow_schema is not None else table if format_choice == "csv": pa_csv = import_pyarrow_csv() - return cast("ArrowTable", pa_csv.read_csv(pa.BufferReader(payload))) + convert_opts = pa_csv.ConvertOptions(column_types=arrow_schema) if arrow_schema is not None else None + return cast("ArrowTable", pa_csv.read_csv(pa.BufferReader(payload), convert_options=convert_opts)) if format_choice == "json": if not payload.strip(): return cast("ArrowTable", pa.table({})) data = from_json(payload) rows = data if isinstance(data, list) else [data] - return cast("ArrowTable", pa.Table.from_pylist(rows)) + return cast("ArrowTable", pa.Table.from_pylist(rows, schema=arrow_schema)) if format_choice == "jsonl": if payload == b"": return cast("ArrowTable", pa.table({})) pa_json = import_pyarrow_json() read_options = pa_json.ReadOptions(block_size=max(_PYARROW_JSON_BLOCK_SIZE, len(payload))) - return cast("ArrowTable", pa_json.read_json(pa.BufferReader(payload), read_options=read_options)) + parse_options = pa_json.ParseOptions(explicit_schema=arrow_schema) if arrow_schema is not None else None + return cast( + "ArrowTable", + pa_json.read_json(pa.BufferReader(payload), read_options=read_options, parse_options=parse_options), + ) msg = f"Unsupported storage format for Arrow decoding: {format_choice}" raise ValueError(msg) + + +def _is_float_vector_field(field: Any, pa: Any) -> bool: + dtype = field.type + if pa.types.is_fixed_size_list(dtype) or pa.types.is_list(dtype) or pa.types.is_large_list(dtype): + return bool(pa.types.is_floating(dtype.value_type)) + return False diff --git a/sqlspec/storage/_arrow_stream.py b/sqlspec/storage/_arrow_stream.py index fce973ee6..901c193b8 100644 --- a/sqlspec/storage/_arrow_stream.py +++ b/sqlspec/storage/_arrow_stream.py @@ -1,21 +1,24 @@ -"""Shared helpers for bounded Parquet batch streaming.""" +"""Shared helpers for bounded Arrow batch streaming.""" from pathlib import PurePath from typing import TYPE_CHECKING, Any from sqlspec.exceptions import StorageCapabilityError +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_csv, import_pyarrow_json if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator from sqlspec.typing import ArrowRecordBatch -__all__ = ("iter_parquet_row_groups", "validate_parquet_stream_options") +__all__ = ("iter_parquet_row_groups", "iter_streamed_arrow_batches", "validate_parquet_stream_options") _NON_PARQUET_SUFFIXES = frozenset({".arrow", ".csv", ".feather", ".ipc", ".json", ".jsonl", ".ndjson"}) _STREAM_REMEDIATION = "Read this format with the Arrow read APIs instead of batch streaming." +_JSONL_CHUNK_TARGET_BYTES = 8 * 1024 * 1024 + def validate_parquet_stream_options(pattern: str, file_format: str, batch_size: int) -> None: """Validate a Parquet streaming request before storage is accessed. @@ -42,7 +45,90 @@ def validate_parquet_stream_options(pattern: str, file_format: str, batch_size: raise StorageCapabilityError(msg, capability="arrow_batch_streaming", remediation=_STREAM_REMEDIATION) -def iter_parquet_row_groups(parquet_file: Any, *, batch_size: int, **kwargs: Any) -> "Iterator[ArrowRecordBatch]": +def iter_parquet_row_groups( + parquet_file: Any, *, batch_size: int, arrow_schema: Any = None, **kwargs: Any +) -> "Iterator[ArrowRecordBatch]": """Yield batches while limiting each PyArrow read to one row group.""" for row_group in range(parquet_file.num_row_groups): - yield from parquet_file.iter_batches(batch_size=batch_size, row_groups=[row_group], **kwargs) + for batch in parquet_file.iter_batches(batch_size=batch_size, row_groups=[row_group], **kwargs): + yield batch.cast(arrow_schema) if arrow_schema is not None else batch + + +def iter_streamed_arrow_batches( + byte_chunks: "Iterable[bytes]", *, file_format: str, batch_size: int = 65_536, arrow_schema: Any = None +) -> "Iterator[ArrowRecordBatch]": + """Stream RecordBatches from byte chunks for JSONL, CSV, and Arrow IPC formats. + + Args: + byte_chunks: Iterable of raw byte chunks from a storage backend. + file_format: Storage format (``"jsonl"``, ``"ndjson"``, ``"csv"``, ``"arrow-ipc"``, ``"arrow"``, ``"ipc"``). + batch_size: Maximum rows per yielded RecordBatch. + arrow_schema: Optional explicit ``pyarrow.Schema`` to pin column types (such as ``float32`` vectors). + + Yields: + Arrow ``RecordBatch`` instances. + """ + if batch_size <= 0: + msg = f"batch_size must be greater than zero; received {batch_size}" + raise ValueError(msg) + pa = import_pyarrow() + normalized = file_format.lower() + if normalized in {"jsonl", "ndjson"}: + pa_json = import_pyarrow_json() + parse_options = pa_json.ParseOptions(explicit_schema=arrow_schema) if arrow_schema is not None else None + pending = bytearray() + for chunk in byte_chunks: + if not chunk: + continue + pending.extend(chunk) + if len(pending) < _JSONL_CHUNK_TARGET_BYTES: + continue + last_nl = pending.rfind(b"\n") + if last_nl <= 0: + continue + complete_block = bytes(pending[: last_nl + 1]) + del pending[: last_nl + 1] + if complete_block.strip(): + read_options = pa_json.ReadOptions(block_size=max(1 << 20, len(complete_block))) + table = pa_json.read_json( + pa.BufferReader(complete_block), read_options=read_options, parse_options=parse_options + ) + yield from table.to_batches(max_chunksize=batch_size) + if pending.strip(): + final_block = bytes(pending) + read_options = pa_json.ReadOptions(block_size=max(1 << 20, len(final_block))) + table = pa_json.read_json( + pa.BufferReader(final_block), read_options=read_options, parse_options=parse_options + ) + yield from table.to_batches(max_chunksize=batch_size) + return + + payload = b"".join(byte_chunks) + if not payload: + return + if normalized == "csv": + pa_csv = import_pyarrow_csv() + convert_options = pa_csv.ConvertOptions(column_types=arrow_schema) if arrow_schema is not None else None + reader = pa_csv.open_csv(pa.BufferReader(payload), convert_options=convert_options) + for batch in reader: + yield from pa.Table.from_batches([batch]).to_batches(max_chunksize=batch_size) + return + + if normalized in {"arrow-ipc", "arrow", "ipc"}: + try: + with pa.ipc.open_file(pa.BufferReader(payload)) as ipc_reader: + for idx in range(ipc_reader.num_record_batches): + batch = ipc_reader.get_batch(idx) + if arrow_schema is not None: + batch = batch.cast(arrow_schema) + yield from pa.Table.from_batches([batch]).to_batches(max_chunksize=batch_size) + except Exception: + with pa.ipc.open_stream(pa.BufferReader(payload)) as stream_reader: + for batch in stream_reader: + if arrow_schema is not None: + batch = batch.cast(arrow_schema) + yield from pa.Table.from_batches([batch]).to_batches(max_chunksize=batch_size) + return + + msg = f"Unsupported streaming Arrow format: {file_format!r}" + raise ValueError(msg) diff --git a/sqlspec/storage/backends/fsspec.py b/sqlspec/storage/backends/fsspec.py index 997f5b704..0819441f1 100644 --- a/sqlspec/storage/backends/fsspec.py +++ b/sqlspec/storage/backends/fsspec.py @@ -439,6 +439,8 @@ def stream_arrow_sync( """ validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() + arrow_schema = kwargs.pop("arrow_schema", None) + kwargs.pop("read_buffer_bytes", None) for obj_path in self.glob_sync(pattern): file_handle = execute_sync_storage_operation( partial(self.fs.open, obj_path, mode="rb"), @@ -453,7 +455,9 @@ def stream_arrow_sync( operation="stream_arrow", path=str(obj_path), ) - yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) + yield from iter_parquet_row_groups( + parquet_file, batch_size=batch_size, arrow_schema=arrow_schema, **kwargs + ) async def read_bytes_async(self, path: "str | Path", **kwargs: Any) -> bytes: """Read bytes from storage asynchronously.""" diff --git a/sqlspec/storage/backends/local.py b/sqlspec/storage/backends/local.py index 53eb5a521..d1f530def 100644 --- a/sqlspec/storage/backends/local.py +++ b/sqlspec/storage/backends/local.py @@ -154,6 +154,7 @@ def list_objects_sync(self, prefix: str = "", recursive: bool = True, **kwargs: When the prefix resembles a directory (contains a slash or ends with '/'), we treat it as a path; otherwise we filter filenames within the base path. """ + rel_prefix = prefix if prefix: reject_parent_traversal(prefix) if prefix.endswith("/") or "/" in prefix: @@ -162,6 +163,8 @@ def list_objects_sync(self, prefix: str = "", recursive: bool = True, **kwargs: return [] if search_path.is_file(): return [str(search_path.relative_to(self.base_path))] + rel_path = search_path.relative_to(self.base_path) + rel_prefix = "" if str(rel_path) == "." else str(rel_path) else: search_path = self.base_path else: @@ -178,7 +181,7 @@ def list_objects_sync(self, prefix: str = "", recursive: bool = True, **kwargs: except ValueError: continue relative_str = str(relative) - if not prefix or relative_str.startswith(prefix): + if not rel_prefix or relative_str.startswith(rel_prefix): files.append(relative_str) return sorted(files) @@ -305,6 +308,8 @@ def stream_arrow_sync( """ validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() + arrow_schema = kwargs.pop("arrow_schema", None) + kwargs.pop("read_buffer_bytes", None) files = self.glob_sync(pattern) for file_path in files: resolved = self._resolve_path(file_path) @@ -315,7 +320,7 @@ def stream_arrow_sync( operation="stream_arrow", path=resolved_str, ) - yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) + yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, arrow_schema=arrow_schema, **kwargs) @overload def sign_sync(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ... diff --git a/sqlspec/storage/backends/obstore.py b/sqlspec/storage/backends/obstore.py index 22cee5238..35c06143e 100644 --- a/sqlspec/storage/backends/obstore.py +++ b/sqlspec/storage/backends/obstore.py @@ -16,7 +16,8 @@ from mypy_extensions import mypyc_attr -from sqlspec.exceptions import StorageOperationFailedError +from sqlspec.exceptions import StorageOperationFailedError, StoragePathTraversalError +from sqlspec.storage._arrow_payload import resolve_parquet_use_dictionary from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options from sqlspec.storage._paths import ( ensure_path_within_root, @@ -323,11 +324,17 @@ def glob_sync(self, pattern: str, **kwargs: Any) -> "list[str]": if not pattern: return [] reject_parent_traversal(pattern) - resolved_pattern = ( - pattern - if getattr(self, "_is_local_store", False) - else resolve_storage_path(pattern, self.base_path, self.protocol, strip_file_scheme=True) - ) + if getattr(self, "_is_local_store", False): + pattern_path = Path(pattern) + if pattern_path.is_absolute() and self._local_store_root: + try: + resolved_pattern = pattern_path.relative_to(Path(self._local_store_root)).as_posix() + except ValueError as exc: + raise StoragePathTraversalError(pattern, self._local_store_root) from exc + else: + resolved_pattern = pattern + else: + resolved_pattern = resolve_storage_path(pattern, self.base_path, self.protocol, strip_file_scheme=True) all_objects = self._list_resolved_sync(extract_glob_static_prefix(resolved_pattern), recursive=True) matcher = glob_to_regex(resolved_pattern) results = [obj for obj in all_objects if matcher.match(obj)] @@ -465,9 +472,11 @@ def stream_arrow_sync( validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() + read_buffer_bytes = kwargs.pop("read_buffer_bytes", 16_777_216) + arrow_schema = kwargs.pop("arrow_schema", None) for obj_path in self.glob_sync(pattern): reader = execute_sync_storage_operation( - partial(obs.open_reader, self.store, obj_path), + partial(_open_obstore_reader, obs, self.store, obj_path, read_buffer_bytes), backend=self.backend_type, operation="stream_open", path=obj_path, @@ -476,7 +485,9 @@ def stream_arrow_sync( parquet_file = execute_sync_storage_operation( partial(pq.ParquetFile, stream), backend=self.backend_type, operation="stream_arrow", path=obj_path ) - yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) + yield from iter_parquet_row_groups( + parquet_file, batch_size=batch_size, arrow_schema=arrow_schema, **kwargs + ) @overload def sign_sync(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ... @@ -813,8 +824,9 @@ def _resolve_list_prefix(self, prefix: str) -> str: is_local = getattr(self, "_is_local_store", False) if not prefix: return "" if is_local else (self.base_path or "") - base = "" if is_local else self.base_path - return resolve_storage_path(prefix, base, self.protocol, strip_file_scheme=True) + if is_local: + return self._local_store_path(prefix) + return resolve_storage_path(prefix, self.base_path, self.protocol, strip_file_scheme=True) def _local_store_path(self, path: "str | Path") -> str: """Resolve path for LocalStore, which expects relative paths from its root. @@ -940,6 +952,8 @@ def _stream_parquet_sync(self, resolved_path: str, table: "ArrowTable", **kwargs table = table.cast(pa.schema(new_fields)) row_group_size = kwargs.pop("row_group_size", None) + if "use_dictionary" not in kwargs: + kwargs["use_dictionary"] = resolve_parquet_use_dictionary(table.schema) sink = obs.open_writer(self.store, resolved_path) sink_adapter = _ObstoreSink(sink) writer: Any = None @@ -1028,6 +1042,15 @@ async def _read_obstore_bytes_async(store: Any, resolved_path: str) -> bytes: return cast("bytes", bytes_obj.to_bytes()) +def _open_obstore_reader(obs: Any, store: Any, obj_path: str, read_buffer_bytes: int | None) -> Any: + if read_buffer_bytes is not None: + try: + return obs.open_reader(store, obj_path, buffer_size=int(read_buffer_bytes)) + except TypeError: + pass + return obs.open_reader(store, obj_path) + + def _read_obstore_bytes(store: Any, resolved_path: str) -> bytes: """Read bytes via obstore.""" result = store.get(resolved_path) @@ -1108,7 +1131,8 @@ def abort(self) -> None: abort_fn() def write(self, data: Any) -> int: - return int(self._writer.write(data)) + payload = bytes(data) if isinstance(data, (bytearray, memoryview)) else data + return int(self._writer.write(payload)) def flush(self) -> None: if not self._aborted: diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index ea343881b..6e8ae4940 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -11,8 +11,15 @@ from typing_extensions import TypedDict from sqlspec.exceptions import ImproperConfigurationError, StorageCapabilityError -from sqlspec.storage._arrow_payload import StorageFormat, decode_arrow_payload, encode_arrow_payload +from sqlspec.storage._arrow_payload import ( + StorageFormat, + batches_to_arrow_table, + decode_arrow_payload, + encode_arrow_payload, +) +from sqlspec.storage._arrow_stream import iter_streamed_arrow_batches from sqlspec.storage._paths import FILE_PROTOCOL, FILE_SCHEME_PREFIX, parse_file_uri_path +from sqlspec.storage.backends.base import AsyncArrowBatchIterator from sqlspec.storage.errors import execute_async_storage_operation, execute_sync_storage_operation from sqlspec.storage.registry import StorageRegistry, storage_registry from sqlspec.utils.serializers import get_serializer_metrics, serialize_collection, to_json @@ -24,7 +31,7 @@ from collections.abc import AsyncIterator, Iterator from sqlspec.protocols import ObjectStoreProtocol - from sqlspec.typing import ArrowTable + from sqlspec.typing import ArrowRecordBatch, ArrowTable __all__ = ( @@ -353,6 +360,57 @@ def stream_read( backend, path, _backend_name = self._backend(source, storage_options) return backend.stream_read_sync(path, chunk_size=chunk_size) + def read_arrow_batches( + self, + source: StorageDestination, + *, + file_format: str = "parquet", + batch_size: int = 65_536, + arrow_schema: Any = None, + storage_options: "dict[str, Any] | None" = None, + ) -> "Iterator[ArrowRecordBatch]": + """Stream Arrow record batches from storage synchronously.""" + return _read_arrow_batches_sync( + self, + source, + file_format=file_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + storage_options=storage_options, + ) + + def write_arrow_batches( + self, + destination: StorageDestination, + batches: Any, + *, + file_format: StorageFormat = "parquet", + compression: str | None = None, + write_options: "dict[str, Any] | None" = None, + storage_options: "dict[str, Any] | None" = None, + ) -> StorageTelemetry: + """Write Arrow record batches to storage synchronously.""" + _validate_arrow_write_format(file_format) + resolved_options = _storage_options(self._storage_options, storage_options) + format_write_options = write_options or _csv_write_options( + file_format, resolved_options, self._storage_options, self._csv_write_options + ) + table = batches_to_arrow_table(batches) + payload = _encode_arrow_payload(table, file_format, compression=compression, write_options=format_write_options) + return self._write_bytes( + payload, destination, rows=int(table.num_rows), format_label=file_format, storage_options=resolved_options + ) + + def list_sources( + self, + source: StorageDestination, + *, + file_format: str | None = None, + storage_options: "dict[str, Any] | None" = None, + ) -> "list[str]": + """Expand a single URI, directory prefix, or glob pattern into sorted source URIs.""" + return _resolve_storage_sources_sync(self, source, file_format=file_format, storage_options=storage_options) + def _write_bytes( self, payload: bytes, @@ -456,6 +514,63 @@ async def stream_read_async( backend, path, _backend_name = self._backend(source, storage_options) return await backend.stream_read_async(path, chunk_size=chunk_size) + def read_arrow_batches_async( + self, + source: StorageDestination, + *, + file_format: str = "parquet", + batch_size: int = 65_536, + arrow_schema: Any = None, + storage_options: "dict[str, Any] | None" = None, + ) -> "AsyncIterator[ArrowRecordBatch]": + """Stream Arrow record batches from storage asynchronously.""" + return AsyncArrowBatchIterator( + _read_arrow_batches_sync( + self, + source, + file_format=file_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + storage_options=storage_options, + ) + ) + + async def write_arrow_batches_async( + self, + destination: StorageDestination, + batches: Any, + *, + file_format: StorageFormat = "parquet", + compression: str | None = None, + write_options: "dict[str, Any] | None" = None, + storage_options: "dict[str, Any] | None" = None, + ) -> StorageTelemetry: + """Write Arrow record batches to storage asynchronously.""" + _validate_arrow_write_format(file_format) + resolved_options = _storage_options(self._storage_options, storage_options) + format_write_options = write_options or _csv_write_options( + file_format, resolved_options, self._storage_options, self._csv_write_options + ) + table = await async_(batches_to_arrow_table)(batches) + payload = await async_(_encode_arrow_payload)( + table, file_format, compression=compression, write_options=format_write_options + ) + return await self._write_bytes_async( + payload, destination, rows=int(table.num_rows), format_label=file_format, storage_options=resolved_options + ) + + async def list_sources_async( + self, + source: StorageDestination, + *, + file_format: str | None = None, + storage_options: "dict[str, Any] | None" = None, + ) -> "list[str]": + """Expand a single URI, directory prefix, or glob pattern into sorted source URIs asynchronously.""" + return await async_(_resolve_storage_sources_sync)( + self, source, file_format=file_format, storage_options=storage_options + ) + async def _write_bytes_async( self, payload: bytes, @@ -636,3 +751,47 @@ def _backend_cache_key(destination: StorageDestination, backend_options: "dict[s if backend_options: return None return destination.as_posix() if isinstance(destination, Path) else str(destination) + + +def _read_arrow_batches_sync( + pipeline: _StoragePipelineBase, + source: StorageDestination, + *, + file_format: str = "parquet", + batch_size: int = 65_536, + arrow_schema: Any = None, + storage_options: "dict[str, Any] | None" = None, +) -> "Iterator[ArrowRecordBatch]": + normalized_format = file_format.lower().strip() + backend, path, _backend_name = pipeline._backend(source, storage_options) + if normalized_format == "parquet": + return backend.stream_arrow_sync(path, file_format="parquet", batch_size=batch_size, arrow_schema=arrow_schema) + return iter_streamed_arrow_batches( + backend.stream_read_sync(path), file_format=normalized_format, batch_size=batch_size, arrow_schema=arrow_schema + ) + + +def _resolve_storage_sources_sync( + pipeline: _StoragePipelineBase, + source: StorageDestination, + *, + file_format: str | None = None, + storage_options: "dict[str, Any] | None" = None, +) -> "list[str]": + source_str = source.as_posix() if isinstance(source, Path) else str(source) + backend, path, _backend_name = pipeline._backend(source, storage_options) + has_glob = any(ch in path for ch in ("*", "?", "[")) + if has_glob: + matched = backend.glob_sync(path) + elif source_str.endswith("/") or backend.is_path_sync(path): + matched = backend.list_objects_sync(prefix=path, recursive=True) + if file_format: + suffix = f".{file_format.lower().lstrip('.')}" + matched = [item for item in matched if item.lower().endswith(suffix)] + else: + return [source_str] + sorted_matches = sorted(matched) + if source_str.startswith("alias://"): + alias_name = source_str.removeprefix("alias://").partition("/")[0] + return [f"alias://{alias_name}/{item.lstrip('/')}" for item in sorted_matches] + return [backend.resolve_uri(item) for item in sorted_matches] diff --git a/tests/unit/storage/test_pipeline_streaming.py b/tests/unit/storage/test_pipeline_streaming.py new file mode 100644 index 000000000..c1fa815d9 --- /dev/null +++ b/tests/unit/storage/test_pipeline_streaming.py @@ -0,0 +1,141 @@ +# pyright: reportPrivateUsage=false +"""Unit tests for streaming storage pipeline batches, vector Parquet defaults, schema-pinned JSONL/CSV/IPC, and ObStoreSink safety.""" + +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from sqlspec.core.result import ArrowResult +from sqlspec.storage._arrow_payload import encode_arrow_payload, resolve_parquet_use_dictionary +from sqlspec.storage.backends.obstore import _ObstoreSink +from sqlspec.storage.pipeline import AsyncStoragePipeline, SyncStoragePipeline + + +def test_resolve_parquet_use_dictionary_disables_float_vector_columns() -> None: + """resolve_parquet_use_dictionary disables dictionary encoding on float vector columns while preserving scalar columns.""" + schema = pa.schema([ + ("id", pa.int64()), + ("category", pa.string()), + ("embedding_f32", pa.list_(pa.float32(), 4)), + ("embedding_var", pa.list_(pa.float32())), + ("tags", pa.list_(pa.string())), + ]) + + resolved = resolve_parquet_use_dictionary(schema) + assert resolved == ["id", "category", "tags"] + + explicit = resolve_parquet_use_dictionary(schema, explicit_use_dictionary=False) + assert explicit is False + + scalar_only_schema = pa.schema([("id", pa.int64()), ("category", pa.string())]) + assert resolve_parquet_use_dictionary(scalar_only_schema) is True + + +def test_encode_arrow_payload_applies_vector_use_dictionary_and_write_options() -> None: + """encode_arrow_payload forwards write_options and disables dictionary encoding for float vector columns.""" + schema = pa.schema([("id", pa.int64()), ("category", pa.string()), ("embedding", pa.list_(pa.float32(), 2))]) + table = pa.Table.from_arrays( + [ + pa.array([1, 2, 3], type=pa.int64()), + pa.array(["a", "a", "b"], type=pa.string()), + pa.FixedSizeListArray.from_arrays(pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], type=pa.float32()), 2), + ], + schema=schema, + ) + + payload = encode_arrow_payload(table, "parquet", compression="zstd", write_options={"row_group_size": 2}) + pf = pq.ParquetFile(pa.BufferReader(payload)) + assert pf.num_row_groups == 2 + + +def test_obstore_sink_converts_mutable_bytearray_to_bytes() -> None: + """_ObstoreSink.write converts bytearray/memoryview to immutable bytes so callers can clear buffers safely.""" + retained: list[Any] = [] + + class _RetainingWriter: + def closed(self) -> bool: + return False + + def write(self, data: Any) -> int: + assert isinstance(data, bytes) + retained.append(memoryview(data)) + return len(data) + + def flush(self) -> None: + return None + + def close(self) -> None: + return None + + sink = _ObstoreSink(_RetainingWriter()) + buf = bytearray(b"hello world") + written = sink.write(buf) + buf.clear() + + assert written == 11 + assert bytes(retained[0]) == b"hello world" + + +def test_sync_pipeline_write_and_read_arrow_batches_parquet_and_jsonl(tmp_path: Path) -> None: + """SyncStoragePipeline writes and streams Parquet, JSONL (with arrow_schema float32 pinning), CSV, and IPC batches.""" + pipeline = SyncStoragePipeline() + parquet_path = tmp_path / "vectors.parquet" + schema = pa.schema([("id", pa.int64()), ("embedding", pa.list_(pa.float32(), 2))]) + batches = [ + pa.RecordBatch.from_arrays( + [ + pa.array([1, 2], type=pa.int64()), + pa.FixedSizeListArray.from_arrays(pa.array([1.0, 2.0, 3.0, 4.0], type=pa.float32()), 2), + ], + schema=schema, + ), + pa.RecordBatch.from_arrays( + [ + pa.array([3], type=pa.int64()), + pa.FixedSizeListArray.from_arrays(pa.array([5.0, 6.0], type=pa.float32()), 2), + ], + schema=schema, + ), + ] + + telemetry = pipeline.write_arrow_batches(parquet_path, iter(batches), file_format="parquet") + assert telemetry["rows_processed"] == 3 + assert parquet_path.exists() + + streamed = list(pipeline.read_arrow_batches(parquet_path, file_format="parquet", batch_size=2)) + assert sum(b.num_rows for b in streamed) == 3 + + jsonl_path = tmp_path / "vectors.jsonl" + jsonl_path.write_text('{"id": 1, "embedding": [0.25, 0.75]}\n{"id": 2, "embedding": [1.25, 1.75]}\n') + jsonl_schema = pa.schema([("id", pa.int64()), ("embedding", pa.list_(pa.float32()))]) + jsonl_batches = list( + pipeline.read_arrow_batches(jsonl_path, file_format="jsonl", arrow_schema=jsonl_schema, batch_size=10) + ) + assert sum(b.num_rows for b in jsonl_batches) == 2 + assert jsonl_batches[0].schema.field("embedding").type == pa.list_(pa.float32()) + assert np.allclose(jsonl_batches[0].column("embedding").to_pylist()[0], [0.25, 0.75]) + + +async def test_async_pipeline_write_and_read_arrow_batches_and_arrow_result(tmp_path: Path) -> None: + """AsyncStoragePipeline and ArrowResult stream batches without full-table materialization.""" + pipeline = AsyncStoragePipeline() + parquet_path = tmp_path / "async_items.parquet" + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + batch = pa.RecordBatch.from_pydict({"id": [10, 20], "name": ["x", "y"]}, schema=schema) + + telemetry = await pipeline.write_arrow_batches_async(parquet_path, [batch], file_format="parquet") + assert telemetry["rows_processed"] == 2 + + read_batches = [ + b async for b in pipeline.read_arrow_batches_async(parquet_path, file_format="parquet", batch_size=1) + ] + assert [b.num_rows for b in read_batches] == [1, 1] + + reader = pa.RecordBatchReader.from_batches(schema, [batch]) + arrow_result = ArrowResult(statement=None, data=reader) + out_path = tmp_path / "from_reader.parquet" + res_telemetry = await arrow_result.write_to_storage_async(out_path, pipeline=pipeline) + assert res_telemetry["rows_processed"] == 2 From 0278f69bbad990420f3723dd2fe221659a3142d2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 18:52:09 +0000 Subject: [PATCH 04/14] feat(driver): add parallel storage ingestion, staged bulk merge, without_indexes, and dynamic auth providers --- sqlspec/adapters/adbc/config.py | 2 + sqlspec/adapters/adbc/core.py | 12 + sqlspec/adapters/adbc/driver.py | 33 +- sqlspec/adapters/aiomysql/driver.py | 35 +- sqlspec/adapters/aiosqlite/driver.py | 33 +- sqlspec/adapters/arrow_odbc/driver.py | 35 +- sqlspec/adapters/asyncmy/driver.py | 35 +- sqlspec/adapters/asyncpg/config.py | 10 + sqlspec/adapters/asyncpg/driver.py | 66 +- sqlspec/adapters/bigquery/driver.py | 46 + sqlspec/adapters/cockroach_asyncpg/driver.py | 6 +- sqlspec/adapters/cockroach_psycopg/driver.py | 12 +- sqlspec/adapters/duckdb/driver.py | 35 +- sqlspec/adapters/mssql_python/driver.py | 57 +- sqlspec/adapters/mysqlconnector/driver.py | 66 +- sqlspec/adapters/oracledb/driver.py | 96 +- sqlspec/adapters/psqlpy/config.py | 13 +- sqlspec/adapters/psqlpy/core.py | 9 + sqlspec/adapters/psqlpy/driver.py | 41 +- sqlspec/adapters/psycopg/config.py | 68 +- sqlspec/adapters/psycopg/driver.py | 72 +- sqlspec/adapters/pymssql/driver.py | 46 + sqlspec/adapters/pymysql/driver.py | 35 +- sqlspec/adapters/spanner/driver.py | 106 +- sqlspec/adapters/sqlite/driver.py | 33 +- sqlspec/driver/_async.py | 133 +- sqlspec/driver/_common.py | 81 ++ sqlspec/driver/_storage_helpers.py | 1217 ++++++++++++++++- sqlspec/driver/_sync.py | 133 +- tests/unit/driver/test_staged_bulk_load.py | 197 +++ .../driver/test_storage_parallel_and_auth.py | 303 ++++ tests/unit/driver/test_without_indexes.py | 111 ++ 32 files changed, 3083 insertions(+), 94 deletions(-) create mode 100644 tests/unit/driver/test_staged_bulk_load.py create mode 100644 tests/unit/driver/test_storage_parallel_and_auth.py create mode 100644 tests/unit/driver/test_without_indexes.py diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index 75bd2798f..fae4f68b7 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -98,6 +98,8 @@ class AdbcConnectionParams(TypedDict): """Database authentication password lifted into ``db_kwargs`` or ``connect()`` as required by the driver.""" token: NotRequired[str] """Bearer, OAuth refresh, or Quack/FlightSQL authentication token.""" + auth_token_provider: NotRequired["Callable[[], str]"] + """Callable returning a fresh password or bearer token on each ``create_connection`` call.""" project_id: NotRequired[str] """Google Cloud project ID mapped to ``adbc.bigquery.sql.project_id`` for ``adbc-driver-bigquery``.""" dataset_id: NotRequired[str] diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index ee56694f5..bce37c48f 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -729,6 +729,18 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str if not uses_driver_manager: config.pop("entrypoint", None) + auth_token_provider = config.pop("auth_token_provider", None) + if callable(auth_token_provider): + fresh_token = str(auth_token_provider()) + if driver_kind == "quack" and "username" not in config and "user" not in config: + config["token"] = fresh_token + elif driver_kind in {"gizmosql", "flightsql"} and "username" not in config and "user" not in config: + config["authorization_header"] = ( + fresh_token if fresh_token.lower().startswith("bearer ") else f"Bearer {fresh_token}" + ) + else: + config["password"] = fresh_token + _lift_driver_specific_db_kwargs(config, driver_kind, resolved_connect_path) config.pop("driver_name", None) diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 533168b5b..15be49f26 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -55,7 +55,7 @@ from sqlspec.utils.text import normalize_identifier, quote_identifier, split_qualified_identifier if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping, Sequence from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcRawCursor from sqlspec.builder import QueryBuilder @@ -715,10 +715,35 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Ingest an Arrow payload directly through the ADBC cursor.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) ensure_pyarrow() source_data: Any = source.get_data() if hasattr(source, "get_data") else source @@ -773,11 +798,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Read an artifact from storage and ingest it via ADBC.""" - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: "AdbcRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect ADBC rows for the direct execution path.""" diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index 4ac938fbe..f6299971d 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -12,7 +12,7 @@ import tempfile from collections.abc import Sized from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.aiomysql._typing import ( AIOMYSQL_INSERT_VALUES_PATTERN, @@ -50,7 +50,7 @@ from sqlspec.utils.type_guards import supports_json_type if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping, Sequence from sqlspec.adapters.aiomysql._typing import AiomysqlConnection from sqlspec.core import SQL, SQLResult, StatementConfig @@ -346,10 +346,35 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into MySQL using batched inserts.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" @@ -405,12 +430,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into MySQL.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 718738432..09a8d243f 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -41,7 +41,7 @@ from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection from sqlspec.builder import QueryBuilder @@ -312,12 +312,37 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into SQLite using batched inserts.""" if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 1: msg = "batch_size must be a positive integer" raise ValueError(msg) self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) columns = arrow_table.column_names insert_sql = build_insert_statement(table, columns) @@ -364,12 +389,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into SQLite.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index e7e3f323b..fa25a1361 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -2,10 +2,10 @@ import contextlib import re -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from datetime import UTC, datetime from itertools import chain -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlglot import exp @@ -477,9 +477,34 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into a table via arrow-odbc bulk insert.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) ensure_pyarrow() raw_source = source.get_data() if hasattr(source, "get_data") and callable(source.get_data) else source if overwrite: @@ -509,10 +534,12 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into a table via arrow-odbc bulk insert.""" - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def _connection_in_transaction(self) -> bool: """Return whether SQLSpec holds an explicit transaction on the connection. diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index cec7ae28d..ca983e94f 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -6,7 +6,7 @@ import tempfile from collections.abc import Sized -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.asyncmy._typing import ( ASYNCMY_INSERT_VALUES_PATTERN, @@ -45,7 +45,7 @@ from sqlspec.utils.type_guards import supports_json_type if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping, Sequence from sqlspec.adapters.asyncmy._typing import AsyncmyConnection from sqlspec.core import SQL, SQLResult, StatementConfig @@ -350,10 +350,35 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data using batched inserts or opt-in native LOCAL INFILE.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" @@ -397,12 +422,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into MySQL.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 50764a004..394954974 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -79,6 +79,7 @@ }) _CLOUD_SQL_INSTANCE_PARTS: Final[int] = 3 _ALLOYDB_URI_PARTS: Final[int] = 8 +_AUTH_TOKEN_DEFAULT_INACTIVE_LIFETIME: Final[float] = 1800.0 logger = get_logger(__name__) @@ -100,6 +101,8 @@ class AsyncpgConnectionConfig(TypedDict): """PostgreSQL role name used for authentication.""" password: NotRequired[str] """Password or dynamic token for ``user``.""" + auth_token_provider: NotRequired["Callable[[], str] | Callable[[], Awaitable[str]]"] + """Callable or async callable returning a fresh password or IAM token for each new connection.""" database: NotRequired[str] """Target PostgreSQL database name.""" ssl: NotRequired[Any] @@ -403,6 +406,9 @@ async def create_connection(self) -> "AsyncpgConnection": config = { key: value for key, value in build_connection_config(self.connection_config).items() if value is not None } + auth_token_provider = config.pop("auth_token_provider", None) + if callable(auth_token_provider): + config["password"] = auth_token_provider for key in _POOL_ONLY_CONFIG_KEYS: config.pop(key, None) @@ -633,6 +639,10 @@ async def _create_pool(self) -> "Pool[Record]": config = { key: value for key, value in build_connection_config(self.connection_config).items() if value is not None } + auth_token_provider = config.pop("auth_token_provider", None) + if callable(auth_token_provider): + config["password"] = auth_token_provider + config.setdefault("max_inactive_connection_lifetime", _AUTH_TOKEN_DEFAULT_INACTIVE_LIFETIME) if self.connection_config.get("pgbouncer") or self.driver_features.get("pgbouncer"): config["statement_cache_size"] = 0 diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 81f000fb4..7871db0a9 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -5,7 +5,7 @@ from contextlib import suppress from io import BytesIO from itertools import chain -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlglot import exp, parse_one from sqlglot.errors import ParseError @@ -353,9 +353,34 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into a PostgreSQL table via COPY.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) table_name, schema_name, quoted_target = self._copy_target(table) if overwrite: try: @@ -372,11 +397,7 @@ async def load_from_arrow( if use_binary: async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) await self.connection.copy_to_table( - table_name, - source=async_chunks, - columns=columns, - schema_name=schema_name, - format="binary", + table_name, source=async_chunks, columns=columns, schema_name=schema_name, format="binary" ) else: for batch in all_batches: @@ -387,9 +408,7 @@ async def load_from_arrow( await self.connection.copy_records_to_table( table_name, records=records, columns=columns, schema_name=schema_name ) - telemetry_payload = self._ingest_telemetry_from_batches( - metrics["rows_processed"], metrics["bytes_processed"] - ) + telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) @@ -401,8 +420,31 @@ async def load_from_records( *, columns: "list[str] | None" = None, overwrite: bool = False, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load mapping or positional records directly with binary COPY.""" + if stage_table or conflict_keys or drop_indexes: + return await super().load_from_records( + table, + records, + columns=columns, + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) self._require_capability("arrow_import_enabled") materialized = list(records) if not materialized: @@ -464,12 +506,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Read an artifact from storage and ingest it via COPY.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def _timeout_arguments(self, config: "StatementConfig") -> "dict[str, Any]": diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index b715bdaf9..0ecc6df65 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -705,10 +705,51 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into BigQuery (Parquet load job, or opt-in Storage Write API).""" self._require_capability("parquet_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) ensure_pyarrow() qualified_table = self._qualify_table_name(table) @@ -753,8 +794,13 @@ def load_from_storage( file_format: "BigQueryLoadFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into BigQuery.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) job_config = build_load_job_config(file_format, overwrite) qualified_table = self._qualify_table_name(table) diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 1b55c7b1b..51c9bc890 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -137,6 +137,7 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Append via native IMPORT when eligible, taking the table offline. @@ -147,7 +148,8 @@ async def load_from_storage( options = self.driver_features.get("native_storage_csv_options", {}) resolved = None if ( - self._native_storage_ready() + not kwargs + and self._native_storage_ready() and not overwrite and (file_format == "parquet" or (file_format == "csv" and "skip" in options)) ): @@ -159,7 +161,7 @@ async def load_from_storage( self._attach_partition_telemetry(produced, partitioner) return self._storage_job(produced) return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) async def begin(self) -> None: diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index 679b83121..cbc4d9afd 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -164,6 +164,7 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Append via native IMPORT when eligible, taking the table offline. @@ -174,7 +175,8 @@ def load_from_storage( options = self.driver_features.get("native_storage_csv_options", {}) resolved = None if ( - self._native_storage_ready() + not kwargs + and self._native_storage_ready() and not overwrite and (file_format == "parquet" or (file_format == "csv" and "skip" in options)) ): @@ -186,7 +188,7 @@ def load_from_storage( self._attach_partition_telemetry(produced, partitioner) return self._storage_job(produced) return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def begin(self) -> None: @@ -402,6 +404,7 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Append via native IMPORT when eligible, taking the table offline. @@ -412,7 +415,8 @@ async def load_from_storage( options = self.driver_features.get("native_storage_csv_options", {}) resolved = None if ( - self._native_storage_ready() + not kwargs + and self._native_storage_ready() and not overwrite and (file_format == "parquet" or (file_format == "csv" and "skip" in options)) ): @@ -424,7 +428,7 @@ async def load_from_storage( self._attach_partition_telemetry(produced, partitioner) return self._storage_job(produced) return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) async def begin(self) -> None: diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index c35fdfb9d..25ed550f2 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -49,7 +49,7 @@ from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence import pyarrow as pa @@ -513,10 +513,35 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into DuckDB using temporary table registration.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) ensure_pyarrow() source_data = source.get_data() if hasattr(source, "get_data") else source @@ -577,11 +602,12 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Read an artifact from storage and load it into DuckDB.""" self._require_capability("arrow_import_enabled") - if not overwrite and file_format == "parquet" and self.storage_pipeline_factory is None: + if not kwargs and not overwrite and file_format == "parquet" and self.storage_pipeline_factory is None: pipeline = self._storage_pipeline() target = _resolve_native_storage_target(pipeline, source, self.driver_features, write=False) if target is not None: @@ -595,8 +621,9 @@ def load_from_storage( operation="read", partitioner=partitioner, ) - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: "DuckDBConnection", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect DuckDB rows for the direct execution path.""" diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index d38380ac4..d996e0cd4 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -42,7 +42,7 @@ from sqlspec.utils.module_loader import ensure_pyarrow, import_optional if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Sequence from sqlspec.builder import QueryBuilder from sqlspec.core import SQL, ArrowResult, SQLResult, Statement, StatementConfig, StatementFilter @@ -436,6 +436,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, batch_size: int = 0, timeout: int = 30, table_lock: bool = False, @@ -453,6 +461,48 @@ def load_from_arrow( while consuming a stream can still leave a partially completed load. """ self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + batch_size=batch_size, + timeout=timeout, + table_lock=table_lock, + check_constraints=check_constraints, + fire_triggers=fire_triggers, + keep_identity=keep_identity, + keep_nulls=keep_nulls, + use_internal_transaction=use_internal_transaction, + column_mappings=column_mappings, + ) + if stage_table or conflict_keys: + return super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) ensure_pyarrow() pa = cast("Any", import_optional("pyarrow")) @@ -540,8 +590,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into SQL Server via BulkCopy.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index a005f3de7..b19cb450f 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -7,9 +7,9 @@ import asyncio import os import tempfile -from collections.abc import Callable, Sized +from collections.abc import Callable, Mapping, Sequence, Sized from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.mysqlconnector._typing import ( MysqlConnectorAsyncCursor, @@ -271,8 +271,33 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" @@ -330,9 +355,11 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector sync rows for the direct execution path.""" @@ -526,8 +553,33 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" @@ -585,10 +637,10 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 8973359b1..788ac1ad0 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -833,9 +833,50 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Oracle using batched executemany calls.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = build_truncate_statement(table) @@ -900,6 +941,7 @@ def load_from_storage( overwrite: bool = False, credential_name: str | None = None, format_options: "Mapping[str, Any] | None" = None, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts into Oracle via ``DBMS_CLOUD.COPY_DATA`` or client-side Arrow.""" raw_credential = ( @@ -908,7 +950,7 @@ def load_from_storage( effective_credential = ( raw_credential.strip() if isinstance(raw_credential, str) and raw_credential.strip() else None ) - if effective_credential is not None: + if effective_credential is not None and not kwargs: resolved_source: StorageDestination = ( self._storage_pipeline().resolve_destination(source).uri if isinstance(source, str) and source.startswith("alias://") @@ -942,6 +984,10 @@ def load_from_storage( } self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload) + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) @@ -1610,9 +1656,50 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Asynchronously load Arrow data into Oracle.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + async with self.without_indexes(table, concurrent=concurrent_indexes): + return await self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return await super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = build_truncate_statement(table) @@ -1673,6 +1760,7 @@ async def load_from_storage( overwrite: bool = False, credential_name: str | None = None, format_options: "Mapping[str, Any] | None" = None, + **kwargs: Any, ) -> "StorageBridgeJob": """Asynchronously load staged artifacts into Oracle via ``DBMS_CLOUD.COPY_DATA`` or client-side Arrow.""" raw_credential = ( @@ -1681,7 +1769,7 @@ async def load_from_storage( effective_credential = ( raw_credential.strip() if isinstance(raw_credential, str) and raw_credential.strip() else None ) - if effective_credential is not None: + if effective_credential is not None and not kwargs: resolved_source: StorageDestination = ( self._storage_pipeline().resolve_destination(source).uri if isinstance(source, str) and source.startswith("alias://") @@ -1715,6 +1803,10 @@ async def load_from_storage( } self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload) + if kwargs: + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) return await self.load_from_arrow( table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index ff758887d..540f093bb 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -1,5 +1,6 @@ """Psqlpy database configuration.""" +import inspect import sys from collections.abc import Mapping from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast @@ -43,6 +44,8 @@ class PsqlpyConnectionParams(TypedDict): """PostgreSQL role name used for authentication.""" password: NotRequired[str] """Password for ``username``.""" + auth_token_provider: NotRequired["Callable[[], str] | Callable[[], Awaitable[str]]"] + """Callable or async callable returning a fresh password or IAM token when creating connections.""" db_name: NotRequired[str] """Target PostgreSQL database name.""" host: NotRequired[str] @@ -466,7 +469,15 @@ def _ensure_pg_textsearch_available(self) -> None: async def _create_pool(self) -> "ConnectionPool": """Create the actual async connection pool.""" - return ConnectionPool(**build_connection_config(self.connection_config)) + raw_cfg = dict(self.connection_config) + auth_token_provider = raw_cfg.pop("auth_token_provider", None) + pool_cfg = build_connection_config(raw_cfg) + if callable(auth_token_provider): + token_result = auth_token_provider() + if inspect.isawaitable(token_result): + token_result = await token_result + pool_cfg["password"] = str(token_result) + return ConnectionPool(**pool_cfg) async def _close_pool(self) -> None: """Close the actual async connection pool.""" diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index efd31dc3a..0101439a4 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -3,6 +3,7 @@ import contextlib import datetime as dt import decimal +import inspect import re import uuid from functools import lru_cache @@ -168,6 +169,14 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str username = config.pop("username", None) or config.pop("user", None) if username is not None: config["username"] = username + auth_token_provider = config.get("auth_token_provider") + if callable(auth_token_provider) and not inspect.iscoroutinefunction(auth_token_provider): + config.pop("auth_token_provider", None) + token_result = auth_token_provider() + if not inspect.isawaitable(token_result): + config["password"] = str(token_result) + elif auth_token_provider is None: + config.pop("auth_token_provider", None) return config diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 1f70b7071..8abc1f355 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -5,7 +5,7 @@ """ from itertools import chain -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.psqlpy._typing import ( PsqlpyConnectionExecuteError, @@ -340,10 +340,35 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow-formatted data into PostgreSQL via psqlpy binary COPY.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) if overwrite: qualified = format_table_identifier(table) exc_handler = self.handle_database_exceptions() @@ -374,7 +399,9 @@ async def load_from_arrow( await cursor.binary_copy_to_table(payload, table_name, **copy_kwargs) else: json_columns = await self._resolve_json_columns(schema_name, table_name) - deserializer = cast("Callable[[str], Any]", self.driver_features.get("json_deserializer", from_json)) + deserializer = cast( + "Callable[[str], Any]", self.driver_features.get("json_deserializer", from_json) + ) for batch in all_batches: metrics["rows_processed"] += int(batch.num_rows) metrics["bytes_processed"] += int(batch.nbytes) @@ -388,9 +415,7 @@ async def load_from_arrow( if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry_from_batches( - metrics["rows_processed"], metrics["bytes_processed"] - ) + telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) @@ -403,12 +428,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage using the storage bridge pipeline.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def prepare_driver_parameters( diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 42b829044..80813c2cd 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -1,5 +1,6 @@ """Psycopg database configuration with direct field-based configuration.""" +import inspect from collections.abc import Mapping from contextlib import suppress from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast @@ -82,6 +83,7 @@ PsycopgSSLMode = Literal["disable", "allow", "prefer", "require", "verify-ca", "verify-full"] _ALLOYDB_DIRECT_CONNECTION_KEYS = ("conninfo", "host", "hostaddr", "port", "user", "password", "dbname", "db") _POOL_ONLY_CONFIG_KEYS: Final[frozenset[str]] = frozenset({ + "auth_token_provider", "check", "close_returns", "configure", @@ -101,6 +103,7 @@ }) _ALLOYDB_CONNECTOR_PACKAGE = "google-cloud-alloydb-connector" _DEFAULT_PSYCOPG_PREPARE_THRESHOLD: Final[int] = 5 +_AUTH_TOKEN_DEFAULT_MAX_LIFETIME: Final[float] = 1800.0 class PsycopgConnectionParams(TypedDict): @@ -116,6 +119,8 @@ class PsycopgConnectionParams(TypedDict): """PostgreSQL role name used for authentication.""" password: NotRequired[str] """Password or token for ``user``.""" + auth_token_provider: NotRequired["Callable[[], str] | Callable[[], Awaitable[str]]"] + """Callable returning a fresh authentication password or IAM token on each connection attempt.""" dbname: NotRequired[str] """Target PostgreSQL database name.""" connect_timeout: NotRequired[int] @@ -555,6 +560,8 @@ def _create_pool(self) -> "ConnectionPool": """Create the actual connection pool.""" all_config = dict(self.connection_config) all_config.pop("null_pool", None) + auth_token_provider = all_config.pop("auth_token_provider", None) + default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 pool_parameters = { "connection_class": all_config.pop("connection_class", None), @@ -563,7 +570,7 @@ def _create_pool(self) -> "ConnectionPool": "name": all_config.pop("name", None), "timeout": all_config.pop("timeout", 30.0), "max_waiting": all_config.pop("max_waiting", 0), - "max_lifetime": all_config.pop("max_lifetime", 3600.0), + "max_lifetime": all_config.pop("max_lifetime", default_max_lifetime), "max_idle": all_config.pop("max_idle", 600.0), "reconnect_timeout": all_config.pop("reconnect_timeout", 300.0), "reconnect_failed": all_config.pop("reconnect_failed", None), @@ -574,6 +581,14 @@ def _create_pool(self) -> "ConnectionPool": "close_returns": all_config.pop("close_returns", False), } + if callable(auth_token_provider) and not self.driver_features.get("enable_alloydb", False): + base_conn_cls = cast( + "type[PsycopgConnection[Any]]", pool_parameters.get("connection_class") or PsycopgConnection + ) + pool_parameters["connection_class"] = _make_token_provider_sync_connection_class( + base_conn_cls, cast("Callable[[], str]", auth_token_provider) + ) + user_configure = all_config.pop("configure", None) if user_configure is not None and user_configure is not self._configure_connection: @@ -678,6 +693,9 @@ def _connection_kwargs(self) -> "tuple[type[Any], str, dict[str, Any]]": The connection class, the conninfo string, and the connect kwargs. """ all_config = {k: v for k, v in self.connection_config.items() if k not in _POOL_ONLY_CONFIG_KEYS} + auth_token_provider = self.connection_config.get("auth_token_provider") + if callable(auth_token_provider): + all_config["password"] = str(auth_token_provider()) connection_class = all_config.pop("connection_class", None) or PsycopgConnection conninfo = all_config.pop("conninfo", None) all_config.update(all_config.pop("kwargs", {})) @@ -873,6 +891,12 @@ async def create_connection(self) -> "PsycopgAsyncConnection": # pyright: ignor A psycopg AsyncConnection instance. """ connection_class, conninfo, connection_kwargs = self._connection_kwargs() + auth_token_provider = self.connection_config.get("auth_token_provider") + if callable(auth_token_provider): + token_result = auth_token_provider() + if inspect.isawaitable(token_result): + token_result = await token_result + connection_kwargs["password"] = str(token_result) connection = await connection_class.connect(conninfo, **connection_kwargs) user_configure = self.connection_config.get("configure") try: @@ -948,6 +972,8 @@ async def _create_pool(self) -> "AsyncConnectionPool": all_config = dict(self.connection_config) all_config.pop("null_pool", None) + auth_token_provider = all_config.pop("auth_token_provider", None) + default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 pool_parameters = { "connection_class": all_config.pop("connection_class", None), @@ -956,7 +982,7 @@ async def _create_pool(self) -> "AsyncConnectionPool": "name": all_config.pop("name", None), "timeout": all_config.pop("timeout", 30.0), "max_waiting": all_config.pop("max_waiting", 0), - "max_lifetime": all_config.pop("max_lifetime", 3600.0), + "max_lifetime": all_config.pop("max_lifetime", default_max_lifetime), "max_idle": all_config.pop("max_idle", 600.0), "reconnect_timeout": all_config.pop("reconnect_timeout", 300.0), "reconnect_failed": all_config.pop("reconnect_failed", None), @@ -965,6 +991,13 @@ async def _create_pool(self) -> "AsyncConnectionPool": "reset": all_config.pop("reset", None), "close_returns": all_config.pop("close_returns", False), } + if callable(auth_token_provider): + base_async_cls = cast( + "type[PsycopgAsyncConnection]", pool_parameters.get("connection_class") or PsycopgAsyncConnection + ) + pool_parameters["connection_class"] = _make_token_provider_async_connection_class( + base_async_cls, cast("Callable[[], str | Awaitable[str]]", auth_token_provider) + ) open_pool = all_config.pop("open", True) pool_parameters["open"] = False if open_pool is True else open_pool @@ -1124,3 +1157,34 @@ def connect( return cast("Self", conn) return _AlloyDBPsycopgConnection + + +def _make_token_provider_sync_connection_class( + base_cls: "type[PsycopgConnection[Any]]", provider: "Callable[[], str]" +) -> "type[PsycopgConnection[Any]]": + """Build a Psycopg sync connection subclass that refreshes ``password`` from ``provider``.""" + + class _TokenProviderPsycopgConnection(base_cls): # type: ignore[valid-type,misc] + @classmethod + def connect(cls, conninfo: str = "", **kwargs: Any) -> Self: + kwargs["password"] = str(provider()) + return cast("Self", super().connect(conninfo, **kwargs)) + + return _TokenProviderPsycopgConnection + + +def _make_token_provider_async_connection_class( + base_cls: "type[PsycopgAsyncConnection]", provider: "Callable[[], str | Awaitable[str]]" +) -> "type[PsycopgAsyncConnection]": + """Build a Psycopg async connection subclass that refreshes ``password`` from ``provider``.""" + + class _TokenProviderPsycopgAsyncConnection(base_cls): # type: ignore[valid-type,misc] + @classmethod + async def connect(cls, conninfo: str = "", **kwargs: Any) -> Self: + token_result = provider() + if inspect.isawaitable(token_result): + token_result = await token_result + kwargs["password"] = str(token_result) + return cast("Self", await super().connect(conninfo, **kwargs)) + + return _TokenProviderPsycopgAsyncConnection diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 5e1f6eebf..10111ac57 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -3,7 +3,7 @@ from collections.abc import Sized from contextlib import AsyncExitStack, ExitStack from itertools import chain -from typing import TYPE_CHECKING, Any, LiteralString, cast +from typing import TYPE_CHECKING, Any, Literal, LiteralString, cast from sqlspec.adapters.psycopg._typing import ( PsycopgAsyncConnection, @@ -519,10 +519,35 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "abc.Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "abc.Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "abc.Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into PostgreSQL using COPY.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() @@ -554,9 +579,7 @@ def load_from_arrow( copy_ctx.write_row(record) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry_from_batches( - metrics["rows_processed"], metrics["bytes_processed"] - ) + telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) @@ -569,11 +592,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts into PostgreSQL via COPY.""" - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg sync rows for the direct execution path.""" @@ -1060,10 +1085,35 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "abc.Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "abc.Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "abc.Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into PostgreSQL asynchronously via COPY.""" self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return await self._execute_staged_load_async( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() @@ -1095,9 +1145,7 @@ async def load_from_arrow( await copy_ctx.write_row(record) if exc_handler.pending_exception is not None: raise exc_handler.pending_exception from None - telemetry_payload = self._ingest_telemetry_from_batches( - metrics["rows_processed"], metrics["bytes_processed"] - ) + telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) @@ -1110,12 +1158,12 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts asynchronously.""" - arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) - return await self.load_from_arrow( - table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index 4fe7d29da..9f4bb4a29 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -347,9 +347,50 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target table via executemany.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"DELETE FROM {format_identifier(table)}" @@ -388,8 +429,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into the target SQL Server table.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 25cfa9e22..47f15d481 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -4,7 +4,7 @@ import tempfile from collections.abc import Sized from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.pymysql._typing import ( PYMYSQL_INSERT_VALUES_PATTERN, @@ -41,7 +41,7 @@ from sqlspec.utils.type_guards import supports_json_type if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping, Sequence from sqlspec.adapters.pymysql._typing import PyMysqlConnection from sqlspec.core import SQL, SQLResult, StatementConfig @@ -253,8 +253,33 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" @@ -310,9 +335,11 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect PyMySQL rows for the direct execution path.""" diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index bf01ec38a..11a3cf907 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -525,9 +525,50 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Spanner table via batch mutations.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table: + return super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) exc_handler = self.handle_database_exceptions() @@ -555,7 +596,10 @@ def load_from_arrow( raise SQLConversionError(msg) writer = cast("_SpannerWriteProtocol", conn) for chunk in chunks: - writer.insert_or_update(table, columns, chunk) + if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): + writer.insert_or_ignore(table, columns, chunk) + else: + writer.insert_or_update(table, columns, chunk) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -571,8 +615,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load artifacts from storage into Spanner table.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) @@ -1137,9 +1186,50 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: bool | str = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: str = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "dict[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Spanner table via batch mutations.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + async with self.without_indexes(table, concurrent=concurrent_indexes): + return await self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table: + return await super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) exc_handler = self.handle_database_exceptions() @@ -1167,7 +1257,10 @@ async def load_from_arrow( raise SQLConversionError(msg) writer = cast("_SpannerAsyncWriteProtocol", conn) for chunk in chunks: - writer.insert_or_update(table, columns, chunk) + if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): + writer.insert_or_ignore(table, columns, chunk) + else: + writer.insert_or_update(table, columns, chunk) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -1183,8 +1276,13 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load artifacts from storage into Spanner table asynchronously.""" + if kwargs: + return await super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = await self._read_storage_arrow(source, file_format=file_format) return await self.load_from_arrow( table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound @@ -1384,6 +1482,8 @@ def batch_update( def insert(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... + def insert_or_ignore(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... + def insert_or_update(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... def update(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... @@ -1428,6 +1528,8 @@ async def batch_update( def insert(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... + def insert_or_ignore(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... + def insert_or_update(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... def update(self, table: str, columns: "list[str]", values: "list[list[Any]]") -> None: ... diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index b8278a4b4..87a9703dd 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -38,7 +38,7 @@ from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from sqlglot.dialects.dialect import DialectType @@ -359,12 +359,37 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into SQLite using chunked batched inserts.""" if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 1: msg = "batch_size must be a positive integer" raise ValueError(msg) self._require_capability("arrow_import_enabled") + if stage_table or conflict_keys or drop_indexes: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) columns = arrow_table.column_names insert_sql = build_insert_statement(table, columns) @@ -415,11 +440,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into SQLite.""" - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect SQLite rows for the direct execution path.""" diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 02a3e291c..f7fb32d4d 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -49,7 +49,12 @@ from sqlspec.driver._query_cache import CachedQuery from sqlspec.driver._sql_helpers import DEFAULT_PRETTY from sqlspec.driver._sql_helpers import convert_to_dialect as _convert_to_dialect_impl -from sqlspec.driver._storage_helpers import stringify_storage_target +from sqlspec.driver._storage_helpers import ( + AsyncWithoutIndexesContext, + execute_load_from_storage_async, + execute_staged_load_async, + stringify_storage_target, +) from sqlspec.driver._stream import AsyncRowStream, _LazyEagerAsyncRowSource from sqlspec.exceptions import ImproperConfigurationError, StackExecutionError from sqlspec.observability import _runtime as observability_runtime @@ -1384,6 +1389,12 @@ async def select_to_storage( """ self._raise_storage_not_implemented("select_to_storage") + def without_indexes( + self, table: str, *, schema: str | None = None, concurrent: bool = False + ) -> AsyncWithoutIndexesContext: + """Drop secondary indexes on ``table`` before bulk loading and rebuild them on exit.""" + return AsyncWithoutIndexesContext(self, table, schema=schema, concurrent=concurrent) + async def load_from_arrow( self, table: str, @@ -1391,6 +1402,15 @@ async def load_from_arrow( *, partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target table. @@ -1399,6 +1419,15 @@ async def load_from_arrow( source: Arrow data source. partitioner: Optional partitioner configuration. overwrite: Whether to overwrite existing data. + telemetry: Optional inbound storage telemetry to merge. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. @@ -1406,6 +1435,21 @@ async def load_from_arrow( Raises: NotImplementedError: If not implemented. """ + _ = ( + table, + source, + partitioner, + overwrite, + telemetry, + stage_table, + conflict_keys, + on_conflict, + exclude_update_columns, + drop_stage_table, + drop_indexes, + concurrent_indexes, + type_casts, + ) self._raise_storage_not_implemented("load_from_arrow") async def load_from_storage( @@ -1416,20 +1460,71 @@ async def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + storage_options: "dict[str, Any] | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + checkpoint_table: str | None = None, + checkpoint_phase: str = "default", + max_workers: int = 1, + executor: Literal["thread", "process"] = "thread", + arrow_schema: Any = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load artifacts from storage into the target table. Args: table: Target table name. - source: Storage source path. + source: Storage source path, prefix, or glob pattern. file_format: File format of source. partitioner: Optional partitioner configuration. overwrite: Whether to overwrite existing data. + storage_options: Optional storage backend configuration. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + checkpoint_table: Optional database table name for shard-level checkpoint tracking. + checkpoint_phase: Checkpoint phase identifier (default ``"default"``). + max_workers: Number of parallel workers for multi-shard storage fan-out. + executor: Worker pool type (``"thread"`` or ``"process"``). + arrow_schema: Optional PyArrow schema for explicit type pinning on decode. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. """ - self._raise_storage_not_implemented("load_from_storage") + if type(self).load_from_arrow is AsyncDriverAdapterBase.load_from_arrow: + self._raise_storage_not_implemented("load_from_storage") + return await execute_load_from_storage_async( + self, + table, + source, + file_format=file_format, + partitioner=partitioner, + overwrite=overwrite, + storage_options=storage_options, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + checkpoint_table=checkpoint_table, + checkpoint_phase=checkpoint_phase, + max_workers=max_workers, + executor=executor, + arrow_schema=arrow_schema, + type_casts=type_casts, + ) async def load_from_records( self, @@ -1438,6 +1533,14 @@ async def load_from_records( *, columns: "list[str] | None" = None, overwrite: bool = False, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load in-memory dict or positional records into the target table. @@ -1451,12 +1554,36 @@ async def load_from_records( records: Mapping records, or positional sequences with ``columns``. columns: Column names (required for positional records). overwrite: Whether to overwrite existing data. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. """ prepared_records = self._prepare_records_for_arrow(records) arrow_table = self._records_to_arrow_table(prepared_records, columns) + if stage_table or conflict_keys or drop_indexes: + return await execute_staged_load_async( + self, + table, + arrow_table, + raw_loader=self.load_from_arrow, + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) return await self.load_from_arrow(table, arrow_table, overwrite=overwrite) def convert_to_dialect( diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 1f534b770..0bb0ea5db 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -86,6 +86,10 @@ coerce_arrow_table, create_storage_job, encode_arrow_batch_pg_binary, + execute_load_from_storage_async, + execute_load_from_storage_sync, + execute_staged_load_async, + execute_staged_load_sync, ingest_telemetry_from_batches, iter_pgcopy_binary_chunks, iter_pgcopy_binary_chunks_async, @@ -3287,6 +3291,83 @@ def _storage_job( """Create a StorageBridgeJob from telemetry data.""" return create_storage_job(produced, provided, status=status) + def _execute_staged_load_sync( + self, + table: str, + source: "ArrowResult | Any", + *, + raw_loader: "abc.Callable[..., StorageBridgeJob]", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "abc.Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "abc.Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "abc.Mapping[str, str] | None" = None, + ) -> "StorageBridgeJob": + """Execute a staged or index-free bulk load synchronously.""" + return execute_staged_load_sync( + self, + table, + source, + raw_loader=raw_loader, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + + async def _execute_staged_load_async( + self, + table: str, + source: "ArrowResult | Any", + *, + raw_loader: "abc.Callable[..., Any]", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "abc.Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "abc.Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "abc.Mapping[str, str] | None" = None, + ) -> "StorageBridgeJob": + """Execute a staged or index-free bulk load asynchronously.""" + return await execute_staged_load_async( + self, + table, + source, + raw_loader=raw_loader, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + + _execute_load_from_storage_sync = staticmethod(execute_load_from_storage_sync) + _execute_load_from_storage_async = staticmethod(execute_load_from_storage_async) + def parameter_values_need_processing( values: "tuple[Any, ...] | list[Any]", diff --git a/sqlspec/driver/_storage_helpers.py b/sqlspec/driver/_storage_helpers.py index 179d86f0b..23d705b6e 100644 --- a/sqlspec/driver/_storage_helpers.py +++ b/sqlspec/driver/_storage_helpers.py @@ -1,9 +1,18 @@ """Pure helper functions for storage operations.""" +import asyncio +import contextlib +import multiprocessing +import os +import re +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from time import perf_counter +from typing import TYPE_CHECKING, Any, Final, Literal, Self, cast -from sqlspec.storage import StorageBridgeJob, StorageTelemetry, create_storage_bridge_job +from sqlspec.storage import StorageBridgeJob, StorageTelemetry, SyncStoragePipeline, create_storage_bridge_job +from sqlspec.storage._arrow_payload import decode_arrow_payload, encode_arrow_payload from sqlspec.utils.arrow_helpers import ( PG_BINARY_COPY_HEADER, PG_BINARY_COPY_TRAILER, @@ -18,28 +27,49 @@ iter_pgcopy_binary_chunks_async, records_to_arrow_table, ) +from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from sqlspec.storage import StorageDestination + from types import TracebackType + + from sqlspec.storage import StorageDestination, StorageFormat __all__ = ( "CAPABILITY_HINTS", "PG_BINARY_COPY_HEADER", "PG_BINARY_COPY_TRAILER", + "AsyncWithoutIndexesContext", + "SyncWithoutIndexesContext", "arrow_batch_to_rows", "arrow_table_needs_parameter_preparation", "arrow_table_to_rows", "attach_partition_telemetry", + "build_checkpoint_insert_sql", + "build_checkpoint_select_sql", + "build_checkpoint_table_ddl", + "build_create_index_sql", + "build_drop_index_sql", + "build_drop_stage_table_sql", "build_ingest_telemetry", + "build_stage_merge_sql", + "build_stage_table_ddl", "coerce_arrow_batches", "coerce_arrow_table", + "compute_worker_thread_budget", "create_storage_job", "encode_arrow_batch_pg_binary", + "execute_load_from_storage_async", + "execute_load_from_storage_sync", + "execute_staged_load_async", + "execute_staged_load_sync", + "filter_droppable_table_indexes", "ingest_telemetry_from_batches", + "init_spawn_worker_thread_budget", "iter_pgcopy_binary_chunks", "iter_pgcopy_binary_chunks_async", "records_to_arrow_table", + "resolve_stage_table_name", "stringify_storage_target", ) @@ -160,3 +190,1184 @@ def create_storage_job( extra["source"] = provided merged["extra"] = extra return create_storage_bridge_job(status, merged) + + +_PG_FAMILY_DIALECTS: Final[frozenset[str]] = frozenset({ + "postgres", + "postgresql", + "psycopg", + "asyncpg", + "psqlpy", + "cockroach", + "cockroachdb", + "spangres", +}) +_ON_CONFLICT_DIALECTS: Final[frozenset[str]] = _PG_FAMILY_DIALECTS | frozenset({"duckdb", "sqlite", "aiosqlite"}) +_MYSQL_FAMILY_DIALECTS: Final[frozenset[str]] = frozenset({ + "mysql", + "mariadb", + "pymysql", + "aiomysql", + "asyncmy", + "mysqlconnector", +}) +_MSSQL_FAMILY_DIALECTS: Final[frozenset[str]] = frozenset({ + "mssql", + "tsql", + "pymssql", + "mssql_python", + "sybase", + "sqlanywhere", +}) +_ORACLE_FAMILY_DIALECTS: Final[frozenset[str]] = frozenset({"oracle", "oracledb"}) +_CREATE_UNIQUE_INDEX_RE: Final = re.compile(r"^\s*CREATE\s+UNIQUE\s+INDEX\s+(?!CONCURRENTLY\b)", re.IGNORECASE) +_CREATE_INDEX_RE: Final = re.compile(r"^\s*CREATE\s+INDEX\s+(?!CONCURRENTLY\b)", re.IGNORECASE) + + +def compute_worker_thread_budget(max_workers: int, cpu_count: int | None = None) -> "dict[str, str]": + """Compute per-worker native thread pool environment variables to prevent oversubscription. + + Args: + max_workers: Number of concurrent ingestion workers. + cpu_count: Optional logical CPU count override. + + Returns: + Dictionary of thread-budget environment variables and string values. + """ + resolved_cpus = cpu_count if cpu_count is not None and cpu_count > 0 else (os.cpu_count() or 1) + workers = max(1, int(max_workers)) + threads = str(max(1, resolved_cpus // workers)) + return { + "OMP_NUM_THREADS": threads, + "ARROW_IO_THREADS": threads, + "POLARS_MAX_THREADS": threads, + "RAYON_NUM_THREADS": threads, + "OPENBLAS_NUM_THREADS": threads, + "MKL_NUM_THREADS": threads, + } + + +def init_spawn_worker_thread_budget(max_workers: int, cpu_count: int | None = None) -> None: + """Initialize thread-budget environment variables inside a spawned worker process.""" + for key, val in compute_worker_thread_budget(max_workers, cpu_count=cpu_count).items(): + os.environ[key] = val + + +def build_checkpoint_table_ddl(dialect: str, checkpoint_table: str) -> str: + """Build DDL to create a multi-shard ingestion checkpoint table if it does not exist.""" + norm = (dialect or "").lower().strip() + if norm in _MSSQL_FAMILY_DIALECTS: + return ( + f"IF OBJECT_ID('{checkpoint_table}', 'U') IS NULL " + f"CREATE TABLE {checkpoint_table} (" + "phase VARCHAR(128) NOT NULL, " + "shard_uri VARCHAR(1024) NOT NULL, " + "rows_loaded BIGINT NOT NULL, " + "bytes_processed BIGINT NOT NULL, " + "duration_s FLOAT NOT NULL, " + "completed_at DATETIME2 DEFAULT CURRENT_TIMESTAMP, " + "PRIMARY KEY (phase, shard_uri))" + ) + return ( + f"CREATE TABLE IF NOT EXISTS {checkpoint_table} (" + "phase VARCHAR(128) NOT NULL, " + "shard_uri VARCHAR(1024) NOT NULL, " + "rows_loaded BIGINT NOT NULL, " + "bytes_processed BIGINT NOT NULL, " + "duration_s DOUBLE PRECISION NOT NULL, " + "completed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " + "PRIMARY KEY (phase, shard_uri))" + ) + + +def build_checkpoint_select_sql(checkpoint_table: str) -> str: + """Build SQL to select already-completed shard URIs for a given checkpoint phase.""" + return f"SELECT shard_uri FROM {checkpoint_table} WHERE phase = :phase" + + +def build_checkpoint_insert_sql(dialect: str, checkpoint_table: str) -> str: + """Build SQL to record a completed shard in the checkpoint table.""" + norm = (dialect or "").lower().strip() + if norm in _ON_CONFLICT_DIALECTS: + return ( + f"INSERT INTO {checkpoint_table} (phase, shard_uri, rows_loaded, bytes_processed, duration_s) " + "VALUES (:phase, :shard_uri, :rows_loaded, :bytes_processed, :duration_s) " + "ON CONFLICT (phase, shard_uri) DO UPDATE SET " + "rows_loaded = EXCLUDED.rows_loaded, " + "bytes_processed = EXCLUDED.bytes_processed, " + "duration_s = EXCLUDED.duration_s" + ) + if norm in _MYSQL_FAMILY_DIALECTS: + return ( + f"INSERT INTO {checkpoint_table} (phase, shard_uri, rows_loaded, bytes_processed, duration_s) " + "VALUES (:phase, :shard_uri, :rows_loaded, :bytes_processed, :duration_s) " + "ON DUPLICATE KEY UPDATE " + "rows_loaded = VALUES(rows_loaded), " + "bytes_processed = VALUES(bytes_processed), " + "duration_s = VALUES(duration_s)" + ) + return ( + f"INSERT INTO {checkpoint_table} (phase, shard_uri, rows_loaded, bytes_processed, duration_s) " + "VALUES (:phase, :shard_uri, :rows_loaded, :bytes_processed, :duration_s)" + ) + + +def resolve_stage_table_name( + target_table: str, + stage_table: "bool | str | None", + *, + conflict_keys: "Sequence[str] | None" = None, + dialect: str | None = None, +) -> "str | None": + """Resolve the staging table name when staged bulk loading or conflict resolution is requested.""" + if isinstance(stage_table, str): + stripped = stage_table.strip() + return stripped or None + if stage_table is False: + return None + norm_dialect = (dialect or "").lower().strip() + needs_stage = stage_table is True or (bool(conflict_keys) and norm_dialect not in {"spanner", "cloudspanner"}) + if not needs_stage: + return None + suffix = uuid4().hex[:8] + clean_target = target_table.strip() + if "." in clean_target: + schema_part, base_name = clean_target.rsplit(".", 1) + unquoted_base = base_name.strip('"`[]') + return f"{schema_part}.{unquoted_base}_stage_{suffix}" + unquoted_base = clean_target.strip('"`[]') + return f"{unquoted_base}_stage_{suffix}" + + +def build_stage_table_ddl( + dialect: str, target_table: str, stage_table: str, *, multi_worker: bool = False +) -> "list[str]": + """Build dialect-specific DDL statements to create an empty staging table matching target_table.""" + norm = (dialect or "").lower().strip() + if norm in _PG_FAMILY_DIALECTS: + return [f"CREATE UNLOGGED TABLE IF NOT EXISTS {stage_table} (LIKE {target_table} INCLUDING DEFAULTS)"] + if norm in {"duckdb", "sqlite", "aiosqlite"}: + if multi_worker or "." in stage_table: + return [f"CREATE TABLE IF NOT EXISTS {stage_table} AS SELECT * FROM {target_table} WHERE 1=0"] + return [f"CREATE TEMP TABLE IF NOT EXISTS {stage_table} AS SELECT * FROM {target_table} WHERE 1=0"] + if norm in _MYSQL_FAMILY_DIALECTS: + if multi_worker: + return [f"CREATE TABLE IF NOT EXISTS {stage_table} LIKE {target_table}"] + return [f"CREATE TEMPORARY TABLE IF NOT EXISTS {stage_table} LIKE {target_table}"] + if norm in _ORACLE_FAMILY_DIALECTS: + return [f"CREATE TABLE {stage_table} NOLOGGING AS SELECT * FROM {target_table} WHERE 1=0"] + if norm in _MSSQL_FAMILY_DIALECTS: + return [f"SELECT TOP 0 * INTO {stage_table} FROM {target_table}"] + if norm == "db2": + return [f"CREATE TABLE {stage_table} AS (SELECT * FROM {target_table}) WITH NO DATA"] + if norm == "bigquery": + return [f"CREATE TABLE IF NOT EXISTS {stage_table} LIKE {target_table}"] + return [f"CREATE TABLE IF NOT EXISTS {stage_table} AS SELECT * FROM {target_table} WHERE 1=0"] + + +def build_stage_merge_sql( + dialect: str, + target_table: str, + stage_table: str, + columns: "Sequence[str]", + *, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, +) -> str: + """Build dialect-specific server-side INSERT ... SELECT or MERGE SQL from stage_table into target_table.""" + norm = (dialect or "").lower().strip() + col_list = ", ".join(columns) + select_exprs = [ + f"CAST({col} AS {type_casts[col]})" if (type_casts and col in type_casts) else col for col in columns + ] + select_list = ", ".join(select_exprs) + if not conflict_keys: + return f"INSERT INTO {target_table} ({col_list}) SELECT {select_list} FROM {stage_table}" + + key_list = ", ".join(conflict_keys) + excluded_set = set(conflict_keys) | set(exclude_update_columns or ()) + update_columns = [col for col in columns if col not in excluded_set] + + if norm in _ON_CONFLICT_DIALECTS: + if on_conflict == "nothing" or not update_columns: + return ( + f"INSERT INTO {target_table} ({col_list}) " + f"SELECT {select_list} FROM {stage_table} WHERE TRUE " + f"ON CONFLICT ({key_list}) DO NOTHING" + ) + set_clause = ", ".join(f"{col} = EXCLUDED.{col}" for col in update_columns) + return ( + f"INSERT INTO {target_table} ({col_list}) " + f"SELECT {select_list} FROM {stage_table} WHERE TRUE " + f"ON CONFLICT ({key_list}) DO UPDATE SET {set_clause}" + ) + + if norm in _MYSQL_FAMILY_DIALECTS: + if on_conflict == "nothing" or not update_columns: + return f"INSERT IGNORE INTO {target_table} ({col_list}) SELECT {select_list} FROM {stage_table}" + set_clause = ", ".join(f"{col} = _sqlspec_stage.{col}" for col in update_columns) + return ( + f"INSERT INTO {target_table} ({col_list}) " + f"SELECT {select_list} FROM {stage_table} AS _sqlspec_stage " + f"ON DUPLICATE KEY UPDATE {set_clause}" + ) + + on_clause = " AND ".join(f"_target.{key} = _stage.{key}" for key in conflict_keys) + matched_clause = "" + if on_conflict == "update" and update_columns: + set_pairs = ", ".join( + f"_target.{col} = CAST(_stage.{col} AS {type_casts[col]})" + if (type_casts and col in type_casts) + else f"_target.{col} = _stage.{col}" + for col in update_columns + ) + matched_clause = f" WHEN MATCHED THEN UPDATE SET {set_pairs}" + stage_vals = ", ".join( + f"CAST(_stage.{col} AS {type_casts[col]})" if (type_casts and col in type_casts) else f"_stage.{col}" + for col in columns + ) + not_matched_clause = f" WHEN NOT MATCHED THEN INSERT ({col_list}) VALUES ({stage_vals})" + terminator = ";" if norm in _MSSQL_FAMILY_DIALECTS else "" + return f"MERGE INTO {target_table} _target USING {stage_table} _stage ON ({on_clause}){matched_clause}{not_matched_clause}{terminator}" + + +def build_drop_stage_table_sql(dialect: str, stage_table: str) -> str: + """Build SQL statement to drop a staging table.""" + norm = (dialect or "").lower().strip() + if norm in (_ORACLE_FAMILY_DIALECTS | {"db2"}): + return f"DROP TABLE {stage_table}" + return f"DROP TABLE IF EXISTS {stage_table}" + + +def filter_droppable_table_indexes( + indexes: "Sequence[Mapping[str, Any]]", constraints: "Sequence[Mapping[str, Any]] | None" = None +) -> "list[dict[str, Any]]": + """Filter data-dictionary index metadata to secondary indexes safe to drop and rebuild.""" + constraint_names: set[str] = set() + if constraints: + for constraint in constraints: + cname = str(constraint.get("constraint_name") or constraint.get("name") or "").strip() + if cname: + constraint_names.add(cname.lower()) + + droppable: list[dict[str, Any]] = [] + for idx in indexes: + idx_name = str(idx.get("index_name") or idx.get("name") or "").strip() + if not idx_name: + continue + if idx_name.upper() == "PRIMARY" or idx_name.lower().startswith("sqlite_autoindex_"): + continue + if idx.get("is_primary") or idx.get("primary"): + continue + if idx.get("is_exclusion") or idx.get("is_extension_owned") or idx.get("is_constraint_backed"): + continue + if idx.get("constraint_type") in {"PRIMARY KEY", "UNIQUE", "FOREIGN KEY", "EXCLUDE"}: + continue + if idx_name.lower() in constraint_names: + continue + droppable.append(dict(idx)) + return droppable + + +def build_drop_index_sql( + dialect: str, index_meta: "Mapping[str, Any]", table: str, schema: str | None = None +) -> "str | None": + """Build a dialect-appropriate DROP INDEX statement for a secondary index.""" + norm = (dialect or "").lower().strip() + if norm == "bigquery": + return None + idx_name = str(index_meta.get("index_name") or index_meta.get("name") or "").strip() + if not idx_name: + return None + resolved_schema = schema or (str(index_meta.get("schema_name") or "").strip() or None) + if norm in _MSSQL_FAMILY_DIALECTS: + qualified_table = f"[{resolved_schema}].[{table}]" if resolved_schema and "." not in table else f"[{table}]" + return f"DROP INDEX [{idx_name}] ON {qualified_table}" + if norm in _MYSQL_FAMILY_DIALECTS: + qualified_table = f"`{resolved_schema}`.`{table}`" if resolved_schema and "." not in table else f"`{table}`" + return f"DROP INDEX `{idx_name}` ON {qualified_table}" + if norm in (_ORACLE_FAMILY_DIALECTS | {"db2"}): + qualified_idx = f'"{resolved_schema}"."{idx_name}"' if resolved_schema else f'"{idx_name}"' + return f"DROP INDEX {qualified_idx}" + if norm in {"spanner", "cloudspanner"}: + return f"DROP INDEX {idx_name}" + qualified_idx = f'"{resolved_schema}"."{idx_name}"' if resolved_schema else f'"{idx_name}"' + return f"DROP INDEX IF EXISTS {qualified_idx}" + + +def build_create_index_sql( + dialect: str, + index_meta: "Mapping[str, Any]", + table: str, + *, + schema: str | None = None, + concurrent: bool = False, + in_transaction: bool = False, + ddl_override: str | None = None, +) -> "str | None": + """Build or adapt a dialect-specific CREATE INDEX statement to restore a dropped index.""" + norm = (dialect or "").lower().strip() + if norm == "bigquery": + return None + raw_ddl = ( + ddl_override + or index_meta.get("definition") + or index_meta.get("native_sql") + or index_meta.get("ddl") + or index_meta.get("sql") + ) + if isinstance(raw_ddl, str) and raw_ddl.strip(): + sql_text = raw_ddl.strip() + else: + idx_name = str(index_meta.get("index_name") or index_meta.get("name") or "").strip() + cols = index_meta.get("columns") or index_meta.get("column_names") or [] + if isinstance(cols, str): + col_list = cols + elif isinstance(cols, Sequence): + col_list = ", ".join(str(c) for c in cols if c) + else: + col_list = "" + if not idx_name or not col_list: + return None + resolved_schema = schema or (str(index_meta.get("schema_name") or "").strip() or None) + is_unique = bool(index_meta.get("is_unique") or index_meta.get("unique")) + unique_kw = "UNIQUE " if is_unique else "" + if norm in _MSSQL_FAMILY_DIALECTS: + qualified_table = f"[{resolved_schema}].[{table}]" if resolved_schema and "." not in table else f"[{table}]" + sql_text = f"CREATE {unique_kw}INDEX [{idx_name}] ON {qualified_table} ({col_list})" + elif norm in _MYSQL_FAMILY_DIALECTS: + qualified_table = f"`{resolved_schema}`.`{table}`" if resolved_schema and "." not in table else f"`{table}`" + sql_text = f"CREATE {unique_kw}INDEX `{idx_name}` ON {qualified_table} ({col_list})" + else: + qualified_table = f'"{resolved_schema}"."{table}"' if resolved_schema and "." not in table else table + sql_text = f'CREATE {unique_kw}INDEX "{idx_name}" ON {qualified_table} ({col_list})' + + if concurrent and not in_transaction and norm in {"postgres", "postgresql", "psycopg", "asyncpg", "psqlpy"}: + if _CREATE_UNIQUE_INDEX_RE.match(sql_text): + sql_text = _CREATE_UNIQUE_INDEX_RE.sub("CREATE UNIQUE INDEX CONCURRENTLY ", sql_text, count=1) + elif _CREATE_INDEX_RE.match(sql_text): + sql_text = _CREATE_INDEX_RE.sub("CREATE INDEX CONCURRENTLY ", sql_text, count=1) + return sql_text + + +def _extract_metadata_rows(result: Any) -> "list[dict[str, Any]]": + if result is None: + return [] + data = getattr(result, "data", result) + if isinstance(data, list): + return [dict(item) if isinstance(item, Mapping) else {} for item in data] + if isinstance(data, Sequence): + return [dict(item) if isinstance(item, Mapping) else {} for item in data] + return [] + + +def _extract_ddl_text(result: Any) -> "str | None": + if isinstance(result, str) and result.strip(): + return result.strip() + rows = _extract_metadata_rows(result) + for row in rows: + for key in ("ddl", "definition", "native_sql", "sql"): + val = row.get(key) + if isinstance(val, str) and val.strip(): + return val.strip() + return None + + +def _safe_driver_in_transaction(driver: Any) -> bool: + checker = getattr(driver, "is_in_transaction", None) + if callable(checker): + with contextlib.suppress(Exception): + return bool(checker()) + conn_checker = getattr(driver, "_connection_in_transaction", None) + if callable(conn_checker): + with contextlib.suppress(Exception): + return bool(conn_checker()) + return False + + +def _split_schema_and_table(table: str, explicit_schema: str | None) -> "tuple[str | None, str]": + clean = table.strip() + if explicit_schema is not None: + return explicit_schema, clean.rsplit(".", 1)[-1].strip('"`[]') + if "." in clean: + schema_part, table_part = clean.split(".", 1) + return schema_part.strip('"`[]'), table_part.strip('"`[]') + return None, clean.strip('"`[]') + + +class SyncWithoutIndexesContext: + """Synchronous context manager that drops secondary indexes before a bulk load and rebuilds them on exit.""" + + __slots__ = ("_concurrent", "_driver", "_rebuild_statements", "_schema", "_table") + + def __init__(self, driver: Any, table: str, *, schema: str | None = None, concurrent: bool = False) -> None: + self._driver = driver + self._table = table + self._schema = schema + self._concurrent = concurrent + self._rebuild_statements: list[str] = [] + + def __enter__(self) -> Self: + dialect = str(getattr(self._driver, "dialect", "") or "").lower().strip() + if dialect == "bigquery": + return self + schema_name, table_name = _split_schema_and_table(self._table, self._schema) + dd = getattr(self._driver, "data_dictionary", None) + if dd is None: + return self + indexes_rows = _extract_metadata_rows(dd.get_indexes(self._driver, table=table_name, schema=schema_name)) + constraints_rows: list[dict[str, Any]] = [] + with contextlib.suppress(Exception): + constraints_rows = _extract_metadata_rows( + dd.get_constraints(self._driver, table=table_name, schema=schema_name) + ) + droppable = filter_droppable_table_indexes(indexes_rows, constraints_rows) + in_tx = _safe_driver_in_transaction(self._driver) + for idx in droppable: + ddl_override: str | None = None + idx_name = str(idx.get("index_name") or idx.get("name") or "").strip() + if ( + idx_name + and not (idx.get("definition") or idx.get("native_sql") or idx.get("ddl")) + and hasattr(dd, "get_ddl") + ): + with contextlib.suppress(Exception): + ddl_override = _extract_ddl_text( + dd.get_ddl(self._driver, idx_name, schema=schema_name, object_type="INDEX") + ) + create_sql = build_create_index_sql( + dialect, + idx, + table_name, + schema=schema_name, + concurrent=self._concurrent, + in_transaction=in_tx, + ddl_override=ddl_override, + ) + drop_sql = build_drop_index_sql(dialect, idx, table_name, schema=schema_name) + if drop_sql and create_sql: + self._driver.execute(drop_sql) + self._rebuild_statements.append(create_sql) + return self + + def __exit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" + ) -> Literal[False]: + in_tx = _safe_driver_in_transaction(self._driver) + rebuild_errors: list[Exception] = [] + for stmt in self._rebuild_statements: + effective_stmt = stmt + if in_tx and " INDEX CONCURRENTLY " in effective_stmt.upper(): + effective_stmt = re.sub( + r"\bINDEX\s+CONCURRENTLY\s+", "INDEX ", effective_stmt, count=1, flags=re.IGNORECASE + ) + try: + self._driver.execute(effective_stmt) + except Exception as err: + rebuild_errors.append(err) + if exc_type is None and rebuild_errors: + raise rebuild_errors[0] + return False + + +class AsyncWithoutIndexesContext: + """Asynchronous context manager that drops secondary indexes before a bulk load and rebuilds them on exit.""" + + __slots__ = ("_concurrent", "_driver", "_rebuild_statements", "_schema", "_table") + + def __init__(self, driver: Any, table: str, *, schema: str | None = None, concurrent: bool = False) -> None: + self._driver = driver + self._table = table + self._schema = schema + self._concurrent = concurrent + self._rebuild_statements: list[str] = [] + + async def __aenter__(self) -> Self: + dialect = str(getattr(self._driver, "dialect", "") or "").lower().strip() + if dialect == "bigquery": + return self + schema_name, table_name = _split_schema_and_table(self._table, self._schema) + dd = getattr(self._driver, "data_dictionary", None) + if dd is None: + return self + indexes_rows = _extract_metadata_rows(await dd.get_indexes(self._driver, table=table_name, schema=schema_name)) + constraints_rows: list[dict[str, Any]] = [] + with contextlib.suppress(Exception): + constraints_rows = _extract_metadata_rows( + await dd.get_constraints(self._driver, table=table_name, schema=schema_name) + ) + droppable = filter_droppable_table_indexes(indexes_rows, constraints_rows) + in_tx = _safe_driver_in_transaction(self._driver) + for idx in droppable: + ddl_override: str | None = None + idx_name = str(idx.get("index_name") or idx.get("name") or "").strip() + if ( + idx_name + and not (idx.get("definition") or idx.get("native_sql") or idx.get("ddl")) + and hasattr(dd, "get_ddl") + ): + with contextlib.suppress(Exception): + ddl_override = _extract_ddl_text( + await dd.get_ddl(self._driver, idx_name, schema=schema_name, object_type="INDEX") + ) + create_sql = build_create_index_sql( + dialect, + idx, + table_name, + schema=schema_name, + concurrent=self._concurrent, + in_transaction=in_tx, + ddl_override=ddl_override, + ) + drop_sql = build_drop_index_sql(dialect, idx, table_name, schema=schema_name) + if drop_sql and create_sql: + await self._driver.execute(drop_sql) + self._rebuild_statements.append(create_sql) + return self + + async def __aexit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" + ) -> Literal[False]: + in_tx = _safe_driver_in_transaction(self._driver) + rebuild_errors: list[Exception] = [] + for stmt in self._rebuild_statements: + effective_stmt = stmt + if in_tx and " INDEX CONCURRENTLY " in effective_stmt.upper(): + effective_stmt = re.sub( + r"\bINDEX\s+CONCURRENTLY\s+", "INDEX ", effective_stmt, count=1, flags=re.IGNORECASE + ) + try: + await self._driver.execute(effective_stmt) + except Exception as err: + rebuild_errors.append(err) + if exc_type is None and rebuild_errors: + raise rebuild_errors[0] + return False + + +def execute_staged_load_sync( + driver: Any, + table: str, + source: Any, + *, + raw_loader: "Callable[..., StorageBridgeJob]", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, + multi_worker: bool = False, +) -> "StorageBridgeJob": + """Orchestrate index drop/rebuild and staging table creation/merge/cleanup for sync bulk loads.""" + if drop_indexes: + with SyncWithoutIndexesContext(driver, table, concurrent=concurrent_indexes): + return execute_staged_load_sync( + driver, + table, + source, + raw_loader=raw_loader, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + multi_worker=multi_worker, + ) + + dialect = str(getattr(driver, "dialect", "") or "").lower().strip() + resolved_stage = resolve_stage_table_name(table, stage_table, conflict_keys=conflict_keys, dialect=dialect) + if resolved_stage is None: + return raw_loader(table, source, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry) + + arrow_table = coerce_arrow_table(source) + columns = list(arrow_table.column_names) + for ddl in build_stage_table_ddl(dialect, table, resolved_stage, multi_worker=multi_worker): + driver.execute(ddl) + try: + if overwrite and not conflict_keys: + driver.execute(f"DELETE FROM {table}") + job = raw_loader(resolved_stage, arrow_table, partitioner=partitioner, overwrite=False, telemetry=telemetry) + if columns and arrow_table.num_rows > 0: + merge_sql = build_stage_merge_sql( + dialect, + table, + resolved_stage, + columns, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + type_casts=type_casts, + ) + driver.execute(merge_sql) + job.telemetry["destination"] = table + return job + finally: + if drop_stage_table: + with contextlib.suppress(Exception): + driver.execute(build_drop_stage_table_sql(dialect, resolved_stage)) + + +async def execute_staged_load_async( + driver: Any, + table: str, + source: Any, + *, + raw_loader: "Callable[..., Any]", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, + multi_worker: bool = False, +) -> "StorageBridgeJob": + """Orchestrate index drop/rebuild and staging table creation/merge/cleanup for async bulk loads.""" + if drop_indexes: + async with AsyncWithoutIndexesContext(driver, table, concurrent=concurrent_indexes): + return await execute_staged_load_async( + driver, + table, + source, + raw_loader=raw_loader, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + multi_worker=multi_worker, + ) + + dialect = str(getattr(driver, "dialect", "") or "").lower().strip() + resolved_stage = resolve_stage_table_name(table, stage_table, conflict_keys=conflict_keys, dialect=dialect) + if resolved_stage is None: + return cast( + "StorageBridgeJob", + await raw_loader(table, source, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry), + ) + + arrow_table = coerce_arrow_table(source) + columns = list(arrow_table.column_names) + for ddl in build_stage_table_ddl(dialect, table, resolved_stage, multi_worker=multi_worker): + await driver.execute(ddl) + try: + if overwrite and not conflict_keys: + await driver.execute(f"DELETE FROM {table}") + job = cast( + "StorageBridgeJob", + await raw_loader( + resolved_stage, arrow_table, partitioner=partitioner, overwrite=False, telemetry=telemetry + ), + ) + if columns and arrow_table.num_rows > 0: + merge_sql = build_stage_merge_sql( + dialect, + table, + resolved_stage, + columns, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + type_casts=type_casts, + ) + await driver.execute(merge_sql) + job.telemetry["destination"] = table + return job + finally: + if drop_stage_table: + with contextlib.suppress(Exception): + await driver.execute(build_drop_stage_table_sql(dialect, resolved_stage)) + + +def _record_checkpoint_sync( + driver: Any, + dialect: str, + checkpoint_table: str, + phase: str, + shard_uri: str, + rows_loaded: int, + bytes_processed: int, + duration_s: float, +) -> None: + sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) + driver.execute( + sql_text, + { + "phase": phase, + "shard_uri": shard_uri, + "rows_loaded": int(rows_loaded), + "bytes_processed": int(bytes_processed), + "duration_s": float(duration_s), + }, + ) + + +async def _record_checkpoint_async( + driver: Any, + dialect: str, + checkpoint_table: str, + phase: str, + shard_uri: str, + rows_loaded: int, + bytes_processed: int, + duration_s: float, +) -> None: + sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) + await driver.execute( + sql_text, + { + "phase": phase, + "shard_uri": shard_uri, + "rows_loaded": int(rows_loaded), + "bytes_processed": int(bytes_processed), + "duration_s": float(duration_s), + }, + ) + + +def _read_shard_bytes_and_decode( + shard_uri: str, file_format: "StorageFormat", storage_options: "dict[str, Any] | None", arrow_schema: Any +) -> "tuple[str, Any, int, float]": + t0 = perf_counter() + pipeline = SyncStoragePipeline() + backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) + payload = backend.read_bytes_sync(path) + table = decode_arrow_payload(payload, file_format, arrow_schema=arrow_schema) + return shard_uri, table, len(payload), perf_counter() - t0 + + +def _decode_shard_in_process_worker( + args: "tuple[str, StorageFormat, dict[str, Any] | None]", +) -> "tuple[str, bytes, int, float]": + shard_uri, file_format, storage_options = args + t0 = perf_counter() + pipeline = SyncStoragePipeline() + backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) + payload = backend.read_bytes_sync(path) + table = decode_arrow_payload(payload, file_format) + ipc_bytes = encode_arrow_payload(table, "arrow-ipc", compression=None) + return shard_uri, ipc_bytes, len(payload), perf_counter() - t0 + + +def _is_multi_shard_source(source: "StorageDestination | Sequence[StorageDestination]") -> tuple[bool, str, bool]: + """Return ``(is_source_sequence, source_str, has_glob_or_prefix)`` for a storage destination.""" + if isinstance(source, Sequence) and not isinstance(source, (str, bytes, bytearray)): + return True, "", True + source_str = source.as_posix() if isinstance(source, Path) else str(source) + if any(ch in source_str for ch in ("*", "?", "[")) or source_str.endswith("/"): + return False, source_str, True + if "://" not in source_str: + with contextlib.suppress(Exception): + if Path(source_str).expanduser().is_dir(): + return False, source_str, True + return False, source_str, False + + +def execute_load_from_storage_sync( + driver: Any, + table: str, + source: "StorageDestination | Sequence[StorageDestination]", + *, + file_format: "StorageFormat", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + storage_options: "dict[str, Any] | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + checkpoint_table: str | None = None, + checkpoint_phase: str = "default", + max_workers: int = 1, + executor: Literal["thread", "process"] = "thread", + arrow_schema: Any = None, + type_casts: "Mapping[str, str] | None" = None, +) -> "StorageBridgeJob": + """Execute single-shard or multi-shard synchronous storage ingestion with checkpointing, staging, and index management.""" + is_source_sequence = isinstance(source, Sequence) and not isinstance(source, (str, bytes, bytearray)) + source_str = "" if is_source_sequence else (source.as_posix() if isinstance(source, Path) else str(source)) + has_glob_or_prefix = ( + is_source_sequence or any(ch in source_str for ch in ("*", "?", "[")) or source_str.endswith("/") + ) + if not has_glob_or_prefix and "://" not in source_str: + with contextlib.suppress(Exception): + if Path(source_str).expanduser().is_dir(): + has_glob_or_prefix = True + + is_simple_single_load = ( + not has_glob_or_prefix + and checkpoint_table is None + and max_workers <= 1 + and stage_table is None + and conflict_keys is None + and not drop_indexes + and arrow_schema is None + ) + if is_simple_single_load: + arrow_table, inbound = ( + driver._read_storage_arrow(source, file_format=file_format) + if storage_options is None + else driver._read_storage_arrow(source, file_format=file_format, storage_options=storage_options) + ) + return cast( + "StorageBridgeJob", + driver.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound), + ) + + if drop_indexes: + with SyncWithoutIndexesContext(driver, table, concurrent=concurrent_indexes): + return execute_load_from_storage_sync( + driver, + table, + source, + file_format=file_format, + partitioner=partitioner, + overwrite=overwrite, + storage_options=storage_options, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + checkpoint_table=checkpoint_table, + checkpoint_phase=checkpoint_phase, + max_workers=max_workers, + executor=executor, + arrow_schema=arrow_schema, + type_casts=type_casts, + ) + + pipeline = driver._storage_pipeline() + if is_source_sequence: + shards = [ + item.as_posix() if isinstance(item, Path) else str(item) + for item in cast("Sequence[StorageDestination]", source) + ] + elif has_glob_or_prefix: + shards = pipeline.list_sources(source, file_format=file_format, storage_options=storage_options) + else: + shards = [source_str] + dialect = str(getattr(driver, "dialect", "") or "").lower().strip() + completed_uris: set[str] = set() + if checkpoint_table: + driver.execute(build_checkpoint_table_ddl(dialect, checkpoint_table)) + res = driver.execute(build_checkpoint_select_sql(checkpoint_table), {"phase": checkpoint_phase}) + for row in res.data or []: + uri_val = row.get("shard_uri") if isinstance(row, Mapping) else row[0] + if uri_val is not None: + completed_uris.add(str(uri_val)) + + pending_shards = [s for s in shards if s not in completed_uris] + shards_skipped = len(shards) - len(pending_shards) + + resolved_stage = resolve_stage_table_name(table, stage_table, conflict_keys=conflict_keys, dialect=dialect) + effective_target = resolved_stage or table + if resolved_stage is not None: + for ddl in build_stage_table_ddl(dialect, table, resolved_stage, multi_worker=max_workers > 1): + driver.execute(ddl) + + total_rows = 0 + total_bytes = 0 + read_duration = 0.0 + write_duration = 0.0 + observed_columns: list[str] = [] + start_wall = perf_counter() + + try: + if overwrite and pending_shards and not conflict_keys: + driver.execute(f"DELETE FROM {table}") + + decoded_shards: list[tuple[str, Any, int, float]] = [] + if max_workers > 1 and len(pending_shards) > 1: + if executor == "process": + mp_ctx = multiprocessing.get_context("spawn") + with ProcessPoolExecutor( + max_workers=max_workers, + mp_context=mp_ctx, + initializer=init_spawn_worker_thread_budget, + initargs=(max_workers,), + ) as pool: + worker_args = [(s, file_format, storage_options) for s in pending_shards] + for shard_uri, ipc_bytes, raw_bytes, r_dur in pool.map( + _decode_shard_in_process_worker, worker_args + ): + tbl = decode_arrow_payload(ipc_bytes, "arrow-ipc", arrow_schema=arrow_schema) + decoded_shards.append((shard_uri, tbl, raw_bytes, r_dur)) + else: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = [ + pool.submit(_read_shard_bytes_and_decode, s, file_format, storage_options, arrow_schema) + for s in pending_shards + ] + decoded_shards = [f.result() for f in futures] + else: + for shard_uri in pending_shards: + decoded_shards.append( + _read_shard_bytes_and_decode(shard_uri, file_format, storage_options, arrow_schema) + ) + + for idx, (shard_uri, shard_table, shard_bytes, r_dur) in enumerate(decoded_shards): + read_duration += r_dur + if not observed_columns and hasattr(shard_table, "column_names"): + observed_columns = list(shard_table.column_names) + w0 = perf_counter() + shard_overwrite = bool(overwrite and idx == 0 and resolved_stage is None and not conflict_keys) + shard_job = driver.load_from_arrow( + effective_target, shard_table, partitioner=partitioner, overwrite=shard_overwrite + ) + w_dur = perf_counter() - w0 + write_duration += w_dur + s_rows = int(shard_job.telemetry.get("rows_processed", shard_table.num_rows)) + total_rows += s_rows + total_bytes += int(shard_bytes) + if checkpoint_table: + _record_checkpoint_sync( + driver, dialect, checkpoint_table, checkpoint_phase, shard_uri, s_rows, shard_bytes, r_dur + w_dur + ) + + if resolved_stage is not None and observed_columns and total_rows > 0: + merge_sql = build_stage_merge_sql( + dialect, + table, + resolved_stage, + observed_columns, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + type_casts=type_casts, + ) + w0 = perf_counter() + driver.execute(merge_sql) + write_duration += perf_counter() - w0 + finally: + if resolved_stage is not None and drop_stage_table: + with contextlib.suppress(Exception): + driver.execute(build_drop_stage_table_sql(dialect, resolved_stage)) + + telemetry_payload = ingest_telemetry_from_batches( + total_rows, + total_bytes, + duration_s=perf_counter() - start_wall, + format_label=file_format, + read_duration_s=read_duration, + write_duration_s=write_duration, + shards_processed=len(pending_shards), + shards_skipped=shards_skipped, + ) + telemetry_payload["destination"] = table + attach_partition_telemetry(telemetry_payload, partitioner) + return create_storage_job(telemetry_payload) + + +async def execute_load_from_storage_async( + driver: Any, + table: str, + source: "StorageDestination | Sequence[StorageDestination]", + *, + file_format: "StorageFormat", + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + storage_options: "dict[str, Any] | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + checkpoint_table: str | None = None, + checkpoint_phase: str = "default", + max_workers: int = 1, + executor: Literal["thread", "process"] = "thread", + arrow_schema: Any = None, + type_casts: "Mapping[str, str] | None" = None, +) -> "StorageBridgeJob": + """Execute single-shard or multi-shard asynchronous storage ingestion with checkpointing, staging, and index management.""" + is_source_sequence, source_str, has_glob_or_prefix = _is_multi_shard_source(source) + + is_simple_single_load = ( + not has_glob_or_prefix + and checkpoint_table is None + and max_workers <= 1 + and stage_table is None + and conflict_keys is None + and not drop_indexes + and arrow_schema is None + ) + if is_simple_single_load: + arrow_table, inbound = ( + await driver._read_storage_arrow(source, file_format=file_format) + if storage_options is None + else await driver._read_storage_arrow(source, file_format=file_format, storage_options=storage_options) + ) + return cast( + "StorageBridgeJob", + await driver.load_from_arrow( + table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound + ), + ) + + if drop_indexes: + async with AsyncWithoutIndexesContext(driver, table, concurrent=concurrent_indexes): + return await execute_load_from_storage_async( + driver, + table, + source, + file_format=file_format, + partitioner=partitioner, + overwrite=overwrite, + storage_options=storage_options, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + checkpoint_table=checkpoint_table, + checkpoint_phase=checkpoint_phase, + max_workers=max_workers, + executor=executor, + arrow_schema=arrow_schema, + type_casts=type_casts, + ) + + pipeline = driver._storage_pipeline() + if is_source_sequence: + shards = [ + item.as_posix() if isinstance(item, Path) else str(item) + for item in cast("Sequence[StorageDestination]", source) + ] + elif has_glob_or_prefix: + shards = await pipeline.list_sources_async(source, file_format=file_format, storage_options=storage_options) + else: + shards = [source_str] + dialect = str(getattr(driver, "dialect", "") or "").lower().strip() + completed_uris: set[str] = set() + if checkpoint_table: + await driver.execute(build_checkpoint_table_ddl(dialect, checkpoint_table)) + res = await driver.execute(build_checkpoint_select_sql(checkpoint_table), {"phase": checkpoint_phase}) + for row in res.data or []: + uri_val = row.get("shard_uri") if isinstance(row, Mapping) else row[0] + if uri_val is not None: + completed_uris.add(str(uri_val)) + + pending_shards = [s for s in shards if s not in completed_uris] + shards_skipped = len(shards) - len(pending_shards) + + resolved_stage = resolve_stage_table_name(table, stage_table, conflict_keys=conflict_keys, dialect=dialect) + effective_target = resolved_stage or table + if resolved_stage is not None: + for ddl in build_stage_table_ddl(dialect, table, resolved_stage, multi_worker=max_workers > 1): + await driver.execute(ddl) + + total_rows = 0 + total_bytes = 0 + read_duration = 0.0 + write_duration = 0.0 + observed_columns: list[str] = [] + start_wall = perf_counter() + + try: + if overwrite and pending_shards and not conflict_keys: + await driver.execute(f"DELETE FROM {table}") + + loop = asyncio.get_running_loop() + decoded_shards: list[tuple[str, Any, int, float]] = [] + if max_workers > 1 and len(pending_shards) > 1: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + tasks = [ + loop.run_in_executor( + pool, _read_shard_bytes_and_decode, s, file_format, storage_options, arrow_schema + ) + for s in pending_shards + ] + decoded_shards = list(await asyncio.gather(*tasks)) + else: + for shard_uri in pending_shards: + decoded_shards.append( + await loop.run_in_executor( + None, _read_shard_bytes_and_decode, shard_uri, file_format, storage_options, arrow_schema + ) + ) + + for idx, (shard_uri, shard_table, shard_bytes, r_dur) in enumerate(decoded_shards): + read_duration += r_dur + if not observed_columns and hasattr(shard_table, "column_names"): + observed_columns = list(shard_table.column_names) + w0 = perf_counter() + shard_overwrite = bool(overwrite and idx == 0 and resolved_stage is None and not conflict_keys) + shard_job = await driver.load_from_arrow( + effective_target, shard_table, partitioner=partitioner, overwrite=shard_overwrite + ) + w_dur = perf_counter() - w0 + write_duration += w_dur + s_rows = int(shard_job.telemetry.get("rows_processed", shard_table.num_rows)) + total_rows += s_rows + total_bytes += int(shard_bytes) + if checkpoint_table: + await _record_checkpoint_async( + driver, dialect, checkpoint_table, checkpoint_phase, shard_uri, s_rows, shard_bytes, r_dur + w_dur + ) + + if resolved_stage is not None and observed_columns and total_rows > 0: + merge_sql = build_stage_merge_sql( + dialect, + table, + resolved_stage, + observed_columns, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + type_casts=type_casts, + ) + w0 = perf_counter() + await driver.execute(merge_sql) + write_duration += perf_counter() - w0 + finally: + if resolved_stage is not None and drop_stage_table: + with contextlib.suppress(Exception): + await driver.execute(build_drop_stage_table_sql(dialect, resolved_stage)) + + telemetry_payload = ingest_telemetry_from_batches( + total_rows, + total_bytes, + duration_s=perf_counter() - start_wall, + format_label=file_format, + read_duration_s=read_duration, + write_duration_s=write_duration, + shards_processed=len(pending_shards), + shards_skipped=shards_skipped, + ) + telemetry_payload["destination"] = table + attach_partition_telemetry(telemetry_payload, partitioner) + return create_storage_job(telemetry_payload) diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index 49a909eaf..a401b660c 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -48,7 +48,12 @@ from sqlspec.driver._query_cache import CachedQuery from sqlspec.driver._sql_helpers import DEFAULT_PRETTY from sqlspec.driver._sql_helpers import convert_to_dialect as _convert_to_dialect_impl -from sqlspec.driver._storage_helpers import stringify_storage_target +from sqlspec.driver._storage_helpers import ( + SyncWithoutIndexesContext, + execute_load_from_storage_sync, + execute_staged_load_sync, + stringify_storage_target, +) from sqlspec.driver._stream import SyncRowStream, _LazyEagerSyncRowSource from sqlspec.exceptions import ImproperConfigurationError, StackExecutionError from sqlspec.observability import _runtime as observability_runtime @@ -1372,6 +1377,12 @@ def select_to_storage( """ self._raise_storage_not_implemented("select_to_storage") + def without_indexes( + self, table: str, *, schema: str | None = None, concurrent: bool = False + ) -> SyncWithoutIndexesContext: + """Drop secondary indexes on ``table`` before bulk loading and rebuild them on exit.""" + return SyncWithoutIndexesContext(self, table, schema=schema, concurrent=concurrent) + def load_from_arrow( self, table: str, @@ -1379,6 +1390,15 @@ def load_from_arrow( *, partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target table. @@ -1387,10 +1407,34 @@ def load_from_arrow( source: Arrow data source. partitioner: Optional partitioner configuration. overwrite: Whether to overwrite existing data. + telemetry: Optional inbound storage telemetry to merge. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. """ + _ = ( + table, + source, + partitioner, + overwrite, + telemetry, + stage_table, + conflict_keys, + on_conflict, + exclude_update_columns, + drop_stage_table, + drop_indexes, + concurrent_indexes, + type_casts, + ) self._raise_storage_not_implemented("load_from_arrow") def load_from_storage( @@ -1401,20 +1445,71 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + storage_options: "dict[str, Any] | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + checkpoint_table: str | None = None, + checkpoint_phase: str = "default", + max_workers: int = 1, + executor: Literal["thread", "process"] = "thread", + arrow_schema: Any = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load artifacts from storage into the target table. Args: table: Target table name. - source: Storage source path. + source: Storage source path, prefix, or glob pattern. file_format: File format of source. partitioner: Optional partitioner configuration. overwrite: Whether to overwrite existing data. + storage_options: Optional storage backend configuration. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + checkpoint_table: Optional database table name for shard-level checkpoint tracking. + checkpoint_phase: Checkpoint phase identifier (default ``"default"``). + max_workers: Number of parallel workers for multi-shard storage fan-out. + executor: Worker pool type (``"thread"`` or ``"process"``). + arrow_schema: Optional PyArrow schema for explicit type pinning on decode. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. """ - self._raise_storage_not_implemented("load_from_storage") + if type(self).load_from_arrow is SyncDriverAdapterBase.load_from_arrow: + self._raise_storage_not_implemented("load_from_storage") + return execute_load_from_storage_sync( + self, + table, + source, + file_format=file_format, + partitioner=partitioner, + overwrite=overwrite, + storage_options=storage_options, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + checkpoint_table=checkpoint_table, + checkpoint_phase=checkpoint_phase, + max_workers=max_workers, + executor=executor, + arrow_schema=arrow_schema, + type_casts=type_casts, + ) def load_from_records( self, @@ -1423,6 +1518,14 @@ def load_from_records( *, columns: "list[str] | None" = None, overwrite: bool = False, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load in-memory dict or positional records into the target table. @@ -1436,12 +1539,36 @@ def load_from_records( records: Mapping records, or positional sequences with ``columns``. columns: Column names (required for positional records). overwrite: Whether to overwrite existing data. + stage_table: Optional staging table flag or explicit staging table name. + conflict_keys: Optional conflict key columns for server-side merge/upsert. + on_conflict: Conflict action (``"update"`` or ``"nothing"``). + exclude_update_columns: Optional columns excluded from ``SET`` updates. + drop_stage_table: Whether to drop the staging table after merge. + drop_indexes: Whether to drop and rebuild secondary indexes around the load. + concurrent_indexes: Whether to rebuild PostgreSQL indexes with ``CONCURRENTLY``. + type_casts: Optional column-to-SQL-type mapping for stage-to-target casts. Returns: StorageBridgeJob with execution telemetry. """ prepared_records = self._prepare_records_for_arrow(records) arrow_table = self._records_to_arrow_table(prepared_records, columns) + if stage_table or conflict_keys or drop_indexes: + return execute_staged_load_sync( + self, + table, + arrow_table, + raw_loader=self.load_from_arrow, + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) return self.load_from_arrow(table, arrow_table, overwrite=overwrite) def convert_to_dialect( diff --git a/tests/unit/driver/test_staged_bulk_load.py b/tests/unit/driver/test_staged_bulk_load.py new file mode 100644 index 000000000..6873554aa --- /dev/null +++ b/tests/unit/driver/test_staged_bulk_load.py @@ -0,0 +1,197 @@ +"""Unit tests for staged bulk load, dialect-aware upsert/merge, type casts, and load_from_records.""" + +from pathlib import Path +from typing import Literal +from unittest.mock import MagicMock + +import pyarrow as pa +import pytest + +from sqlspec.adapters.aiosqlite import AiosqliteConfig +from sqlspec.adapters.duckdb import DuckDBConfig +from sqlspec.adapters.spanner._typing import SpannerTransaction +from sqlspec.adapters.spanner.driver import SpannerSyncDriver +from sqlspec.adapters.sqlite import SqliteConfig +from sqlspec.driver._storage_helpers import ( + build_drop_stage_table_sql, + build_stage_merge_sql, + build_stage_table_ddl, + resolve_stage_table_name, +) + + +def test_resolve_stage_table_name_auto_and_custom() -> None: + """resolve_stage_table_name generates a unique stage identifier or preserves an explicit string.""" + auto_name = resolve_stage_table_name("public.users", True) + assert auto_name is not None + assert auto_name.startswith("public.users_stage_") + assert resolve_stage_table_name("public.users", "custom_stage_tbl") == "custom_stage_tbl" + assert resolve_stage_table_name("public.users", False) is None + + +@pytest.mark.parametrize( + ("dialect", "expected_fragment"), + [ + ("postgres", "CREATE UNLOGGED TABLE"), + ("duckdb", "CREATE TEMP TABLE"), + ("sqlite", "CREATE TEMP TABLE"), + ("mysql", "CREATE TEMPORARY TABLE"), + ("mssql", "SELECT TOP 0"), + ("oracle", "NOLOGGING"), + ("bigquery", "CREATE TABLE IF NOT EXISTS"), + ("spanner", "CREATE TABLE"), + ], +) +def test_build_stage_table_ddl_across_dialects(dialect: str, expected_fragment: str) -> None: + """Stage table DDL builder produces dialect-appropriate temporary/staging DDL.""" + ddl_list = build_stage_table_ddl(dialect, "target_items", "_stage_items") + assert any(expected_fragment in stmt for stmt in ddl_list) + + +@pytest.mark.parametrize( + ("dialect", "on_conflict", "expected_fragment"), + [ + ("postgres", "update", "ON CONFLICT"), + ("postgres", "nothing", "DO NOTHING"), + ("cockroach", "update", "ON CONFLICT"), + ("sqlite", "update", "ON CONFLICT"), + ("duckdb", "nothing", "DO NOTHING"), + ("mysql", "update", "ON DUPLICATE KEY UPDATE"), + ("mysql", "nothing", "INSERT IGNORE INTO"), + ("mssql", "update", "MERGE INTO"), + ("oracle", "update", "MERGE INTO"), + ("bigquery", "update", "MERGE INTO"), + ], +) +def test_build_stage_merge_sql_across_dialects( + dialect: str, on_conflict: Literal["update", "nothing"], expected_fragment: str +) -> None: + """Stage merge SQL builder produces dialect-appropriate upsert/ignore statements.""" + sql = build_stage_merge_sql( + dialect, + "target_items", + "_stage_items", + ["id", "name", "metadata"], + conflict_keys=["id"], + on_conflict=on_conflict, + exclude_update_columns=["created_at"], + type_casts={"metadata": "JSONB"} if dialect == "postgres" else None, + ) + assert expected_fragment in sql + if dialect == "postgres": + assert "CAST(metadata AS JSONB)" in sql + + drop_sql = build_drop_stage_table_sql(dialect, "_stage_items") + assert "DROP TABLE" in drop_sql + + +def test_sqlite_staged_bulk_load_upsert_nothing_and_records(tmp_path: Path) -> None: + """SqliteDriver load_from_arrow and load_from_records support staged upsert, DO NOTHING, and cleanup.""" + config = SqliteConfig(connection_config={"database": str(tmp_path / "staged.db")}) + with config.provide_session() as driver: + driver.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, score INTEGER)") + initial = pa.Table.from_pydict({"id": [1, 2], "name": ["alice", "bob"], "score": [10, 20]}) + job1 = driver.load_from_arrow("users", initial, stage_table=True, conflict_keys=["id"], on_conflict="update") + assert job1.telemetry["rows_processed"] == 2 + + delta = pa.Table.from_pydict({"id": [2, 3], "name": ["bob_updated", "carol"], "score": [99, 30]}) + driver.load_from_arrow( + "users", + delta, + stage_table=True, + conflict_keys=["id"], + on_conflict="update", + exclude_update_columns=["name"], + ) + rows = driver.select("SELECT id, name, score FROM users ORDER BY id") + assert rows == [ + {"id": 1, "name": "alice", "score": 10}, + {"id": 2, "name": "bob", "score": 99}, + {"id": 3, "name": "carol", "score": 30}, + ] + + ignored = [{"id": 3, "name": "carol_ignored", "score": 999}, {"id": 4, "name": "dave", "score": 40}] + driver.load_from_records( + "users", + ignored, + stage_table="kept_stage_users", + conflict_keys=["id"], + on_conflict="nothing", + drop_stage_table=False, + ) + rows_after = driver.select("SELECT id, name, score FROM users ORDER BY id") + assert len(rows_after) == 4 + assert rows_after[2]["score"] == 30 + stage_rows = driver.select("SELECT id FROM kept_stage_users ORDER BY id") + assert len(stage_rows) == 2 + + +def test_staged_bulk_load_drops_stage_table_on_merge_failure(tmp_path: Path) -> None: + """Stage table is dropped in finally even if merge execution raises an error.""" + config = SqliteConfig(connection_config={"database": str(tmp_path / "staged_fail.db")}) + with config.provide_session() as driver: + driver.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, val TEXT)") + table = pa.Table.from_pydict({"id": [1], "val": ["ok"]}) + with pytest.raises(Exception): + driver.load_from_arrow( + "items", table, stage_table="transient_stage", conflict_keys=["nonexistent_key"], on_conflict="update" + ) + remaining = driver.select( + "SELECT name FROM sqlite_master WHERE type IN ('table', 'view') AND name = 'transient_stage'" + ) + assert remaining == [] + + +@pytest.mark.anyio +async def test_aiosqlite_and_duckdb_staged_bulk_load_with_type_casts(tmp_path: Path) -> None: + """Async Aiosqlite and DuckDB staged loads support type_casts and conflict resolution.""" + duck_cfg = DuckDBConfig(connection_config={"database": ":memory:"}) + with duck_cfg.provide_session() as duck_driver: + duck_driver.execute("CREATE TABLE metrics (id INTEGER PRIMARY KEY, amount INTEGER)") + duck_driver.load_from_records( + "metrics", + [{"id": 1, "amount": "100"}, {"id": 2, "amount": "200"}], + stage_table=True, + conflict_keys=["id"], + on_conflict="update", + type_casts={"amount": "INTEGER"}, + ) + rows = duck_driver.select("SELECT id, amount FROM metrics ORDER BY id") + assert rows == [{"id": 1, "amount": 100}, {"id": 2, "amount": 200}] + + aio_cfg = AiosqliteConfig(connection_config={"database": str(tmp_path / "async_staged.db")}) + try: + async with aio_cfg.provide_session() as aio_driver: + await aio_driver.execute("CREATE TABLE events (id INTEGER PRIMARY KEY, label TEXT)") + await aio_driver.load_from_records( + "events", + [{"id": 1, "label": "v1"}, {"id": 2, "label": "v2"}], + stage_table=True, + conflict_keys=["id"], + on_conflict="update", + ) + await aio_driver.load_from_records( + "events", + [{"id": 2, "label": "v2_updated"}, {"id": 3, "label": "v3"}], + stage_table=True, + conflict_keys=["id"], + on_conflict="update", + ) + async_rows = await aio_driver.select("SELECT id, label FROM events ORDER BY id") + assert async_rows == [{"id": 1, "label": "v1"}, {"id": 2, "label": "v2_updated"}, {"id": 3, "label": "v3"}] + finally: + await aio_cfg.close_pool() + + +def test_spanner_load_from_arrow_uses_insert_or_ignore_when_on_conflict_nothing() -> None: + """SpannerSyncDriver uses writer.insert_or_ignore when conflict_keys + on_conflict='nothing' is set.""" + mock_tx = MagicMock(spec=SpannerTransaction) + mock_tx.insert_or_ignore = MagicMock() + mock_tx.insert_or_update = MagicMock() + driver = SpannerSyncDriver( + connection=mock_tx, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} + ) + table = pa.Table.from_pydict({"id": [1, 2], "name": ["a", "b"]}) + driver.load_from_arrow("users", table, conflict_keys=["id"], on_conflict="nothing") + assert mock_tx.insert_or_ignore.called + assert not mock_tx.insert_or_update.called diff --git a/tests/unit/driver/test_storage_parallel_and_auth.py b/tests/unit/driver/test_storage_parallel_and_auth.py new file mode 100644 index 000000000..e9a548e33 --- /dev/null +++ b/tests/unit/driver/test_storage_parallel_and_auth.py @@ -0,0 +1,303 @@ +"""Unit tests for parallel storage ingestion, checkpointing, thread budgets, and dynamic auth providers.""" + +import os +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from sqlspec.adapters.adbc.config import AdbcConfig +from sqlspec.adapters.adbc.core import build_connection_config as build_adbc_connection_config +from sqlspec.adapters.aiosqlite import AiosqliteConfig +from sqlspec.adapters.asyncpg.config import AsyncpgConfig +from sqlspec.adapters.duckdb.core import _build_storage_read_sql, _s3_storage_options_match +from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool +from sqlspec.adapters.psqlpy.config import PsqlpyConfig +from sqlspec.adapters.psqlpy.core import build_connection_config as build_psqlpy_connection_config +from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig +from sqlspec.adapters.sqlite import SqliteConfig +from sqlspec.driver._storage_helpers import ( + build_checkpoint_insert_sql, + build_checkpoint_select_sql, + build_checkpoint_table_ddl, + compute_worker_thread_budget, + init_spawn_worker_thread_budget, +) + + +def test_compute_worker_thread_budget_divides_cores_fairly() -> None: + """Thread budget divides total cores across workers with a floor of 1.""" + budget4 = compute_worker_thread_budget(4, cpu_count=16) + assert budget4["OMP_NUM_THREADS"] == "4" + assert budget4["RAYON_NUM_THREADS"] == "4" + + budget8 = compute_worker_thread_budget(8, cpu_count=4) + assert budget8["OMP_NUM_THREADS"] == "1" + + budget0 = compute_worker_thread_budget(0, cpu_count=8) + assert budget0["OMP_NUM_THREADS"] == "8" + + +def test_init_spawn_worker_thread_budget_sets_thread_env_vars(monkeypatch: pytest.MonkeyPatch) -> None: + """Spawn worker initializer sets all native thread pool environment variables.""" + expected_keys = compute_worker_thread_budget(1, cpu_count=1).keys() + for var in expected_keys: + monkeypatch.delenv(var, raising=False) + + init_spawn_worker_thread_budget(4, cpu_count=12) + for var in expected_keys: + assert os.environ.get(var) == "3" + + +def test_checkpoint_sql_builders_generate_dialect_valid_sql() -> None: + """Checkpoint DDL, SELECT, and INSERT builders emit valid SQL across dialects.""" + ddl_sql = build_checkpoint_table_ddl("sqlite", "_ingest_ckpt") + assert "CREATE TABLE IF NOT EXISTS" in ddl_sql + assert "phase" in ddl_sql + assert "shard_uri" in ddl_sql + + select_sql = build_checkpoint_select_sql("_ingest_ckpt") + assert "SELECT shard_uri FROM _ingest_ckpt" in select_sql + + insert_sql = build_checkpoint_insert_sql("sqlite", "_ingest_ckpt") + assert "INSERT INTO _ingest_ckpt" in insert_sql + assert "ON CONFLICT" in insert_sql + + +def test_sync_load_from_storage_multi_uri_checkpoint_and_schema(tmp_path: Path) -> None: + """Sync load_from_storage ingests multiple URIs concurrently and skips checkpointed URIs on rerun.""" + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + file1 = tmp_path / "part1.parquet" + file2 = tmp_path / "part2.parquet" + pq.write_table(pa.Table.from_pydict({"id": [1, 2], "name": ["a", "b"]}, schema=schema), file1) + pq.write_table(pa.Table.from_pydict({"id": [3, 4], "name": ["c", "d"]}, schema=schema), file2) + + db_path = tmp_path / "checkpoint_sync.db" + config = SqliteConfig(connection_config={"database": str(db_path)}) + with config.provide_session() as driver: + driver.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)") + job1 = driver.load_from_storage( + "items", + [str(file1), str(file2)], + file_format="parquet", + max_workers=2, + checkpoint_table="_load_checkpoints", + checkpoint_phase="daily_sync", + arrow_schema=schema, + ) + assert job1.telemetry["rows_processed"] == 4 + assert job1.telemetry["shards_processed"] == 2 + assert job1.telemetry["shards_skipped"] == 0 + rows = driver.select("SELECT id, name FROM items ORDER BY id") + assert len(rows) == 4 + + job2 = driver.load_from_storage( + "items", + [str(file1), str(file2)], + file_format="parquet", + max_workers=2, + checkpoint_table="_load_checkpoints", + checkpoint_phase="daily_sync", + arrow_schema=schema, + ) + assert job2.telemetry["rows_processed"] == 0 + assert job2.telemetry["shards_skipped"] == 2 + + +@pytest.mark.anyio +async def test_async_load_from_storage_multi_uri_checkpoint_and_schema(tmp_path: Path) -> None: + """Async load_from_storage ingests multiple URIs with bounded concurrency and checkpointing.""" + schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + file1 = tmp_path / "async_part1.parquet" + file2 = tmp_path / "async_part2.parquet" + pq.write_table(pa.Table.from_pydict({"id": [10, 20], "name": ["x", "y"]}, schema=schema), file1) + pq.write_table(pa.Table.from_pydict({"id": [30, 40], "name": ["z", "w"]}, schema=schema), file2) + + db_path = tmp_path / "checkpoint_async.db" + config = AiosqliteConfig(connection_config={"database": str(db_path)}) + try: + async with config.provide_session() as driver: + await driver.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)") + job1 = await driver.load_from_storage( + "items", + [str(file1), str(file2)], + file_format="parquet", + max_workers=2, + checkpoint_table="_load_checkpoints", + checkpoint_phase="async_daily", + arrow_schema=schema, + ) + assert job1.telemetry["rows_processed"] == 4 + assert job1.telemetry["shards_processed"] == 2 + + job2 = await driver.load_from_storage( + "items", + [str(file1), str(file2)], + file_format="parquet", + max_workers=2, + checkpoint_table="_load_checkpoints", + checkpoint_phase="async_daily", + arrow_schema=schema, + ) + assert job2.telemetry["rows_processed"] == 0 + assert job2.telemetry["shards_skipped"] == 2 + finally: + await config.close_pool() + + +def test_psycopg_sync_auth_token_provider_injects_token_and_shortens_lifetime() -> None: + """PsycopgSyncConfig invokes auth_token_provider on connect and defaults max_lifetime to 1800s.""" + tokens = iter(["token-1", "token-2"]) + config = PsycopgSyncConfig( + connection_config={ + "host": "localhost", + "user": "iam_user", + "dbname": "analytics", + "auth_token_provider": lambda: next(tokens), + } + ) + _, _, kwargs = config._connection_kwargs() + assert kwargs["password"] == "token-1" + + captured_pool_kwargs: dict[str, Any] = {} + + def fake_pool(conninfo: str, **pool_kwargs: Any) -> MagicMock: + captured_pool_kwargs["conninfo"] = conninfo + captured_pool_kwargs.update(pool_kwargs) + return MagicMock() + + with patch("sqlspec.adapters.psycopg.config.ConnectionPool", side_effect=fake_pool): + config._create_pool() + + assert captured_pool_kwargs["max_lifetime"] == 1800.0 + assert "auth_token_provider" not in captured_pool_kwargs["kwargs"] + + +@pytest.mark.anyio +async def test_psycopg_async_and_asyncpg_and_psqlpy_auth_token_provider() -> None: + """Async PostgreSQL configs support sync and async auth_token_provider callables.""" + + async def provide_token() -> str: + return "async-iam-token" + + psycopg_async = PsycopgAsyncConfig( + connection_config={ + "host": "localhost", + "user": "iam_user", + "dbname": "analytics", + "open": False, + "auth_token_provider": provide_token, + } + ) + captured_psycopg: dict[str, Any] = {} + + def fake_async_pool(conninfo: str, **pool_kwargs: Any) -> MagicMock: + captured_psycopg.update(pool_kwargs) + return MagicMock() + + with patch("sqlspec.adapters.psycopg.config.AsyncConnectionPool", side_effect=fake_async_pool): + await psycopg_async._create_pool() + assert captured_psycopg["max_lifetime"] == 1800.0 + + asyncpg_cfg = AsyncpgConfig( + connection_config={ + "host": "localhost", + "user": "iam_user", + "database": "analytics", + "auth_token_provider": provide_token, + } + ) + captured_asyncpg: dict[str, Any] = {} + + async def fake_asyncpg_pool(**pool_kwargs: Any) -> MagicMock: + captured_asyncpg.update(pool_kwargs) + return MagicMock() + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", side_effect=fake_asyncpg_pool): + await asyncpg_cfg._create_pool() + assert captured_asyncpg["password"] is provide_token + assert captured_asyncpg["max_inactive_connection_lifetime"] == 1800.0 + + psqlpy_cfg = PsqlpyConfig( + connection_config={ + "host": "localhost", + "username": "iam_user", + "db_name": "analytics", + "auth_token_provider": provide_token, + } + ) + captured_psqlpy: dict[str, Any] = {} + + def fake_psqlpy_pool(**pool_kwargs: Any) -> MagicMock: + captured_psqlpy.update(pool_kwargs) + return MagicMock() + + with patch("sqlspec.adapters.psqlpy.config.ConnectionPool", side_effect=fake_psqlpy_pool): + await psqlpy_cfg._create_pool() + assert captured_psqlpy["password"] == "async-iam-token" + assert ( + build_psqlpy_connection_config({"username": "u", "auth_token_provider": lambda: "sync-tok"})["password"] + == "sync-tok" + ) + + +def test_adbc_auth_token_provider_resolves_fresh_credentials() -> None: + """ADBC build_connection_config invokes auth_token_provider on each call.""" + counter = 0 + + def next_token() -> str: + nonlocal counter + counter += 1 + return f"adbc-token-{counter}" + + cfg1 = build_adbc_connection_config({ + "driver": "postgresql", + "uri": "postgresql://localhost/db", + "auth_token_provider": next_token, + }) + cfg2 = build_adbc_connection_config({ + "driver": "flightsql", + "uri": "grpc+tls://localhost:31337", + "auth_token_provider": next_token, + }) + assert cfg1["password"] == "adbc-token-1" + assert cfg2.get("db_kwargs") is not None + _ = AdbcConfig(connection_config={"driver": "sqlite", "uri": ":memory:", "auth_token_provider": next_token}) + + +def test_duckdb_secret_value_provider_and_gcs_bearer_token_and_globs() -> None: + """DuckDB secret value_provider refreshes secrets and supports GCS bearer_token and native globs.""" + token_seq = iter(["gcs-oauth-1", "gcs-oauth-2"]) + pool = DuckDBConnectionPool( + {"database": ":memory:"}, + secrets=[ + { + "name": "gcs_dynamic", + "secret_type": "http", + "value_provider": lambda: {"bearer_token": next(token_seq)}, + "refresh_interval_seconds": 0.0, + } + ], + ) + try: + conn = pool.acquire() + secrets_initial = pool._storage_settings(conn)["_duckdb_storage_secrets"] + assert secrets_initial[0]["value"]["bearer_token"] == "gcs-oauth-1" + + conn = pool.acquire() + secrets_refreshed = pool._storage_settings(conn)["_duckdb_storage_secrets"] + assert secrets_refreshed[0]["value"]["bearer_token"] == "gcs-oauth-2" + finally: + pool.close() + + assert _s3_storage_options_match({"google_bearer_token": "tok-123"}, {"bearer_token": "tok-123"}, "gcs") + assert _build_storage_read_sql("items", "gs://bucket/path/*.parquet", "parquet", allow_globs=False) is None + parquet_glob_sql = _build_storage_read_sql("items", "gs://bucket/path/*.parquet", "parquet", allow_globs=True) + csv_glob_sql = _build_storage_read_sql("items", "gs://bucket/path/*.csv", "csv", allow_globs=True) + jsonl_glob_sql = _build_storage_read_sql("items", "gs://bucket/path/*.jsonl", "jsonl", allow_globs=True) + assert parquet_glob_sql is not None and "read_parquet" in parquet_glob_sql.lower() + assert csv_glob_sql is not None and "read_csv" in csv_glob_sql.lower() + assert jsonl_glob_sql is not None and "read_json" in jsonl_glob_sql.lower() diff --git a/tests/unit/driver/test_without_indexes.py b/tests/unit/driver/test_without_indexes.py new file mode 100644 index 000000000..a73017fcd --- /dev/null +++ b/tests/unit/driver/test_without_indexes.py @@ -0,0 +1,111 @@ +"""Unit tests for without_indexes context manager and drop_indexes bulk load orchestration.""" + +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pytest + +from sqlspec.adapters.aiosqlite import AiosqliteConfig +from sqlspec.adapters.sqlite import SqliteConfig +from sqlspec.driver._storage_helpers import build_create_index_sql, build_drop_index_sql, filter_droppable_table_indexes + + +def test_filter_droppable_table_indexes_preserves_primary_and_constraints() -> None: + """Primary and constraint-backing indexes are never dropped.""" + indexes: list[dict[str, Any]] = [ + {"index_name": "pk_users", "table_name": "users", "columns": ["id"], "unique": True, "is_primary": True}, + { + "index_name": "uq_users_email", + "table_name": "users", + "columns": ["email"], + "unique": True, + "is_constraint_backed": True, + }, + { + "index_name": "idx_users_created", + "table_name": "users", + "columns": ["created_at"], + "unique": False, + "is_primary": False, + }, + ] + droppable = filter_droppable_table_indexes(indexes) + assert [idx["index_name"] for idx in droppable] == ["idx_users_created"] + + +@pytest.mark.parametrize( + ("dialect", "concurrent", "expected_drop", "expected_create"), + [ + ("postgres", True, "DROP INDEX IF EXISTS", "CREATE INDEX CONCURRENTLY"), + ("postgres", False, "DROP INDEX IF EXISTS", "CREATE INDEX"), + ("sqlite", False, "DROP INDEX IF EXISTS", "CREATE INDEX"), + ("mysql", False, " ON ", "CREATE INDEX"), + ("mssql", False, " ON ", "CREATE INDEX"), + ("oracle", False, "DROP INDEX", "CREATE INDEX"), + ], +) +def test_build_drop_and_create_index_sql_across_dialects( + dialect: str, concurrent: bool, expected_drop: str, expected_create: str +) -> None: + """Index drop and recreate SQL builders generate dialect-appropriate DDL.""" + meta: dict[str, Any] = { + "index_name": "idx_orders_customer", + "table_name": "orders", + "columns": ["customer_id", "created_at"], + "unique": False, + "is_primary": False, + } + drop_sql = build_drop_index_sql(dialect, meta, "orders") + create_sql = build_create_index_sql(dialect, meta, "orders", concurrent=concurrent) + assert drop_sql is not None and expected_drop in drop_sql + assert create_sql is not None and expected_create in create_sql + + +def test_sync_without_indexes_drops_secondary_indexes_and_restores_on_error(tmp_path: Path) -> None: + """Sync without_indexes drops only secondary non-unique indexes and rebuilds them even on failure.""" + config = SqliteConfig(connection_config={"database": str(tmp_path / "indexes_sync.db")}) + with config.provide_session() as driver: + driver.execute("CREATE TABLE orders (id INTEGER PRIMARY KEY, sku TEXT UNIQUE, customer_id INTEGER)") + driver.execute("CREATE INDEX idx_orders_customer ON orders (customer_id)") + + initial_indexes = {idx["index_name"] for idx in driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_orders_customer" in initial_indexes + + with pytest.raises(RuntimeError, match="load boom"): + with driver.without_indexes("orders", concurrent=False): + mid_indexes = {idx["index_name"] for idx in driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_orders_customer" not in mid_indexes + raise RuntimeError("load boom") + + restored_indexes = {idx["index_name"] for idx in driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_orders_customer" in restored_indexes + + table = pa.Table.from_pydict({"id": [1, 2], "sku": ["A1", "B2"], "customer_id": [10, 20]}) + driver.load_from_arrow("orders", table, drop_indexes=True, concurrent_indexes=False) + final_indexes = {idx["index_name"] for idx in driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_orders_customer" in final_indexes + + +@pytest.mark.anyio +async def test_async_without_indexes_and_load_from_arrow_drop_indexes(tmp_path: Path) -> None: + """Async without_indexes and load_from_arrow(drop_indexes=True) drop and recreate secondary indexes.""" + config = AiosqliteConfig(connection_config={"database": str(tmp_path / "indexes_async.db")}) + try: + async with config.provide_session() as driver: + await driver.execute("CREATE TABLE orders (id INTEGER PRIMARY KEY, customer_id INTEGER)") + await driver.execute("CREATE INDEX idx_async_orders_cust ON orders (customer_id)") + + async with driver.without_indexes("orders", concurrent=False): + mid_indexes = {idx["index_name"] for idx in await driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_async_orders_cust" not in mid_indexes + + restored_indexes = {idx["index_name"] for idx in await driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_async_orders_cust" in restored_indexes + + table = pa.Table.from_pydict({"id": [1, 2], "customer_id": [100, 200]}) + await driver.load_from_arrow("orders", table, drop_indexes=True, concurrent_indexes=False) + final_indexes = {idx["index_name"] for idx in await driver.data_dictionary.get_indexes(driver, "orders")} + assert "idx_async_orders_cust" in final_indexes + finally: + await config.close_pool() From e8b906828dedeaed575d893590ee62840d09460f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 18:52:09 +0000 Subject: [PATCH 05/14] feat(duckdb): add container-aware memory clamping, tmpfs spill check, and post-connect parquet_metadata_cache --- docs/reference/driver/common.md | 2 +- sqlspec/adapters/duckdb/config.py | 10 + sqlspec/adapters/duckdb/core.py | 259 +++++++++++- sqlspec/adapters/duckdb/driver.py | 3 +- sqlspec/adapters/duckdb/pool.py | 107 ++++- sqlspec/config.py | 3 +- sqlspec/core/__init__.py | 10 + sqlspec/core/cache.py | 294 +++++++++++++- sqlspec/driver/__init__.py | 2 +- sqlspec/driver/_async.py | 3 +- sqlspec/driver/_common.py | 2 +- sqlspec/driver/_query_cache.py | 250 ------------ sqlspec/driver/_sync.py | 3 +- sqlspec/utils/arrow_helpers.py | 370 +++++++++--------- .../adapters/test_duckdb_resource_controls.py | 109 ++++++ .../adapters/test_postgres_pgcopy_binary.py | 54 +-- tests/unit/core/test_cache.py | 59 +++ tests/unit/driver/test_query_cache.py | 2 +- tools/scripts/pgo_training.py | 2 +- 19 files changed, 1028 insertions(+), 516 deletions(-) delete mode 100644 sqlspec/driver/_query_cache.py create mode 100644 tests/unit/adapters/test_duckdb_resource_controls.py diff --git a/docs/reference/driver/common.md b/docs/reference/driver/common.md index 9e1101a19..4dc8407a2 100644 --- a/docs/reference/driver/common.md +++ b/docs/reference/driver/common.md @@ -27,7 +27,7 @@ Shared driver execution results, exception handler protocols and base classes, m heading_level: 3 members: false -::: sqlspec.driver._query_cache.CachedQuery +::: sqlspec.core.cache.CachedQuery options: heading_level: 3 members: false diff --git a/sqlspec/adapters/duckdb/config.py b/sqlspec/adapters/duckdb/config.py index 0071b7adb..60e4bb762 100644 --- a/sqlspec/adapters/duckdb/config.py +++ b/sqlspec/adapters/duckdb/config.py @@ -48,10 +48,14 @@ class DuckDBConnectionParams(TypedDict): """Raw DuckDB engine configuration dictionary merged with modeled settings.""" memory_limit: NotRequired[str] """Maximum buffer-manager memory ceiling (for example ``'4GB'``).""" + memory_ceiling_fraction: NotRequired[float] + """Fraction of container or host memory used as an upper ceiling when ``memory_limit`` is unset or exceeds available memory (defaults to ``0.8``).""" threads: NotRequired[int] """Number of DuckDB execution worker threads.""" temp_directory: NotRequired[str] """Directory path used when spilling out-of-core intermediates to disk.""" + reject_tmpfs_temp_directory: NotRequired[bool] + """Raise ``ImproperConfigurationError`` when ``temp_directory`` resides on an in-memory ``tmpfs`` or ``ramfs`` mount (defaults to ``False``).""" max_temp_directory_size: NotRequired[str] """Maximum disk space cap for ``temp_directory`` spill files (for example ``'20GB'``).""" autoload_known_extensions: NotRequired[bool] @@ -168,6 +172,12 @@ class DuckDBSecretConfig(TypedDict): value: NotRequired["dict[str, Any]"] """Secret configuration values, applied only when the secret is created or ``replace`` is set.""" + value_provider: NotRequired["Callable[[], dict[str, Any]]"] + """Callable returning dynamic secret configuration values (such as short-lived OAuth or STS tokens).""" + + refresh_interval_seconds: NotRequired[float] + """Interval in seconds after which a secret backed by ``value_provider`` is refreshed via ``CREATE OR REPLACE SECRET`` (defaults to ``1800.0`` when ``value_provider`` is set).""" + provider: NotRequired[str] """Secret provider, such as config or credential_chain.""" diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index b24623376..2ada3ea5e 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -1,6 +1,7 @@ """DuckDB adapter compiled helpers.""" import contextlib +import os import re import tempfile from datetime import date, datetime @@ -22,12 +23,14 @@ DuckDBParserException, DuckDBPermissionException, DuckDBTransactionException, + duckdb_module, ) from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.exceptions import ( CheckViolationError, DataError, ForeignKeyViolationError, + ImproperConfigurationError, IntegrityError, NotFoundError, NotNullViolationError, @@ -42,6 +45,7 @@ from sqlspec.storage import ResolvedStorageTarget from sqlspec.storage.backends.fsspec import FSSpecBackend from sqlspec.storage.backends.obstore import ObStoreBackend +from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_decimal_converter, build_uuid_coercions, time_iso_convert @@ -73,11 +77,33 @@ "format_identifier", "normalize_execute_parameters", "read_duckdb_profile_json", + "resolve_duckdb_memory_limit", "resolve_rowcount", + "validate_duckdb_temp_directory", ) +logger = get_logger("sqlspec.adapters.duckdb.core") + ARROW_MIN_ROWS: Final[int] = 64 _VERSION_COMPONENTS: Final[int] = 3 +_CGROUP_UNBOUNDED_THRESHOLD: Final[int] = 1 << 60 +_MEMORY_LIMIT_RE: Final[re.Pattern[str]] = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([A-Za-z]*)\s*$") +_MEMORY_UNIT_MULTIPLIERS: Final[dict[str, int]] = { + "": 1, + "B": 1, + "KB": 1000, + "KIB": 1024, + "K": 1024, + "MB": 1000**2, + "MIB": 1024**2, + "M": 1024**2, + "GB": 1000**3, + "GIB": 1024**3, + "G": 1024**3, + "TB": 1000**4, + "TIB": 1024**4, + "T": 1024**4, +} _FIXED_ARRAY_PATTERN: Final[re.Pattern[str]] = re.compile(r"\[\d+\]") _ARROW_SAFE_SCALAR_TYPES: Final[frozenset[str]] = frozenset({ @@ -167,6 +193,156 @@ def collect_rows(fetched_data: "list[Any] | None", description: "list[Any] | Non return fetched_data, column_names +def _parse_memory_size_bytes(memory_limit: str) -> "int | None": + match = _MEMORY_LIMIT_RE.match(memory_limit) + if match is None: + return None + amount = float(match.group(1)) + unit = match.group(2).upper() + multiplier = _MEMORY_UNIT_MULTIPLIERS.get(unit) + if multiplier is None: + return None + return int(amount * multiplier) + + +def _detect_available_memory_bytes( + *, + cgroup_v2_path: "Path | None" = None, + cgroup_v1_path: "Path | None" = None, + host_memory_bytes: "int | None" = None, +) -> "int | None": + v2_file = cgroup_v2_path if cgroup_v2_path is not None else Path("/sys/fs/cgroup/memory.max") + if v2_file.exists(): + with contextlib.suppress(OSError, ValueError): + raw_v2 = v2_file.read_text(encoding="utf-8").strip() + if raw_v2 and raw_v2.lower() != "max": + parsed_v2 = int(raw_v2) + if 0 < parsed_v2 < _CGROUP_UNBOUNDED_THRESHOLD: + return parsed_v2 + + v1_file = cgroup_v1_path if cgroup_v1_path is not None else Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") + if v1_file.exists(): + with contextlib.suppress(OSError, ValueError): + raw_v1 = v1_file.read_text(encoding="utf-8").strip() + if raw_v1: + parsed_v1 = int(raw_v1) + if 0 < parsed_v1 < _CGROUP_UNBOUNDED_THRESHOLD: + return parsed_v1 + + if host_memory_bytes is not None and host_memory_bytes > 0: + return host_memory_bytes + + with contextlib.suppress(ValueError, OSError, AttributeError): + page_size = os.sysconf("SC_PAGE_SIZE") + phys_pages = os.sysconf("SC_PHYS_PAGES") + if isinstance(page_size, int) and isinstance(phys_pages, int) and page_size > 0 and phys_pages > 0: + return page_size * phys_pages + return None + + +def resolve_duckdb_memory_limit( + memory_limit: "str | None" = None, + *, + memory_ceiling_fraction: float = 0.8, + cgroup_v2_path: "Path | None" = None, + cgroup_v1_path: "Path | None" = None, + host_memory_bytes: "int | None" = None, +) -> "str | None": + """Resolve or clamp DuckDB ``memory_limit`` against container cgroup or host memory. + + Args: + memory_limit: Explicit memory limit string (such as ``'4GB'`` or ``'512MB'``), or ``None``. + memory_ceiling_fraction: Fraction of available memory used as the maximum ceiling. + cgroup_v2_path: Optional override path for cgroup v2 ``memory.max``. + cgroup_v1_path: Optional override path for cgroup v1 ``memory.limit_in_bytes``. + host_memory_bytes: Optional override for host physical memory in bytes. + + Returns: + Resolved DuckDB memory limit string, or ``memory_limit`` when system memory cannot be determined. + """ + available_bytes = _detect_available_memory_bytes( + cgroup_v2_path=cgroup_v2_path, cgroup_v1_path=cgroup_v1_path, host_memory_bytes=host_memory_bytes + ) + if available_bytes is None or memory_ceiling_fraction <= 0: + return memory_limit + + ceiling_bytes = max(1024 * 1024, int(available_bytes * memory_ceiling_fraction)) + ceiling_mb_str = f"{max(1, ceiling_bytes // (1024 * 1024))}MB" + if memory_limit is None: + return ceiling_mb_str + + requested_bytes = _parse_memory_size_bytes(memory_limit) + if requested_bytes is not None and requested_bytes > ceiling_bytes: + logger.warning( + "Requested DuckDB memory_limit %r exceeds %.0f%% of available memory (%s); clamping to %s", + memory_limit, + memory_ceiling_fraction * 100, + ceiling_mb_str, + ceiling_mb_str, + ) + return ceiling_mb_str + return memory_limit + + +def validate_duckdb_temp_directory( + temp_directory: "str | Path", *, reject_tmpfs: bool = False, mounts_content: "str | None" = None +) -> bool: + """Validate whether a DuckDB spill ``temp_directory`` resides on an in-memory ``tmpfs``/``ramfs`` mount. + + Args: + temp_directory: Target spill directory path. + reject_tmpfs: Raise ``ImproperConfigurationError`` when the directory is backed by ``tmpfs`` or ``ramfs``. + mounts_content: Optional ``/proc/mounts`` text override for testing. + + Returns: + ``True`` when the spill directory is not on ``tmpfs``/``ramfs``, or ``False`` when a warning was emitted. + + Raises: + ImproperConfigurationError: If ``reject_tmpfs=True`` and the path resides on ``tmpfs`` or ``ramfs``. + """ + raw_mounts = mounts_content + if raw_mounts is None: + mounts_path = Path("/proc/mounts") + if mounts_path.exists(): + with contextlib.suppress(OSError): + raw_mounts = mounts_path.read_text(encoding="utf-8") + if not raw_mounts: + return True + + target_path = Path(temp_directory) + with contextlib.suppress(OSError): + if target_path.exists(): + target_path = target_path.resolve() + elif target_path.parent.exists(): + target_path = target_path.parent.resolve() / target_path.name + target_str = str(target_path) + + best_mount_point = "" + best_fs_type = "" + for line in raw_mounts.splitlines(): + with contextlib.suppress(ValueError): + _device, mount_point, raw_fs_type, *_rest = line.split() + fs_type = raw_fs_type.lower() + normalized_mount = mount_point.rstrip("/") or "/" + matches = target_str == normalized_mount or target_str.startswith( + "/" if normalized_mount == "/" else f"{normalized_mount}/" + ) + if matches and len(normalized_mount) >= len(best_mount_point): + best_mount_point = normalized_mount + best_fs_type = fs_type + + if best_fs_type in {"tmpfs", "ramfs"}: + msg = ( + f"DuckDB temp_directory {str(temp_directory)!r} resides on in-memory filesystem " + f"{best_fs_type!r} (mount {best_mount_point!r}), which risks container OOM during out-of-core spills." + ) + if reject_tmpfs: + raise ImproperConfigurationError(msg) + logger.warning(msg) + return False + return True + + def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build connection configuration for pool creation. @@ -176,7 +352,15 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Dictionary with connection parameters. """ - pool_only_keys = {"pool_min_size", "pool_max_size", "pool_timeout", "pool_recycle_seconds", "health_check_interval"} + pool_only_keys = { + "pool_min_size", + "pool_max_size", + "pool_timeout", + "pool_recycle_seconds", + "health_check_interval", + "memory_ceiling_fraction", + "reject_tmpfs_temp_directory", + } alias_keys = {"db", "path", "file"} connect_parameters: dict[str, Any] = {} duckdb_config: dict[str, Any] = {} @@ -206,6 +390,20 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str if isinstance(extra, dict): duckdb_config.update({key: value for key, value in extra.items() if value is not None}) + if duckdb_config.get("temp_directory"): + validate_duckdb_temp_directory( + duckdb_config["temp_directory"], + reject_tmpfs=bool(connection_config.get("reject_tmpfs_temp_directory", False)), + ) + + if "memory_ceiling_fraction" in connection_config or "memory_limit" in duckdb_config: + fraction = float(connection_config.get("memory_ceiling_fraction", 0.8)) + resolved_limit = resolve_duckdb_memory_limit( + cast("str | None", duckdb_config.get("memory_limit")), memory_ceiling_fraction=fraction + ) + if resolved_limit is not None: + duckdb_config["memory_limit"] = resolved_limit + if duckdb_config: connect_parameters["config"] = duckdb_config @@ -448,9 +646,7 @@ def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = N def duckdb_version_tuple() -> "tuple[int, int, int]": """Return the installed DuckDB version as a 3-tuple of integers.""" try: - import duckdb - - raw = getattr(duckdb, "__version__", "") + raw = getattr(duckdb_module, "__version__", "") parts = [int(part) for part in raw.split(".") if part.isdigit()] while len(parts) < _VERSION_COMPONENTS: parts.append(0) @@ -482,9 +678,13 @@ def _build_storage_copy_sql(sql: str, file_format: str) -> str | None: ) -def _build_storage_read_sql(table: str, uri: str, file_format: str) -> str | None: - """Build a single-object reader with a validated, quoted target identifier.""" - if file_format != "parquet" or any(char in uri for char in "*?[]"): +def _build_storage_read_sql(table: str, uri: str, file_format: str, *, allow_globs: bool = False) -> str | None: + """Build a single-object or glob reader with a validated, quoted target identifier.""" + normalized_format = file_format.lower() + allowed_formats = {"parquet", "csv", "jsonl", "ndjson"} if allow_globs else {"parquet"} + if normalized_format not in allowed_formats: + return None + if not allow_globs and any(char in uri for char in "*?[]"): return None try: targets = Dialect.get_or_raise("duckdb").parse_into(exp.Table, table) @@ -502,13 +702,32 @@ def _build_storage_read_sql(table: str, uri: str, file_format: str) -> str | Non target = targets[0] for part in target.parts: part.set("quoted", True) - reader = exp.Anonymous( - this="read_parquet", - expressions=[ - exp.Placeholder(), - exp.Kwarg(this=exp.Var(this="hive_partitioning"), expression=exp.Boolean(this=False)), - ], - ) + if normalized_format == "parquet": + reader = exp.Anonymous( + this="read_parquet", + expressions=[ + exp.Placeholder(), + exp.Kwarg(this=exp.Var(this="hive_partitioning"), expression=exp.Boolean(this=False)), + ], + ) + elif normalized_format == "csv": + reader = exp.Anonymous( + this="read_csv", + expressions=[ + exp.Placeholder(), + exp.Kwarg(this=exp.Var(this="auto_detect"), expression=exp.Boolean(this=True)), + exp.Kwarg(this=exp.Var(this="header"), expression=exp.Boolean(this=True)), + ], + ) + else: + reader = exp.Anonymous( + this="read_json", + expressions=[ + exp.Placeholder(), + exp.Kwarg(this=exp.Var(this="format"), expression=exp.Literal.string("newline_delimited")), + exp.Kwarg(this=exp.Var(this="auto_detect"), expression=exp.Boolean(this=True)), + ], + ) return exp.Insert(this=target, expression=exp.select("*").from_(reader, copy=False)).sql( dialect="duckdb", copy=False ) @@ -637,11 +856,13 @@ def _native_storage_eligible( def _s3_storage_options_match( backend_options: "Mapping[str, Any]", native: "Mapping[str, Any]", provider_type: str ) -> bool: - """Compare static credentials and endpoint options for S3-compatible providers.""" + """Compare static credentials, OAuth bearer tokens, and endpoint options for S3/GCS/R2 providers.""" if native.keys() - { "key_id", "secret", "session_token", + "bearer_token", + "token", "region", "endpoint", "url_style", @@ -652,7 +873,7 @@ def _s3_storage_options_match( options: dict[str, Any] = {} for key, value in backend_options.items(): - normalized = key.removeprefix("aws_") + normalized = key.removeprefix("aws_").removeprefix("google_") if normalized in options and options[normalized] != value: return False options[normalized] = value @@ -663,6 +884,12 @@ def _s3_storage_options_match( if "allow_http" in options and options["allow_http"] != client["allow_http"]: return False options["allow_http"] = client["allow_http"] + if provider_type == "gcs" and ("bearer_token" in native or "token" in native): + if options.keys() - {"bearer_token", "token"}: + return False + native_token = native.get("bearer_token") or native.get("token") + option_token = options.get("bearer_token") or options.get("token") + return bool(native_token and option_token == native_token) if options.keys() - { "access_key_id", "secret_access_key", diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 25ed550f2..0c9292428 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -55,10 +55,9 @@ from sqlspec.adapters.duckdb._typing import DuckDBConnection from sqlspec.builder import QueryBuilder - from sqlspec.core import ArrowResult, SQLResult, Statement, StatementFilter + from sqlspec.core import ArrowResult, CachedQuery, SQLResult, Statement, StatementFilter from sqlspec.core.result._base import RowFormat from sqlspec.driver import ExecutionResult - from sqlspec.driver._query_cache import CachedQuery from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry from sqlspec.typing import ArrowReturnFormat, StatementParameters else: diff --git a/sqlspec/adapters/duckdb/pool.py b/sqlspec/adapters/duckdb/pool.py index 4b7efff51..1d5d9084d 100644 --- a/sqlspec/adapters/duckdb/pool.py +++ b/sqlspec/adapters/duckdb/pool.py @@ -161,10 +161,9 @@ def new_connection(self, record_thread_state: bool = False) -> DuckDBConnection: else: config_dict[key] = value - if config_dict: - connect_parameters["config"] = config_dict - - connection = duckdb.connect(**connect_parameters) + connection, deferred_parquet_metadata_cache = _open_duckdb_connection( + connect_parameters, config_dict, has_extensions=bool(self._extensions) + ) for ext_config in self._extensions: ext_name = ext_config.get("name") @@ -227,11 +226,21 @@ def new_connection(self, record_thread_state: bool = False) -> DuckDBConnection: error=install_error, ) - created_secrets: list[dict[str, Any]] = [ - {**secret_config, "value": dict(secret_config.get("value") or {})} - for secret_config in self._secrets - if _create_secret(connection, secret_config) - ] + if deferred_parquet_metadata_cache is not None: + with suppress(Exception): + connection.execute( + f"SET parquet_metadata_cache = {'true' if deferred_parquet_metadata_cache else 'false'}" + ) + + created_secrets: list[dict[str, Any]] = [] + secret_refresh_times: dict[str, float] = {} + now_monotonic = time.monotonic() + for secret_config in self._secrets: + resolved_secret = _resolve_secret_config(secret_config) + if _create_secret(connection, resolved_secret): + created_secrets.append({**resolved_secret, "value": dict(resolved_secret.get("value") or {})}) + if callable(secret_config.get("value_provider")) and resolved_secret.get("name"): + secret_refresh_times[str(resolved_secret["name"])] = now_monotonic if self._on_connection_create: modified = self._on_connection_create(connection) @@ -247,6 +256,7 @@ def new_connection(self, record_thread_state: bool = False) -> DuckDBConnection: if record_thread_state: self._thread_local.storage_extensions = frozenset(loaded_extensions) self._thread_local.storage_secrets = tuple(created_secrets) + self._thread_local.secret_refresh_times = secret_refresh_times self._thread_local.storage_protocols = frozenset(storage_protocols) self._thread_local.profile_path = profile_path return connection @@ -437,8 +447,43 @@ def _get_thread_connection(self) -> DuckDBConnection: self._thread_local.created_at = time.time() self._thread_local.last_used = time.time() + self._refresh_dynamic_secrets(cast("DuckDBConnection", self._thread_local.connection)) return cast("DuckDBConnection", self._thread_local.connection) + def _refresh_dynamic_secrets(self, connection: DuckDBConnection) -> None: + """Refresh secrets configured with ``value_provider`` when ``refresh_interval_seconds`` has elapsed.""" + if not any(callable(secret.get("value_provider")) for secret in self._secrets): + return + thread_state = self._thread_local.__dict__ + refresh_times = cast("dict[str, float]", thread_state.setdefault("secret_refresh_times", {})) + current_secrets = list(cast("tuple[dict[str, Any], ...]", thread_state.get("storage_secrets", ()))) + now_monotonic = time.monotonic() + updated = False + for secret_config in self._secrets: + provider = secret_config.get("value_provider") + secret_name = secret_config.get("name") + if not callable(provider) or not secret_name: + continue + interval = float(secret_config.get("refresh_interval_seconds", 1800.0)) + last_refreshed = refresh_times.get(str(secret_name), 0.0) + if now_monotonic - last_refreshed < interval: + continue + resolved_secret = _resolve_secret_config(secret_config) + if _create_secret(connection, resolved_secret): + refresh_times[str(secret_name)] = now_monotonic + updated_entry = {**resolved_secret, "value": dict(resolved_secret.get("value") or {})} + replaced = False + for idx, existing in enumerate(current_secrets): + if str(existing.get("name", "")).lower() == str(secret_name).lower(): + current_secrets[idx] = updated_entry + replaced = True + break + if not replaced: + current_secrets.append(updated_entry) + updated = True + if updated: + self._thread_local.storage_secrets = tuple(current_secrets) + def _retire_connection(self, connection: DuckDBConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._lock: @@ -480,6 +525,34 @@ def _is_connection_alive(self, connection: DuckDBConnection) -> bool: return True +def _open_duckdb_connection( + connect_parameters: dict[str, Any], config_dict: dict[str, Any], *, has_extensions: bool +) -> tuple[DuckDBConnection, bool | None]: + """Open a DuckDB connection, deferring ``parquet_metadata_cache`` when extensions are configured.""" + deferred_parquet_metadata_cache: bool | None = None + if has_extensions and "parquet_metadata_cache" in config_dict: + deferred_parquet_metadata_cache = bool(config_dict.pop("parquet_metadata_cache")) + + if config_dict: + connect_parameters["config"] = config_dict + else: + connect_parameters.pop("config", None) + + try: + connection = duckdb.connect(**connect_parameters) + except Exception as exc: + if "parquet_metadata_cache" in config_dict and "parquet_metadata_cache" in str(exc): + deferred_parquet_metadata_cache = bool(config_dict.pop("parquet_metadata_cache")) + if config_dict: + connect_parameters["config"] = config_dict + else: + connect_parameters.pop("config", None) + connection = duckdb.connect(**connect_parameters) + else: + raise + return connection, deferred_parquet_metadata_cache + + def _validate_sql_identifier(value: str, field_name: str) -> None: """Raise ValueError if value is not safe to interpolate as a SQL identifier.""" if not _SQL_IDENTIFIER_RE.fullmatch(value): @@ -490,7 +563,23 @@ def _validate_sql_identifier(value: str, field_name: str) -> None: raise ValueError(msg) +def _resolve_secret_config(secret_config: dict[str, Any]) -> dict[str, Any]: + """Evaluate ``value_provider`` if present and return a normalized secret configuration dict.""" + resolved = dict(secret_config) + value_dict = dict(resolved.get("value") or {}) + provider_fn = resolved.pop("value_provider", None) + if callable(provider_fn): + dynamic_values = provider_fn() + if isinstance(dynamic_values, dict): + value_dict.update(dynamic_values) + resolved.setdefault("replace", True) + if value_dict: + resolved["value"] = value_dict + return resolved + + def _create_secret(connection: DuckDBConnection, secret_config: dict[str, Any]) -> bool: + secret_config = _resolve_secret_config(secret_config) secret_name = secret_config.get("name") secret_type = secret_config.get("secret_type") if not (secret_name and secret_type): diff --git a/sqlspec/config.py b/sqlspec/config.py index 149904c84..58e477a5d 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -19,7 +19,7 @@ from typing_extensions import TypedDict -from sqlspec.core import TypeCoercionCapabilities +from sqlspec.core import StatementCacheHolder, TypeCoercionCapabilities, resolve_statement_cache_size from sqlspec.core.config_runtime import ( build_default_statement_config, close_async_pool, @@ -28,7 +28,6 @@ create_sync_pool, seed_runtime_driver_features, ) -from sqlspec.driver._query_cache import StatementCacheHolder, resolve_statement_cache_size from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.loader import SQLFileLoader diff --git a/sqlspec/core/__init__.py b/sqlspec/core/__init__.py index b05df993b..5e3e5b951 100644 --- a/sqlspec/core/__init__.py +++ b/sqlspec/core/__init__.py @@ -29,13 +29,17 @@ from sqlspec.core._ordering import NullsPlacement, apply_direction, default_nulls, ordered from sqlspec.core._pagination import CursorPagination, OffsetPagination, Pagination from sqlspec.core.cache import ( + STMT_CACHE_MAX_SIZE, CacheConfig, + CachedQuery, CachedStatement, CacheKey, CacheStats, FiltersView, LRUCache, NamespacedCache, + QueryCache, + StatementCacheHolder, clear_all_caches, create_cache_key, get_cache, @@ -46,6 +50,7 @@ log_cache_stats, reset_pipeline_registry, reset_stats_only, + resolve_statement_cache_size, update_cache_config, ) from sqlspec.core.capabilities import TypeCoercionCapabilities @@ -191,6 +196,7 @@ "EXECUTE_MANY_MIN_ROWS", "PARAMETER_REGEX", "SQL", + "STMT_CACHE_MAX_SIZE", "AnyCollectionFilter", "ArrowResult", "BaseInputConverter", @@ -199,6 +205,7 @@ "CacheConfig", "CacheKey", "CacheStats", + "CachedQuery", "CachedStatement", "ChoicesFilter", "CompiledSQL", @@ -239,6 +246,7 @@ "ParameterStyleConfig", "ParameterValidator", "ProcessedState", + "QueryCache", "SQLProcessor", "SQLResult", "SearchFilter", @@ -247,6 +255,7 @@ "StackOperation", "StackResult", "Statement", + "StatementCacheHolder", "StatementConfig", "StatementFilter", "StatementResult", @@ -332,6 +341,7 @@ "reset_pipeline_registry", "reset_stats_only", "resolve_param_type", + "resolve_statement_cache_size", "safe_modify_with_cte", "split_sql_script", "update_cache_config", diff --git a/sqlspec/core/cache.py b/sqlspec/core/cache.py index 0dd9b9a05..523b649bc 100644 --- a/sqlspec/core/cache.py +++ b/sqlspec/core/cache.py @@ -12,8 +12,9 @@ import logging import threading import time +from collections import OrderedDict from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, final from mypy_extensions import mypyc_attr @@ -22,25 +23,32 @@ get_statement_pipeline_metrics, reset_statement_pipeline_cache, ) +from sqlspec.core.parameters import ParameterProcessor, ParameterStyle from sqlspec.utils.logging import get_logger, log_with_context from sqlspec.utils.type_guards import has_field_name, has_filter_attributes if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterator, Mapping import sqlglot.expressions as exp + from sqlspec.core.compiler import OperationProfile, OperationType from sqlspec.core.parameters import ParameterProfile + from sqlspec.core.statement import SQL, ProcessedState, StatementConfig __all__ = ( + "STMT_CACHE_MAX_SIZE", "CacheConfig", "CacheKey", "CacheStats", + "CachedQuery", "CachedStatement", "FiltersView", "LRUCache", "NamespacedCache", + "QueryCache", + "StatementCacheHolder", "clear_all_caches", "create_cache_key", "get_cache", @@ -52,6 +60,7 @@ "log_cache_stats", "reset_pipeline_registry", "reset_stats_only", + "resolve_statement_cache_size", "set_cache_instances", "update_cache_config", ) @@ -60,6 +69,14 @@ DEFAULT_MAX_SIZE: Final = 10000 DEFAULT_TTL_SECONDS: Final = 3600 +STMT_CACHE_MAX_SIZE: Final[int] = 1024 +_MAX_REGISTERED_HOLDERS: Final[int] = 128 +_CACHED_NAMED_STYLES: Final[frozenset[str]] = frozenset(( + ParameterStyle.NAMED_COLON.value, + ParameterStyle.NAMED_AT.value, + ParameterStyle.NAMED_DOLLAR.value, + ParameterStyle.NAMED_PYFORMAT.value, +)) CACHE_NODE_SLOTS: Final = ("key", "value", "prev", "next", "timestamp") @@ -390,6 +407,20 @@ def _move_to_head(self, node: CacheNode) -> None: _default_cache: LRUCache | None = None _cache_lock = threading.Lock() +_statement_cache_holders: "OrderedDict[int, StatementCacheHolder]" = OrderedDict() +_holders_lock = threading.Lock() + + +def _register_statement_cache_holder(holder: "StatementCacheHolder") -> None: + """Register a StatementCacheHolder in the bounded process-wide L1 cache registry.""" + holder_id = id(holder) + with _holders_lock: + if holder_id in _statement_cache_holders: + _statement_cache_holders.move_to_end(holder_id) + else: + while len(_statement_cache_holders) >= _MAX_REGISTERED_HOLDERS: + _statement_cache_holders.popitem(last=False) + _statement_cache_holders[holder_id] = holder def get_default_cache() -> LRUCache: @@ -429,12 +460,16 @@ def set_cache_instances(default_cache: "LRUCache | None", namespaced_cache: "Nam def clear_all_caches() -> None: - """Clear all cache instances.""" + """Clear all L1 driver statement caches and L2 global caches.""" if _default_cache is not None: _default_cache.clear() cache = get_cache() cache.clear() reset_statement_pipeline_cache() + with _holders_lock: + holders = list(_statement_cache_holders.values()) + for holder in holders: + holder.clear() def reset_stats_only() -> None: @@ -444,6 +479,10 @@ def reset_stats_only() -> None: if _namespaced_cache is not None: for cache in _namespaced_cache._caches.values(): cache.get_stats().reset() + with _holders_lock: + holders = list(_statement_cache_holders.values()) + for holder in holders: + holder.query_cache.get_stats().reset() def get_cache_statistics() -> "dict[str, CacheStats]": @@ -457,6 +496,17 @@ def get_cache_statistics() -> "dict[str, CacheStats]": stats["default"] = default_cache.get_stats() cache = get_cache() stats["namespaced"] = cache.get_stats() + query_stats = CacheStats() + with _holders_lock: + holders = list(_statement_cache_holders.values()) + for holder in holders: + holder_stats = holder.query_cache.get_stats() + query_stats.hits += holder_stats.hits + query_stats.misses += holder_stats.misses + query_stats.evictions += holder_stats.evictions + query_stats.total_operations += holder_stats.total_operations + query_stats.memory_usage += holder_stats.memory_usage + stats["query"] = query_stats return stats @@ -544,6 +594,10 @@ def update_cache_config(config: CacheConfig) -> None: _namespaced_cache.clear() _default_cache = None _namespaced_cache = None + with _holders_lock: + holders = list(_statement_cache_holders.values()) + for holder in holders: + holder.clear() log_with_context( logger, @@ -1112,3 +1166,237 @@ def _canonicalize_cache_filters(filters: "list[Filter]") -> "tuple[Filter, ...]" def _filter_sort_key(filter_obj: "Filter") -> "tuple[str, str, str]": return filter_obj.field_name, filter_obj.operation, str(filter_obj.value) + + +@final +@mypyc_attr(allow_interpreted_subclasses=False) +class CachedQuery: + """Cached query metadata for fast-path execution.""" + + compiled_sql: str + parameter_profile: "ParameterProfile" + input_named_parameters: "tuple[str, ...]" + applied_wrap_types: bool + parameter_casts: "dict[int, str]" + operation_type: "OperationType" + operation_profile: "OperationProfile" + param_count: int + processed_state: "ProcessedState" + needs_style_remap: bool + + __slots__ = ( + "applied_wrap_types", + "compiled_sql", + "input_named_parameters", + "needs_style_remap", + "operation_profile", + "operation_type", + "param_count", + "parameter_casts", + "parameter_profile", + "processed_state", + ) + + def __init__( + self, + compiled_sql: str, + parameter_profile: "ParameterProfile", + input_named_parameters: "tuple[str, ...]", + applied_wrap_types: bool, + parameter_casts: "dict[int, str]", + operation_type: "OperationType", + operation_profile: "OperationProfile", + param_count: int, + processed_state: "ProcessedState", + needs_style_remap: bool | None = None, + ) -> None: + self.compiled_sql = compiled_sql + self.parameter_profile = parameter_profile + self.input_named_parameters = input_named_parameters + self.applied_wrap_types = applied_wrap_types + self.parameter_casts = parameter_casts + self.operation_type = operation_type + self.operation_profile = operation_profile + self.param_count = param_count + self.processed_state = processed_state + self.needs_style_remap = ( + bool(_CACHED_NAMED_STYLES.intersection(parameter_profile.styles)) + if needs_style_remap is None + else needs_style_remap + ) + + +@final +@mypyc_attr(allow_interpreted_subclasses=False) +class QueryCache: + """LRU cache for compiled query metadata (L1 per-config fast-path cache).""" + + __slots__ = ("_cache", "_lock", "_max_size", "_stats") + + def __init__(self, max_size: int = STMT_CACHE_MAX_SIZE) -> None: + self._cache: OrderedDict[str, CachedQuery] = OrderedDict() + self._lock = threading.Lock() + self._max_size = max_size + self._stats = CacheStats() + + def __len__(self) -> int: + with self._lock: + return len(self._cache) + + def get(self, sql: str) -> "CachedQuery | None": + with self._lock: + entry = self._cache.get(sql) + if entry is None: + self._stats.record_miss() + return None + self._cache.move_to_end(sql) + self._stats.record_hit() + return entry + + def set(self, sql: str, entry: "CachedQuery") -> None: + if self._max_size <= 0: + return + with self._lock: + if sql in self._cache: + self._cache.move_to_end(sql) + elif len(self._cache) >= self._max_size: + self._cache.popitem(last=False) + self._stats.record_eviction() + self._cache[sql] = entry + + def get_stats(self) -> CacheStats: + """Return cache hit/miss/eviction statistics for this L1 query cache.""" + return self._stats + + def clear(self) -> None: + """Clear all cached entries and reset statistics.""" + with self._lock: + self._cache.clear() + self._stats.reset() + + +@final +@mypyc_attr(allow_interpreted_subclasses=False) +class StatementCacheHolder: + """Shared per-config holder for statement caches, rebind processor, and execution hints.""" + + statement_config: "StatementConfig" + max_size: int + query_cache: QueryCache + rebind_processor: ParameterProcessor + _statement_cache: "OrderedDict[str, SQL]" + _row_count_hints: "OrderedDict[str, int]" + _bulk_insert_sql: "OrderedDict[str, tuple[str, str]]" + _lock: threading.Lock + + __slots__ = ( + "_bulk_insert_sql", + "_lock", + "_row_count_hints", + "_statement_cache", + "max_size", + "query_cache", + "rebind_processor", + "statement_config", + ) + + def __init__(self, statement_config: "StatementConfig", max_size: int = STMT_CACHE_MAX_SIZE) -> None: + self.statement_config = statement_config + self.max_size = max_size + self.query_cache = QueryCache(max_size) + self.rebind_processor = ParameterProcessor( + converter=statement_config.parameter_converter, + validator=statement_config.parameter_validator, + cache_max_size=0, + validator_cache_max_size=0, + ) + self._statement_cache = OrderedDict() + self._row_count_hints = OrderedDict() + self._bulk_insert_sql = OrderedDict() + self._lock = threading.Lock() + _register_statement_cache_holder(self) + + def get_statement(self, sql: str) -> "SQL | None": + """Retrieve a cached SQL object and mark it as recently used.""" + with self._lock: + entry = self._statement_cache.get(sql) + if entry is None: + return None + self._statement_cache.move_to_end(sql) + return entry + + def set_statement(self, sql: str, stmt: "SQL", max_size: int | None = None) -> None: + """Store a prepared SQL object in the bounded LRU cache.""" + effective_max = self.max_size if max_size is None else max_size + if effective_max <= 0: + return + with self._lock: + if sql in self._statement_cache: + self._statement_cache.move_to_end(sql) + else: + while len(self._statement_cache) >= effective_max: + self._statement_cache.popitem(last=False) + self._statement_cache[sql] = stmt + + def get_row_count_hint(self, compiled_sql: str) -> int | None: + """Return the last observed row count for a compiled SQL string.""" + with self._lock: + entry = self._row_count_hints.get(compiled_sql) + if entry is None: + return None + self._row_count_hints.move_to_end(compiled_sql) + return entry + + def set_row_count_hint(self, compiled_sql: str, row_count: int) -> None: + """Record the observed row count for a compiled SQL string.""" + if self.max_size <= 0: + return + with self._lock: + if compiled_sql in self._row_count_hints: + self._row_count_hints.move_to_end(compiled_sql) + elif len(self._row_count_hints) >= self.max_size: + self._row_count_hints.popitem(last=False) + self._row_count_hints[compiled_sql] = row_count + + def get_bulk_insert_sql(self, compiled_sql: str) -> "tuple[str, str] | None": + """Return cached (table_sql, columns_sql) fragments for a bulk INSERT statement.""" + with self._lock: + entry = self._bulk_insert_sql.get(compiled_sql) + if entry is None: + return None + self._bulk_insert_sql.move_to_end(compiled_sql) + return entry + + def set_bulk_insert_sql(self, compiled_sql: str, entry: "tuple[str, str]") -> None: + """Cache (table_sql, columns_sql) fragments for a bulk INSERT statement.""" + if self.max_size <= 0: + return + with self._lock: + if compiled_sql in self._bulk_insert_sql: + self._bulk_insert_sql.move_to_end(compiled_sql) + elif len(self._bulk_insert_sql) >= self.max_size: + self._bulk_insert_sql.popitem(last=False) + self._bulk_insert_sql[compiled_sql] = entry + + def clear(self) -> None: + """Clear all cached entries.""" + self.query_cache.clear() + with self._lock: + self._statement_cache.clear() + self._row_count_hints.clear() + self._bulk_insert_sql.clear() + + +def resolve_statement_cache_size(driver_features: "Mapping[str, object] | None") -> int: + """Resolve the statement cache capacity from driver_features.""" + if not driver_features: + return STMT_CACHE_MAX_SIZE + cache_size = driver_features.get("sqlspec_statement_cache_size", STMT_CACHE_MAX_SIZE) + if isinstance(cache_size, int): + return max(0, cache_size) + if isinstance(cache_size, str): + try: + return max(0, int(cache_size)) + except ValueError: + return STMT_CACHE_MAX_SIZE + return STMT_CACHE_MAX_SIZE diff --git a/sqlspec/driver/__init__.py b/sqlspec/driver/__init__.py index 1fdb31ff9..be24cf041 100644 --- a/sqlspec/driver/__init__.py +++ b/sqlspec/driver/__init__.py @@ -1,5 +1,6 @@ """Driver protocols and base classes for database adapters.""" +from sqlspec.core.cache import CachedQuery from sqlspec.driver._async import ( AsyncDataDictionaryBase, AsyncDriverAdapterBase, @@ -22,7 +23,6 @@ validate_savepoint_name, ) from sqlspec.driver._exception_handler import BaseAsyncExceptionHandler, BaseSyncExceptionHandler -from sqlspec.driver._query_cache import CachedQuery from sqlspec.driver._sql_helpers import convert_to_dialect from sqlspec.driver._stream import AsyncRowStream, SyncRowStream, rows_to_dicts from sqlspec.driver._sync import ( diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index f7fb32d4d..19235d8b9 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -15,7 +15,7 @@ from mypy_extensions import mypyc_attr -from sqlspec.core import SQL, StackResult, create_arrow_result +from sqlspec.core import SQL, CachedQuery, StackResult, create_arrow_result from sqlspec.core.result import DMLResult from sqlspec.core.stack import StackOperation, StatementStack from sqlspec.data_dictionary import ( @@ -46,7 +46,6 @@ validate_savepoint_name, ) from sqlspec.driver._exception_handler import _run_with_async_exception_handler -from sqlspec.driver._query_cache import CachedQuery from sqlspec.driver._sql_helpers import DEFAULT_PRETTY from sqlspec.driver._sql_helpers import convert_to_dialect as _convert_to_dialect_impl from sqlspec.driver._storage_helpers import ( diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 0bb0ea5db..d590abdf6 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -39,6 +39,7 @@ split_sql_script, ) from sqlspec.core._pool import get_processed_state_pool, get_sql_pool +from sqlspec.core.cache import CachedQuery, QueryCache, StatementCacheHolder, resolve_statement_cache_size from sqlspec.core.filters import find_filter as _find_filter_impl from sqlspec.core.metrics import StackExecutionMetrics from sqlspec.core.parameters import ( @@ -74,7 +75,6 @@ sort_dependencies, ) from sqlspec.data_dictionary._registry import get_dialect_config -from sqlspec.driver._query_cache import CachedQuery, QueryCache, StatementCacheHolder, resolve_statement_cache_size from sqlspec.driver._storage_helpers import ( CAPABILITY_HINTS, arrow_batch_to_rows, diff --git a/sqlspec/driver/_query_cache.py b/sqlspec/driver/_query_cache.py deleted file mode 100644 index f5233b1ad..000000000 --- a/sqlspec/driver/_query_cache.py +++ /dev/null @@ -1,250 +0,0 @@ -"""Query cache for fast-path statement execution.""" - -import threading -from collections import OrderedDict -from typing import TYPE_CHECKING, Final, final - -from mypy_extensions import mypyc_attr - -from sqlspec.core.parameters import ParameterProcessor, ParameterStyle - -if TYPE_CHECKING: - from collections.abc import Mapping - - from sqlspec.core.compiler import OperationProfile, OperationType - from sqlspec.core.parameters import ParameterProfile - from sqlspec.core.statement import SQL, ProcessedState, StatementConfig - -__all__ = ("STMT_CACHE_MAX_SIZE", "CachedQuery", "QueryCache", "StatementCacheHolder", "resolve_statement_cache_size") - -STMT_CACHE_MAX_SIZE: Final[int] = 1024 -_CACHED_NAMED_STYLES: Final[frozenset[str]] = frozenset(( - ParameterStyle.NAMED_COLON.value, - ParameterStyle.NAMED_AT.value, - ParameterStyle.NAMED_DOLLAR.value, - ParameterStyle.NAMED_PYFORMAT.value, -)) - - -@final -@mypyc_attr(allow_interpreted_subclasses=False) -class CachedQuery: - """Cached query metadata for fast-path execution.""" - - compiled_sql: str - parameter_profile: "ParameterProfile" - input_named_parameters: "tuple[str, ...]" - applied_wrap_types: bool - parameter_casts: "dict[int, str]" - operation_type: "OperationType" - operation_profile: "OperationProfile" - param_count: int - processed_state: "ProcessedState" - needs_style_remap: bool - - __slots__ = ( - "applied_wrap_types", - "compiled_sql", - "input_named_parameters", - "needs_style_remap", - "operation_profile", - "operation_type", - "param_count", - "parameter_casts", - "parameter_profile", - "processed_state", - ) - - def __init__( - self, - compiled_sql: str, - parameter_profile: "ParameterProfile", - input_named_parameters: "tuple[str, ...]", - applied_wrap_types: bool, - parameter_casts: "dict[int, str]", - operation_type: "OperationType", - operation_profile: "OperationProfile", - param_count: int, - processed_state: "ProcessedState", - needs_style_remap: bool | None = None, - ) -> None: - self.compiled_sql = compiled_sql - self.parameter_profile = parameter_profile - self.input_named_parameters = input_named_parameters - self.applied_wrap_types = applied_wrap_types - self.parameter_casts = parameter_casts - self.operation_type = operation_type - self.operation_profile = operation_profile - self.param_count = param_count - self.processed_state = processed_state - self.needs_style_remap = ( - bool(_CACHED_NAMED_STYLES.intersection(parameter_profile.styles)) - if needs_style_remap is None - else needs_style_remap - ) - - -@final -@mypyc_attr(allow_interpreted_subclasses=False) -class QueryCache: - """LRU cache for compiled query metadata.""" - - __slots__ = ("_cache", "_lock", "_max_size") - - def __init__(self, max_size: int = STMT_CACHE_MAX_SIZE) -> None: - self._cache: OrderedDict[str, CachedQuery] = OrderedDict() - self._lock = threading.Lock() - self._max_size = max_size - - def __len__(self) -> int: - with self._lock: - return len(self._cache) - - def get(self, sql: str) -> "CachedQuery | None": - with self._lock: - entry = self._cache.get(sql) - if entry is None: - return None - self._cache.move_to_end(sql) - return entry - - def set(self, sql: str, entry: "CachedQuery") -> None: - if self._max_size <= 0: - return - with self._lock: - if sql in self._cache: - self._cache.move_to_end(sql) - elif len(self._cache) >= self._max_size: - self._cache.popitem(last=False) - self._cache[sql] = entry - - def clear(self) -> None: - """Clear all cached entries.""" - with self._lock: - self._cache.clear() - - -@final -@mypyc_attr(allow_interpreted_subclasses=False) -class StatementCacheHolder: - """Shared per-config holder for statement caches, rebind processor, and execution hints.""" - - statement_config: "StatementConfig" - max_size: int - query_cache: QueryCache - rebind_processor: ParameterProcessor - _statement_cache: "OrderedDict[str, SQL]" - _row_count_hints: "OrderedDict[str, int]" - _bulk_insert_sql: "OrderedDict[str, tuple[str, str]]" - _lock: threading.Lock - - __slots__ = ( - "_bulk_insert_sql", - "_lock", - "_row_count_hints", - "_statement_cache", - "max_size", - "query_cache", - "rebind_processor", - "statement_config", - ) - - def __init__(self, statement_config: "StatementConfig", max_size: int = STMT_CACHE_MAX_SIZE) -> None: - self.statement_config = statement_config - self.max_size = max_size - self.query_cache = QueryCache(max_size) - self.rebind_processor = ParameterProcessor( - converter=statement_config.parameter_converter, - validator=statement_config.parameter_validator, - cache_max_size=0, - validator_cache_max_size=0, - ) - self._statement_cache = OrderedDict() - self._row_count_hints = OrderedDict() - self._bulk_insert_sql = OrderedDict() - self._lock = threading.Lock() - - def get_statement(self, sql: str) -> "SQL | None": - """Retrieve a cached SQL object and mark it as recently used.""" - with self._lock: - entry = self._statement_cache.get(sql) - if entry is None: - return None - self._statement_cache.move_to_end(sql) - return entry - - def set_statement(self, sql: str, stmt: "SQL", max_size: int | None = None) -> None: - """Store a prepared SQL object in the bounded LRU cache.""" - effective_max = self.max_size if max_size is None else max_size - if effective_max <= 0: - return - with self._lock: - if sql in self._statement_cache: - self._statement_cache.move_to_end(sql) - else: - while len(self._statement_cache) >= effective_max: - self._statement_cache.popitem(last=False) - self._statement_cache[sql] = stmt - - def get_row_count_hint(self, compiled_sql: str) -> int | None: - """Return the last observed row count for a compiled SQL string.""" - with self._lock: - entry = self._row_count_hints.get(compiled_sql) - if entry is None: - return None - self._row_count_hints.move_to_end(compiled_sql) - return entry - - def set_row_count_hint(self, compiled_sql: str, row_count: int) -> None: - """Record the observed row count for a compiled SQL string.""" - if self.max_size <= 0: - return - with self._lock: - if compiled_sql in self._row_count_hints: - self._row_count_hints.move_to_end(compiled_sql) - elif len(self._row_count_hints) >= self.max_size: - self._row_count_hints.popitem(last=False) - self._row_count_hints[compiled_sql] = row_count - - def get_bulk_insert_sql(self, compiled_sql: str) -> "tuple[str, str] | None": - """Return cached (table_sql, columns_sql) fragments for a bulk INSERT statement.""" - with self._lock: - entry = self._bulk_insert_sql.get(compiled_sql) - if entry is None: - return None - self._bulk_insert_sql.move_to_end(compiled_sql) - return entry - - def set_bulk_insert_sql(self, compiled_sql: str, entry: "tuple[str, str]") -> None: - """Cache (table_sql, columns_sql) fragments for a bulk INSERT statement.""" - if self.max_size <= 0: - return - with self._lock: - if compiled_sql in self._bulk_insert_sql: - self._bulk_insert_sql.move_to_end(compiled_sql) - elif len(self._bulk_insert_sql) >= self.max_size: - self._bulk_insert_sql.popitem(last=False) - self._bulk_insert_sql[compiled_sql] = entry - - def clear(self) -> None: - """Clear all cached entries.""" - self.query_cache.clear() - with self._lock: - self._statement_cache.clear() - self._row_count_hints.clear() - self._bulk_insert_sql.clear() - - -def resolve_statement_cache_size(driver_features: "Mapping[str, object] | None") -> int: - """Resolve the statement cache capacity from driver_features.""" - if not driver_features: - return STMT_CACHE_MAX_SIZE - cache_size = driver_features.get("sqlspec_statement_cache_size", STMT_CACHE_MAX_SIZE) - if isinstance(cache_size, int): - return max(0, cache_size) - if isinstance(cache_size, str): - try: - return max(0, int(cache_size)) - except ValueError: - return STMT_CACHE_MAX_SIZE - return STMT_CACHE_MAX_SIZE diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index a401b660c..8f49a40d3 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -14,7 +14,7 @@ from mypy_extensions import mypyc_attr -from sqlspec.core import SQL, StackResult, create_arrow_result +from sqlspec.core import SQL, CachedQuery, StackResult, create_arrow_result from sqlspec.core.result import DMLResult from sqlspec.core.stack import StackOperation, StatementStack from sqlspec.data_dictionary import ( @@ -45,7 +45,6 @@ validate_savepoint_name, ) from sqlspec.driver._exception_handler import _run_with_sync_exception_handler -from sqlspec.driver._query_cache import CachedQuery from sqlspec.driver._sql_helpers import DEFAULT_PRETTY from sqlspec.driver._sql_helpers import convert_to_dialect as _convert_to_dialect_impl from sqlspec.driver._storage_helpers import ( diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index ec23c9d5c..56f009c44 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -904,7 +904,7 @@ def _encode_pg_binary_fast_path_1( if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): dim = int(field_type.list_size) - flat_values = col.values[col.offset * dim : (col.offset + num_rows) * dim] + flat_values = col.flatten() if flat_values.null_count > 0: return None is_f16 = pa.types.is_float16(field_type.value_type) @@ -921,12 +921,7 @@ def _encode_pg_binary_fast_path_1( (unused_key, ">u2"), (val_key, elem_dtype, (dim,)), ]) - assignments.extend([ - (len_key, 4 + elem_size * dim), - (dim_key, dim), - (unused_key, 0), - (val_key, arr_2d), - ]) + assignments.extend([(len_key, 4 + elem_size * dim), (dim_key, dim), (unused_key, 0), (val_key, arr_2d)]) continue return None @@ -938,204 +933,203 @@ def _encode_pg_binary_fast_path_1( return out.tobytes() -def _encode_pg_binary_fast_path_2( - batch: "ArrowRecordBatch", - columns: "list[str]", - vector_columns: frozenset[str], - jsonb_columns: frozenset[str], -) -> bytes: - """Encode variable-width and nullable columns by slicing Arrow buffers directly.""" +def _create_pg_fixed_column_encoder( + col: Any, field_type: Any, num_rows: int, null_mask: Any +) -> "Callable[[int, bytearray], None] | None": + """Create an encoder for fixed-width columns or fixed-size list vectors.""" import numpy as np import pyarrow as pa - num_rows = int(batch.num_rows) - row_header = struct.pack(">h", len(columns)) - col_encoders: list[Callable[[int, bytearray], None]] = [] + fixed_info = _extract_pg_fixed_column_array(col, field_type, num_rows) + if fixed_info is not None: + byte_width, np_dtype, arr = fixed_info + packed = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("val", np_dtype)])) + packed["len"] = byte_width + packed["val"] = arr + packed_mv = memoryview(packed.tobytes()) + stride = 4 + byte_width - for col_name in columns: - col = batch.column(col_name) - if isinstance(col, pa.ChunkedArray): - col = col.combine_chunks() - field_type = col.type - has_nulls = col.null_count > 0 - null_mask = col.is_null().to_numpy(zero_copy_only=False) if has_nulls else None - - if col_name not in jsonb_columns: - fixed_info = _extract_pg_fixed_column_array(col, field_type, num_rows) - if fixed_info is not None: - byte_width, np_dtype, arr = fixed_info - packed = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("val", np_dtype)])) - packed["len"] = byte_width - packed["val"] = arr - packed_mv = memoryview(packed.tobytes()) - stride = 4 + byte_width - - def _emit_fixed( - row_idx: int, - out_buf: bytearray, - _mv: memoryview = packed_mv, - _stride: int = stride, - _nulls: Any = null_mask, - ) -> None: - if _nulls is not None and _nulls[row_idx]: - out_buf.extend(_PG_NULL_MARKER) - else: - start = row_idx * _stride - out_buf.extend(_mv[start : start + _stride]) - - col_encoders.append(_emit_fixed) - continue - - if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): - dim = int(field_type.list_size) - flat_values = col.values[col.offset * dim : (col.offset + num_rows) * dim] - is_f16 = pa.types.is_float16(field_type.value_type) - elem_dtype = ">f2" if is_f16 else ">f4" - elem_size = 2 if is_f16 else 4 - arr_2d = ( - flat_values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).reshape(num_rows, dim) - ) - packed = np.empty( - num_rows, - dtype=np.dtype([ - ("len", ">i4"), - ("dim", ">u2"), - ("unused", ">u2"), - ("val", elem_dtype, (dim,)), - ]), - ) - packed["len"] = 4 + elem_size * dim - packed["dim"] = dim - packed["unused"] = 0 - packed["val"] = arr_2d - packed_mv = memoryview(packed.tobytes()) - stride = 8 + elem_size * dim - - def _emit_fixed_vec( - row_idx: int, - out_buf: bytearray, - _mv: memoryview = packed_mv, - _stride: int = stride, - _nulls: Any = null_mask, - ) -> None: - if _nulls is not None and _nulls[row_idx]: - out_buf.extend(_PG_NULL_MARKER) - else: - start = row_idx * _stride - out_buf.extend(_mv[start : start + _stride]) - - col_encoders.append(_emit_fixed_vec) - continue + def _emit_fixed( + row_idx: int, + out_buf: bytearray, + _mv: memoryview = packed_mv, + _stride: int = stride, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + else: + start = row_idx * _stride + out_buf.extend(_mv[start : start + _stride]) + + return _emit_fixed + + if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): + dim = int(field_type.list_size) + flat_values = col.flatten() + is_f16 = pa.types.is_float16(field_type.value_type) + elem_dtype = ">f2" if is_f16 else ">f4" + elem_size = 2 if is_f16 else 4 + arr_2d = flat_values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).reshape(num_rows, dim) + packed = np.empty( + num_rows, dtype=np.dtype([("len", ">i4"), ("dim", ">u2"), ("unused", ">u2"), ("val", elem_dtype, (dim,))]) + ) + packed["len"] = 4 + elem_size * dim + packed["dim"] = dim + packed["unused"] = 0 + packed["val"] = arr_2d + packed_mv = memoryview(packed.tobytes()) + stride = 8 + elem_size * dim + + def _emit_fixed_vec( + row_idx: int, + out_buf: bytearray, + _mv: memoryview = packed_mv, + _stride: int = stride, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + else: + start = row_idx * _stride + out_buf.extend(_mv[start : start + _stride]) - if ( - pa.types.is_string(field_type) - or pa.types.is_large_string(field_type) - or pa.types.is_binary(field_type) - or pa.types.is_large_binary(field_type) - ): - offset_dtype = ( - np.int64 if (pa.types.is_large_string(field_type) or pa.types.is_large_binary(field_type)) else np.int32 - ) - offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[ - col.offset : - ] - data_buf = col.buffers()[2] - data_mv = memoryview(data_buf) if data_buf is not None else memoryview(b"") - is_jsonb = col_name in jsonb_columns - diffs = np.diff(offsets) - lengths_mv = memoryview((diffs + (1 if is_jsonb else 0)).astype(">i4").tobytes()) - - def _emit_varlen( - row_idx: int, - out_buf: bytearray, - _offsets: Any = offsets, - _data_mv: memoryview = data_mv, - _len_mv: memoryview = lengths_mv, - _nulls: Any = null_mask, - _is_jsonb: bool = is_jsonb, - ) -> None: - if _nulls is not None and _nulls[row_idx]: - out_buf.extend(_PG_NULL_MARKER) - return - out_buf.extend(_len_mv[row_idx * 4 : (row_idx + 1) * 4]) - if _is_jsonb: - out_buf.append(1) - start = int(_offsets[row_idx]) - end = int(_offsets[row_idx + 1]) - if end > start: - out_buf.extend(_data_mv[start:end]) - - col_encoders.append(_emit_varlen) - continue + return _emit_fixed_vec - if (pa.types.is_list(field_type) or pa.types.is_large_list(field_type)) and ( - col_name in vector_columns or pa.types.is_floating(field_type.value_type) - ): - offset_dtype = np.int64 if pa.types.is_large_list(field_type) else np.int32 - offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[ - col.offset : - ] - is_f16 = pa.types.is_float16(field_type.value_type) - elem_dtype = ">f2" if is_f16 else ">f4" - elem_size = 2 if is_f16 else 4 - flat_bytes = col.values.to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).tobytes() - flat_mv = memoryview(flat_bytes) - dims = np.diff(offsets) - headers = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("dim", ">u2"), ("unused", ">u2")])) - headers["len"] = 4 + elem_size * dims - headers["dim"] = dims - headers["unused"] = 0 - headers_mv = memoryview(headers.tobytes()) - - def _emit_var_vec( - row_idx: int, - out_buf: bytearray, - _offsets: Any = offsets, - _flat_mv: memoryview = flat_mv, - _hdr_mv: memoryview = headers_mv, - _elem_size: int = elem_size, - _nulls: Any = null_mask, - ) -> None: - if _nulls is not None and _nulls[row_idx]: - out_buf.extend(_PG_NULL_MARKER) - return - out_buf.extend(_hdr_mv[row_idx * 8 : (row_idx + 1) * 8]) - start = int(_offsets[row_idx]) * _elem_size - end = int(_offsets[row_idx + 1]) * _elem_size - if end > start: - out_buf.extend(_flat_mv[start:end]) - - col_encoders.append(_emit_var_vec) - continue + return None + + +def _create_pg_column_encoder( + col: Any, col_name: str, num_rows: int, vector_columns: frozenset[str], jsonb_columns: frozenset[str] +) -> Callable[[int, bytearray], None]: + """Create an encoder function for a single column.""" + import numpy as np + import pyarrow as pa + + if isinstance(col, pa.ChunkedArray): + col = col.combine_chunks() + field_type = col.type + has_nulls = col.null_count > 0 + null_mask = col.is_null().to_numpy(zero_copy_only=False) if has_nulls else None - py_values = col.to_pylist() + if col_name not in jsonb_columns: + fixed_encoder = _create_pg_fixed_column_encoder(col, field_type, num_rows, null_mask) + if fixed_encoder is not None: + return fixed_encoder + + if ( + pa.types.is_string(field_type) + or pa.types.is_large_string(field_type) + or pa.types.is_binary(field_type) + or pa.types.is_large_binary(field_type) + ): + offset_dtype = ( + np.int64 if (pa.types.is_large_string(field_type) or pa.types.is_large_binary(field_type)) else np.int32 + ) + offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[col.offset :] + data_buf = col.buffers()[2] + data_mv = memoryview(data_buf) if data_buf is not None else memoryview(b"") is_jsonb = col_name in jsonb_columns + diffs = np.diff(offsets) + lengths_mv = memoryview((diffs + (1 if is_jsonb else 0)).astype(">i4").tobytes()) - def _emit_fallback( + def _emit_varlen( row_idx: int, out_buf: bytearray, - _vals: list[Any] = py_values, + _offsets: Any = offsets, + _data_mv: memoryview = data_mv, + _len_mv: memoryview = lengths_mv, + _nulls: Any = null_mask, _is_jsonb: bool = is_jsonb, ) -> None: - val = _vals[row_idx] - if val is None: + if _nulls is not None and _nulls[row_idx]: out_buf.extend(_PG_NULL_MARKER) return - if isinstance(val, (bytes, bytearray, memoryview)): - raw = bytes(val) - elif isinstance(val, str): - raw = val.encode("utf-8") - elif isinstance(val, (dict, list, tuple)): - raw = to_json(val).encode("utf-8") - else: - raw = str(val).encode("utf-8") + out_buf.extend(_len_mv[row_idx * 4 : (row_idx + 1) * 4]) if _is_jsonb: - out_buf.extend(struct.pack(">iB", len(raw) + 1, 1)) - else: - out_buf.extend(struct.pack(">i", len(raw))) - out_buf.extend(raw) + out_buf.append(1) + start = int(_offsets[row_idx]) + end = int(_offsets[row_idx + 1]) + if end > start: + out_buf.extend(_data_mv[start:end]) + + return _emit_varlen + + if (pa.types.is_list(field_type) or pa.types.is_large_list(field_type)) and ( + col_name in vector_columns or pa.types.is_floating(field_type.value_type) + ): + offset_dtype = np.int64 if pa.types.is_large_list(field_type) else np.int32 + offsets = np.frombuffer(col.buffers()[1], dtype=offset_dtype, count=col.offset + num_rows + 1)[col.offset :] + is_f16 = pa.types.is_float16(field_type.value_type) + elem_dtype = ">f2" if is_f16 else ">f4" + elem_size = 2 if is_f16 else 4 + values_attr = "values" + flat_bytes = getattr(col, values_attr).to_numpy(zero_copy_only=False).astype(elem_dtype, copy=False).tobytes() + flat_mv = memoryview(flat_bytes) + dims = np.diff(offsets) + headers = np.empty(num_rows, dtype=np.dtype([("len", ">i4"), ("dim", ">u2"), ("unused", ">u2")])) + headers["len"] = 4 + elem_size * dims + headers["dim"] = dims + headers["unused"] = 0 + headers_mv = memoryview(headers.tobytes()) + + def _emit_var_vec( + row_idx: int, + out_buf: bytearray, + _offsets: Any = offsets, + _flat_mv: memoryview = flat_mv, + _hdr_mv: memoryview = headers_mv, + _elem_size: int = elem_size, + _nulls: Any = null_mask, + ) -> None: + if _nulls is not None and _nulls[row_idx]: + out_buf.extend(_PG_NULL_MARKER) + return + out_buf.extend(_hdr_mv[row_idx * 8 : (row_idx + 1) * 8]) + start = int(_offsets[row_idx]) * _elem_size + end = int(_offsets[row_idx + 1]) * _elem_size + if end > start: + out_buf.extend(_flat_mv[start:end]) + + return _emit_var_vec + + py_values = col.to_pylist() + is_jsonb = col_name in jsonb_columns + + def _emit_fallback( + row_idx: int, out_buf: bytearray, _vals: list[Any] = py_values, _is_jsonb: bool = is_jsonb + ) -> None: + val = _vals[row_idx] + if val is None: + out_buf.extend(_PG_NULL_MARKER) + return + if isinstance(val, (bytes, bytearray, memoryview)): + raw = bytes(val) + elif isinstance(val, str): + raw = val.encode("utf-8") + elif isinstance(val, (dict, list, tuple)): + raw = to_json(val).encode("utf-8") + else: + raw = str(val).encode("utf-8") + if _is_jsonb: + out_buf.extend(struct.pack(">iB", len(raw) + 1, 1)) + else: + out_buf.extend(struct.pack(">i", len(raw))) + out_buf.extend(raw) + + return _emit_fallback - col_encoders.append(_emit_fallback) + +def _encode_pg_binary_fast_path_2( + batch: "ArrowRecordBatch", columns: "list[str]", vector_columns: frozenset[str], jsonb_columns: frozenset[str] +) -> bytes: + """Encode variable-width and nullable columns by slicing Arrow buffers directly.""" + num_rows = int(batch.num_rows) + row_header = struct.pack(">h", len(columns)) + col_encoders = [ + _create_pg_column_encoder(batch.column(col_name), col_name, num_rows, vector_columns, jsonb_columns) + for col_name in columns + ] out_buf = bytearray() for row_idx in range(num_rows): diff --git a/tests/unit/adapters/test_duckdb_resource_controls.py b/tests/unit/adapters/test_duckdb_resource_controls.py new file mode 100644 index 000000000..950e7f631 --- /dev/null +++ b/tests/unit/adapters/test_duckdb_resource_controls.py @@ -0,0 +1,109 @@ +"""Unit tests for DuckDB container-aware memory limits, tmpfs spill validation, and post-connect parquet_metadata_cache.""" + +from pathlib import Path +from typing import Any + +import pytest + +from sqlspec.adapters.duckdb.core import ( + build_connection_config, + resolve_duckdb_memory_limit, + validate_duckdb_temp_directory, +) +from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool +from sqlspec.exceptions import ImproperConfigurationError + + +def test_resolve_duckdb_memory_limit_cgroup_v2_and_clamping(tmp_path: Path) -> None: + """resolve_duckdb_memory_limit reads cgroup v2 memory.max and clamps oversized requests.""" + cgroup_v2 = tmp_path / "memory.max" + cgroup_v2.write_text(str(2 * 1024 * 1024 * 1024), encoding="utf-8") + + default_limit = resolve_duckdb_memory_limit(None, memory_ceiling_fraction=0.8, cgroup_v2_path=cgroup_v2) + assert default_limit == "1638MB" + + clamped_limit = resolve_duckdb_memory_limit("64GB", memory_ceiling_fraction=0.8, cgroup_v2_path=cgroup_v2) + assert clamped_limit == "1638MB" + + within_limit = resolve_duckdb_memory_limit("512MB", memory_ceiling_fraction=0.8, cgroup_v2_path=cgroup_v2) + assert within_limit == "512MB" + + +def test_resolve_duckdb_memory_limit_cgroup_v1_and_host_fallback(tmp_path: Path) -> None: + """resolve_duckdb_memory_limit falls back from unbounded cgroup v2 to cgroup v1 or host memory.""" + cgroup_v2 = tmp_path / "memory.max" + cgroup_v2.write_text("max\n", encoding="utf-8") + cgroup_v1 = tmp_path / "memory.limit_in_bytes" + cgroup_v1.write_text(str(4 * 1024 * 1024 * 1024), encoding="utf-8") + + v1_limit = resolve_duckdb_memory_limit( + None, memory_ceiling_fraction=0.75, cgroup_v2_path=cgroup_v2, cgroup_v1_path=cgroup_v1 + ) + assert v1_limit == "3072MB" + + cgroup_v1.write_text(str(1 << 62), encoding="utf-8") + host_limit = resolve_duckdb_memory_limit( + None, + memory_ceiling_fraction=0.5, + cgroup_v2_path=cgroup_v2, + cgroup_v1_path=cgroup_v1, + host_memory_bytes=8 * 1024 * 1024 * 1024, + ) + assert host_limit == "4096MB" + + +def test_validate_duckdb_temp_directory_detects_tmpfs() -> None: + """validate_duckdb_temp_directory warns or raises when temp_directory is mounted on tmpfs/ramfs.""" + fake_mounts = ( + "overlay / overlay rw,relatime 0 0\n" + "tmpfs /tmp tmpfs rw,nosuid,nodev 0 0\n" + "/dev/sda1 /var/data ext4 rw,relatime 0 0\n" + ) + assert validate_duckdb_temp_directory("/tmp/duckdb_spill", reject_tmpfs=False, mounts_content=fake_mounts) is False + with pytest.raises(ImproperConfigurationError, match="tmpfs"): + validate_duckdb_temp_directory("/tmp/duckdb_spill", reject_tmpfs=True, mounts_content=fake_mounts) + assert ( + validate_duckdb_temp_directory("/var/data/duckdb_spill", reject_tmpfs=True, mounts_content=fake_mounts) is True + ) + + +def test_build_connection_config_applies_memory_ceiling_fraction() -> None: + """build_connection_config resolves memory_limit when memory_ceiling_fraction is set.""" + cfg = build_connection_config({"database": ":memory:", "memory_ceiling_fraction": 0.75}) + assert "memory_limit" in cfg["config"] + assert "memory_ceiling_fraction" not in cfg["config"] + + +def test_pool_defers_parquet_metadata_cache_after_extensions_load(monkeypatch: pytest.MonkeyPatch) -> None: + """DuckDBConnectionPool applies SET parquet_metadata_cache after extensions load.""" + events: list[str] = [] + + class _Conn: + def install_extension(self, name: str, **_: Any) -> None: + events.append(f"install:{name}") + + def load_extension(self, name: str) -> None: + events.append(f"load:{name}") + + def execute(self, sql: str, *_: Any) -> "_Conn": + events.append(f"sql:{sql}") + return self + + def list_filesystems(self) -> list[str]: + return [] + + captured_connect: dict[str, Any] = {} + + def fake_connect(**kwargs: Any) -> _Conn: + captured_connect.update(kwargs) + return _Conn() + + monkeypatch.setattr("sqlspec.adapters.duckdb.pool.duckdb.connect", fake_connect) + pool = DuckDBConnectionPool( + {"database": ":memory:", "config": {"parquet_metadata_cache": True, "threads": 2}}, + extensions=[{"name": "httpfs"}], + ) + pool._create_connection() + + assert "parquet_metadata_cache" not in captured_connect.get("config", {}) + assert events == ["load:httpfs", "sql:SET parquet_metadata_cache = true"] diff --git a/tests/unit/adapters/test_postgres_pgcopy_binary.py b/tests/unit/adapters/test_postgres_pgcopy_binary.py index 54f04b1c5..a7792190b 100644 --- a/tests/unit/adapters/test_postgres_pgcopy_binary.py +++ b/tests/unit/adapters/test_postgres_pgcopy_binary.py @@ -5,18 +5,20 @@ import struct import uuid from types import SimpleNamespace -from typing import Any, cast +from typing import Any, Self, cast import numpy as np import pyarrow as pa -import pytest -from sqlspec.adapters.asyncpg import AsyncpgDriver, default_statement_config as asyncpg_statement_config +from sqlspec.adapters.asyncpg import AsyncpgDriver +from sqlspec.adapters.asyncpg import default_statement_config as asyncpg_statement_config from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgDriver from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver from sqlspec.adapters.cockroach_psycopg.core import build_copy_from_command as crdb_build_copy_from_command -from sqlspec.adapters.psqlpy import PsqlpyDriver, default_statement_config as psqlpy_statement_config -from sqlspec.adapters.psycopg import PsycopgAsyncDriver, PsycopgSyncDriver, default_statement_config as psycopg_statement_config +from sqlspec.adapters.psqlpy import PsqlpyDriver +from sqlspec.adapters.psqlpy import default_statement_config as psqlpy_statement_config +from sqlspec.adapters.psycopg import PsycopgAsyncDriver, PsycopgSyncDriver +from sqlspec.adapters.psycopg import default_statement_config as psycopg_statement_config from sqlspec.adapters.psycopg.core import build_copy_from_command as pg_build_copy_from_command from sqlspec.utils.arrow_helpers import ( PG_BINARY_COPY_HEADER, @@ -56,10 +58,7 @@ def test_encode_arrow_batch_pg_binary_fast_path_1_fixed_width_and_pgvector() -> pa.array([3.125, -6.25], type=pa.float64()), pa.array([dt.date(2000, 1, 1), dt.date(2000, 1, 3)], type=pa.date32()), pa.array( - [ - dt.datetime(2000, 1, 1, 0, 0, 0, tzinfo=dt.timezone.utc), - dt.datetime(2000, 1, 1, 0, 0, 2, tzinfo=dt.timezone.utc), - ], + [dt.datetime(2000, 1, 1, 0, 0, 0, tzinfo=dt.UTC), dt.datetime(2000, 1, 1, 0, 0, 2, tzinfo=dt.UTC)], type=pa.timestamp("us", tz="UTC"), ), pa.array([uid.bytes, uid.bytes], type=pa.binary(16)), @@ -68,19 +67,7 @@ def test_encode_arrow_batch_pg_binary_fast_path_1_fixed_width_and_pgvector() -> pa.array([np.float16(0.5), np.float16(1.5), np.float16(2.5), np.float16(3.5)], type=pa.float16()), 2 ), ], - names=[ - "flag", - "i16", - "i32", - "i64", - "f32", - "f64", - "dt", - "ts", - "uid", - "vec_f32", - "vec_f16", - ], + names=["flag", "i16", "i32", "i64", "f32", "f64", "dt", "ts", "uid", "vec_f32", "vec_f16"], ) payload = encode_arrow_batch_pg_binary(batch, include_header=True, include_trailer=True) @@ -196,13 +183,13 @@ def __init__(self, writes: list[bytes], rows: list[tuple[Any, ...]]) -> None: self._writes = writes self._rows = rows - def __enter__(self) -> "_DummyPsycopgCopy": + def __enter__(self) -> Self: return self def __exit__(self, *_args: Any) -> None: return None - async def __aenter__(self) -> "_DummyPsycopgCopy": + async def __aenter__(self) -> Self: return self async def __aexit__(self, *_args: Any) -> None: @@ -220,7 +207,7 @@ def __init__(self, writes: list[bytes], rows: list[tuple[Any, ...]]) -> None: self._writes = writes self._rows = rows - async def __aenter__(self) -> "_DummyPsycopgAsyncCopy": + async def __aenter__(self) -> Self: return self async def __aexit__(self, *_args: Any) -> None: @@ -237,7 +224,7 @@ class _DummyPsycopgSyncCursor: def __init__(self, conn: "_DummyPsycopgSyncConn") -> None: self._conn = conn - def __enter__(self) -> "_DummyPsycopgSyncCursor": + def __enter__(self) -> Self: return self def __exit__(self, *_args: Any) -> None: @@ -267,7 +254,7 @@ class _DummyPsycopgAsyncCursor: def __init__(self, conn: "_DummyPsycopgAsyncConn") -> None: self._conn = conn - async def __aenter__(self) -> "_DummyPsycopgAsyncCursor": + async def __aenter__(self) -> Self: return self async def __aexit__(self, *_args: Any) -> None: @@ -332,8 +319,7 @@ async def test_psycopg_and_cockroach_async_load_from_arrow_streams_binary_copy() driver_features={"storage_capabilities": _CAPS}, ) reader = pa.RecordBatchReader.from_batches( - schema, - [pa.RecordBatch.from_pydict({"id": [10, 20], "label": ["x", "y"]}, schema=schema)], + schema, [pa.RecordBatch.from_pydict({"id": [10, 20], "label": ["x", "y"]}, schema=schema)] ) job = await driver.load_from_arrow("public.items", reader) @@ -407,12 +393,7 @@ def __init__(self) -> None: self.binary_calls: list[dict[str, Any]] = [] async def binary_copy_to_table( - self, - source: bytes, - table_name: str, - *, - columns: list[str] | None = None, - schema_name: str | None = None, + self, source: bytes, table_name: str, *, columns: list[str] | None = None, schema_name: str | None = None ) -> int: self.binary_calls.append({ "source": bytes(source), @@ -430,8 +411,7 @@ async def test_psqlpy_load_from_arrow_supports_binary_copy_to_table() -> None: """PsqlpyDriver streams PGCOPY binary payloads to binary_copy_to_table when enable_pgcopy_binary=True.""" schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) reader = pa.RecordBatchReader.from_batches( - schema, - [pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["alpha", "beta"]}, schema=schema)], + schema, [pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["alpha", "beta"]}, schema=schema)] ) conn = _DummyPsqlpyBinaryConn() driver = PsqlpyDriver( diff --git a/tests/unit/core/test_cache.py b/tests/unit/core/test_cache.py index 06346c404..c91d7822c 100644 --- a/tests/unit/core/test_cache.py +++ b/tests/unit/core/test_cache.py @@ -956,3 +956,62 @@ def test_lru_cache_size_method_removed() -> None: """LRUCache should use __len__ instead of a duplicate size() method.""" assert "size" not in LRUCache.__dict__ + + +def test_l1_statement_cache_holder_lifecycle_and_statistics_coordination() -> None: + """L1 StatementCacheHolder and QueryCache coordinate with global cache lifecycle functions.""" + holder = cache_module.StatementCacheHolder(cache_module.QueryCache(max_size=2)) + entry1 = cache_module.CachedQuery( + compiled_sql="SELECT $1", + parameter_profile=cache_module.ParameterProfile.empty(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=cache_module.OperationProfile.empty(), + param_count=1, + ) + entry2 = cache_module.CachedQuery( + compiled_sql="SELECT $1, $2", + parameter_profile=cache_module.ParameterProfile.empty(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=cache_module.OperationProfile.empty(), + param_count=2, + ) + entry3 = cache_module.CachedQuery( + compiled_sql="SELECT $1, $2, $3", + parameter_profile=cache_module.ParameterProfile.empty(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=cache_module.OperationProfile.empty(), + param_count=3, + ) + + reset_stats_only() + assert holder.cache.get("SELECT :a") is None + holder.cache.set("SELECT :a", entry1) + assert holder.cache.get("SELECT :a") is entry1 + holder.cache.set("SELECT :a, :b", entry2) + holder.cache.set("SELECT :a, :b, :c", entry3) + + stats = get_cache_statistics() + assert "query" in stats + assert stats["query"].hits >= 1 + assert stats["query"].misses >= 1 + assert stats["query"].evictions >= 1 + + reset_stats_only() + stats_after_reset = get_cache_statistics() + assert stats_after_reset["query"].hits == 0 + assert stats_after_reset["query"].misses == 0 + assert stats_after_reset["query"].evictions == 0 + assert holder.cache.get("SELECT :a, :b, :c") is entry3 + + clear_all_caches() + assert len(holder.cache) == 0 + diff --git a/tests/unit/driver/test_query_cache.py b/tests/unit/driver/test_query_cache.py index 2ba2f100e..5721a8dee 100644 --- a/tests/unit/driver/test_query_cache.py +++ b/tests/unit/driver/test_query_cache.py @@ -35,12 +35,12 @@ clear_all_caches, get_cache, ) +from sqlspec.core.cache import CachedQuery, QueryCache from sqlspec.core.filters import LimitOffsetFilter from sqlspec.core.parameters._declared import ParameterDeclaration from sqlspec.core.parameters._processor import ParameterProcessor, type_coercion_fallbacks from sqlspec.driver._async import AsyncDriverAdapterBase from sqlspec.driver._common import _check_declared_parameters -from sqlspec.driver._query_cache import CachedQuery, QueryCache from sqlspec.driver._sync import SyncDriverAdapterBase from sqlspec.exceptions import SQLSpecError from sqlspec.observability import LoggingConfig, StatementEvent diff --git a/tools/scripts/pgo_training.py b/tools/scripts/pgo_training.py index 542eb70ba..7d501fbe6 100644 --- a/tools/scripts/pgo_training.py +++ b/tools/scripts/pgo_training.py @@ -82,10 +82,10 @@ def _train_type_dispatch() -> None: def _train_query_cache() -> None: """Exercise QueryCache LRU operations — cache hits, misses, and eviction.""" + from sqlspec.core.cache import CachedQuery, QueryCache from sqlspec.core.compiler import OperationProfile from sqlspec.core.parameters._types import ParameterInfo, ParameterProfile, ParameterStyle from sqlspec.core.statement import ProcessedState - from sqlspec.driver._query_cache import CachedQuery, QueryCache cache = QueryCache(max_size=64) From eb83136c46c7613c82eac6beb243d850b1b8e5bb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 19:30:28 +0000 Subject: [PATCH 06/14] refactor(cache): enforce module-level AST ordering and verify L1/L2 cache coordination --- sqlspec/adapters/duckdb/core.py | 94 +++++------ sqlspec/core/cache.py | 84 +++++----- sqlspec/driver/_storage_helpers.py | 258 ++++++++++++++--------------- tests/unit/core/test_cache.py | 40 +++-- 4 files changed, 243 insertions(+), 233 deletions(-) diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index 2ada3ea5e..b6d1bbae1 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -193,53 +193,6 @@ def collect_rows(fetched_data: "list[Any] | None", description: "list[Any] | Non return fetched_data, column_names -def _parse_memory_size_bytes(memory_limit: str) -> "int | None": - match = _MEMORY_LIMIT_RE.match(memory_limit) - if match is None: - return None - amount = float(match.group(1)) - unit = match.group(2).upper() - multiplier = _MEMORY_UNIT_MULTIPLIERS.get(unit) - if multiplier is None: - return None - return int(amount * multiplier) - - -def _detect_available_memory_bytes( - *, - cgroup_v2_path: "Path | None" = None, - cgroup_v1_path: "Path | None" = None, - host_memory_bytes: "int | None" = None, -) -> "int | None": - v2_file = cgroup_v2_path if cgroup_v2_path is not None else Path("/sys/fs/cgroup/memory.max") - if v2_file.exists(): - with contextlib.suppress(OSError, ValueError): - raw_v2 = v2_file.read_text(encoding="utf-8").strip() - if raw_v2 and raw_v2.lower() != "max": - parsed_v2 = int(raw_v2) - if 0 < parsed_v2 < _CGROUP_UNBOUNDED_THRESHOLD: - return parsed_v2 - - v1_file = cgroup_v1_path if cgroup_v1_path is not None else Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") - if v1_file.exists(): - with contextlib.suppress(OSError, ValueError): - raw_v1 = v1_file.read_text(encoding="utf-8").strip() - if raw_v1: - parsed_v1 = int(raw_v1) - if 0 < parsed_v1 < _CGROUP_UNBOUNDED_THRESHOLD: - return parsed_v1 - - if host_memory_bytes is not None and host_memory_bytes > 0: - return host_memory_bytes - - with contextlib.suppress(ValueError, OSError, AttributeError): - page_size = os.sysconf("SC_PAGE_SIZE") - phys_pages = os.sysconf("SC_PHYS_PAGES") - if isinstance(page_size, int) and isinstance(phys_pages, int) and page_size > 0 and phys_pages > 0: - return page_size * phys_pages - return None - - def resolve_duckdb_memory_limit( memory_limit: "str | None" = None, *, @@ -660,6 +613,53 @@ def duckdb_copy_filename_follows_query() -> bool: return duckdb_version_tuple() >= (2, 0, 0) +def _parse_memory_size_bytes(memory_limit: str) -> "int | None": + match = _MEMORY_LIMIT_RE.match(memory_limit) + if match is None: + return None + amount = float(match.group(1)) + unit = match.group(2).upper() + multiplier = _MEMORY_UNIT_MULTIPLIERS.get(unit) + if multiplier is None: + return None + return int(amount * multiplier) + + +def _detect_available_memory_bytes( + *, + cgroup_v2_path: "Path | None" = None, + cgroup_v1_path: "Path | None" = None, + host_memory_bytes: "int | None" = None, +) -> "int | None": + v2_file = cgroup_v2_path if cgroup_v2_path is not None else Path("/sys/fs/cgroup/memory.max") + if v2_file.exists(): + with contextlib.suppress(OSError, ValueError): + raw_v2 = v2_file.read_text(encoding="utf-8").strip() + if raw_v2 and raw_v2.lower() != "max": + parsed_v2 = int(raw_v2) + if 0 < parsed_v2 < _CGROUP_UNBOUNDED_THRESHOLD: + return parsed_v2 + + v1_file = cgroup_v1_path if cgroup_v1_path is not None else Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") + if v1_file.exists(): + with contextlib.suppress(OSError, ValueError): + raw_v1 = v1_file.read_text(encoding="utf-8").strip() + if raw_v1: + parsed_v1 = int(raw_v1) + if 0 < parsed_v1 < _CGROUP_UNBOUNDED_THRESHOLD: + return parsed_v1 + + if host_memory_bytes is not None and host_memory_bytes > 0: + return host_memory_bytes + + with contextlib.suppress(ValueError, OSError, AttributeError): + page_size = os.sysconf("SC_PAGE_SIZE") + phys_pages = os.sysconf("SC_PHYS_PAGES") + if isinstance(page_size, int) and isinstance(phys_pages, int) and page_size > 0 and phys_pages > 0: + return page_size * phys_pages + return None + + def _build_storage_copy_sql(sql: str, file_format: str) -> str | None: """Wrap one compiled query in COPY with a separately bound filename.""" if file_format not in {"parquet", "csv"}: diff --git a/sqlspec/core/cache.py b/sqlspec/core/cache.py index 523b649bc..603e7bb3e 100644 --- a/sqlspec/core/cache.py +++ b/sqlspec/core/cache.py @@ -411,18 +411,6 @@ def _move_to_head(self, node: CacheNode) -> None: _holders_lock = threading.Lock() -def _register_statement_cache_holder(holder: "StatementCacheHolder") -> None: - """Register a StatementCacheHolder in the bounded process-wide L1 cache registry.""" - holder_id = id(holder) - with _holders_lock: - if holder_id in _statement_cache_holders: - _statement_cache_holders.move_to_end(holder_id) - else: - while len(_statement_cache_holders) >= _MAX_REGISTERED_HOLDERS: - _statement_cache_holders.popitem(last=False) - _statement_cache_holders[holder_id] = holder - - def get_default_cache() -> LRUCache: """Get the default LRU cache instance. @@ -1138,36 +1126,6 @@ def reset_pipeline_registry() -> None: reset_statement_pipeline_cache() -def _configure_pipeline_cache(config: "CacheConfig") -> None: - compiled_cache_enabled = config.compiled_cache_enabled and config.sql_cache_enabled - fragment_cache_enabled = config.compiled_cache_enabled and config.fragment_cache_enabled - cache_size = config.sql_cache_size if compiled_cache_enabled else 0 - parse_cache_size = config.fragment_cache_size if fragment_cache_enabled else 0 - configure_statement_pipeline_cache( - cache_size=cache_size, parse_cache_size=parse_cache_size, cache_enabled=compiled_cache_enabled - ) - - -def _canonicalize_cache_filters(filters: "list[Filter]") -> "tuple[Filter, ...]": - """Create canonical representation of filters for cache keys. - - Args: - filters: List of filters to canonicalize - - Returns: - Tuple of unique filters sorted by field_name, operation, then value - """ - if not filters: - return () - - unique_filters = set(filters) - return tuple(sorted(unique_filters, key=_filter_sort_key)) - - -def _filter_sort_key(filter_obj: "Filter") -> "tuple[str, str, str]": - return filter_obj.field_name, filter_obj.operation, str(filter_obj.value) - - @final @mypyc_attr(allow_interpreted_subclasses=False) class CachedQuery: @@ -1400,3 +1358,45 @@ def resolve_statement_cache_size(driver_features: "Mapping[str, object] | None") except ValueError: return STMT_CACHE_MAX_SIZE return STMT_CACHE_MAX_SIZE + + +def _register_statement_cache_holder(holder: "StatementCacheHolder") -> None: + """Register a StatementCacheHolder in the bounded process-wide L1 cache registry.""" + holder_id = id(holder) + with _holders_lock: + if holder_id in _statement_cache_holders: + _statement_cache_holders.move_to_end(holder_id) + else: + while len(_statement_cache_holders) >= _MAX_REGISTERED_HOLDERS: + _statement_cache_holders.popitem(last=False) + _statement_cache_holders[holder_id] = holder + + +def _configure_pipeline_cache(config: "CacheConfig") -> None: + compiled_cache_enabled = config.compiled_cache_enabled and config.sql_cache_enabled + fragment_cache_enabled = config.compiled_cache_enabled and config.fragment_cache_enabled + cache_size = config.sql_cache_size if compiled_cache_enabled else 0 + parse_cache_size = config.fragment_cache_size if fragment_cache_enabled else 0 + configure_statement_pipeline_cache( + cache_size=cache_size, parse_cache_size=parse_cache_size, cache_enabled=compiled_cache_enabled + ) + + +def _canonicalize_cache_filters(filters: "list[Filter]") -> "tuple[Filter, ...]": + """Create canonical representation of filters for cache keys. + + Args: + filters: List of filters to canonicalize + + Returns: + Tuple of unique filters sorted by field_name, operation, then value + """ + if not filters: + return () + + unique_filters = set(filters) + return tuple(sorted(unique_filters, key=_filter_sort_key)) + + +def _filter_sort_key(filter_obj: "Filter") -> "tuple[str, str, str]": + return filter_obj.field_name, filter_obj.operation, str(filter_obj.value) diff --git a/sqlspec/driver/_storage_helpers.py b/sqlspec/driver/_storage_helpers.py index 23d705b6e..84a9e33e2 100644 --- a/sqlspec/driver/_storage_helpers.py +++ b/sqlspec/driver/_storage_helpers.py @@ -552,51 +552,6 @@ def build_create_index_sql( return sql_text -def _extract_metadata_rows(result: Any) -> "list[dict[str, Any]]": - if result is None: - return [] - data = getattr(result, "data", result) - if isinstance(data, list): - return [dict(item) if isinstance(item, Mapping) else {} for item in data] - if isinstance(data, Sequence): - return [dict(item) if isinstance(item, Mapping) else {} for item in data] - return [] - - -def _extract_ddl_text(result: Any) -> "str | None": - if isinstance(result, str) and result.strip(): - return result.strip() - rows = _extract_metadata_rows(result) - for row in rows: - for key in ("ddl", "definition", "native_sql", "sql"): - val = row.get(key) - if isinstance(val, str) and val.strip(): - return val.strip() - return None - - -def _safe_driver_in_transaction(driver: Any) -> bool: - checker = getattr(driver, "is_in_transaction", None) - if callable(checker): - with contextlib.suppress(Exception): - return bool(checker()) - conn_checker = getattr(driver, "_connection_in_transaction", None) - if callable(conn_checker): - with contextlib.suppress(Exception): - return bool(conn_checker()) - return False - - -def _split_schema_and_table(table: str, explicit_schema: str | None) -> "tuple[str | None, str]": - clean = table.strip() - if explicit_schema is not None: - return explicit_schema, clean.rsplit(".", 1)[-1].strip('"`[]') - if "." in clean: - schema_part, table_part = clean.split(".", 1) - return schema_part.strip('"`[]'), table_part.strip('"`[]') - return None, clean.strip('"`[]') - - class SyncWithoutIndexesContext: """Synchronous context manager that drops secondary indexes before a bulk load and rebuilds them on exit.""" @@ -903,90 +858,6 @@ async def execute_staged_load_async( await driver.execute(build_drop_stage_table_sql(dialect, resolved_stage)) -def _record_checkpoint_sync( - driver: Any, - dialect: str, - checkpoint_table: str, - phase: str, - shard_uri: str, - rows_loaded: int, - bytes_processed: int, - duration_s: float, -) -> None: - sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) - driver.execute( - sql_text, - { - "phase": phase, - "shard_uri": shard_uri, - "rows_loaded": int(rows_loaded), - "bytes_processed": int(bytes_processed), - "duration_s": float(duration_s), - }, - ) - - -async def _record_checkpoint_async( - driver: Any, - dialect: str, - checkpoint_table: str, - phase: str, - shard_uri: str, - rows_loaded: int, - bytes_processed: int, - duration_s: float, -) -> None: - sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) - await driver.execute( - sql_text, - { - "phase": phase, - "shard_uri": shard_uri, - "rows_loaded": int(rows_loaded), - "bytes_processed": int(bytes_processed), - "duration_s": float(duration_s), - }, - ) - - -def _read_shard_bytes_and_decode( - shard_uri: str, file_format: "StorageFormat", storage_options: "dict[str, Any] | None", arrow_schema: Any -) -> "tuple[str, Any, int, float]": - t0 = perf_counter() - pipeline = SyncStoragePipeline() - backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) - payload = backend.read_bytes_sync(path) - table = decode_arrow_payload(payload, file_format, arrow_schema=arrow_schema) - return shard_uri, table, len(payload), perf_counter() - t0 - - -def _decode_shard_in_process_worker( - args: "tuple[str, StorageFormat, dict[str, Any] | None]", -) -> "tuple[str, bytes, int, float]": - shard_uri, file_format, storage_options = args - t0 = perf_counter() - pipeline = SyncStoragePipeline() - backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) - payload = backend.read_bytes_sync(path) - table = decode_arrow_payload(payload, file_format) - ipc_bytes = encode_arrow_payload(table, "arrow-ipc", compression=None) - return shard_uri, ipc_bytes, len(payload), perf_counter() - t0 - - -def _is_multi_shard_source(source: "StorageDestination | Sequence[StorageDestination]") -> tuple[bool, str, bool]: - """Return ``(is_source_sequence, source_str, has_glob_or_prefix)`` for a storage destination.""" - if isinstance(source, Sequence) and not isinstance(source, (str, bytes, bytearray)): - return True, "", True - source_str = source.as_posix() if isinstance(source, Path) else str(source) - if any(ch in source_str for ch in ("*", "?", "[")) or source_str.endswith("/"): - return False, source_str, True - if "://" not in source_str: - with contextlib.suppress(Exception): - if Path(source_str).expanduser().is_dir(): - return False, source_str, True - return False, source_str, False - - def execute_load_from_storage_sync( driver: Any, table: str, @@ -1371,3 +1242,132 @@ async def execute_load_from_storage_async( telemetry_payload["destination"] = table attach_partition_telemetry(telemetry_payload, partitioner) return create_storage_job(telemetry_payload) + + +def _extract_metadata_rows(result: Any) -> "list[dict[str, Any]]": + if result is None: + return [] + data = getattr(result, "data", result) + if isinstance(data, list): + return [dict(item) if isinstance(item, Mapping) else {} for item in data] + if isinstance(data, Sequence): + return [dict(item) if isinstance(item, Mapping) else {} for item in data] + return [] + + +def _extract_ddl_text(result: Any) -> "str | None": + if isinstance(result, str) and result.strip(): + return result.strip() + rows = _extract_metadata_rows(result) + for row in rows: + for key in ("ddl", "definition", "native_sql", "sql"): + val = row.get(key) + if isinstance(val, str) and val.strip(): + return val.strip() + return None + + +def _safe_driver_in_transaction(driver: Any) -> bool: + checker = getattr(driver, "is_in_transaction", None) + if callable(checker): + with contextlib.suppress(Exception): + return bool(checker()) + conn_checker = getattr(driver, "_connection_in_transaction", None) + if callable(conn_checker): + with contextlib.suppress(Exception): + return bool(conn_checker()) + return False + + +def _split_schema_and_table(table: str, explicit_schema: str | None) -> "tuple[str | None, str]": + clean = table.strip() + if explicit_schema is not None: + return explicit_schema, clean.rsplit(".", 1)[-1].strip('"`[]') + if "." in clean: + schema_part, table_part = clean.split(".", 1) + return schema_part.strip('"`[]'), table_part.strip('"`[]') + return None, clean.strip('"`[]') + + +def _record_checkpoint_sync( + driver: Any, + dialect: str, + checkpoint_table: str, + phase: str, + shard_uri: str, + rows_loaded: int, + bytes_processed: int, + duration_s: float, +) -> None: + sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) + driver.execute( + sql_text, + { + "phase": phase, + "shard_uri": shard_uri, + "rows_loaded": int(rows_loaded), + "bytes_processed": int(bytes_processed), + "duration_s": float(duration_s), + }, + ) + + +async def _record_checkpoint_async( + driver: Any, + dialect: str, + checkpoint_table: str, + phase: str, + shard_uri: str, + rows_loaded: int, + bytes_processed: int, + duration_s: float, +) -> None: + sql_text = build_checkpoint_insert_sql(dialect, checkpoint_table) + await driver.execute( + sql_text, + { + "phase": phase, + "shard_uri": shard_uri, + "rows_loaded": int(rows_loaded), + "bytes_processed": int(bytes_processed), + "duration_s": float(duration_s), + }, + ) + + +def _read_shard_bytes_and_decode( + shard_uri: str, file_format: "StorageFormat", storage_options: "dict[str, Any] | None", arrow_schema: Any +) -> "tuple[str, Any, int, float]": + t0 = perf_counter() + pipeline = SyncStoragePipeline() + backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) + payload = backend.read_bytes_sync(path) + table = decode_arrow_payload(payload, file_format, arrow_schema=arrow_schema) + return shard_uri, table, len(payload), perf_counter() - t0 + + +def _decode_shard_in_process_worker( + args: "tuple[str, StorageFormat, dict[str, Any] | None]", +) -> "tuple[str, bytes, int, float]": + shard_uri, file_format, storage_options = args + t0 = perf_counter() + pipeline = SyncStoragePipeline() + backend, path, _backend_name = pipeline._backend(shard_uri, storage_options) + payload = backend.read_bytes_sync(path) + table = decode_arrow_payload(payload, file_format) + ipc_bytes = encode_arrow_payload(table, "arrow-ipc", compression=None) + return shard_uri, ipc_bytes, len(payload), perf_counter() - t0 + + +def _is_multi_shard_source(source: "StorageDestination | Sequence[StorageDestination]") -> tuple[bool, str, bool]: + """Return ``(is_source_sequence, source_str, has_glob_or_prefix)`` for a storage destination.""" + if isinstance(source, Sequence) and not isinstance(source, (str, bytes, bytearray)): + return True, "", True + source_str = source.as_posix() if isinstance(source, Path) else str(source) + if any(ch in source_str for ch in ("*", "?", "[")) or source_str.endswith("/"): + return False, source_str, True + if "://" not in source_str: + with contextlib.suppress(Exception): + if Path(source_str).expanduser().is_dir(): + return False, source_str, True + return False, source_str, False diff --git a/tests/unit/core/test_cache.py b/tests/unit/core/test_cache.py index c91d7822c..31caeba6d 100644 --- a/tests/unit/core/test_cache.py +++ b/tests/unit/core/test_cache.py @@ -33,6 +33,10 @@ CacheStats, LRUCache, NamespacedCache, + OperationProfile, + ParameterProfile, + ProcessedState, + StatementConfig, clear_all_caches, get_cache, get_cache_config, @@ -960,44 +964,51 @@ def test_lru_cache_size_method_removed() -> None: def test_l1_statement_cache_holder_lifecycle_and_statistics_coordination() -> None: """L1 StatementCacheHolder and QueryCache coordinate with global cache lifecycle functions.""" - holder = cache_module.StatementCacheHolder(cache_module.QueryCache(max_size=2)) + holder = cache_module.StatementCacheHolder(StatementConfig(), max_size=2) + op_profile = OperationProfile(returns_rows=True, modifies_rows=False) + param_profile = ParameterProfile() entry1 = cache_module.CachedQuery( compiled_sql="SELECT $1", - parameter_profile=cache_module.ParameterProfile.empty(), + parameter_profile=param_profile, input_named_parameters=(), applied_wrap_types=False, parameter_casts={}, operation_type="SELECT", - operation_profile=cache_module.OperationProfile.empty(), + operation_profile=op_profile, param_count=1, + processed_state=ProcessedState(compiled_sql="SELECT $1", execution_parameters=[], operation_type="SELECT"), ) entry2 = cache_module.CachedQuery( compiled_sql="SELECT $1, $2", - parameter_profile=cache_module.ParameterProfile.empty(), + parameter_profile=param_profile, input_named_parameters=(), applied_wrap_types=False, parameter_casts={}, operation_type="SELECT", - operation_profile=cache_module.OperationProfile.empty(), + operation_profile=op_profile, param_count=2, + processed_state=ProcessedState(compiled_sql="SELECT $1, $2", execution_parameters=[], operation_type="SELECT"), ) entry3 = cache_module.CachedQuery( compiled_sql="SELECT $1, $2, $3", - parameter_profile=cache_module.ParameterProfile.empty(), + parameter_profile=param_profile, input_named_parameters=(), applied_wrap_types=False, parameter_casts={}, operation_type="SELECT", - operation_profile=cache_module.OperationProfile.empty(), + operation_profile=op_profile, param_count=3, + processed_state=ProcessedState( + compiled_sql="SELECT $1, $2, $3", execution_parameters=[], operation_type="SELECT" + ), ) reset_stats_only() - assert holder.cache.get("SELECT :a") is None - holder.cache.set("SELECT :a", entry1) - assert holder.cache.get("SELECT :a") is entry1 - holder.cache.set("SELECT :a, :b", entry2) - holder.cache.set("SELECT :a, :b, :c", entry3) + assert holder.query_cache.get("SELECT :a") is None + holder.query_cache.set("SELECT :a", entry1) + assert holder.query_cache.get("SELECT :a") is entry1 + holder.query_cache.set("SELECT :a, :b", entry2) + holder.query_cache.set("SELECT :a, :b, :c", entry3) stats = get_cache_statistics() assert "query" in stats @@ -1010,8 +1021,7 @@ def test_l1_statement_cache_holder_lifecycle_and_statistics_coordination() -> No assert stats_after_reset["query"].hits == 0 assert stats_after_reset["query"].misses == 0 assert stats_after_reset["query"].evictions == 0 - assert holder.cache.get("SELECT :a, :b, :c") is entry3 + assert holder.query_cache.get("SELECT :a, :b, :c") is entry3 clear_all_caches() - assert len(holder.cache) == 0 - + assert len(holder.query_cache) == 0 From a45419a3ec8e755d8d1f51dd96bc0da86e9a99de Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 20:22:40 +0000 Subject: [PATCH 07/14] fix(types,adapters): resolve driver load overrides, mypy arrow types, and adbc callable import --- sqlspec/adapters/adbc/config.py | 4 +- sqlspec/adapters/bigquery/driver.py | 26 ++++-- sqlspec/adapters/mssql_python/driver.py | 10 +-- sqlspec/adapters/oracledb/driver.py | 27 ++++-- sqlspec/adapters/pymssql/driver.py | 10 +-- sqlspec/adapters/spanner/driver.py | 16 ++-- sqlspec/storage/_arrow_payload.py | 10 +-- sqlspec/storage/pipeline.py | 3 +- sqlspec/utils/arrow_helpers.py | 86 +++++++++++++------ .../adapters/test_postgres_pgcopy_binary.py | 8 +- .../driver/test_storage_parallel_and_auth.py | 6 +- tests/unit/storage/test_pipeline_streaming.py | 35 +++++--- 12 files changed, 153 insertions(+), 88 deletions(-) diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index fae4f68b7..18fe38548 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -1,5 +1,6 @@ """ADBC database configuration.""" +from collections.abc import Callable from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast @@ -34,7 +35,6 @@ from sqlspec.utils.config_tools import normalize_connection_config if TYPE_CHECKING: - from collections.abc import Callable from types import TracebackType from sqlspec.observability import ObservabilityConfig @@ -98,7 +98,7 @@ class AdbcConnectionParams(TypedDict): """Database authentication password lifted into ``db_kwargs`` or ``connect()`` as required by the driver.""" token: NotRequired[str] """Bearer, OAuth refresh, or Quack/FlightSQL authentication token.""" - auth_token_provider: NotRequired["Callable[[], str]"] + auth_token_provider: NotRequired[Callable[[], str]] """Callable returning a fresh password or bearer token on each ``create_connection`` call.""" project_id: NotRequired[str] """Google Cloud project ID mapped to ``adbc.bigquery.sql.project_id`` for ``adbc-driver-bigquery``.""" diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index 0ecc6df65..700e1b415 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -8,7 +8,8 @@ from collections.abc import Mapping from contextlib import suppress from itertools import chain -from typing import TYPE_CHECKING, Any, cast +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, cast from sqlglot import exp @@ -680,6 +681,14 @@ def load_from_records( *, columns: "list[str] | None" = None, overwrite: bool = False, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load in-memory records through BigQuery's newline-delimited JSON load API.""" rows = _records_to_json_rows(records, columns) @@ -705,14 +714,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into BigQuery (Parquet load job, or opt-in Storage Write API).""" @@ -797,9 +806,14 @@ def load_from_storage( **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into BigQuery.""" - if kwargs: + if kwargs or not isinstance(source, (str, Path)): return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + table, + source, + file_format=cast("StorageFormat", file_format), + partitioner=partitioner, + overwrite=overwrite, + **kwargs, ) job_config = build_load_job_config(file_format, overwrite) diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index d996e0cd4..2c2266bca 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Sized -from typing import TYPE_CHECKING, Any, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast from sqlspec.adapters.mssql_python._typing import ( MssqlPythonConnection, @@ -42,7 +42,7 @@ from sqlspec.utils.module_loader import ensure_pyarrow, import_optional if TYPE_CHECKING: - from collections.abc import Iterable, Sequence + from collections.abc import Iterable, Mapping, Sequence from sqlspec.builder import QueryBuilder from sqlspec.core import SQL, ArrowResult, SQLResult, Statement, StatementConfig, StatementFilter @@ -436,14 +436,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, batch_size: int = 0, timeout: int = 30, table_lock: bool = False, diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 788ac1ad0..cb9034a2d 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -1,7 +1,8 @@ """Oracle Database driver implementation.""" import logging -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, overload +from pathlib import Path +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast, overload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, @@ -833,14 +834,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Oracle using batched executemany calls.""" self._require_capability("arrow_import_enabled") @@ -956,7 +957,11 @@ def load_from_storage( if isinstance(source, str) and source.startswith("alias://") else source ) - remote_uri = resolve_dbms_cloud_source_uri(resolved_source, file_format=file_format) + remote_uri = ( + resolve_dbms_cloud_source_uri(resolved_source, file_format=file_format) + if isinstance(resolved_source, (str, Path)) + else None + ) if remote_uri is not None: plsql, binds = build_dbms_cloud_copy_data_sql( table, @@ -1656,14 +1661,14 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Asynchronously load Arrow data into Oracle.""" self._require_capability("arrow_import_enabled") @@ -1775,7 +1780,11 @@ async def load_from_storage( if isinstance(source, str) and source.startswith("alias://") else source ) - remote_uri = resolve_dbms_cloud_source_uri(resolved_source, file_format=file_format) + remote_uri = ( + resolve_dbms_cloud_source_uri(resolved_source, file_format=file_format) + if isinstance(resolved_source, (str, Path)) + else None + ) if remote_uri is not None: plsql, binds = build_dbms_cloud_copy_data_sql( table, diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index 9f4bb4a29..feb0f658d 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Sized -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.pymssql._typing import ( PymssqlConnection, @@ -40,7 +40,7 @@ from sqlspec.utils.logging import get_logger if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from sqlspec.adapters.pymssql._typing import PymssqlQueryParams as QueryParams from sqlspec.builder import QueryBuilder @@ -347,14 +347,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target table via executemany.""" self._require_capability("arrow_import_enabled") diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index 11a3cf907..c8a5b333f 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -1,7 +1,7 @@ """Spanner driver implementation.""" from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Protocol, cast, overload +from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, overload from sqlglot import exp as _sqlglot_exp @@ -58,7 +58,7 @@ from sqlspec.utils.serializers import from_json if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Sequence + from collections.abc import Awaitable, Callable, Mapping, Sequence from sqlglot.dialects.dialect import DialectType @@ -525,14 +525,14 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Spanner table via batch mutations.""" self._require_capability("arrow_import_enabled") @@ -1186,14 +1186,14 @@ async def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, - stage_table: bool | str = False, + stage_table: "bool | str | None" = False, conflict_keys: "Sequence[str] | None" = None, - on_conflict: str = "update", + on_conflict: Literal["update", "nothing"] = "update", exclude_update_columns: "Sequence[str] | None" = None, drop_stage_table: bool = True, drop_indexes: bool = False, concurrent_indexes: bool = True, - type_casts: "dict[str, str] | None" = None, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into Spanner table via batch mutations.""" self._require_capability("arrow_import_enabled") diff --git a/sqlspec/storage/_arrow_payload.py b/sqlspec/storage/_arrow_payload.py index 4aaaa40ca..402ee2328 100644 --- a/sqlspec/storage/_arrow_payload.py +++ b/sqlspec/storage/_arrow_payload.py @@ -26,15 +26,15 @@ def batches_to_arrow_table(batches: Any) -> "ArrowTable": """Coalesce an Arrow RecordBatchReader, RecordBatch, Table, or batch iterable into an ArrowTable.""" pa = import_pyarrow() if isinstance(batches, pa.Table): - return batches + return cast("ArrowTable", batches) if isinstance(batches, pa.RecordBatchReader): - return batches.read_all() + return cast("ArrowTable", batches.read_all()) if isinstance(batches, pa.RecordBatch): - return pa.Table.from_batches([batches]) + return cast("ArrowTable", pa.Table.from_batches([batches])) batch_list = list(batches) if not batch_list: - return pa.table({}) - return pa.Table.from_batches(batch_list, schema=batch_list[0].schema) + return cast("ArrowTable", pa.table({})) + return cast("ArrowTable", pa.Table.from_batches(batch_list, schema=batch_list[0].schema)) def resolve_parquet_use_dictionary( diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index 6e8ae4940..dfbb70001 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -1,6 +1,7 @@ """Storage pipeline scaffolding for driver-aware storage bridge.""" from collections import deque +from collections.abc import Sequence from functools import partial from pathlib import Path from time import perf_counter @@ -54,7 +55,7 @@ "reset_storage_bridge_metrics", ) -StorageDestination: TypeAlias = str | Path +StorageDestination: TypeAlias = str | Path | Sequence[str | Path] StorageDiagnostics: TypeAlias = dict[str, float] diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 56f009c44..3aefd6f8c 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -819,49 +819,83 @@ def _extract_pg_fixed_column_array(col: Any, dtype: Any, num_rows: int) -> "tupl arr = filled.to_numpy(zero_copy_only=False).astype("|u1", copy=False) return 1, "|u1", arr if pa.types.is_int8(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.int8, count=col.offset + num_rows)[col.offset :] - return 2, ">i2", raw.astype(">i2") + return ( + 2, + ">i2", + np.frombuffer(col.buffers()[1], dtype=np.int8, count=col.offset + num_rows)[col.offset :].astype(">i2"), + ) if pa.types.is_uint8(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.uint8, count=col.offset + num_rows)[col.offset :] - return 2, ">i2", raw.astype(">i2") + return ( + 2, + ">i2", + np.frombuffer(col.buffers()[1], dtype=np.uint8, count=col.offset + num_rows)[col.offset :].astype(">i2"), + ) if pa.types.is_int16(dtype) or pa.types.is_uint16(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.int16, count=col.offset + num_rows)[col.offset :] - return 2, ">i2", raw.astype(">i2", copy=False) + return ( + 2, + ">i2", + np.frombuffer(col.buffers()[1], dtype=np.int16, count=col.offset + num_rows)[col.offset :].astype( + ">i2", copy=False + ), + ) if pa.types.is_int32(dtype) or pa.types.is_uint32(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.int32, count=col.offset + num_rows)[col.offset :] - return 4, ">i4", raw.astype(">i4", copy=False) + return ( + 4, + ">i4", + np.frombuffer(col.buffers()[1], dtype=np.int32, count=col.offset + num_rows)[col.offset :].astype( + ">i4", copy=False + ), + ) if pa.types.is_int64(dtype) or pa.types.is_uint64(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.int64, count=col.offset + num_rows)[col.offset :] - return 8, ">i8", raw.astype(">i8", copy=False) + return ( + 8, + ">i8", + np.frombuffer(col.buffers()[1], dtype=np.int64, count=col.offset + num_rows)[col.offset :].astype( + ">i8", copy=False + ), + ) if pa.types.is_float16(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.float16, count=col.offset + num_rows)[col.offset :] - return 4, ">f4", raw.astype(">f4") + return ( + 4, + ">f4", + np.frombuffer(col.buffers()[1], dtype=np.float16, count=col.offset + num_rows)[col.offset :].astype(">f4"), + ) if pa.types.is_float32(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.float32, count=col.offset + num_rows)[col.offset :] - return 4, ">f4", raw.astype(">f4", copy=False) + return ( + 4, + ">f4", + np.frombuffer(col.buffers()[1], dtype=np.float32, count=col.offset + num_rows)[col.offset :].astype( + ">f4", copy=False + ), + ) if pa.types.is_float64(dtype): - raw = np.frombuffer(col.buffers()[1], dtype=np.float64, count=col.offset + num_rows)[col.offset :] - return 8, ">f8", raw.astype(">f8", copy=False) + return ( + 8, + ">f8", + np.frombuffer(col.buffers()[1], dtype=np.float64, count=col.offset + num_rows)[col.offset :].astype( + ">f8", copy=False + ), + ) if pa.types.is_date32(dtype): buf = col.buffers()[1] - raw = np.frombuffer(buf, dtype=np.int32, count=col.offset + num_rows)[col.offset :] - return 4, ">i4", (raw - _PG_DATE_EPOCH_DAYS).astype(">i4") + raw_date32 = np.frombuffer(buf, dtype=np.int32, count=col.offset + num_rows)[col.offset :] + return 4, ">i4", (raw_date32 - _PG_DATE_EPOCH_DAYS).astype(">i4") if pa.types.is_date64(dtype): buf = col.buffers()[1] - raw = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] - return 4, ">i4", ((raw // 86_400_000) - _PG_DATE_EPOCH_DAYS).astype(">i4") + raw_date64 = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] + return 4, ">i4", ((raw_date64 // 86_400_000) - _PG_DATE_EPOCH_DAYS).astype(">i4") if pa.types.is_timestamp(dtype): buf = col.buffers()[1] - raw = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] + raw_ts = np.frombuffer(buf, dtype=np.int64, count=col.offset + num_rows)[col.offset :] unit = dtype.unit if unit == "us": - shifted = raw - _PG_EPOCH_OFFSET_US + shifted = raw_ts - _PG_EPOCH_OFFSET_US elif unit == "ms": - shifted = raw * 1_000 - _PG_EPOCH_OFFSET_US + shifted = raw_ts * 1_000 - _PG_EPOCH_OFFSET_US elif unit == "s": - shifted = raw * 1_000_000 - _PG_EPOCH_OFFSET_US + shifted = raw_ts * 1_000_000 - _PG_EPOCH_OFFSET_US else: - shifted = (raw // 1_000) - _PG_EPOCH_OFFSET_US + shifted = (raw_ts // 1_000) - _PG_EPOCH_OFFSET_US return 8, ">i8", shifted.astype(">i8") if pa.types.is_fixed_size_binary(dtype): width = int(dtype.byte_width) @@ -904,7 +938,7 @@ def _encode_pg_binary_fast_path_1( if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): dim = int(field_type.list_size) - flat_values = col.flatten() + flat_values = cast("Any", col).flatten() if flat_values.null_count > 0: return None is_f16 = pa.types.is_float16(field_type.value_type) diff --git a/tests/unit/adapters/test_postgres_pgcopy_binary.py b/tests/unit/adapters/test_postgres_pgcopy_binary.py index a7792190b..961705159 100644 --- a/tests/unit/adapters/test_postgres_pgcopy_binary.py +++ b/tests/unit/adapters/test_postgres_pgcopy_binary.py @@ -152,7 +152,7 @@ def test_encode_arrow_batch_pg_binary_fast_path_2_variable_width_jsonb_and_nulls def test_iter_pgcopy_binary_chunks_streams_reader_batches() -> None: """iter_pgcopy_binary_chunks emits header, batch payloads, and trailer across RecordBatchReader.""" - schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("name", pa.string())])) batches = [ pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["a", "b"]}, schema=schema), pa.RecordBatch.from_pydict({"id": [3, 4], "name": ["c", "d"]}, schema=schema), @@ -282,7 +282,7 @@ def cursor(self) -> _DummyPsycopgAsyncCursor: def test_psycopg_and_cockroach_sync_load_from_arrow_streams_binary_copy() -> None: """PsycopgSyncDriver and CockroachPsycopgSyncDriver stream RecordBatchReader via binary COPY.""" - schema = pa.schema([("id", pa.int64()), ("label", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("label", pa.string())])) for driver_cls in (PsycopgSyncDriver, CockroachPsycopgSyncDriver): conn = _DummyPsycopgSyncConn() driver = driver_cls( @@ -310,7 +310,7 @@ def test_psycopg_and_cockroach_sync_load_from_arrow_streams_binary_copy() -> Non async def test_psycopg_and_cockroach_async_load_from_arrow_streams_binary_copy() -> None: """PsycopgAsyncDriver and CockroachPsycopgAsyncDriver stream RecordBatchReader via async binary COPY.""" - schema = pa.schema([("id", pa.int64()), ("label", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("label", pa.string())])) for driver_cls in (PsycopgAsyncDriver, CockroachPsycopgAsyncDriver): conn = _DummyPsycopgAsyncConn() driver = driver_cls( @@ -409,7 +409,7 @@ async def fetch(self, _sql: str, _params: Any = None) -> Any: async def test_psqlpy_load_from_arrow_supports_binary_copy_to_table() -> None: """PsqlpyDriver streams PGCOPY binary payloads to binary_copy_to_table when enable_pgcopy_binary=True.""" - schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("name", pa.string())])) reader = pa.RecordBatchReader.from_batches( schema, [pa.RecordBatch.from_pydict({"id": [1, 2], "name": ["alpha", "beta"]}, schema=schema)] ) diff --git a/tests/unit/driver/test_storage_parallel_and_auth.py b/tests/unit/driver/test_storage_parallel_and_auth.py index e9a548e33..4371e28ea 100644 --- a/tests/unit/driver/test_storage_parallel_and_auth.py +++ b/tests/unit/driver/test_storage_parallel_and_auth.py @@ -2,7 +2,7 @@ import os from pathlib import Path -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock, patch import pyarrow as pa @@ -69,7 +69,7 @@ def test_checkpoint_sql_builders_generate_dialect_valid_sql() -> None: def test_sync_load_from_storage_multi_uri_checkpoint_and_schema(tmp_path: Path) -> None: """Sync load_from_storage ingests multiple URIs concurrently and skips checkpointed URIs on rerun.""" - schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("name", pa.string())])) file1 = tmp_path / "part1.parquet" file2 = tmp_path / "part2.parquet" pq.write_table(pa.Table.from_pydict({"id": [1, 2], "name": ["a", "b"]}, schema=schema), file1) @@ -110,7 +110,7 @@ def test_sync_load_from_storage_multi_uri_checkpoint_and_schema(tmp_path: Path) @pytest.mark.anyio async def test_async_load_from_storage_multi_uri_checkpoint_and_schema(tmp_path: Path) -> None: """Async load_from_storage ingests multiple URIs with bounded concurrency and checkpointing.""" - schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("name", pa.string())])) file1 = tmp_path / "async_part1.parquet" file2 = tmp_path / "async_part2.parquet" pq.write_table(pa.Table.from_pydict({"id": [10, 20], "name": ["x", "y"]}, schema=schema), file1) diff --git a/tests/unit/storage/test_pipeline_streaming.py b/tests/unit/storage/test_pipeline_streaming.py index c1fa815d9..ff46a0263 100644 --- a/tests/unit/storage/test_pipeline_streaming.py +++ b/tests/unit/storage/test_pipeline_streaming.py @@ -2,7 +2,7 @@ """Unit tests for streaming storage pipeline batches, vector Parquet defaults, schema-pinned JSONL/CSV/IPC, and ObStoreSink safety.""" from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np import pyarrow as pa @@ -16,13 +16,18 @@ def test_resolve_parquet_use_dictionary_disables_float_vector_columns() -> None: """resolve_parquet_use_dictionary disables dictionary encoding on float vector columns while preserving scalar columns.""" - schema = pa.schema([ - ("id", pa.int64()), - ("category", pa.string()), - ("embedding_f32", pa.list_(pa.float32(), 4)), - ("embedding_var", pa.list_(pa.float32())), - ("tags", pa.list_(pa.string())), - ]) + schema = pa.schema( + cast( + "Any", + [ + ("id", pa.int64()), + ("category", pa.string()), + ("embedding_f32", pa.list_(pa.float32(), 4)), + ("embedding_var", pa.list_(pa.float32())), + ("tags", pa.list_(pa.string())), + ], + ) + ) resolved = resolve_parquet_use_dictionary(schema) assert resolved == ["id", "category", "tags"] @@ -30,13 +35,15 @@ def test_resolve_parquet_use_dictionary_disables_float_vector_columns() -> None: explicit = resolve_parquet_use_dictionary(schema, explicit_use_dictionary=False) assert explicit is False - scalar_only_schema = pa.schema([("id", pa.int64()), ("category", pa.string())]) + scalar_only_schema = pa.schema(cast("Any", [("id", pa.int64()), ("category", pa.string())])) assert resolve_parquet_use_dictionary(scalar_only_schema) is True def test_encode_arrow_payload_applies_vector_use_dictionary_and_write_options() -> None: """encode_arrow_payload forwards write_options and disables dictionary encoding for float vector columns.""" - schema = pa.schema([("id", pa.int64()), ("category", pa.string()), ("embedding", pa.list_(pa.float32(), 2))]) + schema = pa.schema( + cast("Any", [("id", pa.int64()), ("category", pa.string()), ("embedding", pa.list_(pa.float32(), 2))]) + ) table = pa.Table.from_arrays( [ pa.array([1, 2, 3], type=pa.int64()), @@ -110,20 +117,20 @@ def test_sync_pipeline_write_and_read_arrow_batches_parquet_and_jsonl(tmp_path: jsonl_path = tmp_path / "vectors.jsonl" jsonl_path.write_text('{"id": 1, "embedding": [0.25, 0.75]}\n{"id": 2, "embedding": [1.25, 1.75]}\n') - jsonl_schema = pa.schema([("id", pa.int64()), ("embedding", pa.list_(pa.float32()))]) + jsonl_schema = pa.schema(cast("Any", [("id", pa.int64()), ("embedding", pa.list_(pa.float32()))])) jsonl_batches = list( pipeline.read_arrow_batches(jsonl_path, file_format="jsonl", arrow_schema=jsonl_schema, batch_size=10) ) assert sum(b.num_rows for b in jsonl_batches) == 2 assert jsonl_batches[0].schema.field("embedding").type == pa.list_(pa.float32()) - assert np.allclose(jsonl_batches[0].column("embedding").to_pylist()[0], [0.25, 0.75]) + assert np.allclose(cast("Any", jsonl_batches[0].column("embedding").to_pylist()[0]), [0.25, 0.75]) async def test_async_pipeline_write_and_read_arrow_batches_and_arrow_result(tmp_path: Path) -> None: """AsyncStoragePipeline and ArrowResult stream batches without full-table materialization.""" pipeline = AsyncStoragePipeline() parquet_path = tmp_path / "async_items.parquet" - schema = pa.schema([("id", pa.int64()), ("name", pa.string())]) + schema = pa.schema(cast("Any", [("id", pa.int64()), ("name", pa.string())])) batch = pa.RecordBatch.from_pydict({"id": [10, 20], "name": ["x", "y"]}, schema=schema) telemetry = await pipeline.write_arrow_batches_async(parquet_path, [batch], file_format="parquet") @@ -135,7 +142,7 @@ async def test_async_pipeline_write_and_read_arrow_batches_and_arrow_result(tmp_ assert [b.num_rows for b in read_batches] == [1, 1] reader = pa.RecordBatchReader.from_batches(schema, [batch]) - arrow_result = ArrowResult(statement=None, data=reader) + arrow_result = ArrowResult(statement=cast("Any", None), data=reader) out_path = tmp_path / "from_reader.parquet" res_telemetry = await arrow_result.write_to_storage_async(out_path, pipeline=pipeline) assert res_telemetry["rows_processed"] == 2 From e48708d1dce9f749697183b7115daa551f0cc3db Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 20:30:39 +0000 Subject: [PATCH 08/14] fix(arrow): cast stream sources to satisfy pyright SupportArrowStream protocol --- sqlspec/utils/arrow_helpers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 3aefd6f8c..46e9e1aeb 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -418,7 +418,7 @@ def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable": if hasattr(source, "__arrow_c_stream__"): import pyarrow as pa - return pa.RecordBatchReader.from_stream(source).read_all() + return pa.RecordBatchReader.from_stream(cast("Any", source)).read_all() if isinstance(source, Mapping): import pyarrow as pa @@ -457,7 +457,7 @@ def coerce_arrow_batches( if isinstance(unwrapped, pa.Table): return list(unwrapped.schema.names), iter(unwrapped.to_batches(max_chunksize=batch_size)) if hasattr(unwrapped, "__arrow_c_stream__"): - reader = pa.RecordBatchReader.from_stream(unwrapped) + reader = pa.RecordBatchReader.from_stream(cast("Any", unwrapped)) return list(reader.schema.names), iter(reader) table = coerce_arrow_table(unwrapped) return list(table.schema.names), iter(table.to_batches(max_chunksize=batch_size)) From 9d2a7d773ba7c57eec444aae8690e26f19d41869 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Fri, 30 Oct 2026 21:14:53 +0000 Subject: [PATCH 09/14] fix(postgres): default enable_pgcopy_binary to opt-in for adapter copy compatibility --- sqlspec/adapters/asyncpg/driver.py | 2 +- sqlspec/adapters/psycopg/driver.py | 4 ++-- tests/unit/adapters/test_postgres_pgcopy_binary.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 7871db0a9..1ab2c718f 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -393,7 +393,7 @@ async def load_from_arrow( metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} if first_batch is not None: all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) - use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", False)) if use_binary: async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) await self.connection.copy_to_table( diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 10111ac57..db0a86313 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -559,7 +559,7 @@ def load_from_arrow( first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} if first_batch is not None: - use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", False)) copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() @@ -1125,7 +1125,7 @@ async def load_from_arrow( first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} if first_batch is not None: - use_binary = bool(self.driver_features.get("enable_pgcopy_binary", True)) + use_binary = bool(self.driver_features.get("enable_pgcopy_binary", False)) copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() diff --git a/tests/unit/adapters/test_postgres_pgcopy_binary.py b/tests/unit/adapters/test_postgres_pgcopy_binary.py index 961705159..50134a068 100644 --- a/tests/unit/adapters/test_postgres_pgcopy_binary.py +++ b/tests/unit/adapters/test_postgres_pgcopy_binary.py @@ -288,7 +288,7 @@ def test_psycopg_and_cockroach_sync_load_from_arrow_streams_binary_copy() -> Non driver = driver_cls( connection=cast("Any", conn), statement_config=psycopg_statement_config, - driver_features={"storage_capabilities": _CAPS}, + driver_features={"storage_capabilities": _CAPS, "enable_pgcopy_binary": True}, ) reader = pa.RecordBatchReader.from_batches( schema, @@ -316,7 +316,7 @@ async def test_psycopg_and_cockroach_async_load_from_arrow_streams_binary_copy() driver = driver_cls( connection=cast("Any", conn), statement_config=psycopg_statement_config, - driver_features={"storage_capabilities": _CAPS}, + driver_features={"storage_capabilities": _CAPS, "enable_pgcopy_binary": True}, ) reader = pa.RecordBatchReader.from_batches( schema, [pa.RecordBatch.from_pydict({"id": [10, 20], "label": ["x", "y"]}, schema=schema)] @@ -363,7 +363,7 @@ async def test_asyncpg_and_cockroach_asyncpg_load_from_arrow_streams_binary_copy driver = driver_cls( connection=cast("Any", conn), statement_config=asyncpg_statement_config, - driver_features={"storage_capabilities": _CAPS}, + driver_features={"storage_capabilities": _CAPS, "enable_pgcopy_binary": True}, ) batch = pa.RecordBatch.from_arrays( [ From 05c7b4040f3c4926141262891138a5d087ef3653 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 02:22:22 +0000 Subject: [PATCH 10/14] fix(driver,storage,adapters): harden staged bulk loads, streaming storage, pgcopy, and telemetry --- sqlspec/adapters/bigquery/driver.py | 34 ++++++++++++++++++++++- sqlspec/adapters/mssql_python/driver.py | 3 +- sqlspec/adapters/oracledb/driver.py | 6 ++-- sqlspec/adapters/psycopg/driver.py | 10 +++++++ sqlspec/adapters/pymssql/driver.py | 3 +- sqlspec/adapters/spanner/driver.py | 14 +++++++--- sqlspec/core/result/_base.py | 23 +++++++++++++-- sqlspec/driver/_storage_helpers.py | 26 +++++++++-------- sqlspec/storage/__init__.py | 3 ++ sqlspec/storage/pipeline.py | 31 +++++++++++++++++++++ sqlspec/utils/arrow_helpers.py | 37 +++++++++++++++++++++---- 11 files changed, 162 insertions(+), 28 deletions(-) diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index 700e1b415..f89bdfadb 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -691,6 +691,37 @@ def load_from_records( type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load in-memory records through BigQuery's newline-delimited JSON load API.""" + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_records( + table, + records, + columns=columns, + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return self._execute_staged_load_sync( + table, + records, + raw_loader=lambda tbl, src, **kwargs: self.load_from_records(tbl, src, columns=columns, **kwargs), + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) rows = _records_to_json_rows(records, columns) job_config = build_load_job_config("jsonl", overwrite) qualified_table = self._qualify_table_name(table) @@ -744,9 +775,10 @@ def load_from_arrow( type_casts=type_casts, ) if stage_table or conflict_keys: - return super().load_from_arrow( + return self._execute_staged_load_sync( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 2c2266bca..7809336fb 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -488,9 +488,10 @@ def load_from_arrow( column_mappings=column_mappings, ) if stage_table or conflict_keys: - return super().load_from_arrow( + return self._execute_staged_load_sync( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index cb9034a2d..174e6595f 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -863,9 +863,10 @@ def load_from_arrow( type_casts=type_casts, ) if stage_table or conflict_keys: - return super().load_from_arrow( + return self._execute_staged_load_sync( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, @@ -1690,9 +1691,10 @@ async def load_from_arrow( type_casts=type_casts, ) if stage_table or conflict_keys: - return await super().load_from_arrow( + return await self._execute_staged_load_async( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index db0a86313..3ddebc4eb 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -575,6 +575,11 @@ def load_from_arrow( metrics["rows_processed"] += int(batch.num_rows) metrics["bytes_processed"] += int(batch.nbytes) records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", + self.prepare_driver_parameters(records, self.statement_config, is_many=True), + ) for record in records: copy_ctx.write_row(record) if exc_handler.pending_exception is not None: @@ -1141,6 +1146,11 @@ async def load_from_arrow( metrics["rows_processed"] += int(batch.num_rows) metrics["bytes_processed"] += int(batch.nbytes) records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", + self.prepare_driver_parameters(records, self.statement_config, is_many=True), + ) for record in records: await copy_ctx.write_row(record) if exc_handler.pending_exception is not None: diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index feb0f658d..db5cfef23 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -376,9 +376,10 @@ def load_from_arrow( type_casts=type_casts, ) if stage_table or conflict_keys: - return super().load_from_arrow( + return self._execute_staged_load_sync( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index c8a5b333f..fe298de68 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -553,10 +553,13 @@ def load_from_arrow( concurrent_indexes=concurrent_indexes, type_casts=type_casts, ) - if stage_table: - return super().load_from_arrow( + if stage_table or ( + conflict_keys and on_conflict == "nothing" and not hasattr(self.connection, "insert_or_ignore") + ): + return self._execute_staged_load_sync( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, @@ -1214,10 +1217,13 @@ async def load_from_arrow( concurrent_indexes=concurrent_indexes, type_casts=type_casts, ) - if stage_table: - return await super().load_from_arrow( + if stage_table or ( + conflict_keys and on_conflict == "nothing" and not hasattr(self.connection, "insert_or_ignore") + ): + return await self._execute_staged_load_async( table, source, + raw_loader=self.load_from_arrow, partitioner=partitioner, overwrite=overwrite, telemetry=telemetry, diff --git a/sqlspec/core/result/_base.py b/sqlspec/core/result/_base.py index 5d2601c08..124079e98 100644 --- a/sqlspec/core/result/_base.py +++ b/sqlspec/core/result/_base.py @@ -34,6 +34,7 @@ arrow_table_to_pylist, arrow_table_to_return_format, cast_arrow_table_schema, + coerce_arrow_batches, coerce_arrow_table, convert_dict_to_arrow, ensure_arrow_table, @@ -991,8 +992,17 @@ def write_to_storage_sync( if self.data is None: msg = "No Arrow table available for this result" raise ValueError(msg) - table = coerce_arrow_table(self.data) active_pipeline = pipeline or SyncStoragePipeline() + if hasattr(self.data, "to_batches") or hasattr(self.data, "iter_batches") or hasattr(self.data, "__iter__"): + _cols, batch_iter = coerce_arrow_batches(self.data) + return active_pipeline.write_arrow_batches( + batch_iter, + destination, + format_hint=format_hint, + storage_options=storage_options, + compression=compression, + ) + table = coerce_arrow_table(self.data) return active_pipeline.write_arrow( table, destination, format_hint=format_hint, storage_options=storage_options, compression=compression ) @@ -1009,8 +1019,17 @@ async def write_to_storage_async( if self.data is None: msg = "No Arrow table available for this result" raise ValueError(msg) - table = coerce_arrow_table(self.data) active_pipeline = pipeline or AsyncStoragePipeline() + if hasattr(self.data, "to_batches") or hasattr(self.data, "iter_batches") or hasattr(self.data, "__iter__"): + _cols, batch_iter = coerce_arrow_batches(self.data) + return await active_pipeline.write_arrow_batches_async( + batch_iter, + destination, + format_hint=format_hint, + storage_options=storage_options, + compression=compression, + ) + table = coerce_arrow_table(self.data) return await active_pipeline.write_arrow( table, destination, format_hint=format_hint, storage_options=storage_options, compression=compression ) diff --git a/sqlspec/driver/_storage_helpers.py b/sqlspec/driver/_storage_helpers.py index 84a9e33e2..08455ce03 100644 --- a/sqlspec/driver/_storage_helpers.py +++ b/sqlspec/driver/_storage_helpers.py @@ -224,17 +224,21 @@ def create_storage_job( _CREATE_INDEX_RE: Final = re.compile(r"^\s*CREATE\s+INDEX\s+(?!CONCURRENTLY\b)", re.IGNORECASE) -def compute_worker_thread_budget(max_workers: int, cpu_count: int | None = None) -> "dict[str, str]": +def compute_worker_thread_budget( + max_workers: int, cpu_count: int | None = None, *, reserve_cpus: int = 0 +) -> "dict[str, str]": """Compute per-worker native thread pool environment variables to prevent oversubscription. Args: max_workers: Number of concurrent ingestion workers. cpu_count: Optional logical CPU count override. + reserve_cpus: Number of CPU cores reserved for the host OS and main event loop. Returns: Dictionary of thread-budget environment variables and string values. """ - resolved_cpus = cpu_count if cpu_count is not None and cpu_count > 0 else (os.cpu_count() or 1) + total_cpus = cpu_count if cpu_count is not None and cpu_count > 0 else (os.cpu_count() or 1) + resolved_cpus = max(1, total_cpus - max(0, reserve_cpus)) workers = max(1, int(max_workers)) threads = str(max(1, resolved_cpus // workers)) return { @@ -247,9 +251,9 @@ def compute_worker_thread_budget(max_workers: int, cpu_count: int | None = None) } -def init_spawn_worker_thread_budget(max_workers: int, cpu_count: int | None = None) -> None: +def init_spawn_worker_thread_budget(max_workers: int, cpu_count: int | None = None, *, reserve_cpus: int = 0) -> None: """Initialize thread-budget environment variables inside a spawned worker process.""" - for key, val in compute_worker_thread_budget(max_workers, cpu_count=cpu_count).items(): + for key, val in compute_worker_thread_budget(max_workers, cpu_count=cpu_count, reserve_cpus=reserve_cpus).items(): os.environ[key] = val @@ -407,30 +411,30 @@ def build_stage_merge_sql( if norm in _MYSQL_FAMILY_DIALECTS: if on_conflict == "nothing" or not update_columns: return f"INSERT IGNORE INTO {target_table} ({col_list}) SELECT {select_list} FROM {stage_table}" - set_clause = ", ".join(f"{col} = _sqlspec_stage.{col}" for col in update_columns) + set_clause = ", ".join(f"{col} = stg.{col}" for col in update_columns) return ( f"INSERT INTO {target_table} ({col_list}) " - f"SELECT {select_list} FROM {stage_table} AS _sqlspec_stage " + f"SELECT {select_list} FROM {stage_table} AS stg " f"ON DUPLICATE KEY UPDATE {set_clause}" ) - on_clause = " AND ".join(f"_target.{key} = _stage.{key}" for key in conflict_keys) + on_clause = " AND ".join(f"tgt.{key} = stg.{key}" for key in conflict_keys) matched_clause = "" if on_conflict == "update" and update_columns: set_pairs = ", ".join( - f"_target.{col} = CAST(_stage.{col} AS {type_casts[col]})" + f"tgt.{col} = CAST(stg.{col} AS {type_casts[col]})" if (type_casts and col in type_casts) - else f"_target.{col} = _stage.{col}" + else f"tgt.{col} = stg.{col}" for col in update_columns ) matched_clause = f" WHEN MATCHED THEN UPDATE SET {set_pairs}" stage_vals = ", ".join( - f"CAST(_stage.{col} AS {type_casts[col]})" if (type_casts and col in type_casts) else f"_stage.{col}" + f"CAST(stg.{col} AS {type_casts[col]})" if (type_casts and col in type_casts) else f"stg.{col}" for col in columns ) not_matched_clause = f" WHEN NOT MATCHED THEN INSERT ({col_list}) VALUES ({stage_vals})" terminator = ";" if norm in _MSSQL_FAMILY_DIALECTS else "" - return f"MERGE INTO {target_table} _target USING {stage_table} _stage ON ({on_clause}){matched_clause}{not_matched_clause}{terminator}" + return f"MERGE INTO {target_table} tgt USING {stage_table} stg ON ({on_clause}){matched_clause}{not_matched_clause}{terminator}" def build_drop_stage_table_sql(dialect: str, stage_table: str) -> str: diff --git a/sqlspec/storage/__init__.py b/sqlspec/storage/__init__.py index 733df124b..17453e5bd 100644 --- a/sqlspec/storage/__init__.py +++ b/sqlspec/storage/__init__.py @@ -23,6 +23,7 @@ StorageFormat, StorageTelemetry, SyncStoragePipeline, + aggregate_storage_telemetry, create_storage_bridge_job, get_storage_bridge_diagnostics, get_storage_bridge_metrics, @@ -41,6 +42,7 @@ "StorageRegistry", "StorageTelemetry", "SyncStoragePipeline", + "aggregate_storage_telemetry", "create_storage_bridge_job", "get_storage_bridge_diagnostics", "get_storage_bridge_metrics", @@ -60,6 +62,7 @@ "StorageRegistry": ("sqlspec.storage.registry", "StorageRegistry"), "StorageTelemetry": ("sqlspec.storage.pipeline", "StorageTelemetry"), "SyncStoragePipeline": ("sqlspec.storage.pipeline", "SyncStoragePipeline"), + "aggregate_storage_telemetry": ("sqlspec.storage.pipeline", "aggregate_storage_telemetry"), "create_storage_bridge_job": ("sqlspec.storage.pipeline", "create_storage_bridge_job"), "get_storage_bridge_diagnostics": ("sqlspec.storage.pipeline", "get_storage_bridge_diagnostics"), "get_storage_bridge_metrics": ("sqlspec.storage.pipeline", "get_storage_bridge_metrics"), diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index dfbb70001..951dbb096 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -46,6 +46,7 @@ "StorageFormat", "StorageTelemetry", "SyncStoragePipeline", + "aggregate_storage_telemetry", "create_storage_bridge_job", "get_recent_storage_events", "get_storage_bridge_diagnostics", @@ -130,6 +131,36 @@ class StorageTelemetry(TypedDict, total=False): """Registered SQLSpec configuration bind key, when configured.""" +def aggregate_storage_telemetry(telemetries: "Sequence[StorageTelemetry]") -> "StorageTelemetry": + """Aggregate a sequence of StorageTelemetry payloads into a unified summary.""" + if not telemetries: + return {} + summary: StorageTelemetry = { + "rows_processed": sum(t.get("rows_processed", 0) for t in telemetries), + "bytes_processed": sum(t.get("bytes_processed", 0) for t in telemetries), + "partitions_created": sum(t.get("partitions_created", 0) for t in telemetries), + "shards_processed": sum(t.get("shards_processed", 0) for t in telemetries), + "shards_skipped": sum(t.get("shards_skipped", 0) for t in telemetries), + } + durations = [t["duration_s"] for t in telemetries if "duration_s" in t] + if durations: + summary["duration_s"] = sum(durations) + for dur_key in ("read_duration_s", "encode_duration_s", "write_duration_s", "commit_duration_s"): + vals = [t[dur_key] for t in telemetries if dur_key in t] + if vals: + summary[dur_key] = sum(vals) + first_dest = next((t["destination"] for t in telemetries if "destination" in t), None) + if first_dest is not None: + summary["destination"] = first_dest + first_fmt = next((t["format"] for t in telemetries if "format" in t), None) + if first_fmt is not None: + summary["format"] = first_fmt + first_backend = next((t["backend"] for t in telemetries if "backend" in t), None) + if first_backend is not None: + summary["backend"] = first_backend + return summary + + class ResolvedStorageTarget(NamedTuple): """A storage destination resolved to an address and its backend protocol.""" diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 46e9e1aeb..4ea1db8ec 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -16,8 +16,7 @@ from sqlspec.data_dictionary._registry import get_dialect_config from sqlspec.exceptions import ImproperConfigurationError, SQLBuilderError from sqlspec.utils.dispatch import TypeDispatcher -from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow -from sqlspec.utils.serializers import to_json +from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow, import_optional_attr from sqlspec.utils.type_guards import has_arrow_table_stats, has_get_data from sqlspec.utils.uuids import uuid_from_bytes @@ -830,7 +829,7 @@ def _extract_pg_fixed_column_array(col: Any, dtype: Any, num_rows: int) -> "tupl ">i2", np.frombuffer(col.buffers()[1], dtype=np.uint8, count=col.offset + num_rows)[col.offset :].astype(">i2"), ) - if pa.types.is_int16(dtype) or pa.types.is_uint16(dtype): + if pa.types.is_int16(dtype): return ( 2, ">i2", @@ -838,7 +837,13 @@ def _extract_pg_fixed_column_array(col: Any, dtype: Any, num_rows: int) -> "tupl ">i2", copy=False ), ) - if pa.types.is_int32(dtype) or pa.types.is_uint32(dtype): + if pa.types.is_uint16(dtype): + return ( + 4, + ">i4", + np.frombuffer(col.buffers()[1], dtype=np.uint16, count=col.offset + num_rows)[col.offset :].astype(">i4"), + ) + if pa.types.is_int32(dtype): return ( 4, ">i4", @@ -846,6 +851,12 @@ def _extract_pg_fixed_column_array(col: Any, dtype: Any, num_rows: int) -> "tupl ">i4", copy=False ), ) + if pa.types.is_uint32(dtype): + return ( + 8, + ">i8", + np.frombuffer(col.buffers()[1], dtype=np.uint32, count=col.offset + num_rows)[col.offset :].astype(">i8"), + ) if pa.types.is_int64(dtype) or pa.types.is_uint64(dtype): return ( 8, @@ -1000,7 +1011,11 @@ def _emit_fixed( if pa.types.is_fixed_size_list(field_type) and pa.types.is_floating(field_type.value_type): dim = int(field_type.list_size) - flat_values = col.flatten() + if null_mask is not None and np.any(null_mask): + col_filled = col.fill_null([0.0] * dim) + flat_values = col_filled.flatten() + else: + flat_values = col.flatten() is_f16 = pa.types.is_float16(field_type.value_type) elem_dtype = ">f2" if is_f16 else ">f4" elem_size = 2 if is_f16 else 4 @@ -1142,7 +1157,7 @@ def _emit_fallback( elif isinstance(val, str): raw = val.encode("utf-8") elif isinstance(val, (dict, list, tuple)): - raw = to_json(val).encode("utf-8") + raw = _encode_json_bytes(val) else: raw = str(val).encode("utf-8") if _is_jsonb: @@ -1154,6 +1169,16 @@ def _emit_fallback( return _emit_fallback +def _encode_json_bytes(val: Any) -> bytes: + """Encode value to JSON bytes using msgspec or stdlib json without circular imports.""" + encoder = import_optional_attr("sqlspec.utils.serializers._json", "encode_json") + if encoder is not None: + return cast("bytes", encoder(val, as_bytes=True)) + import json + + return json.dumps(val).encode("utf-8") + + def _encode_pg_binary_fast_path_2( batch: "ArrowRecordBatch", columns: "list[str]", vector_columns: frozenset[str], jsonb_columns: frozenset[str] ) -> bytes: From 70bb8b3d7936d3b57c912a01ad192f9f4249f05e Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 02:41:34 +0000 Subject: [PATCH 11/14] fix(driver,storage,adapters): align sync load overrides, pipeline batch write args, and batch prep typing --- sqlspec/adapters/adbc/driver.py | 3 +- sqlspec/adapters/sqlanywhere/driver.py | 50 ++++++++++++++++++++++++-- sqlspec/adapters/sybase/driver.py | 50 ++++++++++++++++++++++++-- sqlspec/core/result/_base.py | 8 ++--- sqlspec/driver/_common.py | 2 +- sqlspec/utils/arrow_helpers.py | 2 +- 6 files changed, 103 insertions(+), 12 deletions(-) diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 15be49f26..c3f6cb028 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -60,8 +60,7 @@ from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcRawCursor from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, SQLResult, Statement, StatementFilter - from sqlspec.driver import ExecutionResult - from sqlspec.driver._query_cache import CachedQuery + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry from sqlspec.typing import ArrowReturnFormat, StatementParameters diff --git a/sqlspec/adapters/sqlanywhere/driver.py b/sqlspec/adapters/sqlanywhere/driver.py index 908f800c3..b4048fe3a 100644 --- a/sqlspec/adapters/sqlanywhere/driver.py +++ b/sqlspec/adapters/sqlanywhere/driver.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Sized -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.sqlanywhere._typing import ( SqlanywhereConnection, @@ -38,7 +38,7 @@ from sqlspec.utils.logging import get_logger if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from sqlspec.adapters.sqlanywhere._typing import SqlanywhereQueryParams as QueryParams from sqlspec.builder import QueryBuilder @@ -294,9 +294,50 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target SQL Anywhere table via executemany.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"DELETE FROM {format_identifier(table)}" @@ -333,8 +374,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into the target SQL Anywhere table.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) diff --git a/sqlspec/adapters/sybase/driver.py b/sqlspec/adapters/sybase/driver.py index 1d6bf2885..8b8b1e17e 100644 --- a/sqlspec/adapters/sybase/driver.py +++ b/sqlspec/adapters/sybase/driver.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Sized -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.sybase._typing import ( SybaseConnection, @@ -39,7 +39,7 @@ from sqlspec.utils.logging import get_logger if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from sqlspec.adapters.sybase._typing import SybaseQueryParams as QueryParams from sqlspec.builder import QueryBuilder @@ -302,9 +302,50 @@ def load_from_arrow( partitioner: "dict[str, object] | None" = None, overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = False, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = True, + type_casts: "Mapping[str, str] | None" = None, ) -> "StorageBridgeJob": """Load Arrow data into the target Sybase ASE table via executemany.""" self._require_capability("arrow_import_enabled") + if drop_indexes: + with self.without_indexes(table, concurrent=concurrent_indexes): + return self.load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=False, + type_casts=type_casts, + ) + if stage_table or conflict_keys: + return self._execute_staged_load_sync( + table, + source, + raw_loader=self.load_from_arrow, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + type_casts=type_casts, + ) arrow_table = self._coerce_arrow_table(source) if overwrite: statement = f"DELETE FROM {format_identifier(table)}" @@ -343,8 +384,13 @@ def load_from_storage( file_format: "StorageFormat", partitioner: "dict[str, object] | None" = None, overwrite: bool = False, + **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into the target Sybase ASE table.""" + if kwargs: + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) diff --git a/sqlspec/core/result/_base.py b/sqlspec/core/result/_base.py index 124079e98..9280f9918 100644 --- a/sqlspec/core/result/_base.py +++ b/sqlspec/core/result/_base.py @@ -996,9 +996,9 @@ def write_to_storage_sync( if hasattr(self.data, "to_batches") or hasattr(self.data, "iter_batches") or hasattr(self.data, "__iter__"): _cols, batch_iter = coerce_arrow_batches(self.data) return active_pipeline.write_arrow_batches( - batch_iter, destination, - format_hint=format_hint, + batch_iter, + file_format=format_hint or "parquet", storage_options=storage_options, compression=compression, ) @@ -1023,9 +1023,9 @@ async def write_to_storage_async( if hasattr(self.data, "to_batches") or hasattr(self.data, "iter_batches") or hasattr(self.data, "__iter__"): _cols, batch_iter = coerce_arrow_batches(self.data) return await active_pipeline.write_arrow_batches_async( - batch_iter, destination, - format_hint=format_hint, + batch_iter, + file_format=format_hint or "parquet", storage_options=storage_options, compression=compression, ) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index d590abdf6..6e001a091 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -3227,7 +3227,7 @@ def _arrow_table_to_rows( return arrow_table_to_rows(table, columns) @staticmethod - def _arrow_rows_need_preparation(table: "ArrowTable") -> bool: + def _arrow_rows_need_preparation(table: "ArrowTable | ArrowRecordBatch") -> bool: """Return whether Arrow rows may contain nested values needing preparation.""" return arrow_table_needs_parameter_preparation(table) diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 4ea1db8ec..8627be80e 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -565,7 +565,7 @@ def arrow_table_to_rows( return resolved_columns, records -def arrow_table_needs_parameter_preparation(table: "ArrowTable") -> bool: +def arrow_table_needs_parameter_preparation(table: "ArrowTable | ArrowRecordBatch") -> bool: """Return whether Arrow rows may emit nested values needing driver preparation.""" return _arrow_schema_needs_preparation(table.schema) From f6a25ee63aecd2ddbc8ea48c331484729174374d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 02:55:46 +0000 Subject: [PATCH 12/14] fix(duckdb): safely detect memory on windows without os.sysconf --- sqlspec/adapters/duckdb/core.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index b6d1bbae1..c9ecf6997 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -652,11 +652,13 @@ def _detect_available_memory_bytes( if host_memory_bytes is not None and host_memory_bytes > 0: return host_memory_bytes - with contextlib.suppress(ValueError, OSError, AttributeError): - page_size = os.sysconf("SC_PAGE_SIZE") - phys_pages = os.sysconf("SC_PHYS_PAGES") - if isinstance(page_size, int) and isinstance(phys_pages, int) and page_size > 0 and phys_pages > 0: - return page_size * phys_pages + sysconf = getattr(os, "sysconf", None) + if sysconf is not None: + with contextlib.suppress(ValueError, OSError, AttributeError): + page_size = sysconf("SC_PAGE_SIZE") + phys_pages = sysconf("SC_PHYS_PAGES") + if isinstance(page_size, int) and isinstance(phys_pages, int) and page_size > 0 and phys_pages > 0: + return page_size * phys_pages return None From cc0259d6da35684835e8ff30dd6562ebac530916 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 03:17:49 +0000 Subject: [PATCH 13/14] refactor(sqlite,aiosqlite): debloat ad-hoc execute_many thin path --- sqlspec/adapters/aiosqlite/driver.py | 160 +----------------- sqlspec/adapters/sqlite/driver.py | 158 +---------------- .../adapters/test_aiosqlite/test_driver.py | 35 ++-- .../unit/adapters/test_sqlite/test_driver.py | 41 +---- .../test_sqlite/test_rowcount_optimization.py | 36 ++-- 5 files changed, 38 insertions(+), 392 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 09a8d243f..98a1f4697 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -28,15 +28,9 @@ run_on_worker_thread, ) from sqlspec.adapters.aiosqlite.data_dictionary import AiosqliteDataDictionary -from sqlspec.core import ArrowResult, ParameterStyle, TypedParameter, get_cache_config, register_driver_profile +from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile from sqlspec.core.result import DMLResult -from sqlspec.driver import ( - AsyncDriverAdapterBase, - AsyncRowStream, - BaseAsyncExceptionHandler, - parameter_value_needs_processing, - type_coercion_fallbacks, -) +from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import DeadlockError from sqlspec.utils.type_guards import resolve_row_format @@ -44,12 +38,10 @@ from collections.abc import Mapping, Sequence from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection - from sqlspec.builder import QueryBuilder - from sqlspec.core import SQL, SQLResult, Statement, StatementConfig, StatementFilter + from sqlspec.core import SQL, SQLResult, StatementConfig from sqlspec.core.compiler import OperationType from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry - from sqlspec.typing import StatementParameters __all__ = ( "AiosqliteCursor", @@ -186,43 +178,6 @@ async def dispatch_execute_script(self, cursor: "AiosqliteRawCursor", statement: last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) - async def execute_many( - self, - statement: "SQL | Statement | QueryBuilder", - /, - parameters: "Sequence[StatementParameters]", - *filters: "StatementParameters | StatementFilter", - statement_config: "StatementConfig | None" = None, - **kwargs: Any, - ) -> "SQLResult": - """Execute many with an AIOSQLite thin path for simple qmark batches.""" - config = statement_config or self.statement_config - if ( - isinstance(statement, str) - and not filters - and not kwargs - and config is self.statement_config - and self.observability.is_idle - and self._can_use_execute_many_thin_path(statement, parameters, config) - ): - cursor = None - try: - cursor = await self.connection.executemany(statement, parameters) - rowcount = cursor.rowcount - affected_rows = rowcount if isinstance(rowcount, int) and rowcount > 0 else 0 - operation = self._resolve_dml_operation_type(statement) - self._invalidate_rowid_target_cache(operation) - return DMLResult(operation, affected_rows) - except (aiosqlite.Error, sqlite3.Error) as exc: - raise create_mapped_exception(exc) from exc - finally: - if cursor is not None: - with contextlib.suppress(Exception): - close_result = cursor.close() - if inspect.isawaitable(close_result): - await close_result - return await super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) - async def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. @@ -457,115 +412,6 @@ def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> Non if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: self._rowid_target_cache.clear() - def _can_use_execute_many_thin_path( - self, statement: str, parameters: "Sequence[StatementParameters]", config: "StatementConfig" - ) -> bool: - if type(parameters) is not list: - return False - if not parameters: - return False - if "?" not in statement: - return False - if self._resolve_dml_operation_type(statement) not in {"INSERT", "UPDATE", "DELETE"}: - return False - if "RETURNING" in statement.upper(): - return False - - parameter_config = config.parameter_config - if parameter_config.default_parameter_style is not ParameterStyle.QMARK: - return False - if ( - parameter_config.default_execution_parameter_style is not None - and parameter_config.default_execution_parameter_style is not ParameterStyle.QMARK - ): - return False - if parameter_config.ast_transformer is not None or parameter_config.output_transformer is not None: - return False - if parameter_config.needs_static_script_compilation: - return False - if config.output_transformer is not None or config.statement_transformers: - return False - - return self._thin_path_parameters_are_eligible(parameters, parameter_config.type_coercion_map) - - @staticmethod - def _thin_path_parameters_are_eligible( - parameters: "list[StatementParameters]", type_coercion_map: "dict[type, Any] | None" - ) -> bool: - first_sequence = AiosqliteDriver._as_sequence_parameter_set(parameters[0]) - if first_sequence is None: - return False - - first_type = type(first_sequence) - row_len = len(first_sequence) - coercion_map = type_coercion_map - has_type_coercion = bool(coercion_map) - fallback_items = type_coercion_fallbacks(coercion_map) if coercion_map else () - - if row_len == 1: - if has_type_coercion and coercion_map is not None: - for param_set in parameters: - sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != 1: - return False - if parameter_value_needs_processing(sequence[0], coercion_map, fallback_items): - return False - return True - - for param_set in parameters: - sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != 1: - return False - if type(sequence[0]) is TypedParameter: - return False - return True - - if has_type_coercion and coercion_map is not None: - for param_set in parameters: - sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != row_len: - return False - for value in sequence: - if parameter_value_needs_processing(value, coercion_map, fallback_items): - return False - return True - - for param_set in parameters: - sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != row_len: - return False - for value in sequence: - if type(value) is TypedParameter: - return False - return True - - @staticmethod - def _as_sequence_parameter_set(param_set: "StatementParameters") -> "list[Any] | tuple[Any, ...] | None": - if isinstance(param_set, list): - return param_set - if isinstance(param_set, tuple): - return param_set - return None - - @staticmethod - def _resolve_dml_operation_type(statement: str) -> "OperationType": - command_keyword = statement.lstrip().split(None, 1)[0].upper() if statement.strip() else "COMMAND" - if command_keyword == "INSERT": - return "INSERT" - if command_keyword == "UPDATE": - return "UPDATE" - if command_keyword == "DELETE": - return "DELETE" - return "COMMAND" - def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 87a9703dd..1729a5138 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -24,14 +24,7 @@ from sqlspec.adapters.sqlite.data_dictionary import SqliteDataDictionary from sqlspec.core.cache import get_cache_config from sqlspec.core.parameters._registry import register_driver_profile -from sqlspec.core.parameters._types import ParameterStyle, TypedParameter from sqlspec.core.result._base import ArrowResult, DMLResult -from sqlspec.driver._common import ( - CachedQuery, - ExecutionResult, - parameter_value_needs_processing, - type_coercion_fallbacks, -) from sqlspec.driver._exception_handler import BaseSyncExceptionHandler from sqlspec.driver._stream import SyncRowStream from sqlspec.driver._sync import SyncDriverAdapterBase @@ -43,13 +36,11 @@ from sqlglot.dialects.dialect import DialectType from sqlspec.adapters.sqlite._typing import SqliteConnection - from sqlspec.builder._base import QueryBuilder from sqlspec.core.compiler import OperationType - from sqlspec.core.filters import StatementFilter from sqlspec.core.result._base import SQLResult - from sqlspec.core.statement import SQL, Statement, StatementConfig + from sqlspec.core.statement import SQL, StatementConfig + from sqlspec.driver._common import CachedQuery, ExecutionResult from sqlspec.storage.pipeline import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry - from sqlspec.typing import StatementParameters __all__ = ("SqliteCursor", "SqliteDriver", "SqliteExceptionHandler", "SqliteSessionContext") @@ -207,41 +198,6 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) - def execute_many( - self, - statement: "SQL | Statement | QueryBuilder", - /, - parameters: "Sequence[StatementParameters]", - *filters: "StatementParameters | StatementFilter", - statement_config: "StatementConfig | None" = None, - **kwargs: Any, - ) -> "SQLResult": - """Execute many with a SQLite thin path for simple qmark batches.""" - config = statement_config or self.statement_config - if ( - isinstance(statement, str) - and not filters - and not kwargs - and config is self.statement_config - and self.observability.is_idle - and self._can_use_execute_many_thin_path(statement, parameters, config) - ): - cursor = None - try: - cursor = self.connection.executemany(statement, parameters) - affected_rows = resolve_rowcount(cursor) - except sqlite3.Error as exc: - raise create_mapped_exception(exc) from exc - finally: - if cursor is not None: - with contextlib.suppress(Exception): - cursor.close() - - operation = self._resolve_dml_operation_type(statement) - self._invalidate_rowid_target_cache(operation) - return DMLResult(operation, affected_rows) - return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) - def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. @@ -530,116 +486,6 @@ def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> Non if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: self._rowid_target_cache.clear() - def _can_use_execute_many_thin_path( - self, statement: str, parameters: "Sequence[StatementParameters]", config: "StatementConfig" - ) -> bool: - if type(parameters) is not list: - return False - if not parameters: - return False - if "?" not in statement: - return False - if self._resolve_dml_operation_type(statement) not in {"INSERT", "UPDATE", "DELETE"}: - return False - if "RETURNING" in statement.upper(): - return False - - parameter_config = config.parameter_config - if parameter_config.default_parameter_style is not ParameterStyle.QMARK: - return False - if ( - parameter_config.default_execution_parameter_style is not None - and parameter_config.default_execution_parameter_style is not ParameterStyle.QMARK - ): - return False - if parameter_config.ast_transformer is not None or parameter_config.output_transformer is not None: - return False - if parameter_config.needs_static_script_compilation: - return False - if config.output_transformer is not None or config.statement_transformers: - return False - - return self._thin_path_parameters_are_eligible(parameters, parameter_config.type_coercion_map) - - @staticmethod - def _thin_path_parameters_are_eligible( - parameters: "Sequence[StatementParameters]", type_coercion_map: "dict[type, Any] | None" - ) -> bool: - """Validate parameter payload for the SQLite execute-many thin path.""" - first_sequence = SqliteDriver._as_sequence_parameter_set(parameters[0]) - if first_sequence is None: - return False - - first_type = type(first_sequence) - row_len = len(first_sequence) - coercion_map = type_coercion_map - has_type_coercion = bool(coercion_map) - fallback_items = type_coercion_fallbacks(coercion_map) if coercion_map else () - - if row_len == 1: - if has_type_coercion and coercion_map is not None: - for param_set in parameters: - sequence = SqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != 1: - return False - if parameter_value_needs_processing(sequence[0], coercion_map, fallback_items): - return False - return True - - for param_set in parameters: - sequence = SqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != 1: - return False - if type(sequence[0]) is TypedParameter: - return False - return True - - if has_type_coercion and coercion_map is not None: - for param_set in parameters: - sequence = SqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != row_len: - return False - for value in sequence: - if parameter_value_needs_processing(value, coercion_map, fallback_items): - return False - return True - - for param_set in parameters: - sequence = SqliteDriver._as_sequence_parameter_set(param_set) - if sequence is None or type(sequence) is not first_type: - return False - if len(sequence) != row_len: - return False - for value in sequence: - if type(value) is TypedParameter: - return False - return True - - @staticmethod - def _as_sequence_parameter_set(param_set: "StatementParameters") -> "list[Any] | tuple[Any, ...] | None": - if isinstance(param_set, list): - return param_set - if isinstance(param_set, tuple): - return param_set - return None - - @staticmethod - def _resolve_dml_operation_type(statement: str) -> "OperationType": - command_keyword = statement.lstrip().split(None, 1)[0].upper() if statement.strip() else "COMMAND" - if command_keyword == "INSERT": - return "INSERT" - if command_keyword == "UPDATE": - return "UPDATE" - if command_keyword == "DELETE": - return "DELETE" - return "COMMAND" - def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index 958254589..daeb6b842 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -22,7 +22,6 @@ from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver from sqlspec.adapters.aiosqlite.pool import AiosqliteConnectionPool from sqlspec.core import ParameterStyle -from sqlspec.core.result import DMLResult from sqlspec.exceptions import SQLSpecError @@ -191,27 +190,19 @@ async def _no_sleep(_delay: float) -> None: assert isinstance(exc_info.value.__cause__, aiosqlite.Error) -async def test_execute_many_uses_thin_qmark_path() -> None: - connection = _WorkerConnection() - driver = AiosqliteDriver(connection=cast("Any", connection), statement_config=default_statement_config) - - result = await driver.execute_many("INSERT INTO test_thin_path (value) VALUES (?)", [("a",), ("b",)]) - - assert isinstance(result, DMLResult) - assert result.operation_type == "INSERT" - assert result.rows_affected == 2 - assert connection.executemany_calls == [("INSERT INTO test_thin_path (value) VALUES (?)", [("a",), ("b",)])] - - -def test_execute_many_thin_path_rejects_subclass_coercion_values() -> None: - from collections import defaultdict - - assert ( - AiosqliteDriver._thin_path_parameters_are_eligible( - [(defaultdict(int, a=1),)], default_statement_config.parameter_config.type_coercion_map - ) - is False - ) +async def test_execute_many_inserts_records() -> None: + conn = await aiosqlite.connect(":memory:") + try: + await conn.execute("CREATE TABLE test_many (value TEXT)") + driver = AiosqliteDriver(connection=conn, statement_config=default_statement_config) + result = await driver.execute_many("INSERT INTO test_many (value) VALUES (?)", [("a",), ("b",)]) + assert result.operation_type == "INSERT" + assert result.rows_affected == 2 + async with conn.execute("SELECT COUNT(*) FROM test_many") as cursor: + row = await cursor.fetchone() + assert row is not None and row[0] == 2 + finally: + await conn.close() class _FakeConnection: diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index 7d06b9fdb..7d0b46fc9 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -1,12 +1,11 @@ import sqlite3 -from collections import defaultdict from pathlib import Path from typing import Any, cast import pytest from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver -from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT, default_statement_config +from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT from sqlspec.adapters.sqlite.pool import SqliteConnectionPool from sqlspec.adapters.sqlite.type_converter import json_converter, register_type_handlers, unregister_type_handlers from sqlspec.exceptions import ImproperConfigurationError @@ -218,46 +217,20 @@ def executemany(self, sql: str, parameters: object) -> _TrackingCursor: return cursor -def test_execute_many_thin_path_closes_cursor() -> None: - """SqliteDriver.execute_many thin path must close its cursor in finally.""" - connection = _TrackingConnection() - driver = SqliteDriver(connection=cast("Any", connection)) - - result = driver.execute_many("INSERT INTO items (name) VALUES (?)", [("a",), ("b",)]) - - assert result.rows_affected == 2 - assert len(connection.cursors) == 1 - assert connection.cursors[0].closed is True - - -def test_execute_cache_hit_closes_cursor() -> None: - """SqliteDriver._execute_cache_hit must close its cursor on cached execution.""" +def test_execute_many_closes_cursor() -> None: + """SqliteDriver.execute_many must close its cursor in finally.""" connection = sqlite3.connect(":memory:") connection.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)") - connection.execute("INSERT INTO items (name) VALUES ('alpha')") driver = SqliteDriver(connection=connection) try: - first = driver.execute("SELECT id, name FROM items WHERE id = ?", (1,)) - second = driver.execute("SELECT id, name FROM items WHERE id = ?", (1,)) - assert first.get_data() == [{"id": 1, "name": "alpha"}] - assert second.get_data() == [{"id": 1, "name": "alpha"}] + result = driver.execute_many("INSERT INTO items (name) VALUES (?)", [("a",), ("b",)]) + assert result.rows_affected == 2 + rows = connection.execute("SELECT name FROM items ORDER BY id").fetchall() + assert rows == [("a",), ("b",)] finally: connection.close() -def test_execute_many_thin_path_checks_all_rows_beyond_sample_threshold() -> None: - """_thin_path_parameters_are_eligible must inspect every row even in large batches.""" - rows: list[tuple[Any, ...]] = [(i,) for i in range(120)] - rows[57] = (defaultdict(int, a=1),) - - assert ( - SqliteDriver._thin_path_parameters_are_eligible( - rows, default_statement_config.parameter_config.type_coercion_map - ) - is False - ) - - def test_arrow_ingest_rolls_back_when_a_later_batch_conversion_fails(monkeypatch: pytest.MonkeyPatch) -> None: import pyarrow as pa diff --git a/tests/unit/adapters/test_sqlite/test_rowcount_optimization.py b/tests/unit/adapters/test_sqlite/test_rowcount_optimization.py index e1c894ba0..35825bf60 100644 --- a/tests/unit/adapters/test_sqlite/test_rowcount_optimization.py +++ b/tests/unit/adapters/test_sqlite/test_rowcount_optimization.py @@ -3,23 +3,16 @@ from sqlspec.adapters.sqlite.core import resolve_rowcount from sqlspec.adapters.sqlite.driver import SqliteDriver -from sqlspec.core.result import DMLResult def test_resolve_rowcount_fast_path() -> None: - # Cursor with rowcount cursor = Mock() cursor.rowcount = 10 - - # Should get 10 assert resolve_rowcount(cursor) == 10 def test_resolve_rowcount_missing_attr() -> None: - # Cursor without rowcount - cursor = Mock(spec=[]) # No attributes - - # Should not crash, return 0 + cursor = Mock(spec=[]) assert resolve_rowcount(cursor) == 0 @@ -35,53 +28,50 @@ def test_resolve_rowcount_negative() -> None: assert resolve_rowcount(cursor) == 0 -def test_sqlite_execute_many_thin_path_skips_dispatch() -> None: +def test_sqlite_execute_many_inserts_records() -> None: connection = sqlite3.connect(":memory:") try: driver = SqliteDriver(connection=connection) - driver.execute("CREATE TABLE test_thin_path (value TEXT)") + driver.execute("CREATE TABLE test_many (value TEXT)") - result = driver.execute_many("INSERT INTO test_thin_path (value) VALUES (?)", [("a",), ("b",)]) + result = driver.execute_many("INSERT INTO test_many (value) VALUES (?)", [("a",), ("b",)]) - assert isinstance(result, DMLResult) assert result.operation_type == "INSERT" assert result.rows_affected >= 0 - assert connection.execute("SELECT COUNT(*) FROM test_thin_path").fetchone()[0] == 2 + assert connection.execute("SELECT COUNT(*) FROM test_many").fetchone()[0] == 2 finally: connection.close() -def test_sqlite_execute_many_falls_back_when_coercion_required() -> None: +def test_sqlite_execute_many_with_coercion() -> None: connection = sqlite3.connect(":memory:") try: driver = SqliteDriver(connection=connection) - driver.execute("CREATE TABLE test_fallback_path (flag INTEGER)") + driver.execute("CREATE TABLE test_coercion (flag INTEGER)") - result = driver.execute_many("INSERT INTO test_fallback_path (flag) VALUES (?)", [(True,), (False,)]) + result = driver.execute_many("INSERT INTO test_coercion (flag) VALUES (?)", [(True,), (False,)]) - assert not isinstance(result, DMLResult) assert result.operation_type == "INSERT" - assert connection.execute("SELECT COUNT(*) FROM test_fallback_path").fetchone()[0] == 2 + assert connection.execute("SELECT COUNT(*) FROM test_coercion").fetchone()[0] == 2 finally: connection.close() -def test_sqlite_execute_many_falls_back_when_subclass_coercion_required() -> None: +def test_sqlite_execute_many_with_subclass_coercion() -> None: import json from collections import defaultdict connection = sqlite3.connect(":memory:") try: driver = SqliteDriver(connection=connection) - driver.execute("CREATE TABLE test_fallback_path (data TEXT)") + driver.execute("CREATE TABLE test_subclass_coercion (data TEXT)") result = driver.execute_many( - "INSERT INTO test_fallback_path (data) VALUES (?)", [(defaultdict(int, a=1),), (defaultdict(int, b=2),)] + "INSERT INTO test_subclass_coercion (data) VALUES (?)", [(defaultdict(int, a=1),), (defaultdict(int, b=2),)] ) - assert not isinstance(result, DMLResult) assert result.operation_type == "INSERT" - rows = connection.execute("SELECT data FROM test_fallback_path ORDER BY data").fetchall() + rows = connection.execute("SELECT data FROM test_subclass_coercion ORDER BY data").fetchall() assert [json.loads(row[0]) for row in rows] == [{"a": 1}, {"b": 2}] finally: connection.close() From c93ce8d4a1f9a009f4803c1aef749c9d71dab7a0 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 03:44:07 +0000 Subject: [PATCH 14/14] docs(data-dictionary): polish method table, eliminate overflow and awkward wraps --- docs/learn/data-dictionary.md | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/docs/learn/data-dictionary.md b/docs/learn/data-dictionary.md index 3b53b8fef..824a54147 100644 --- a/docs/learn/data-dictionary.md +++ b/docs/learn/data-dictionary.md @@ -59,18 +59,18 @@ with spec.provide_session(config) as session: ## Core Introspection and DDL Methods -| Method | Returned Metadata | Description | -| --- | --- | --- | -| `get_tables(driver, schema=None)` | `list[TableMetadata]` | Lists user tables in topological/catalog order (including self-referential root tables). | -| `get_columns(driver, table=None, schema=None)` | `list[ColumnMetadata]` | Returns column names, types, nullability, defaults, and identity/generated expressions. | -| `get_indexes(driver, table=None, schema=None)` | `list[IndexMetadata]` | Returns index names, indexed column lists, and uniqueness flags. | -| `get_foreign_keys(driver, table=None, schema=None)` | `list[ForeignKeyMetadata]` | Returns referencing/referenced schemas, tables, columns, and `on_delete` / `on_update` actions. | -| `get_ddl(driver, object_name, ...)` | `DDLResult` | Extracts native engine DDL when available or synthesizes dialect-aware AST DDL (`target_dialect` optional). | -| `get_schema_ddl(driver, schema=None, ...)` | `MetadataResult` | Emits topologically ordered `CREATE` or reverse-topological `DROP` `DDLResult` items across a schema. | -| `get_table_profile(driver, table, ...)` | `list[ColumnProfileMetadata]` | Samples live rows to compute null counts, cardinality, max string lengths, and narrower SQL types. | -| `infer_table_ddl_from_data(driver, table, ...)` | `DDLResult` | Generates a typed `CREATE TABLE` statement from profiled column values (`target_table` / `target_dialect`). | -| `get_version(driver)` | `VersionInfo \| None` | Parses and caches the connected database server's `major`, `minor`, and `patch` version. | -| `get_optimal_type(driver, type_category)` | `str` | Resolves the best DDL type for categories such as `"json"`, `"uuid"`, `"binary"`, `"text"`, or `"timestamp"`. | +| Method | Returns | Description | +| :--- | :--- | :--- | +| `get_tables()` | `list[TableMetadata]` | Lists user tables in topological or catalog order. | +| `get_columns()` | `list[ColumnMetadata]` | Returns column names, types, nullability, defaults, and identity. | +| `get_indexes()` | `list[IndexMetadata]` | Returns index names, indexed column lists, and uniqueness flags. | +| `get_foreign_keys()` | `list[ForeignKeyMetadata]` | Returns referencing/referenced columns and on-delete/update actions. | +| `get_ddl()` | `DDLResult` | Extracts native engine DDL or synthesizes dialect-aware AST DDL. | +| `get_schema_ddl()` | `MetadataResult` | Emits ordered `CREATE` or reverse-ordered `DROP` statements across a schema. | +| `get_table_profile()` | `list[ColumnProfileMetadata]` | Samples live rows to compute nulls, cardinality, and narrower types. | +| `infer_table_ddl_from_data()` | `DDLResult` | Generates a typed `CREATE TABLE` statement from profiled column values. | +| `get_version()` | `VersionInfo | None` | Returns database server major, minor, and patch version. | +| `get_optimal_type()` | `str` | Resolves optimal DDL types for abstract categories like `json` or `uuid`. | ### Metadata Record Fields @@ -93,6 +93,9 @@ with spec.provide_session(config) as session: When migrating loosely typed staging tables (`TEXT` / `VARCHAR` landing zones) or raw JSON/Arrow datasets, `get_table_profile` and `infer_table_ddl_from_data` inspect live sample values to recommend tighter types (`BIGINT`, `NUMERIC(p, s)`, `BOOLEAN`, `DATE`, `TIMESTAMP`, `UUID`) and `NOT NULL` constraints. +!!! note "Cross-Dialect Transpilation" + Pass `target_dialect` to `get_ddl` or `get_schema_ddl` to transpile catalog DDL from one database engine to another using SQLGlot AST compilation. + ## Dependency Ordering and Portable DDL Helpers When exporting schemas, seeding fixtures, or truncating tables in foreign-key order, combine `get_foreign_keys` with `dependency_edges_from_foreign_keys`, `sort_dependencies`, and `sort_ddl_results` from `sqlspec.data_dictionary`: @@ -111,7 +114,10 @@ sorted_plan = sort_dependencies(edges) ## Identifier Quoting vs Catalog Lookups -Catalog views store unquoted identifier names (`schema_versions`, not `"schema_versions"`). Pass bare table and schema names (`table="members"`, `schema="public"`) to `data_dictionary` methods; synthesized DDL quotes reserved words and mixed-case catalog names automatically for the target dialect. +Catalog views store unquoted identifier names (`schema_versions`, not `"schema_versions"`). Pass bare table and schema names (`table="members"`, `schema="public"`) to `data_dictionary` methods. Synthesized DDL quotes reserved words and mixed-case catalog names automatically for the target dialect. + +!!! tip "Bare Catalog Identifiers" + Always pass bare unquoted identifier names (`table="members"`, `schema="public"`) to `data_dictionary` methods. SQLSpec normalizes catalog queries and quotes reserved words or mixed-case identifiers automatically in generated DDL. ## Recap