From 264f9cb846e6a9d6aa5db84f0daf7c328398547b Mon Sep 17 00:00:00 2001 From: Luccas Gomes Date: Mon, 13 Jul 2026 11:27:32 -0300 Subject: [PATCH] Add gzip format for datastore dumps --- datastore/api/endpoints/dump.py | 24 +++++-- .../engines/bigquery/backend.py | 11 ++-- datastore/services/dump.py | 40 +++++++++++- tests/test_datastore_dump.py | 64 ++++++++++++++++++- 4 files changed, 127 insertions(+), 12 deletions(-) diff --git a/datastore/api/endpoints/dump.py b/datastore/api/endpoints/dump.py index b319ede..a7b56c4 100644 --- a/datastore/api/endpoints/dump.py +++ b/datastore/api/endpoints/dump.py @@ -23,22 +23,34 @@ from datastore.api.context import Context from datastore.api.responses import ERROR_RESPONSES from datastore.infrastructure.engines import get_datastore_engine -from datastore.services.dump import stream_csv_shards, stream_ndjson_shards +from datastore.services.dump import ( + stream_csv_shards, + stream_gzip_csv_shards, + stream_ndjson_shards, +) -DumpFormat = Literal["csv", "ndjson", "parquet"] +DumpFormat = Literal["csv", "gzip", "ndjson", "parquet"] _MEDIA_TYPE: dict[str, str] = { "csv": "text/csv", + "gzip": "application/gzip", "ndjson": "application/x-ndjson", "parquet": "application/vnd.apache.parquet", } +_DOWNLOAD_EXT: dict[str, str] = { + "csv": "csv", + "gzip": "csv.gz", + "ndjson": "ndjson", + "parquet": "parquet", +} + router = APIRouter(tags=["Datastore Download"], responses=ERROR_RESPONSES) @router.get( "/datastore/dump/{resource_id}", - summary="Download an entire table (CSV / NDJSON / Parquet)", + summary="Download an entire table (CSV / gzip CSV / NDJSON / Parquet)", responses={ 302: {"description": "Single-shard export — redirect to a signed GCS URL."}, 200: {"description": "Multi-shard export — streamed CSV / NDJSON body."}, @@ -49,7 +61,7 @@ async def dump( resource_id: str, fmt: Annotated[DumpFormat, Query(alias="format")] = "csv", ): - """Download an entire resource as `csv` (default), `ndjson`, or `parquet`. + """Download an entire resource as `csv` (default), `gzip`, `ndjson`, or `parquet`. Small exports redirect (302) straight to a signed GCS URL; large ones stream a concatenated body. Select the format with `?format=`. @@ -64,6 +76,8 @@ async def dump( if fmt == "csv": body = stream_csv_shards(urls) + elif fmt == "gzip": + body = stream_gzip_csv_shards(urls) elif fmt == "ndjson": body = stream_ndjson_shards(urls) else: # pragma: no cover — Parquet never returns >1 shard @@ -74,7 +88,7 @@ async def dump( media_type=_MEDIA_TYPE[fmt], headers={ "Content-Disposition": ( - f'attachment; filename="{resource_id}.{fmt}"' + f'attachment; filename="{resource_id}.{_DOWNLOAD_EXT[fmt]}"' ), }, ) diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index d93086e..58b8a8c 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -935,7 +935,7 @@ 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/NDJSON: wildcard URI → BigQuery shards above 1 GB. + - CSV/gzip/NDJSON: wildcard URI → BigQuery shards above 1 GB. - Parquet: single-file URI; >1 GB → 413, switch format. - Cache key = `table.modified`; unchanged tables skip the extract. - Older revisions are GC'd on cache miss. @@ -1009,8 +1009,11 @@ async def _list(b: Any, p: str) -> list[Any]: # `header=true` is the documented default for CSV but some # client versions / project configs treat it as false; be # explicit so the column names always land in shard 0. - # NDJSON / Parquet ignore the option. - extra_opts = ", header=true" if fmt == "csv" else "" + # NDJSON / Parquet ignore the option. `format=gzip` is CSV + # with BigQuery-side GZIP compression. + extra_opts = ", header=true" if fmt in {"csv", "gzip"} else "" + if fmt == "gzip": + extra_opts += ", compression='GZIP'" sql = ( f"EXPORT DATA OPTIONS(" f"uri='{uri}', format='{_FMT[fmt]['bq']}', overwrite=true" @@ -1309,6 +1312,7 @@ def _translate_bigquery_error( # extension on the GCS object so clients see the file type they expect. _FMT: dict[str, dict[str, str]] = { "csv": {"ext": "csv", "bq": "CSV"}, + "gzip": {"ext": "csv.gz", "bq": "CSV"}, "ndjson": {"ext": "json", "bq": "JSON"}, "parquet": {"ext": "parquet", "bq": "PARQUET"}, } @@ -1338,4 +1342,3 @@ def _is_export_too_large(exc: BaseException) -> bool: """ msg = str(exc).lower() return "single uri" in msg or "wildcard" in msg - diff --git a/datastore/services/dump.py b/datastore/services/dump.py index c9fb703..9607ba6 100644 --- a/datastore/services/dump.py +++ b/datastore/services/dump.py @@ -21,6 +21,7 @@ from __future__ import annotations +import zlib from collections.abc import AsyncIterator import httpx @@ -63,6 +64,43 @@ async def stream_ndjson_shards(urls: list[str]) -> AsyncIterator[bytes]: yield chunk +async def stream_gzip_csv_shards(urls: list[str]) -> AsyncIterator[bytes]: + """Stream-concat gzip-compressed CSV shards into one gzip file. + + BigQuery emits one CSV header per exported shard. For gzip output we + decompress each shard, strip duplicate headers after shard 0, and + recompress the combined CSV as a single gzip stream. + """ + compressor = zlib.compressobj(wbits=16 + zlib.MAX_WBITS) + async with httpx.AsyncClient(timeout=_TIMEOUT) as client: + for i, url in enumerate(urls): + async with client.stream("GET", url) as resp: + resp.raise_for_status() + csv_chunks = _decompress_gzip(resp.aiter_bytes(_CHUNK_BYTES)) + if i > 0: + csv_chunks = _skip_first_line(csv_chunks) + async for chunk in csv_chunks: + compressed = compressor.compress(chunk) + if compressed: + yield compressed + tail = compressor.flush() + if tail: + yield tail + + +async def _decompress_gzip( + chunks: AsyncIterator[bytes], +) -> AsyncIterator[bytes]: + decompressor = zlib.decompressobj(wbits=16 + zlib.MAX_WBITS) + async for chunk in chunks: + data = decompressor.decompress(chunk) + if data: + yield data + tail = decompressor.flush() + if tail: + yield tail + + async def _skip_first_line( chunks: AsyncIterator[bytes], ) -> AsyncIterator[bytes]: @@ -80,5 +118,3 @@ async def _skip_first_line( break async for chunk in chunks: yield chunk - - diff --git a/tests/test_datastore_dump.py b/tests/test_datastore_dump.py index 2180cd4..fbc0e0a 100644 --- a/tests/test_datastore_dump.py +++ b/tests/test_datastore_dump.py @@ -11,6 +11,7 @@ from __future__ import annotations +import gzip from collections.abc import AsyncIterator from typing import Any from unittest.mock import MagicMock, patch @@ -51,7 +52,7 @@ def test_single_shard_returns_302(client: TestClient) -> None: assert response.content == b"" -@pytest.mark.parametrize("fmt", ["csv", "ndjson", "parquet"]) +@pytest.mark.parametrize("fmt", ["csv", "gzip", "ndjson", "parquet"]) def test_each_format_supports_single_shard_redirect( fmt: str, client: TestClient, ) -> None: @@ -111,6 +112,35 @@ def test_multi_shard_ndjson_pure_byte_concat(client: TestClient) -> None: ) +def test_multi_shard_gzip_stream_concat_dedups_header( + client: TestClient, +) -> None: + """Gzip CSV shards are decompressed, header-deduped, then emitted + as one gzip-compressed CSV file.""" + shards = { + "url-1": gzip.compress(b"col1,col2\na,1\nb,2\n"), + "url-2": gzip.compress(b"col1,col2\nc,3\nd,4\n"), + "url-3": gzip.compress(b"col1,col2\ne,5\n"), + } + + with _patch_dump(list(shards.keys())), _patch_httpx_stream(shards): + response = client.get( + DUMP_URL, params={"format": "gzip"}, follow_redirects=False, + ) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("application/gzip") + assert response.headers["content-disposition"] == ( + 'attachment; filename="balancing_auction_results_2025.csv.gz"' + ) + assert gzip.decompress(response.content).decode().splitlines() == [ + "col1,col2", + "a,1", "b,2", + "c,3", "d,4", + "e,5", + ] + + # --- error paths ---------------------------------------------------------- @@ -188,6 +218,13 @@ def test_build_export_select_parquet_returns_star() -> None: assert _build_export_select(schema, fmt="parquet") == "*" +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( + schema, fmt="csv", + ) + + def _bq_field(name: str, field_type: str) -> Any: f = MagicMock() f.name = name @@ -434,6 +471,31 @@ def test_dump_cache_miss_submits_extract_then_returns_urls() -> None: assert backend.client.query.call_count == 1 +def test_dump_gzip_export_uses_csv_with_gzip_compression() -> None: + """`format=gzip` is a CSV export with BigQuery-side gzip compression + and `.csv.gz` object names.""" + import asyncio + + new_blob = _fake_blob("dumps/res-1/gzip/_000.csv.gz", "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", "gzip")) + + assert urls == ["https://fresh"] + sql = backend.client.query.call_args.args[0] + assert "format='CSV'" in sql + assert "header=true" in sql + assert "compression='GZIP'" in sql + assert "_*.csv.gz" in sql + + 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