From fcb8fc9089c448ae87edee3d0e3eb6e3a807483a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 1 Aug 2026 20:15:35 +0000 Subject: [PATCH 1/2] fix: validate storage payload formats --- docs/reference/storage.rst | 15 +++ sqlspec/storage/_arrow_payload.py | 11 +- sqlspec/storage/_utils.py | 21 +++- sqlspec/storage/pipeline.py | 22 +++- tests/unit/storage/test_payload_codecs.py | 129 ++++++++++++++++++++++ 5 files changed, 190 insertions(+), 8 deletions(-) create mode 100644 tests/unit/storage/test_payload_codecs.py diff --git a/docs/reference/storage.rst b/docs/reference/storage.rst index ae4f0bd2e..6326826ff 100644 --- a/docs/reference/storage.rst +++ b/docs/reference/storage.rst @@ -6,6 +6,21 @@ Storage abstraction layer with multiple backend support (local filesystem, fsspec, obstore), configuration-based registration, and Arrow table import/export with CSV format support. +Write and read formats +====================== + +Row-oriented writes accept only ``json`` and newline-delimited ``jsonl``. +Arrow-table writes accept only ``parquet``, ``arrow-ipc``, and ``csv``. The +pipeline rejects mismatched formats before encoding or storage I/O, so it +cannot write one payload type under another format label. Read APIs retain the +full format set because they decode all five formats into Arrow tables. + +JSONL reads use PyArrow's native JSON reader. Its type inference applies to the +result, including conversion of date-like strings to Arrow timestamps. This +avoids Python per-line decoding and ``Table.from_pylist()`` copies. It does not +make ``load_from_storage()`` bounded-memory: that API reads the complete object +payload before decoding it. + Pipelines ========= diff --git a/sqlspec/storage/_arrow_payload.py b/sqlspec/storage/_arrow_payload.py index 605ffaf8d..12ade7b56 100644 --- a/sqlspec/storage/_arrow_payload.py +++ b/sqlspec/storage/_arrow_payload.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast -from sqlspec.storage._utils import import_pyarrow, import_pyarrow_csv, import_pyarrow_parquet +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_csv, import_pyarrow_json, import_pyarrow_parquet from sqlspec.utils.serializers import from_json if TYPE_CHECKING: @@ -56,13 +56,14 @@ def decode_arrow_payload(payload: bytes, format_choice: StorageFormat) -> "Arrow if format_choice == "csv": pa_csv = import_pyarrow_csv() return cast("ArrowTable", pa_csv.read_csv(pa.BufferReader(payload))) - text_payload = payload.decode() if format_choice == "json": - data = from_json(text_payload) + data = from_json(payload.decode()) rows = data if isinstance(data, list) else [data] return cast("ArrowTable", pa.Table.from_pylist(rows)) if format_choice == "jsonl": - rows = [from_json(line) for line in text_payload.splitlines() if line.strip()] - return cast("ArrowTable", pa.Table.from_pylist(rows)) + if payload == b"": + return cast("ArrowTable", pa.table({})) + pa_json = import_pyarrow_json() + return cast("ArrowTable", pa_json.read_json(pa.BufferReader(payload))) msg = f"Unsupported storage format for Arrow decoding: {format_choice}" raise ValueError(msg) diff --git a/sqlspec/storage/_utils.py b/sqlspec/storage/_utils.py index fbdc3f399..8daf73bc3 100644 --- a/sqlspec/storage/_utils.py +++ b/sqlspec/storage/_utils.py @@ -6,7 +6,13 @@ from sqlspec.utils.logging import get_logger, log_with_context from sqlspec.utils.module_loader import ensure_pyarrow -__all__ = ("_log_storage_event", "import_pyarrow", "import_pyarrow_csv", "import_pyarrow_parquet") +__all__ = ( + "_log_storage_event", + "import_pyarrow", + "import_pyarrow_csv", + "import_pyarrow_json", + "import_pyarrow_parquet", +) logger = get_logger(__name__) @@ -50,6 +56,19 @@ def import_pyarrow_csv() -> "Any": return pa_csv +def import_pyarrow_json() -> "Any": + """Import PyArrow JSON module with optional dependency guard. + + Returns: + PyArrow JSON module. + """ + + ensure_pyarrow() + import pyarrow.json as pa_json + + return pa_json + + def _log_storage_event( event: str, *, diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index cae77781b..e201f13df 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -144,6 +144,8 @@ def reset(self) -> None: _METRICS = _StorageBridgeMetrics() _RECENT_STORAGE_EVENTS: "deque[StorageTelemetry]" = deque(maxlen=25) _EMPTY_STORAGE_OPTIONS: dict[str, Any] = {} +_ARROW_WRITE_FORMATS = frozenset({"parquet", "arrow-ipc", "csv"}) +_ROW_WRITE_FORMATS = frozenset({"json", "jsonl"}) def _storage_options(default_options: "dict[str, Any]", storage_options: "dict[str, Any] | None") -> "dict[str, Any]": @@ -225,6 +227,18 @@ def _encode_row_payload(rows: "list[Any]", format_hint: StorageFormat) -> bytes: return bytes(buffer) +def _validate_arrow_write_format(format_choice: StorageFormat) -> None: + if format_choice not in _ARROW_WRITE_FORMATS: + msg = "Arrow storage writes support only Parquet, Arrow IPC, and CSV formats" + raise ValueError(msg) + + +def _validate_row_write_format(format_choice: StorageFormat) -> None: + if format_choice not in _ROW_WRITE_FORMATS: + msg = "Row storage writes support only JSON and JSONL formats" + raise ValueError(msg) + + def _encode_arrow_payload( table: "ArrowTable", format_choice: StorageFormat, @@ -353,8 +367,9 @@ def write_rows( ) -> StorageTelemetry: """Write dictionary rows to storage using cached serializers.""" - serialized = serialize_collection(rows) format_choice = format_hint or "jsonl" + _validate_row_write_format(format_choice) + serialized = serialize_collection(rows) payload = _encode_row_payload(serialized, format_choice) resolved_options = _storage_options(self._storage_options, storage_options) return self._write_bytes( @@ -373,6 +388,7 @@ def write_arrow( """Write an Arrow table to storage using zero-copy buffers.""" format_choice = format_hint or "parquet" + _validate_arrow_write_format(format_choice) resolved_options = _storage_options(self._storage_options, storage_options) format_write_options = _csv_write_options( format_choice, resolved_options, self._storage_options, self._csv_write_options @@ -485,8 +501,9 @@ async def write_rows( format_hint: StorageFormat | None = None, storage_options: "dict[str, Any] | None" = None, ) -> StorageTelemetry: - serialized = serialize_collection(rows) format_choice = format_hint or "jsonl" + _validate_row_write_format(format_choice) + serialized = serialize_collection(rows) payload = await async_(_encode_row_payload)(serialized, format_choice) resolved_options = _storage_options(self._storage_options, storage_options) return await self._write_bytes_async( @@ -503,6 +520,7 @@ async def write_arrow( compression: str | None = None, ) -> StorageTelemetry: format_choice = format_hint or "parquet" + _validate_arrow_write_format(format_choice) resolved_options = _storage_options(self._storage_options, storage_options) format_write_options = _csv_write_options( format_choice, resolved_options, self._storage_options, self._csv_write_options diff --git a/tests/unit/storage/test_payload_codecs.py b/tests/unit/storage/test_payload_codecs.py new file mode 100644 index 000000000..884a86b54 --- /dev/null +++ b/tests/unit/storage/test_payload_codecs.py @@ -0,0 +1,129 @@ +# pyright: reportPrivateUsage=false +"""Focused storage payload codec and write-format contract tests.""" + +from typing import Any + +import pyarrow as pa +import pytest + +from sqlspec.storage._arrow_payload import decode_arrow_payload, encode_arrow_payload +from sqlspec.storage.pipeline import ( + AsyncStoragePipeline, + SyncStoragePipeline, + get_recent_storage_events, + reset_storage_bridge_events, +) + + +class _TrackingBackend: + backend_type = "tracking" + + def __init__(self) -> None: + self.writes: list[tuple[str, bytes]] = [] + + def write_bytes_sync(self, path: str, payload: bytes) -> None: + self.writes.append((path, payload)) + + async def write_bytes_async(self, path: str, payload: bytes) -> None: + self.writes.append((path, payload)) + + +def _install_backend(monkeypatch: pytest.MonkeyPatch, pipeline_type: type[Any]) -> _TrackingBackend: + backend = _TrackingBackend() + monkeypatch.setattr(pipeline_type, "_backend", lambda *_args, **_kwargs: (backend, "payload", backend.backend_type)) + return backend + + +def test_decode_jsonl_uses_native_timestamp_inference() -> None: + payload = b'{"id":1,"created_at":"2026-08-01T12:34:56"}\n' + + table = decode_arrow_payload(payload, "jsonl") + + assert table.schema.field("created_at").type == pa.timestamp("s") + assert table.to_pylist() == [{"id": 1, "created_at": table.column("created_at")[0].as_py()}] + + +def test_decode_empty_jsonl_preserves_empty_table() -> None: + table = decode_arrow_payload(b"", "jsonl") + + assert table.equals(pa.table({})) + + +def test_decode_whitespace_only_jsonl_is_delegated_to_pyarrow() -> None: + table = decode_arrow_payload(b" \n\t", "jsonl") + + assert table.equals(pa.table({})) + + +def test_json_array_decode_keeps_existing_shape() -> None: + table = decode_arrow_payload(b'[{"id":1},{"id":2}]', "json") + + assert table.to_pylist() == [{"id": 1}, {"id": 2}] + + +@pytest.mark.parametrize("format_choice", ["parquet", "arrow-ipc", "csv"]) +def test_arrow_payload_round_trip_has_expected_signature(format_choice: str) -> None: + table = pa.table({"id": [1, 2], "name": ["alpha", "beta"]}) + + payload = encode_arrow_payload(table, format_choice, compression=None) # type: ignore[arg-type] + restored = decode_arrow_payload(payload, format_choice) # type: ignore[arg-type] + + signatures = {"parquet": b"PAR1", "arrow-ipc": b"ARROW1", "csv": b'"id","name"'} + assert payload.startswith(signatures[format_choice]) + assert restored.to_pylist() == table.to_pylist() + + +@pytest.mark.parametrize("format_hint", ["parquet", "arrow-ipc", "csv"]) +def test_sync_row_write_rejects_arrow_formats_before_encoding_or_io( + monkeypatch: pytest.MonkeyPatch, format_hint: str +) -> None: + pipeline = SyncStoragePipeline() + backend = _install_backend(monkeypatch, SyncStoragePipeline) + reset_storage_bridge_events() + + with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"): + pipeline.write_rows([{"id": 1}], "payload", format_hint=format_hint) # type: ignore[arg-type] + + assert backend.writes == [] + assert get_recent_storage_events() == [] + + +@pytest.mark.parametrize("format_hint", ["json", "jsonl"]) +def test_sync_arrow_write_rejects_row_formats_before_encoding_or_io( + monkeypatch: pytest.MonkeyPatch, format_hint: str +) -> None: + pipeline = SyncStoragePipeline() + backend = _install_backend(monkeypatch, SyncStoragePipeline) + reset_storage_bridge_events() + + with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"): + pipeline.write_arrow(pa.table({"id": [1]}), "payload", format_hint=format_hint) # type: ignore[arg-type] + + assert backend.writes == [] + assert get_recent_storage_events() == [] + + +@pytest.mark.parametrize("format_hint", ["parquet", "arrow-ipc", "csv"]) +async def test_async_row_write_rejects_arrow_formats_before_io( + monkeypatch: pytest.MonkeyPatch, format_hint: str +) -> None: + pipeline = AsyncStoragePipeline() + backend = _install_backend(monkeypatch, AsyncStoragePipeline) + + with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"): + await pipeline.write_rows([{"id": 1}], "payload", format_hint=format_hint) # type: ignore[arg-type] + + assert backend.writes == [] + + +@pytest.mark.parametrize("format_hint", ["json", "jsonl"]) +async def test_async_arrow_write_rejects_row_formats_before_io( + monkeypatch: pytest.MonkeyPatch, format_hint: str +) -> None: + pipeline = AsyncStoragePipeline() + backend = _install_backend(monkeypatch, AsyncStoragePipeline) + + with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"): + await pipeline.write_arrow(pa.table({"id": [1]}), "payload", format_hint=format_hint) # type: ignore[arg-type] + + assert backend.writes == [] From be17fdbb4f68637cb9c8974aebd523b8cb4d8237 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 1 Aug 2026 20:20:19 +0000 Subject: [PATCH 2/2] test: prove invalid storage writes have no side effects --- tests/unit/storage/test_payload_codecs.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/storage/test_payload_codecs.py b/tests/unit/storage/test_payload_codecs.py index 884a86b54..033738deb 100644 --- a/tests/unit/storage/test_payload_codecs.py +++ b/tests/unit/storage/test_payload_codecs.py @@ -6,6 +6,7 @@ import pyarrow as pa import pytest +import sqlspec.storage.pipeline as storage_pipeline from sqlspec.storage._arrow_payload import decode_arrow_payload, encode_arrow_payload from sqlspec.storage.pipeline import ( AsyncStoragePipeline, @@ -79,6 +80,8 @@ def test_sync_row_write_rejects_arrow_formats_before_encoding_or_io( ) -> None: pipeline = SyncStoragePipeline() backend = _install_backend(monkeypatch, SyncStoragePipeline) + monkeypatch.setattr(storage_pipeline, "serialize_collection", pytest.fail) + monkeypatch.setattr(storage_pipeline, "_encode_row_payload", pytest.fail) reset_storage_bridge_events() with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"): @@ -94,6 +97,7 @@ def test_sync_arrow_write_rejects_row_formats_before_encoding_or_io( ) -> None: pipeline = SyncStoragePipeline() backend = _install_backend(monkeypatch, SyncStoragePipeline) + monkeypatch.setattr(storage_pipeline, "_encode_arrow_payload", pytest.fail) reset_storage_bridge_events() with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"): @@ -109,11 +113,15 @@ async def test_async_row_write_rejects_arrow_formats_before_io( ) -> None: pipeline = AsyncStoragePipeline() backend = _install_backend(monkeypatch, AsyncStoragePipeline) + monkeypatch.setattr(storage_pipeline, "serialize_collection", pytest.fail) + monkeypatch.setattr(storage_pipeline, "_encode_row_payload", pytest.fail) + reset_storage_bridge_events() with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"): await pipeline.write_rows([{"id": 1}], "payload", format_hint=format_hint) # type: ignore[arg-type] assert backend.writes == [] + assert get_recent_storage_events() == [] @pytest.mark.parametrize("format_hint", ["json", "jsonl"]) @@ -122,8 +130,11 @@ async def test_async_arrow_write_rejects_row_formats_before_io( ) -> None: pipeline = AsyncStoragePipeline() backend = _install_backend(monkeypatch, AsyncStoragePipeline) + monkeypatch.setattr(storage_pipeline, "_encode_arrow_payload", pytest.fail) + reset_storage_bridge_events() with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"): await pipeline.write_arrow(pa.table({"id": [1]}), "payload", format_hint=format_hint) # type: ignore[arg-type] assert backend.writes == [] + assert get_recent_storage_events() == []