From d988a249382c23830ec7c50b2438d6017b793574 Mon Sep 17 00:00:00 2001 From: Luccas Gomes Date: Tue, 14 Jul 2026 15:50:13 -0300 Subject: [PATCH] Fix BigQuery parquet dump exports --- datastore/api/endpoints/dump.py | 6 +- .../engines/bigquery/backend.py | 39 +++++++++---- tests/test_datastore_dump.py | 56 +++++++++++++++++++ 3 files changed, 86 insertions(+), 15 deletions(-) diff --git a/datastore/api/endpoints/dump.py b/datastore/api/endpoints/dump.py index a7b56c4..04e1a47 100644 --- a/datastore/api/endpoints/dump.py +++ b/datastore/api/endpoints/dump.py @@ -2,14 +2,14 @@ Behaviour by shard count (decided by BigQuery from the export size): - - **1 shard** (≤ 1 GB, or any-size Parquet): 302 redirect to the + - **1 shard** (≤ 1 GB, including Parquet): 302 redirect to the GCS signed URL. Zero server bandwidth — bytes go GCS → client. - **N shards** (>1 GB CSV/NDJSON): `StreamingResponse` over `services.dump.stream_*_shards`, which pulls each shard from GCS via async httpx and byte-forwards (CSV header-dedup; NDJSON pure concat). Memory ≈ one chunk in flight; no threadpool consumption. -Parquet >1 GB is refused upstream with 413 (parquet shards can't be +Multi-shard Parquet is refused with 413 (parquet shards can't be byte-concatenated). Caller picks CSV/NDJSON. """ @@ -80,7 +80,7 @@ async def dump( body = stream_gzip_csv_shards(urls) elif fmt == "ndjson": body = stream_ndjson_shards(urls) - else: # pragma: no cover — Parquet never returns >1 shard + else: # pragma: no cover — the engine rejects multi-shard Parquet raise RuntimeError(f"unexpected multi-shard format: {fmt}") return StreamingResponse( diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index ffeb439..f817033 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -940,8 +940,9 @@ def _count_rows(self, resource_id: str) -> int: async def dump(self, resource_id: str, fmt: str) -> list[str]: """Submit `EXPORT DATA`; poll non-blockingly; return signed URLs. - - CSV/gzip/NDJSON: wildcard URI → BigQuery shards above 1 GB. - - Parquet: single-file URI; >1 GB → 413, switch format. + - All formats use a wildcard URI because BigQuery SQL + `EXPORT DATA` requires one; Parquet must still produce one + shard because parquet shards cannot be byte-concatenated. - Cache key = `table.modified`; unchanged tables skip the extract. - Older revisions are GC'd on cache miss. - All BQ + GCS calls are offloaded via `asyncio.to_thread`; the @@ -996,11 +997,7 @@ async def dump(self, resource_id: str, fmt: str) -> list[str]: ) ext = _FMT[fmt]["ext"] prefix = f"dumps/{resource_id}/{fmt}/{rev}" - uri = ( - f"gs://{bucket}/{prefix}.{ext}" - if fmt == "parquet" - else f"gs://{bucket}/{prefix}_*.{ext}" - ) + uri = f"gs://{bucket}/{prefix}_*.{ext}" async def _list(b: Any, p: str) -> list[Any]: return sorted( @@ -1101,6 +1098,14 @@ def _gc() -> int: # (signing needs IAM signBlob under workload identity). blobs = await _list(rw_gcs, prefix) + if fmt == "parquet" and len(blobs) > 1: + raise PayloadTooLargeError( + f"resource {resource_id!r} exported as multiple parquet " + "shards; single-file download isn't possible. Try " + "`format=csv` or `format=ndjson` for sharded multi-file " + "downloads instead." + ) + expiry = timedelta( hours=getattr(self.config, "BIGQUERY_EXPORT_URL_EXPIRY_HOURS", 1), ) @@ -1342,13 +1347,23 @@ def _translate_bigquery_error( def _build_export_select(schema: Any, fmt: str) -> str: """SELECT column list for EXPORT DATA. - Parquet preserves native logical types → `*` is enough. For CSV / - NDJSON, every column goes through `format_select_column` (in - `bigquery/lib.py`) — the same helper `datastore_search` uses — so a - given column renders identically in a dump and in a search response. + Parquet preserves native logical types, but BigQuery cannot export + native JSON columns to Parquet. Keep `*` for the common no-JSON path; + otherwise cast only JSON columns to strings. For CSV / NDJSON, every + column goes through `format_select_column` (in `bigquery/lib.py`) — + the same helper `datastore_search` uses — so a given column renders + identically in a dump and in a search response. """ if fmt == "parquet": - return "*" + fields = list(schema) + if not any((f.field_type or "").upper() == "JSON" for f in fields): + return "*" + return ", ".join( + f"TO_JSON_STRING(`{f.name}`) AS `{f.name}`" + if (f.field_type or "").upper() == "JSON" + else f"`{f.name}`" + for f in fields + ) return ", ".join( format_select_column(f.name, f.field_type) for f in schema ) diff --git a/tests/test_datastore_dump.py b/tests/test_datastore_dump.py index fbc0e0a..47756ca 100644 --- a/tests/test_datastore_dump.py +++ b/tests/test_datastore_dump.py @@ -218,6 +218,18 @@ def test_build_export_select_parquet_returns_star() -> None: assert _build_export_select(schema, fmt="parquet") == "*" +def test_build_export_select_parquet_casts_json_columns() -> None: + schema = [ + _bq_field("id", "INT64"), + _bq_field("bidder_metadata", "JSON"), + _bq_field("delivery_start", "TIMESTAMP"), + ] + assert _build_export_select(schema, fmt="parquet") == ( + "`id`, TO_JSON_STRING(`bidder_metadata`) AS `bidder_metadata`, " + "`delivery_start`" + ) + + def test_build_export_select_gzip_matches_csv_formatting() -> None: schema = [_bq_field("delivery_start", "TIMESTAMP")] assert _build_export_select(schema, fmt="gzip") == _build_export_select( @@ -496,6 +508,50 @@ def test_dump_gzip_export_uses_csv_with_gzip_compression() -> None: assert "_*.csv.gz" in sql +def test_dump_parquet_export_uses_wildcard_uri() -> None: + """BigQuery SQL EXPORT DATA requires a wildcard URI even for + formats where this endpoint only accepts a single shard. + """ + import asyncio + + new_blob = _fake_blob("dumps/res-1/parquet/_000.parquet", "https://fresh") + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], [new_blob], [new_blob]] + + job = MagicMock() + job.state = "DONE" + job.error_result = None + backend.client.query.return_value = job + + urls = asyncio.run(backend.dump("res-1", "parquet")) + + assert urls == ["https://fresh"] + sql = backend.client.query.call_args.args[0] + assert "format='PARQUET'" in sql + assert "_*.parquet" in sql + + +def test_dump_multi_shard_parquet_returns_413() -> None: + import asyncio + + blobs = [ + _fake_blob("dumps/res-1/parquet/_000.parquet"), + _fake_blob("dumps/res-1/parquet/_001.parquet"), + ] + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], blobs, blobs] + + job = MagicMock() + job.state = "DONE" + job.error_result = None + backend.client.query.return_value = job + + with pytest.raises(PayloadTooLargeError, match="multiple parquet shards"): + asyncio.run(backend.dump("res-1", "parquet")) + + def test_dump_cache_miss_deletes_older_revisions() -> None: """After a successful extract on cache miss, blobs from any older revision under `dumps///` should be deleted to keep