Add gzip format for datastore dumps

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]}"'
             ),
         },
     )

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
-

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
-
-

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/<rev>_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/<rid>/<fmt>/` should be deleted to keep