From 20ffd964271440bb95a6178bf7540df557f5e240 Mon Sep 17 00:00:00 2001 From: Luccas Gomes Date: Fri, 6 Nov 2026 08:05:17 -0300 Subject: [PATCH] feat(dump): optional fields param to download a column subset GET /dump/{resource_id}?fields=a,b (or a JSON array) exports only those columns, in the given order, for every row. Unknown names are a 400. Each column subset is cached and GC-swept under its own dumps//fields-// prefix, so it never serves or deletes the whole-table file. Parquet subsets always use an explicit column list instead of SELECT *. Downloads get a _selected_columns filename suffix. --- API.md | 6 + datastore/api/endpoints/dump.py | 29 ++++- datastore/infrastructure/engines/base.py | 11 +- .../engines/bigquery/backend.py | 11 +- .../infrastructure/engines/bigquery/export.py | 47 +++++-- tests/test_datastore_dump.py | 121 +++++++++++++++++- 6 files changed, 205 insertions(+), 20 deletions(-) diff --git a/API.md b/API.md index a71ccce..c2686f9 100644 --- a/API.md +++ b/API.md @@ -392,6 +392,12 @@ row stats — a column-level metadata catalog without a side store. Download an entire resource. Pick the format with `?format=csv` (default), `gzip`, `ndjson`, or `parquet`. +To leave columns out, list the ones to keep with `?fields=` — a comma-separated +string (`fields=auctionID,deliveryStart`) or a JSON array of strings. Columns +come out in the order given; every row is still downloaded. An unknown column +name is a `400`. The file name gets a `_selected_columns` suffix, and each +column subset is cached separately from the whole-table file. + This route returns a **file**, never the JSON envelope — the only JSON it can produce is an error. The content type comes from the signed GCS URL and follows `format`: diff --git a/datastore/api/endpoints/dump.py b/datastore/api/endpoints/dump.py index b5b306d..bc0e0b9 100644 --- a/datastore/api/endpoints/dump.py +++ b/datastore/api/endpoints/dump.py @@ -26,11 +26,11 @@ DocumentFormat, DumpFormat, ) -from datastore.core.exceptions import ServerError +from datastore.core.exceptions import ServerError, ValidationError from datastore.core.helper import content_disposition, dump_filename_base from datastore.infrastructure.engines import get_datastore_engine from datastore.schemas.request import DatastoreDumpSQLRequest -from datastore.schemas.validators import ResourceId +from datastore.schemas.validators import ResourceId, to_csv_list from datastore.services.metadata import ( json_bytes, table_schema_csv, @@ -164,19 +164,38 @@ async def dump( context: Context, resource_id: ResourceId, fmt: Annotated[DumpFormat, Query(alias="format")] = "csv", + fields: Annotated[ + str | None, + Query( + description=( + "Columns to download, in this order. Accepts a comma-separated " + 'string (`a,b,c`) or a JSON array of strings (`["a","b"]`). ' + "Default: every column. Every row is still downloaded." + ), + examples=["auctionID,deliveryStart,clearingPrice"], + ), + ] = None, ): - """Download an entire resource; pick the format with `?format=`. + """Download an entire resource; pick the format with `?format=` and, + optionally, the columns with `?fields=`. The file is named after the uploaded file for `url_type: upload` resources, else the resource's `download_filename`; when neither is - available, after the resource id.""" + available, after the resource id. A `fields` download gets a + `_selected_columns` suffix.""" + try: + field_names = to_csv_list(fields) + except ValueError as e: + raise ValidationError(f"fields: {e}") from e authorized = await context.authorize(resource_id=resource_id, permission="read") # Live resource: the cached one still holds the pre-replace url. resource = await context.fresh_resource(resource_id, authorized.get("resource")) filename_base = dump_filename_base(resource, DUMP_EXTENSIONS[fmt]) or resource_id + if field_names: + filename_base = f"{filename_base}_selected_columns" engine = get_datastore_engine(context, mode="ro") - urls = await engine.dump(resource_id, fmt, filename_base=filename_base) + urls = await engine.dump(resource_id, fmt, filename_base=filename_base, fields=field_names) return download_response(request, urls, fmt, filename_base) diff --git a/datastore/infrastructure/engines/base.py b/datastore/infrastructure/engines/base.py index d567733..4821456 100644 --- a/datastore/infrastructure/engines/base.py +++ b/datastore/infrastructure/engines/base.py @@ -127,10 +127,17 @@ def info(self, resource_id: str) -> InfoResult: @abstractmethod async def dump( - self, resource_id: str, fmt: str, *, filename_base: str | None = None + self, + resource_id: str, + fmt: str, + *, + filename_base: str | None = None, + fields: list[str] | None = None, ) -> list[str]: """Download a table as CSV/NDJSON/Parquet. The file is named - `filename_base` (default: `resource_id`) plus the format's extension.""" + `filename_base` (default: `resource_id`) plus the format's extension. + `fields` limits the download to those columns, in that order + (default: every column).""" @abstractmethod async def dump_sql( diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index 7aff47a..99012aa 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -976,10 +976,17 @@ def _count_rows(self, resource_id: str) -> int: # ----- downloads (pipeline lives in bigquery/export.py) --------------- async def dump( - self, resource_id: str, fmt: str, *, filename_base: str | None = None + self, + resource_id: str, + fmt: str, + *, + filename_base: str | None = None, + fields: list[str] | None = None, ) -> list[str]: """Whole-table download → signed URL(s). See `export.py`.""" - return await bq_export.dump(self, resource_id, fmt, filename_base=filename_base) + return await bq_export.dump( + self, resource_id, fmt, filename_base=filename_base, fields=fields + ) async def dump_sql( self, diff --git a/datastore/infrastructure/engines/bigquery/export.py b/datastore/infrastructure/engines/bigquery/export.py index ad0bc3a..f6f61ce 100644 --- a/datastore/infrastructure/engines/bigquery/export.py +++ b/datastore/infrastructure/engines/bigquery/export.py @@ -113,11 +113,21 @@ def _dry_run_error(e: Exception) -> ValidationError: async def dump( - backend: Any, resource_id: str, fmt: str, *, filename_base: str | None = None + backend: Any, + resource_id: str, + fmt: str, + *, + filename_base: str | None = None, + fields: list[str] | None = None, ) -> list[str]: """Whole-table download → signed URLs. Cache key `dumps////` with `rev` = `table.modified`. The name - only lives in the signed URL, so it doesn't affect the cache.""" + only lives in the signed URL, so it doesn't affect the cache. + + `fields` narrows the download to those columns, in the given order + (every row is still exported). A column subset is cached apart from + the whole table, under `dumps//fields-///`, + with `fhash` = hash of the column list.""" if backend.client is None: return [] @@ -128,17 +138,23 @@ async def dump( # read and every request exports. cacheable = table.modified is not None rev = f"{int(table.modified.timestamp() * 1_000_000):x}" if cacheable else uuid4().hex[:12] - prefix = f"{_EXPORT_ROOT}/{resource_id}/{fmt}/{rev}/" source = f"`{backend.config.BIGQUERY_PROJECT}.{backend.config.BIGQUERY_DATASET}.{resource_id}`" - columns = _dump_column_order(table.schema) + if fields: + columns = _select_dump_columns(table.schema, fields) + fhash = hashlib.sha256("\0".join(f.name for f in columns).encode()).hexdigest()[:16] + key = f"{_EXPORT_ROOT}/{resource_id}/fields-{fhash}/{fmt}/" + else: + columns = _dump_column_order(table.schema) + key = f"{_EXPORT_ROOT}/{resource_id}/{fmt}/" + prefix = f"{key}{rev}/" async def build_export_sql(uri: str) -> tuple[str, bytes | None]: # Schema is already known from get_table — no dry run needed. sql = _export_data_sql( uri, fmt, - _export_select_list(columns, fmt), + _export_select_list(columns, fmt, allow_star=not fields), source=source, suffix=" ORDER BY `_id`", ) @@ -152,7 +168,7 @@ async def build_export_sql(uri: str) -> tuple[str, bytes | None]: fmt=fmt, cacheable=cacheable, build_export_sql=build_export_sql, - sweep_prefix=f"{_EXPORT_ROOT}/{resource_id}/{fmt}/", + sweep_prefix=key, filename_base=filename_base or resource_id, what=f"resource {resource_id!r}", ) @@ -634,7 +650,19 @@ def _dump_column_order(schema: Any) -> list[Any]: ] -def _export_select_list(schema: Any, fmt: str) -> str: +def _select_dump_columns(schema: Any, fields: list[str]) -> list[Any]: + """The table columns named by `fields`, in `fields` order, duplicates + dropped. An unknown name is the caller's mistake → 400.""" + by_name = {f.name: f for f in schema} + unknown = [name for name in fields if name not in by_name] + if unknown: + raise ValidationError( + "unknown field(s): " + ", ".join(repr(name) for name in unknown), + ) + return [by_name[name] for name in dict.fromkeys(fields)] + + +def _export_select_list(schema: Any, fmt: str, *, allow_star: bool = True) -> str: """SELECT column list for an export. Every format applies the same `format_select_column` casts as @@ -646,11 +674,12 @@ def _export_select_list(schema: Any, fmt: str) -> str: while csv/ndjson rendered a whole-second ISO string - the same row came back as `...T07:08:35` in csv and `...T07:08:35.908256` in parquet. `*` is still emitted when there is nothing to cast, which - keeps the common case free of a column list. + keeps the common case free of a column list. `allow_star=False` + (a column subset) always gets the explicit list. """ fields = list(schema) if fmt == "parquet": - if not any( + if allow_star and not any( (f.field_type or "").upper() in ("JSON", "TIMESTAMP", "DATETIME") for f in fields ): return "*" diff --git a/tests/test_datastore_dump.py b/tests/test_datastore_dump.py index 6226a1f..63a94e3 100644 --- a/tests/test_datastore_dump.py +++ b/tests/test_datastore_dump.py @@ -19,12 +19,14 @@ import httpx import pytest from datastore.core.constants import DUMP_PREFIX +from datastore.core.exceptions import ValidationError from datastore.infrastructure.engines.bigquery import BigQueryBackend from datastore.infrastructure.engines.bigquery.export import ( _csv_header_bytes, _dump_column_order, _export_select_list, _is_export_too_large, + _select_dump_columns, ) from fastapi.testclient import TestClient @@ -38,10 +40,22 @@ def _patch_dump(urls_or_exc: list[str] | Exception, calls: list[dict[str, Any]] arguments are appended to `calls` when given.""" async def fake( - self: BigQueryBackend, resource_id: str, fmt: str, *, filename_base: str | None = None + self: BigQueryBackend, + resource_id: str, + fmt: str, + *, + filename_base: str | None = None, + fields: list[str] | None = None, ) -> list[str]: if calls is not None: - calls.append({"resource_id": resource_id, "fmt": fmt, "filename_base": filename_base}) + calls.append( + { + "resource_id": resource_id, + "fmt": fmt, + "filename_base": filename_base, + "fields": fields, + } + ) if isinstance(urls_or_exc, Exception): raise urls_or_exc return urls_or_exc @@ -166,6 +180,7 @@ def test_dump_uses_download_filename(client: TestClient, fake_ckan: FakeCKAN) -> "resource_id": "balancing_auction_results_2025", "fmt": "gzip", "filename_base": "Auction Results", + "fields": None, } @@ -236,6 +251,40 @@ def test_sharded_parquet_zip_uses_download_filename( # --- error paths ---------------------------------------------------------- +@pytest.mark.parametrize( + "fields", + ["auctionID,deliveryStart", " auctionID , deliveryStart ", '["auctionID","deliveryStart"]'], +) +def test_dump_passes_fields_and_suffixes_filename(client: TestClient, fields: str) -> None: + """`?fields=` reaches the engine as a list (comma-separated or JSON + array) and the file is marked as a column subset.""" + calls: list[dict[str, Any]] = [] + with _patch_dump(["https://signed"], calls): + response = client.get(DUMP_URL, params={"fields": fields}, follow_redirects=False) + assert response.status_code == 302 + assert calls[0]["fields"] == ["auctionID", "deliveryStart"] + assert calls[0]["filename_base"].endswith("_selected_columns") + + +@pytest.mark.parametrize("params", [{}, {"fields": ""}]) +def test_dump_without_fields_downloads_every_column( + client: TestClient, params: dict[str, str] +) -> None: + calls: list[dict[str, Any]] = [] + with _patch_dump(["https://signed"], calls): + client.get(DUMP_URL, params=params, follow_redirects=False) + assert calls[0]["fields"] is None + assert not calls[0]["filename_base"].endswith("_selected_columns") + + +def test_dump_with_malformed_fields_returns_validation_error(client: TestClient) -> None: + calls: list[dict[str, Any]] = [] + with _patch_dump(["https://signed"], calls): + response = client.get(DUMP_URL, params={"fields": "[not json"}) + assert response.status_code == 400 + assert calls == [] + + def test_unknown_format_returns_validation_error(client: TestClient) -> None: response = client.get(DUMP_URL, params={"format": "xml"}) assert response.status_code == 400 @@ -623,4 +672,72 @@ def test_dump_cache_key_changes_when_table_modified_advances() -> None: assert first_prefix != second_prefix, "table.modified change must produce a different cache key" +def test_select_dump_columns_keeps_requested_order_and_drops_duplicates() -> None: + schema = [_bq_field("_id", "INT64"), _bq_field("a", "STRING"), _bq_field("b", "INT64")] + assert [f.name for f in _select_dump_columns(schema, ["b", "_id", "b"])] == ["b", "_id"] + + +def test_select_dump_columns_rejects_unknown_names() -> None: + schema = [_bq_field("a", "STRING")] + with pytest.raises(ValidationError, match="'nope'"): + _select_dump_columns(schema, ["a", "nope"]) + + +def test_build_export_select_parquet_subset_never_returns_star() -> None: + schema = [_bq_field("id", "INT64")] + assert _export_select_list(schema, "parquet", allow_star=False) == "`id`" + + +def test_dump_with_fields_exports_only_those_columns_under_own_cache_key() -> None: + """A column subset selects just those columns, keeps every row + (ORDER BY `_id`, no WHERE), and caches/sweeps under its own + `fields-` prefix so it never serves or deletes the whole-table file.""" + import asyncio + + backend, storage_client = _engine_with_storage([]) + table = backend.client.get_table.return_value + table.schema = [ + _bq_field("_id", "INT64"), + _bq_field("_updated_at", "TIMESTAMP"), + _bq_field("a", "STRING"), + _bq_field("b", "INT64"), + ] + bucket_obj = storage_client.bucket.return_value + new_blob = _fake_blob("dumps/res-1/fields-x/csv/_000.csv") + bucket_obj.list_blobs.side_effect = [[], [new_blob], [new_blob]] + backend.client.query.return_value = MagicMock() + + async def go() -> None: + await backend.dump("res-1", "csv", fields=["b", "a"]) + if backend._cleanup_tasks: + await asyncio.gather(*list(backend._cleanup_tasks)) + + asyncio.run(go()) + + sql = backend.client.query.call_args.args[0] + select = sql.split(" AS SELECT ", 1)[1] + assert select.startswith("`b`, `a` FROM ") + assert "WHERE" not in select + assert select.endswith("ORDER BY `_id`") + prefixes = [c.kwargs["prefix"] for c in bucket_obj.list_blobs.call_args_list] + assert all(p.startswith("dumps/res-1/fields-") for p in prefixes) + assert not any(p.startswith("dumps/res-1/csv/") for p in prefixes) + # The header names only the selected columns. + composed = bucket_obj.blob.return_value + assert any( + call.args and call.args[0] == b"b,a\n" + for call in composed.upload_from_string.call_args_list + ) + + +def test_dump_with_unknown_field_returns_validation_error_before_export() -> None: + import asyncio + + backend, _ = _engine_with_storage([]) + backend.client.get_table.return_value.schema = [_bq_field("a", "STRING")] + with pytest.raises(ValidationError): + asyncio.run(backend.dump("res-1", "csv", fields=["nope"])) + backend.client.query.assert_not_called() + + # --- test infrastructure --------------------------------------------------