diff --git a/API.md b/API.md index b32e330..fa87dfe 100644 --- a/API.md +++ b/API.md @@ -71,6 +71,7 @@ token (except under `anonymous`). | POST | `/api/3/action/datastore_delete` | Delete rows, drop columns, or drop the table | | GET | `/api/3/action/datastore_search` | Search a resource (streaming) | | GET | `/api/3/action/datastore_search_sql` | Run a read-only SQL `SELECT` (streaming) | +| GET | `/datastore/dump/query` | Download the result of a SQL `SELECT` as a file | | GET | `/api/3/action/datastore_info` | Schema + row stats for a resource | | GET | `/datastore/dump/{resource_id}` | Download a whole resource (CSV/NDJSON/Parquet) | | GET | `/` · `/health` · `/ready` | Welcome / liveness / readiness | @@ -306,6 +307,7 @@ GET /api/3/action/datastore_search Run a single read-only `SELECT` / `WITH` statement and stream the result. Tables are referenced by `resource_id`; each is authorized individually, and functions are checked against the engine's allow-list. Include a `LIMIT` (required). +To export the result as a file instead, use [`GET /datastore/dump/query`](#get-datastoredumpquery). ### Query parameters @@ -380,17 +382,74 @@ row stats — a column-level metadata catalog without a side store. ## `GET /datastore/dump/{resource_id}` Download an entire resource. Pick the format with `?format=csv` (default), -`ndjson`, or `parquet`. - -- **Small export (single shard):** `302` redirect to a signed GCS URL (bytes go - straight from storage to the client). -- **Large export (multiple shards, CSV/NDJSON):** a streamed body - (`200`, ~constant memory). -- Parquet over the single-shard limit returns `413` — switch to CSV/NDJSON. +`gzip`, `ndjson`, or `parquet`. + +- **csv / gzip / ndjson** — `302` redirect to a signed GCS URL, at any size. + Shards from a large export are stitched into one object server-side, so the + bytes go straight from storage to the client (resumable, no server + bandwidth, no server CPU). `gzip` is compressed by BigQuery at export time. +- **parquet** — `302` when the export is a single file. Parquet shards can't be + merged (footer + magic bytes), so a sharded export instead returns `200` with + a **zip** of the parts (`Content-Type: application/zip`, members named + `_NN.parquet`). Entries are stored, not deflated — parquet is + already compressed. This is the one download the server streams itself, so + it has no `Content-Length` and can't be resumed; unzip before querying. Requires `read` permission on the resource and a configured export bucket (`BIGQUERY_EXPORT_BUCKET`). +`query` is a **reserved name** on this route — `/datastore/dump/query` is the SQL +download endpoint below, so a resource literally named `query` can't be dumped +by this URL. + +--- + +## `GET /datastore/dump/query` + +Download the result of a **SQL `SELECT`** as a single file — filtered +downloads at any size. Same validation as `datastore_search_sql` (single +`SELECT`/`WITH`, per-table auth, function allow-list); the response is the +file itself, not the CKAN envelope. + +### Query parameters + +| Name | Type | Notes | +|---|---|---| +| `sql` | string | A single `SELECT`/`WITH`; no multi-statement, no DML/DDL. `LIMIT` optional. | +| `format` | string | `csv` (default) \| `gzip` (gzipped CSV) \| `ndjson` \| `parquet`. | + +### Example + +```http +GET /datastore/dump/query + ?sql=SELECT * FROM "c6153a74-43cb-4edf-8bdf-bb664feca937" WHERE accepted = true + &format=csv +``` + +### Response + +Identical to `/datastore/dump/{resource_id}` above: + +- **csv / gzip / ndjson** — `302` to a signed GCS URL at any size (shards are + composed into one object). The URL expires after + `BIGQUERY_EXPORT_URL_EXPIRY_HOURS` (default 1h). +- **parquet** — `302` for one file; a sharded export returns `200` with a + streamed zip of the parts, members named `query_NN.parquet`. + +### Rules (vs the JSON API) + +- `LIMIT` is **optional** — absent exports the full result set; present it is + honored as written, and the `SEARCH_RESULT_ROWS_MAX` cap does not apply + (`OFFSET` without `LIMIT` is rejected). +- Repeated downloads of the same SQL are served from a GCS cache until any + referenced table changes; SQL calling non-deterministic functions + (`now()`, `current_date`, …) re-exports on every request. +- Row order in the file follows the SQL's `ORDER BY` when it sorts on output + columns (BigQuery preserves it across shards, which are composed in order). Without one, results fall + back to `_id` order when the query outputs `_id` (e.g. `SELECT *`) — the + same order the JSON API pages in; otherwise order is undefined. +- Requires a configured export bucket (`BIGQUERY_EXPORT_BUCKET`). + --- ## Health diff --git a/CLAUDE.md b/CLAUDE.md index 62c3fc7..574813e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -10,7 +10,7 @@ storage backend (BigQuery Datastore or Ducklake as future support). - **Pluggable storage backend** selected by `DATASTORE_ENGINE` (`bigquery` today; `ducklake` planned). - **Pluggable auth** selected by `AUTH_TYPE` (`ckan` / `jwt` / `anonymous`). Provider lives in `datastore/auth//`; only the CKAN provider touches the network, and its TTL cache is local to that provider. - **Standalone-capable** — runs without an upstream CKAN under `AUTH_TYPE=anonymous` or `AUTH_TYPE=jwt`. CKAN is only required when `AUTH_TYPE=ckan`. -- **Streaming search responses** (peak memory ≈ 1 row) for `datastore_search` / `datastore_search_sql`. +- **Streaming search responses** (peak memory ≈ 1 row) for `datastore_search` / `datastore_search_sql`; the sharded-parquet download streams a zip at ≈ 1 chunk. - Strict request validation, structured CKAN-shaped error envelopes. @@ -187,7 +187,8 @@ datastore-api/ │ │ │ # dispatch, pagination links, │ │ │ # function allow-list) │ │ └── streaming.py # streaming.py – byte-yielding writers -│ │ # (objects/lists/csv/tsv) +│ │ # (objects/lists/csv/tsv +│ │ # + zip_archive_writer) │ │ │ │ ── 6. INFRASTRUCTURE (adapters to the outside world) ─ │ └── infrastructure/ @@ -208,6 +209,8 @@ datastore-api/ │ | | # registry imports `Backend`, so the │ | | # concrete class name is engine-private. │ | ├── backend.py # DatastoreBackend subclass (placeholder) +│ | ├── export.py # Download pipeline (dump/dump_sql): +│ | | # cache → EXPORT DATA → compose → sign │ | ├── client.py # google-cloud-bigquery `Client` construction │ | ├── lib.py # Backend-specific helpers (optional) │ | └── allowed_functions.txt # Per-engine datastore_search_sql @@ -423,6 +426,7 @@ Each endpoint takes a single `ContextDep`. The handler calls `context.authorize( | POST | `/api/3/action/datastore_delete` | **implemented** | `DatastoreDeleteRequest` | `DatastoreDeleteResponse` | | GET | `/api/3/action/datastore_search` | **implemented** (streaming) | `DatastoreSearchRequest` | `DatastoreSearchResponse` | | GET | `/api/3/action/datastore_search_sql` | **implemented** (streaming) | `DatastoreSearchSQLRequest` | `DatastoreSearchResponse` | +| GET | `/datastore/dump/query` | **implemented** | `sql=`, `format=csv\|gzip\|ndjson\|parquet` | 302 → GCS *or* streaming body (see §5.3) | | GET | `/api/3/action/datastore_info` | **implemented** | `DatastoreInfoRequest` | `DatastoreInfoResponse` | | GET | `/datastore/dump/{resource_id}` | **implemented** | `format=csv\|ndjson\|parquet` | 302 → GCS *or* streaming body (see §5.3) | @@ -437,26 +441,35 @@ The BigQuery engine is wired end-to-end: DDL, MERGE-based upsert, DML delete, pa ### 5.3 `GET /datastore/dump/{resource_id}` -Full-table download, **one URL → one file** from the caller's point of view. Bytes never sit in API memory — small dumps redirect to GCS, large dumps stream-concat through async httpx. +Full-table download, **one URL → one file** from the caller's point of +view. Bytes never pass through API memory — the one exception is a +sharded parquet export, which is zipped on the way out. -Pipeline: +Pipeline (`_prepare_download` in [bigquery/export.py](datastore/infrastructure/engines/bigquery/export.py)): -1. **Resolve cache key** — read `table.modified` from BigQuery, compute `rev = hex(microsec_epoch(modified))`, prefix becomes `dumps///`. -2. **GCS cache lookup** — `list_blobs(prefix=…)`. Non-empty → skip steps 3-5; log `cache HIT`. -3. **Submit `EXPORT DATA`** — Parquet → single-file URI `gs:///.parquet`; CSV/NDJSON → wildcard URI `gs:///_*.` so BigQuery shards >1 GB exports. The SELECT casts `TIMESTAMP` + `DATETIME` columns to ISO 8601 for CSV/NDJSON; Parquet keeps native types. -4. **Poll non-blockingly** — `await asyncio.to_thread(job.reload)` + `await asyncio.sleep(_DUMP_POLL_INTERVAL_SECONDS)` between iterations. Worker thread is held only during the brief `reload` call, not the wait. -5. **GC stale revisions** — after a successful extract, sweep `dumps///` and delete any blob that doesn't start with the current `prefix`. Best-effort, failures logged. Storage stays bounded to one rev per `(rid, fmt)`. -6. **Sign URLs** — V4 signed URLs with `response-content-disposition: attachment; filename="."` (single shard) or `_NN.` (multi-shard, 1-indexed, zero-padded). Signing offloaded to a thread (IAM round-trip under workload identity). -7. **Return**: - - 1 URL → `RedirectResponse(302)`. Bytes flow GCS → client; server is **out of the byte path**. - - N URLs → `StreamingResponse` over `services.dump.stream_*_shards` (async httpx, 64 KiB chunks, serial shard walk, CSV header-dedup via `_skip_first_line`, NDJSON pure byte-concat). +1. **Resolve cache key** — read `table.modified`, `rev = hex(microsec_epoch(modified))`, prefix `dumps////`. Everything this service writes lives under the single `dumps/` root — table dumps keyed on the resource id, query dumps on a SQL hash — so one lifecycle rule covers it all. Every request for that (resource, format, table version) shares the **revision directory**; each individual export writes into its own `/` beneath it: -Per-stream resource profile (multi-shard branch): ~64 KiB resident memory, **0** worker threads, byte-copy CPU only, async cancellation propagates from client disconnect → httpx → GCS connection released. +``` +dumps//// ← cache key, shared by all requests + └── / ← one export's private scratch + ├── part_*. + ├── data. ← composed (csv/gzip, ndjson >1 shard) + └── _SUCCESS ← written last; publishes the attempt +``` +2. **Cache lookup** (`_complete_attempt`) — one `list_blobs(prefix=…)`, grouped by attempt directory; serve the newest attempt carrying `_SUCCESS`. Attempts without it are in-flight or abandoned: never served, never deleted on the response path. Within the winner a composed `data.` is the whole download; `part_*` shards beside it are leftovers awaiting background deletion. Parquet never composes, so its shards *are* the download. +3. **Submit `EXPORT DATA`** — wildcard URI `gs:////part_*.`. The wildcard is mandatory (BigQuery rejects a bare URI — *"Option 'uri' value must be a wild card URI"*) and shards by write parallelism, not size, so a 40 MB result routinely lands as several files. `job.result()` waits, holding one worker thread for the export's duration. CSV **and gzip** export **header-less** (gzip adds `compression='GZIP'`, so BigQuery does the compressing); the SELECT casts `TIMESTAMP` + `DATETIME` to ISO 8601 for CSV/NDJSON and `TO_JSON_STRING` for JSON columns on Parquet. +4. **Compose** — csv/gzip/ndjson shards are stitched into ONE object server-side with GCS `compose` (≤32 sources per call, chained beyond that). csv and gzip compose a synthesized header member in first — plain bytes for csv, a gzip member for gzip — so the result carries exactly one header. (Verified against the real bucket: BigQuery's gzip objects carry **no** `Content-Encoding`, so compose byte-concats them and multi-member gzip decompresses as one file.) Parquet is never composable (footer + magic bytes) and stays as shards. Compose sources are the explicit shard list from **this attempt only** — never a fresh listing — so nothing foreign can be swept into the output. +5. **Publish** — write the zero-byte `/_SUCCESS`. This is the commit point: the attempt is unreadable before it, readable after, so no caller ever sees a half-written export. Written *after* the compose, so it can never become file content. +6. **Cleanup — in the background** (`_cleanup_in_background`; never on the response path): the compose source shards are deleted, and superseded attempts + revisions under `dumps///` are swept (`_delete_old_cache` — this request's own attempt is never touched, and the sweep is **always age-gated by the signed-URL expiry** so a sibling attempt that may still be exporting, or whose URLs are live, survives). +7. **Sign URLs** — V4 with `response-content-disposition: attachment; filename="."` (one file) or `_NN.` (multi-file parquet, 1-indexed). Signing is offloaded to a thread (IAM round-trip under workload identity). +8. **Return** (`download_response` in [api/endpoints/dump.py](datastore/api/endpoints/dump.py)): + - 1 URL → `RedirectResponse(302)`. Bytes flow GCS → client; the server is **out of the byte path** for every format, and downloads are resumable. + - N URLs (sharded parquet only) → `200` + **a streamed zip** of the parts (`zip_archive_writer` in [services/streaming.py](datastore/services/streaming.py)). The API fetches each signed URL over `app.state.http` and frames it into the archive a chunk at a time, so one export is always one file at one URL. Entries are `ZIP_STORED` (parquet is already compressed; deflating costs CPU for no size win) and `force_zip64=True` (member sizes aren't known up front). This is the **only** path where the server carries the bytes — no `Content-Length`, no range support, and a fetch failure mid-archive truncates a response that already returned 200. Errors: -- Parquet >1 GB → `EXPORT DATA` job fails with a "single URI / wildcard" message; classifier in `_is_export_too_large` flips it to `PayloadTooLargeError` (413). Caller switches to `format=csv` or `format=ndjson`. -- Any other BigQuery / GCS failure → `ServerError` (500) with the upstream message. +- Any BigQuery / GCS failure → `ServerError` (500) with the upstream message. A ">1 GB single URI" failure is classified by `_is_export_too_large` into `PayloadTooLargeError` (413) — defensive only, since the wildcard URI means BigQuery shards instead of failing. - `BIGQUERY_EXPORT_BUCKET` unset → `ServerError` at request time (the lifespan doesn't fail-fast because dump is an optional capability). +- **Concurrency.** Requests for the same (query, table version) share a revision directory but export into **separate attempt directories**, so they can never overwrite each other's objects, and neither is servable until its own `_SUCCESS` lands. The residual cost is that both run an export — duplicate billed scans, bounded and rare. Fixing that needs single-flighting (an in-process lock, or a deterministic BigQuery job id so the loser waits on the winner's job); it is a cost optimisation, not a correctness one, and is deliberately not implemented. Required IAM. Dump follows a strict **ro for reading, rw for writing/updating** model — see [bigquery/client.py](datastore/infrastructure/engines/bigquery/client.py) `load_credentials` + `_build_bq_client` / `_build_storage_client` on the backend: @@ -466,17 +479,33 @@ Required IAM. Dump follows a strict **ro for reading, rw for writing/updating** | `list_blobs` cache lookup | RO GCS | Reading GCS objects. | | `client.query("EXPORT DATA …")` | RW BQ (built on demand) | BigQuery writes shards to GCS under this SA's identity — it's a write op even though the SQL surface is `SELECT`. | | Post-extract `list_blobs` refresh | RW GCS | Blobs are passed straight to `generate_signed_url` next; we want them bound to the rw client. | +| `compose` (csv/ndjson) | RW GCS | Server-side concat into one object: reads each source's metadata (`storage.objects.get` — **`list` alone is not enough**) and creates the composite. | +| `upload_from_string` (csv header member) | RW GCS | The one-row header composed in front of the header-less csv shards. | | `delete` (GC) | RW GCS | Writing/deleting objects. | | `generate_signed_url` | RW GCS | Under workload identity this calls IAM `signBlob`, which typically only the rw SA holds via `iam.serviceAccountTokenCreator`. | Concrete perm sets: - **RO SA** (`BIGQUERY_CREDENTIALS_RO`) — `bigquery.tables.get` + `storage.objects.list`. -- **RW SA** (`BIGQUERY_CREDENTIALS`) — `bigquery.jobs.create` + `bigquery.tables.export` + `bigquery.tables.getData` + `storage.objects.{create,list,delete}` + `iam.serviceAccountTokenCreator`. +- **RW SA** (`BIGQUERY_CREDENTIALS`) — `bigquery.jobs.create` + `bigquery.tables.export` + `bigquery.tables.getData` + `storage.objects.{create,get,list,delete}` + `iam.serviceAccountTokenCreator`. **`get` is required by `compose`** (it reads each source object); a role with only create/list/delete 403s on the compose step. A single SA works if both perm sets land on the same identity — `BIGQUERY_CREDENTIALS_RO` empty falls through to ADC; same env var can drive both. `_build_bq_client` and `_build_storage_client` on the backend are deliberately small + stub-friendly so tests inject mocks without monkey-patching `google.cloud.*` globally. -A 24h object-lifecycle rule on the bucket is a useful belt-and-braces: the engine GCs older revs already, but lifecycle catches anything stranded by a crashed dump. +A 24h object-lifecycle rule on the bucket is **required** in practice: the engine GCs older revs already, but lifecycle is the only thing that cleans abandoned `dumps//` prefixes (SQL downloads whose query is never re-issued — see below) and anything stranded by a crashed dump. + +### SQL download (`GET /datastore/dump/query`) + +`GET /datastore/dump/query?sql=&format=csv|gzip|ndjson|parquet` exports the result of an arbitrary vetted SELECT through the same pipeline as `/datastore/dump/{resource_id}` — engine method `dump_sql` in [bigquery/export.py](datastore/infrastructure/engines/bigquery/export.py), response shaping shared via `download_response` in [api/endpoints/dump.py](datastore/api/endpoints/dump.py) (302 for the composed file · gzip streamed · JSON URL list for multi-file parquet). Same SQL validation + per-table auth as `datastore_search_sql` (`DatastoreDumpSQLRequest` subclasses its request schema); the action API itself stays pure JSON envelope. The route is declared before `/datastore/dump/{resource_id}`, making `query` a reserved resource name on the dump family. + +Deltas vs the whole-table dump: + +- **LIMIT is optional, uncapped.** `datastore_search_sql` requires a LIMIT literal; the dump request schema relaxes it (`parse_sql_pagination(require_limit=…)` via the `_REQUIRE_LIMIT` class flag), honors a present LIMIT as written, and `SEARCH_RESULT_ROWS_MAX` does not apply. OFFSET without LIMIT is rejected. +- **Cache key** = `dumps////`: `qhash` = sha256 of the qualified SQL, `rev` = sha256 over every referenced table's `(rid, modified)` pair — any table change → new rev. Query dumps share the `dumps/` root with table dumps; the identity segment is a 16-hex hash rather than a resource id, so the two only collide if a table is literally named like one. +- **Non-deterministic SQL bypasses the cache.** Queries calling `now()`, `current_date`, … (`_NON_DETERMINISTIC_SQL_FUNCTIONS`) skip the lookup and export under a fresh uuid rev per run. +- **RO dry run → RW export.** The user SQL is dry-run on the RO client first (free; clean 400 on SQL that doesn't compile; yields the output schema for the same per-format casts `dump()` uses — ISO timestamps for CSV/NDJSON, `TO_JSON_STRING` for JSON→parquet). The `EXPORT DATA` itself must run under the RW SA (it writes GCS objects); containment = single-statement/SELECT-only schema validation + per-table authorize + function allow-list + the user SQL riding in subquery position (`AS SELECT … FROM ()`). +- **Age-gated GC.** Stale revisions under `dumps///` are deleted only once older than the signed-URL expiry, so a re-export can't kill shards whose URLs are still live. +- **Row order** (`_outer_order_by`): BigQuery ignores a subquery's ORDER BY without LIMIT, so ordering lives on the **outer** exported query — a user's top-level `ORDER BY` is hoisted there when its keys are output columns; with no ORDER BY, `ORDER BY _id` is applied when `_id` is in the output (mirrors JSON mode's `default_order_by`). Otherwise the file is unordered. BigQuery preserves outer ORDER BY globally across shards; shards concat in name order. +- Every cache miss is a **billed query** (EXPORT DATA never uses BigQuery's result cache); `maximum_bytes_billed` is a possible future cost cap. The GCS client is built with the same credentials as the BigQuery client for the active engine mode (`load_credentials(config, mode)` in [bigquery/client.py](datastore/infrastructure/engines/bigquery/client.py)). Without this shim, a service-account JSON loaded via `BIGQUERY_CREDENTIALS_RO` would drive BigQuery but `storage.Client(...)` would silently fall back to ADC — a near-invisible identity split. Workload identity / `GOOGLE_APPLICATION_CREDENTIALS`-style setups still work because `load_credentials` returns `None` for ADC and the storage client follows the same default-credentials path. @@ -763,7 +792,7 @@ Optional fields appear in `result` only when requested: ### 6.4 `GET /api/3/datastore_search_sql` -**Query params**: `sql` (required), `limit` (default 32000). +**Query params**: `sql` (required; must carry a `LIMIT` literal). To export the result as a file instead of the JSON envelope, use `GET /datastore/dump/query?sql=…&format=…` (LIMIT optional + uncapped there — see §5.3 "SQL download"). **Example request — daily clearing-price summary** ``` diff --git a/README.md b/README.md index 2298bb6..dec9f97 100644 --- a/README.md +++ b/README.md @@ -31,8 +31,18 @@ Copy [.env.example](.env.example) and fill it in. The essentials: | `AUTH_TYPE` | `ckan` | Auth provider: `ckan` · `jwt` · `anonymous` | | `CKAN_URL` | — | CKAN base URL (required when `AUTH_TYPE=ckan`) | | `BIGQUERY_PROJECT` / `BIGQUERY_DATASET` | — | Required when `DATASTORE_ENGINE=bigquery` | +| `BIGQUERY_EXPORT_BUCKET` | — | GCS bucket for downloads (`/datastore/dump/{resource_id}`, `/datastore/dump/query`) | | `REDIS_URL` | — | Cache backend; empty → in-process cache | +**Note on the export bucket:** everything the service writes lives under +a single `dumps/` prefix — `dumps//…` for whole-table downloads +and `dumps//…` for SQL downloads — so one lifecycle rule covers +everything. Configure a ~24h rule on +`BIGQUERY_EXPORT_BUCKET`: the engine garbage-collects superseded export +revisions itself, but abandoned query-download prefixes (queries never +re-issued) and files stranded by a crashed export are only ever cleaned by +the lifecycle rule. + ## Documentation - **[API.md](API.md)** — full API reference (endpoints, request/response, examples). diff --git a/datastore/api/endpoints/datastore.py b/datastore/api/endpoints/datastore.py index c1a9ece..f700296 100644 --- a/datastore/api/endpoints/datastore.py +++ b/datastore/api/endpoints/datastore.py @@ -180,7 +180,8 @@ async def datastore_search_sql( params: Annotated[DatastoreSearchSQLRequest, Query()], ): """`GET /api/3/datastore_search_sql` — execute a raw SQL SELECT and stream. - Accepts a single `sql` query parameter; + Accepts a single `sql` query parameter. To export the result as a + file instead, use `GET /datastore/dump/query`. """ for resource_id in params.resource_ids: await context.authorize(resource_id=resource_id, permission="read") diff --git a/datastore/api/endpoints/dump.py b/datastore/api/endpoints/dump.py index 04e1a47..06c4eb6 100644 --- a/datastore/api/endpoints/dump.py +++ b/datastore/api/endpoints/dump.py @@ -1,94 +1,135 @@ -"""`GET /datastore/dump/{resource_id}` — single download for a table. +"""Download endpoints: `/datastore/dump/{resource_id}` + `/datastore/dump/query`. -Behaviour by shard count (decided by BigQuery from the export size): - - - **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. - -Multi-shard Parquet is refused with 413 (parquet shards can't be -byte-concatenated). Caller picks CSV/NDJSON. +csv / gzip / ndjson shards are composed into one GCS object, so those +always redirect — the server never touches the bytes. Parquet can't be +composed (footer + magic bytes), so a parquet export that shards is +zipped on the way out: the API fetches the parts and streams one +archive, which keeps every download a single file at one URL. """ from __future__ import annotations -from typing import Annotated, Literal +from typing import Annotated from fastapi import APIRouter, Query +from starlette.requests import Request from starlette.responses import RedirectResponse, StreamingResponse from datastore.api.context import Context from datastore.api.responses import ERROR_RESPONSES +from datastore.core.constants import DUMP_EXTENSIONS, DumpFormat +from datastore.core.exceptions import ServerError from datastore.infrastructure.engines import get_datastore_engine -from datastore.services.dump import ( - stream_csv_shards, - stream_gzip_csv_shards, - stream_ndjson_shards, -) +from datastore.schemas.request import DatastoreDumpSQLRequest +from datastore.services.read import dump_sql_datastore +from datastore.services.streaming import zip_archive_writer -DumpFormat = Literal["csv", "gzip", "ndjson", "parquet"] +router = APIRouter(tags=["Datastore Download"], responses=ERROR_RESPONSES) -_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", -} +def download_response( + request: Request, + urls: list[str], + fmt: DumpFormat, + filename_base: str, +) -> RedirectResponse | StreamingResponse: + """Shape the engine's signed URL(s) into a response: + + - one file → 302 to the signed URL; bytes go GCS → client and + the download is resumable + - many files → 200 streaming one zip of the parts (parquet only — + every other format composes to a single object) + - none → 500; the engine isn't configured + + The zip is the only path where the server carries the bytes, and it + exists so a caller always gets one file from one URL. `Content- + Length` is unknowable mid-stream, so the response is chunked: no + progress bar, and no resuming a dropped connection. + """ + if not urls: + raise ServerError( + "export produced no downloadable files " + "(datastore engine is not configured)" + ) + if len(urls) == 1: + return RedirectResponse(url=urls[0], status_code=302) -router = APIRouter(tags=["Datastore Download"], responses=ERROR_RESPONSES) + ext = DUMP_EXTENSIONS[fmt] + members = [ + (f"{filename_base}_{i + 1:02d}.{ext}", url) + for i, url in enumerate(urls) + ] + return StreamingResponse( + zip_archive_writer(request.app.state.http, members), + media_type="application/zip", + headers={ + "Content-Disposition": ( + f'attachment; filename="{filename_base}.zip"' + ), + }, + ) + + +@router.get( + "/datastore/dump/query", + summary="Download the result of a SQL SELECT (CSV / gzip CSV / NDJSON / Parquet)", + responses={ + 302: {"description": "Redirect to the signed download URL."}, + 200: { + "description": ( + "Sharded parquet export — one streamed zip of the parts." + ), + "content": {"application/zip": {}}, + }, + }, +) +async def dump_sql( + request: Request, + context: Context, + params: Annotated[DatastoreDumpSQLRequest, Query()], +): + """Download a vetted SQL SELECT's result as one file (no envelope). + + Same validation as `datastore_search_sql`; `LIMIT` optional and + uncapped. Declared before `/{resource_id}` → `query` is reserved. + """ + for resource_id in params.resource_ids: + await context.authorize(resource_id=resource_id, permission="read") + + urls = await dump_sql_datastore( + context, + { + "sql": params.sql, + "fmt": params.format, + "resource_ids": params.resource_ids, + "function_names": params.function_names, + }, + ) + return download_response(request, urls, params.format, "query") @router.get( "/datastore/dump/{resource_id}", 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."}, + 302: {"description": "Redirect to the signed Download URL."}, + 200: { + "description": ( + "Multi-file parquet export — one streamed zip." + ), + "content": {"application/zip": {}}, + }, }, ) async def dump( + request: Request, context: Context, resource_id: str, fmt: Annotated[DumpFormat, Query(alias="format")] = "csv", ): - """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=`. - """ + """Download an entire resource; pick the format with `?format=`.""" await context.authorize(resource_id=resource_id, permission="read") engine = get_datastore_engine(context, mode="ro") urls = await engine.dump(resource_id, fmt) - - if len(urls) == 1: - return RedirectResponse(url=urls[0], status_code=302) - - 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 — the engine rejects multi-shard Parquet - raise RuntimeError(f"unexpected multi-shard format: {fmt}") - - return StreamingResponse( - body, - media_type=_MEDIA_TYPE[fmt], - headers={ - "Content-Disposition": ( - f'attachment; filename="{resource_id}.{_DOWNLOAD_EXT[fmt]}"' - ), - }, - ) + return download_response(request, urls, fmt, resource_id) diff --git a/datastore/core/constants.py b/datastore/core/constants.py index 6a8ceae..e8d19b2 100644 --- a/datastore/core/constants.py +++ b/datastore/core/constants.py @@ -1,5 +1,23 @@ from __future__ import annotations +from typing import Literal + +# Download formats served by the export pipeline (`/datastore/dump/…` and +# `datastore_search_sql?download=…`). Lives here — not in `api/` — because +# both the request schemas (pydantic layer) and the endpoints (starlette +# layer) need it, and `schemas/` must not import from `api/`. +DUMP_FORMATS: tuple[str, ...] = ("csv", "gzip", "ndjson", "parquet") +DumpFormat = Literal["csv", "gzip", "ndjson", "parquet"] + +# File extension per dump format — names the download and, for a +# multi-file export, each member inside the zip. +DUMP_EXTENSIONS: dict[str, str] = { + "csv": "csv", + "gzip": "csv.gz", + "ndjson": "json", + "parquet": "parquet", +} + POSTGRES_TYPES: dict[str, str] = { # integer "int2": "int2", diff --git a/datastore/infrastructure/engines/base.py b/datastore/infrastructure/engines/base.py index e638cef..73a2c20 100644 --- a/datastore/infrastructure/engines/base.py +++ b/datastore/infrastructure/engines/base.py @@ -130,7 +130,27 @@ def info(self, resource_id: str) -> InfoResult: @abstractmethod async def dump(self, resource_id: str, fmt: str) -> list[str]: - """Download a table as CSV/NDJSON/Parquet. + """Download a table as CSV/NDJSON/Parquet. + """ + + @abstractmethod + async def dump_sql( + self, + sql: str, + fmt: str, + *, + resource_ids: list[str], + function_names: list[str], + ) -> list[str]: + """Export the result of a vetted SELECT as `fmt`; return signed + URLs — normally one (engines may merge shards into a single + object), more only for formats that cannot be merged. + + `resource_ids` / `function_names` are the names already parsed + (and authorized / allow-listed) by the schema + service layers — + passed in so engines don't re-parse the SQL. Engines may use + `function_names` to decide result cacheability (non-deterministic + functions make a cached export stale by definition). """ @abstractmethod diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index f817033..ad912e0 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -16,6 +16,10 @@ build on those two. 4. CKAN action methods (`create`, `upsert`, `search`, `search_sql`, `delete`, `info`, `get_columns`, `healthcheck`). + +Downloads (`dump` / `dump_sql` — the cache → export → compose → sign +pipeline) live in [export.py](export.py); the methods here are one-line +delegates into that module. """ from __future__ import annotations @@ -27,7 +31,6 @@ from datastore.core.config import Config from datastore.core.exceptions import ( NotFoundError, - PayloadTooLargeError, ServerError, ValidationError, ) @@ -37,6 +40,7 @@ SearchResult, WriteResult, ) +from datastore.infrastructure.engines.bigquery import export as bq_export from datastore.infrastructure.engines.bigquery.lib import ( PK_CONFLICT_SENTINEL, SYSTEM_COLUMN_NAMES, @@ -45,7 +49,6 @@ default_order_by, delete_sql, drop_columns_sql, - format_select_column, insert_guarded_sql, insert_sql, merge_sql, @@ -82,6 +85,9 @@ def __init__( self.config = config self.client: Any = None self._health: tuple[float, bool] | None = None + # In-flight fire-and-forget cleanup tasks (compose part deletion, + # revision GC) — tracked so tests can await them deterministically. + self._cleanup_tasks: set[Any] = set() def initialize(self) -> None: """Build the BigQuery client when project + dataset are configured. @@ -937,199 +943,29 @@ def _count_rows(self, resource_id: str) -> int: return 0 return int(rows[0]["n"]) - async def dump(self, resource_id: str, fmt: str) -> list[str]: - """Submit `EXPORT DATA`; poll non-blockingly; return signed URLs. - - - 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 - poll loop releases the worker between `job.reload` calls. - """ - import asyncio - from datetime import timedelta - from uuid import uuid4 - - if self.client is None: - return [] - - bucket = ( - getattr(self.config, "BIGQUERY_EXPORT_BUCKET", "") or "" - ).strip() - if not bucket: - raise ServerError( - "BIGQUERY_EXPORT_BUCKET is not configured — " - "/datastore/dump cannot run without an export bucket." - ) - - from google.cloud import bigquery - - # Clients: ro for reads (BQ get_table, GCS list); rw for the - # rest (BQ EXPORT DATA writes shards under its identity; GCS - # delete + sign). One bucket handle per client. - rw_bq = self._build_bq_client("rw") - ro_gcs = self._build_storage_client("ro").bucket(bucket) - rw_gcs = self._build_storage_client("rw").bucket(bucket) - - table_ref = bigquery.TableReference.from_string( - f"{self.config.BIGQUERY_PROJECT}" - f".{self.config.BIGQUERY_DATASET}.{resource_id}" - ) - from google.api_core.exceptions import NotFound - - try: - table = await asyncio.to_thread(self.client.get_table, table_ref) - except NotFound as e: - raise NotFoundError( - f"resource {resource_id!r} is not declared; nothing to dump" - ) from e - except Exception as e: - raise ServerError( - f"BigQuery get_table failed for resource {resource_id!r}: {e}" - ) from e - - rev = ( - f"{int(table.modified.timestamp() * 1_000_000):x}" - if table.modified is not None - else uuid4().hex[:12] - ) - ext = _FMT[fmt]["ext"] - prefix = f"dumps/{resource_id}/{fmt}/{rev}" - uri = f"gs://{bucket}/{prefix}_*.{ext}" - - async def _list(b: Any, p: str) -> list[Any]: - return sorted( - await asyncio.to_thread(lambda: list(b.list_blobs(prefix=p))), - key=lambda x: x.name, - ) - - blobs = await _list(ro_gcs, prefix) - - if not blobs: - # `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. `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" - f"{extra_opts}" - ") AS " - f"SELECT {_build_export_select(table.schema, fmt)} FROM " - f"`{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}` " - f"ORDER BY `_id`" - ) - try: - job = await asyncio.to_thread(rw_bq.query, sql) - except Exception as e: - raise ServerError( - f"BigQuery EXPORT DATA submit failed for resource " - f"{resource_id!r}: {e}" - ) from e - - while True: - await asyncio.to_thread(job.reload) - if job.state == "DONE": - break - await asyncio.sleep(_DUMP_POLL_INTERVAL_SECONDS) - - if job.error_result: - err_msg = (job.error_result or {}).get("message", "") - if _is_export_too_large(RuntimeError(err_msg)): - raise PayloadTooLargeError( - f"resource {resource_id!r} exceeds 1 GB after export " - f"as {fmt!r}; single-file download isn't possible. " - "Try `format=csv` or `format=ndjson` for sharded " - "multi-file downloads instead." - ) - raise ServerError( - f"BigQuery EXPORT DATA failed for resource " - f"{resource_id!r}: {err_msg}" - ) - - log.info( - "BigQuery dump cache MISS: resource=%s format=%s rev=%s", - resource_id, fmt, rev, - ) - blobs = await _list(rw_gcs, prefix) - if not blobs: - raise ServerError( - f"BigQuery EXPORT DATA wrote no shards for resource " - f"{resource_id!r}; check job logs." - ) + # ----- downloads (pipeline lives in bigquery/export.py) --------------- - # GC stale revisions under dumps///. Best-effort. - def _gc() -> int: - deleted = 0 - for old in rw_gcs.list_blobs(prefix=f"dumps/{resource_id}/{fmt}/"): - if old.name.startswith(prefix): - continue - try: - old.delete() - deleted += 1 - except Exception as gc_err: # noqa: BLE001 - log.warning("dump GC: failed to delete %s: %s", old.name, gc_err) - return deleted - - try: - gc_count = await asyncio.to_thread(_gc) - if gc_count: - log.info( - "BigQuery dump GC: resource=%s format=%s removed=%d", - resource_id, fmt, gc_count, - ) - except Exception as gc_err: # noqa: BLE001 - log.warning( - "BigQuery dump GC failed for resource=%s format=%s: %s", - resource_id, fmt, gc_err, - ) - else: - log.info( - "BigQuery dump cache HIT: resource=%s format=%s rev=%s shards=%d", - resource_id, fmt, rev, len(blobs), - ) - # Re-fetch via rw so the blobs we sign carry rw credentials - # (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." - ) + async def dump(self, resource_id: str, fmt: str) -> list[str]: + """Whole-table download → signed URL(s). See `export.py`.""" + return await bq_export.dump(self, resource_id, fmt) - expiry = timedelta( - hours=getattr(self.config, "BIGQUERY_EXPORT_URL_EXPIRY_HOURS", 1), + async def dump_sql( + self, + sql: str, + fmt: str, + *, + resource_ids: list[str], + function_names: list[str], + ) -> list[str]: + """SQL-result download → signed URL(s). See `export.py`.""" + return await bq_export.dump_sql( + self, + sql, + fmt, + resource_ids=resource_ids, + function_names=function_names, ) - def _sign_all() -> list[str]: - out: list[str] = [] - for i, blob in enumerate(blobs): - filename = ( - f"{resource_id}.{ext}" - if len(blobs) == 1 - else f"{resource_id}_{i + 1:02d}.{ext}" - ) - out.append( - blob.generate_signed_url( - version="v4", - expiration=expiry, - method="GET", - response_disposition=f'attachment; filename="{filename}"', - ) - ) - return out - - return await asyncio.to_thread(_sign_all) - def _build_bq_client(self, mode: str) -> Any: """Construct an on-demand BigQuery client for `mode` ("ro" / "rw"). @@ -1322,59 +1158,3 @@ def _translate_bigquery_error( "json": "object", "bytes": "string", } - - -# --- EXPORT DATA helpers ----------------------------------------------------- - -# Seconds between BigQuery job-status polls during a dump. Each poll -# is a quick metadata HTTP call (~tens of ms); between polls the worker -# thread is released so other requests can run. Bumping this down makes -# small jobs complete faster, bumping it up means fewer reload calls -# per job — 1 s is a safe middle. -_DUMP_POLL_INTERVAL_SECONDS = 1.0 - -# Per-format filename extension + BigQuery EXPORT DATA `format` value. -# BigQuery writes newline-delimited JSON to `.json` files; we keep that -# 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"}, -} - - -def _build_export_select(schema: Any, fmt: str) -> str: - """SELECT column list for EXPORT DATA. - - 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": - 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 - ) - - -def _is_export_too_large(exc: BaseException) -> bool: - """Does this BigQuery error look like ">1 GB single-file rejection"? - - BigQuery's exact wording shifts across SDK versions; both phrasings - we've seen contain `single URI` or `wildcard`. False negatives just - surface as a generic 500 instead of 413 — annoying but not silent. - """ - msg = str(exc).lower() - return "single uri" in msg or "wildcard" in msg diff --git a/datastore/infrastructure/engines/bigquery/export.py b/datastore/infrastructure/engines/bigquery/export.py new file mode 100644 index 0000000..1b5a1b3 --- /dev/null +++ b/datastore/infrastructure/engines/bigquery/export.py @@ -0,0 +1,706 @@ +"""BigQuery download pipeline — `/datastore/dump/{rid}` + `/datastore/dump/query`. + +Read top-down: entry points, then the workflow they share, then the +helpers each step uses. + + 1. ENTRY POINTS `dump` (whole table) · `dump_sql` (vetted SELECT) + — derive the cache key + export SQL, nothing else + 2. WORKFLOW `_prepare_download`: + HIT? newest attempt with a `_SUCCESS` marker + └─ yes ──▶ sign → 302 (no BigQuery) + EXPORT EXPORT DATA → shards, own attempt dir + COMPOSE csv/ndjson shards → ONE object + PUBLISH write `_SUCCESS` + GC old attempts/revisions — BACKGROUND + SIGN V4 URL(s) → 302 + 3. STEP HELPERS one per workflow step, in the order used + 4. SQL BUILDERS pure text/bytes helpers + 5. BACKEND small accessors over the engine instance + +Functions take the `BigQueryBackend` instance (`backend`) explicitly; +`google.cloud` imports stay lazy (optional dep). +""" + +from __future__ import annotations + +import asyncio +import gzip +import hashlib +import logging +import time +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta, timezone +from typing import Any +from uuid import uuid4 + +from datastore.core.exceptions import ( + NotFoundError, + PayloadTooLargeError, + ServerError, + ValidationError, +) +from datastore.infrastructure.engines.bigquery.lib import ( + format_select_column, + qualify_table_refs, +) + +log = logging.getLogger(__name__) + +# Per-format file extension + BigQuery EXPORT DATA `format` value. +_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"}, +} + +# Content type stamped on the composed object. +_CONTENT_TYPE = { + "csv": "text/csv", + "gzip": "application/gzip", + "ndjson": "application/x-ndjson", +} + +# GCS `compose` accepts ≤32 sources per call; more are chained. +_COMPOSE_MAX_SOURCES = 32 + +# Written last, once the bytes are final. An attempt without it is +# invisible — mid-export or abandoned, either way not servable. +_SUCCESS = "_SUCCESS" + +# Bucket root. Everything this service writes lives under one prefix, so +# a lifecycle rule can target it directly. Table dumps key on the +# resource id, query dumps on a SQL hash — both sit directly beneath. +_EXPORT_ROOT = "dumps" + +# SQL functions whose results change per run → the cache must be bypassed. +_NON_DETERMINISTIC_SQL_FUNCTIONS = frozenset( + { + "now", + "current_date", + "current_datetime", + "current_time", + "current_timestamp", + "rand", + "generate_uuid", + "session_user", + "current_user", + } +) + +# Formats that survive a single-file download — named in 413 messages. +_FORMAT_HINT = "`format=csv` or `format=ndjson`" + + +# ============================================================================ +# 1. ENTRY POINTS +# ============================================================================ + + +async def dump(backend: Any, resource_id: str, fmt: str) -> list[str]: + """Whole-table download → signed URLs. Cache key + `dumps////` with `rev` = `table.modified`.""" + if backend.client is None: + return [] + + bucket = _get_export_bucket(backend) + table = await _get_table(backend, resource_id) + + # No `modified` → no stable version to key on, so the cache can't be + # 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}" + f".{backend.config.BIGQUERY_DATASET}.{resource_id}`" + ) + + 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(table.schema, fmt), + source=source, + suffix=" ORDER BY `_id`", + ) + header = ( + _csv_header_bytes(table.schema, fmt) + if fmt in ("csv", "gzip") + else None + ) + return sql, header + + return await _prepare_download( + backend, + bucket=bucket, + prefix=prefix, + fmt=fmt, + cacheable=cacheable, + build_export_sql=build_export_sql, + sweep_prefix=f"{_EXPORT_ROOT}/{resource_id}/{fmt}/", + filename_base=resource_id, + what=f"resource {resource_id!r}", + ) + + +async def dump_sql( + backend: Any, + sql: str, + fmt: str, + *, + resource_ids: list[str], + function_names: list[str], +) -> list[str]: + """SQL-result download → signed URLs. + + Cache key `dumps////`: `qhash` = hash of the + qualified SQL, `rev` = hash of every referenced table's `modified`. + Non-deterministic SQL (`now()`, …) is never cached (uuid rev). A + free RO dry run validates the SQL and yields the output schema for + the per-format casts; the export wraps the user SQL as a subquery. + """ + if backend.client is None: + return [] + if backend.mode != "ro": + raise ServerError( + "datastore_search_sql download must run on a read-only " + "engine; got mode=" + repr(backend.mode) + ) + + from google.cloud import bigquery + + bucket = _get_export_bucket(backend) + + try: + qualified_sql = qualify_table_refs( + sql, + project=backend.config.BIGQUERY_PROJECT, + dataset=backend.config.BIGQUERY_DATASET, + ) + except Exception as e: + raise ServerError(f"failed to qualify table references in SQL: {e}") from e + + tables = {rid: await _get_table(backend, rid) for rid in resource_ids} + + # Deterministic SQL over tables with known versions → stable rev. + cacheable = not (set(function_names) & _NON_DETERMINISTIC_SQL_FUNCTIONS) and all( + t.modified is not None for t in tables.values() + ) + if cacheable: + pairs = sorted( + (rid, int(t.modified.timestamp() * 1_000_000)) for rid, t in tables.items() + ) + rev = hashlib.sha256( + "|".join(f"{rid}:{us}" for rid, us in pairs).encode() + ).hexdigest()[:16] + else: + rev = uuid4().hex[:12] + + qhash = hashlib.sha256(qualified_sql.encode()).hexdigest()[:16] + prefix = f"{_EXPORT_ROOT}/{qhash}/{fmt}/{rev}/" + + async def build_export_sql(uri: str) -> tuple[str, bytes | None]: + # RO dry run (free): validates the SQL + yields the output + # schema. Deferred — a cache hit never pays for it. + dry_cfg = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False) + try: + dry_job = await asyncio.to_thread( + backend.client.query, + qualified_sql, + job_config=dry_cfg, + ) + except Exception as e: + raise ValidationError(f"sql failed BigQuery validation: {e}") from e + export_sql = _export_data_sql( + uri, + fmt, + _export_select_list(dry_job.schema, fmt), + source=f"({qualified_sql})", + suffix=_outer_order_by(qualified_sql, dry_job.schema), + ) + header = ( + _csv_header_bytes(dry_job.schema, fmt) + if fmt in ("csv", "gzip") + else None + ) + return export_sql, header + + return await _prepare_download( + backend, + bucket=bucket, + prefix=prefix, + fmt=fmt, + cacheable=cacheable, + build_export_sql=build_export_sql, + sweep_prefix=f"{_EXPORT_ROOT}/{qhash}/{fmt}/", + filename_base=f"query_{qhash[:8]}", + what=f"sql query {qhash[:8]}", + ) + + +# ============================================================================ +# 2. WORKFLOW +# ============================================================================ + + +async def _prepare_download( + backend: Any, + *, + bucket: str, + prefix: str, + fmt: str, + cacheable: bool, + build_export_sql: Callable[[str], Awaitable[tuple[str, bytes | None]]], + sweep_prefix: str, + filename_base: str, + what: str, +) -> list[str]: + """The download workflow shared by `dump` and `dump_sql`: + + list prefix → an attempt with `_SUCCESS`? sign its files + → otherwise export · compose · publish · GC + + Every export writes into its own `/`, and an attempt + is invisible until its `_SUCCESS` lands. Two requests for the same + (query, table version) therefore cost a duplicate export at worst — + neither can overwrite the other's objects or read a half-written one. + """ + ro_gcs = backend._build_storage_client("ro").bucket(bucket) + rw_gcs = backend._build_storage_client("rw").bucket(bucket) + + blobs: list[Any] = [] + if cacheable and _complete_attempt(await _list_files_sorted(ro_gcs, prefix), fmt): + # Re-list via rw: signing needs rw creds, and a blob signs through + # whichever client listed it. + blobs = _complete_attempt(await _list_files_sorted(rw_gcs, prefix), fmt) or [] + + if blobs: + log.info("BigQuery export cache HIT: %s prefix=%s", what, prefix) + else: + attempt = f"{prefix}{uuid4().hex[:12]}/" + rw_bq = backend._build_bq_client("rw") + export_sql, header_bytes = await build_export_sql( + f"gs://{bucket}/{attempt}part_*.{_FMT[fmt]['ext']}" + ) + await _run_export_job(rw_bq, export_sql, what=what, fmt=fmt) + log.info("BigQuery export cache MISS: %s attempt=%s", what, attempt) + blobs = await _list_files_sorted(rw_gcs, attempt) + if not blobs: + raise ServerError( + f"BigQuery EXPORT DATA wrote no shards for {what}; " "check job logs." + ) + + blobs = await _compose_single_file( + backend, rw_gcs, attempt, fmt, blobs, header_bytes + ) + + # Publishes the attempt. Written after the compose, so it can + # never be a compose source. + await asyncio.to_thread( + rw_gcs.blob(f"{attempt}{_SUCCESS}").upload_from_string, b"" + ) + + def gc() -> None: + removed = _delete_old_cache( + rw_gcs, + sweep_prefix=sweep_prefix, + keep_prefix=attempt, + min_age=_url_expiry(backend), + ) + if removed: + log.info("BigQuery export GC: %s removed=%d", what, removed) + + _cleanup_in_background(backend, gc, what=f"old cache {what}") + + return await _signed_urls( + backend, + blobs, + filename_base=filename_base, + ext=_FMT[fmt]["ext"], + ) + + +# ============================================================================ +# 3. STEP HELPERS (in workflow order) +# ============================================================================ + + +async def _list_files_sorted(bucket: Any, prefix: str) -> list[Any]: + """`list_blobs(prefix=…)` off-thread, name-sorted (= shard order).""" + return sorted( + await asyncio.to_thread( + lambda: list(bucket.list_blobs(prefix=prefix)), + ), + key=lambda x: x.name, + ) + + +def _complete_attempt(blobs: list[Any], fmt: str) -> list[Any] | None: + """Files of the newest attempt carrying `_SUCCESS`, else `None`. + + A listing of `` spans every attempt made for this (query, + table version). Only a published one counts, so an in-flight or + abandoned attempt can't be served or mistaken for a finished one. + + Within the winner, the composed `data.` is the whole download + when it exists; `part_*` shards beside it are leftovers awaiting + background deletion. Parquet never composes, so its shards *are* the + download. + """ + attempts: dict[str, list[Any]] = {} + for blob in blobs: + attempts.setdefault(blob.name.rsplit("/", 1)[0], []).append(blob) + + published = { + name: objs + for name, objs in attempts.items() + if any(o.name.endswith(f"/{_SUCCESS}") for o in objs) + } + if not published: + return None + + newest = max(published, key=lambda a: _published_at(published[a])) + files = [o for o in published[newest] if not o.name.endswith(f"/{_SUCCESS}")] + composed = [o for o in files if o.name.endswith(f"/data.{_FMT[fmt]['ext']}")] + return composed or files + + +def _published_at(objs: list[Any]) -> datetime: + """When an attempt's `_SUCCESS` landed — picks between attempts that + both finished.""" + epoch = datetime.min.replace(tzinfo=timezone.utc) + for obj in objs: + if obj.name.endswith(f"/{_SUCCESS}"): + created = getattr(obj, "time_created", None) + return created if isinstance(created, datetime) else epoch + return epoch + + +async def _run_export_job( + rw_bq: Any, + export_sql: str, + *, + what: str, + fmt: str, +) -> None: + """Run `EXPORT DATA` to completion; >1 GB → 413, other errors → 500. + + `job.result()` does the waiting (the client polls with backoff), so + it holds one worker thread for the length of the export. + """ + t0 = time.perf_counter() + try: + job = await asyncio.to_thread(rw_bq.query, export_sql) + await asyncio.to_thread(job.result) + except Exception as e: + if _is_export_too_large(e): + raise PayloadTooLargeError( + f"{what} exceeds 1 GB after export as {fmt!r}; " + "single-file download isn't possible. " + f"Try {_FORMAT_HINT} for sharded multi-file downloads " + "instead." + ) from e + raise ServerError(f"BigQuery EXPORT DATA failed for {what}: {e}") from e + + log.debug( + "BigQuery export: %s in %.2fs", + what, + time.perf_counter() - t0, + ) + + +async def _compose_single_file( + backend: Any, + rw_gcs: Any, + attempt: str, + fmt: str, + shards: list[Any], + header_bytes: bytes | None, +) -> list[Any]: + """Compose csv/ndjson shards (+ csv header) into ONE GCS object → + `[composite]` → the endpoint 302s a single URL. csv always composes + (to inject the header); ndjson only when >1 shard; gzip/parquet pass + through. Source shards are deleted in the background afterwards.""" + composable = fmt in ("csv", "gzip") or (fmt == "ndjson" and len(shards) > 1) + if not composable: + return shards # lone ndjson shard / gzip / parquet + + ext = _FMT[fmt]["ext"] + content_type = _CONTENT_TYPE[fmt] + + def compose() -> tuple[Any, list[Any]]: + sources = list(shards) + extras: list[Any] = [] + + if header_bytes is not None: # csv header member, sorts first + header = rw_gcs.blob(f"{attempt}header.{ext}") + header.upload_from_string(header_bytes, content_type=content_type) + sources = [header, *sources] + extras.append(header) + + composite = rw_gcs.blob(f"{attempt}data.{ext}") + composite.content_type = content_type + composite.compose(sources[:_COMPOSE_MAX_SOURCES]) + i = _COMPOSE_MAX_SOURCES + while i < len(sources): + # `composite` itself counts as one source → add up to 31 more. + batch = sources[i : i + _COMPOSE_MAX_SOURCES - 1] + composite.compose([composite, *batch]) + i += _COMPOSE_MAX_SOURCES - 1 + return composite, [*shards, *extras] + + t0 = time.perf_counter() + composite, parts = await asyncio.to_thread(compose) + log.debug( + "BigQuery compose: %s <- %d shard(s) in %.0fms", + attempt, + len(shards), + (time.perf_counter() - t0) * 1000, + ) + + # The composite is standalone → its source shards are garbage. + # Deleted off the response path. (Old revisions: `_delete_old_cache`.) + _cleanup_in_background( + backend, + lambda: _delete_blobs(parts, "compose sources"), + what=f"source shards {attempt}", + ) + return [composite] + + +def _delete_old_cache( + rw_gcs: Any, + *, + sweep_prefix: str, + keep_prefix: str, + min_age: timedelta | None = None, +) -> int: + """Delete older attempts and revisions under `sweep_prefix`; this + request's own attempt (`keep_prefix`) is never touched. `min_age` + spares anything younger than the signed-URL lifetime — covering both + a sibling attempt still exporting and one whose URLs are live.""" + cutoff = datetime.now(timezone.utc) - min_age if min_age else None + + def is_stale(blob: Any) -> bool: + if blob.name.startswith(keep_prefix): + return False # this request's own attempt + created = getattr(blob, "time_created", None) + if cutoff is not None and created is not None and created > cutoff: + return False # URLs may still be live + return True + + stale = [b for b in rw_gcs.list_blobs(prefix=sweep_prefix) if is_stale(b)] + return _delete_blobs(stale, "old cache") + + +def _cleanup_in_background( + backend: Any, + fn: Callable[[], None], + *, + what: str, +) -> None: + """Fire-and-forget GCS cleanup on a worker thread — the response + never waits on deletes. The task is held on the backend (a cached + singleton) because the loop keeps only weak refs; tests await it.""" + + async def run() -> None: + try: + await asyncio.to_thread(fn) + except Exception as e: # noqa: BLE001 + log.warning("background cleanup failed (%s): %s", what, e) + + task = asyncio.get_running_loop().create_task(run()) + backend._cleanup_tasks.add(task) + task.add_done_callback(backend._cleanup_tasks.discard) + + +def _delete_blobs(blobs: list[Any], why: str) -> int: + """Delete blobs, logging (never raising) individual failures — every + delete here is best-effort cleanup. Sync: run via `to_thread`.""" + deleted = 0 + for blob in blobs: + try: + blob.delete() + deleted += 1 + except Exception as e: # noqa: BLE001 + log.warning( + "%s: failed to delete %s: %s", + why, + getattr(blob, "name", "?"), + e, + ) + return deleted + + +async def _signed_urls( + backend: Any, + blobs: list[Any], + *, + filename_base: str, + ext: str, +) -> list[str]: + """V4-sign each blob with an attachment filename (`.`, + or `_NN.` when there are several).""" + expiry = _url_expiry(backend) + + def sign_all() -> list[str]: + out: list[str] = [] + for i, blob in enumerate(blobs): + filename = ( + f"{filename_base}.{ext}" + if len(blobs) == 1 + else f"{filename_base}_{i + 1:02d}.{ext}" + ) + out.append( + blob.generate_signed_url( + version="v4", + expiration=expiry, + method="GET", + response_disposition=f'attachment; filename="{filename}"', + ) + ) + return out + + return await asyncio.to_thread(sign_all) + + +# ============================================================================ +# 4. SQL / BYTES BUILDERS (pure) +# ============================================================================ + + +def _export_data_sql( + uri: str, + fmt: str, + select_list: str, + source: str, + suffix: str = "", +) -> str: + """`EXPORT DATA` statement text. csv exports header-less (the compose + step injects one `_csv_header_bytes` member); gzip keeps per-shard + headers (streamed with dedup).""" + if fmt == "csv": + extra_opts = ", header=false" + elif fmt == "gzip": + extra_opts = ", header=false, compression='GZIP'" + else: + extra_opts = "" + return ( + f"EXPORT DATA OPTIONS(" + f"uri='{uri}', format='{_FMT[fmt]['bq']}', overwrite=true" + f"{extra_opts}" + ") AS " + f"SELECT {select_list} FROM {source}{suffix}" + ) + + +def _export_select_list(schema: Any, fmt: str) -> str: + """SELECT column list: parquet casts JSON columns to strings (BQ + can't export native JSON to parquet); csv/ndjson use the same + `format_select_column` casts as `datastore_search`.""" + if fmt == "parquet": + 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) + + +def _csv_header_bytes(schema: Any, fmt: str) -> bytes: + """The one CSV header row composed in front of the header-less + shards. gzip-compressed for `gzip` so it concatenates with the + gzip shards as another member (column names never need quoting).""" + row = (",".join(f.name for f in schema) + "\n").encode() + return gzip.compress(row) if fmt == "gzip" else row + + +def _outer_order_by(sql: str, schema: Any) -> str: + """Outer `ORDER BY` for a wrapped SQL export (BigQuery ignores a + subquery's ORDER BY without LIMIT). Hoists the user's ORDER BY when + its keys are output columns; else `ORDER BY _id` when `_id` is in + the output; else "" (unordered).""" + import sqlglot + from sqlglot import expressions as exp + + out_cols = {f.name for f in schema} + try: + tree = sqlglot.parse_one(sql, dialect="bigquery") + except Exception: + return "" + + order = tree.args.get("order") + if order is None: + return " ORDER BY `_id`" if "_id" in out_cols else "" + + keys: list[str] = [] + for ordered in order.expressions: + col = ordered.this + if not isinstance(col, exp.Column) or col.name not in out_cols: + return "" + hoisted = ordered.copy() + hoisted.set("this", exp.column(col.name, quoted=True)) + keys.append(hoisted.sql(dialect="bigquery")) + return " ORDER BY " + ", ".join(keys) + + +def _is_export_too_large(exc: BaseException) -> bool: + """Does this BigQuery error look like ">1 GB single-file rejection"?""" + msg = str(exc).lower() + return "single uri" in msg or "wildcard" in msg + + +# ============================================================================ +# 5. BACKEND ACCESSORS +# ============================================================================ + + +def _get_export_bucket(backend: Any) -> str: + """Validated GCS export bucket name; `ServerError` when unset.""" + bucket = (getattr(backend.config, "BIGQUERY_EXPORT_BUCKET", "") or "").strip() + if not bucket: + raise ServerError( + "BIGQUERY_EXPORT_BUCKET is not configured — " + "/datastore/dump cannot run without an export bucket." + ) + return bucket + + +def _url_expiry(backend: Any) -> timedelta: + """Signed-URL lifetime (also the GC age gate for uuid revs).""" + return timedelta( + hours=getattr(backend.config, "BIGQUERY_EXPORT_URL_EXPIRY_HOURS", 1), + ) + + +async def _get_table(backend: Any, resource_id: str) -> Any: + """Fetch table metadata; NotFound → 404, anything else → 500.""" + from google.api_core.exceptions import NotFound + from google.cloud import bigquery + + table_ref = bigquery.TableReference.from_string( + f"{backend.config.BIGQUERY_PROJECT}" + f".{backend.config.BIGQUERY_DATASET}.{resource_id}" + ) + try: + return await asyncio.to_thread(backend.client.get_table, table_ref) + except NotFound as e: + raise NotFoundError( + f"resource {resource_id!r} is not declared; nothing to dump" + ) from e + except Exception as e: + raise ServerError( + f"BigQuery get_table failed for resource {resource_id!r}: {e}" + ) from e diff --git a/datastore/schemas/request.py b/datastore/schemas/request.py index 030d48b..29bd092 100644 --- a/datastore/schemas/request.py +++ b/datastore/schemas/request.py @@ -1,7 +1,7 @@ from __future__ import annotations import re -from typing import Annotated, Any, Literal +from typing import Annotated, Any, ClassVar, Literal from pydantic import ( BaseModel, @@ -12,6 +12,7 @@ model_validator, ) +from datastore.core.constants import DumpFormat from datastore.schemas.validators import ( FieldSpec, StringOrList, @@ -327,6 +328,10 @@ class DatastoreSearchSQLRequest(BaseModel): model_config = ConfigDict(extra="forbid") + # Subclass hook: `DatastoreDumpSQLRequest` flips this to False — + # a download exports the full result set, so LIMIT is optional there. + _REQUIRE_LIMIT: ClassVar[bool] = True + sql: str = Field( description=( "A single read-only `SELECT` / `WITH` statement. Reference " @@ -344,7 +349,7 @@ class DatastoreSearchSQLRequest(BaseModel): # the read-only properties below give callers a clean attribute. _resource_ids: list[str] = PrivateAttr(default_factory=list) _function_names: list[str] = PrivateAttr(default_factory=list) - _limit: int = PrivateAttr(default=0) + _limit: int | None = PrivateAttr(default=None) _offset: int = PrivateAttr(default=0) @property @@ -360,8 +365,10 @@ def function_names(self) -> list[str]: return self._function_names @property - def limit(self) -> int: - """`LIMIT` literal parsed from the SQL — required.""" + def limit(self) -> int | None: + """`LIMIT` literal parsed from the SQL. Required here (never + `None`); `None` only on `DatastoreDumpSQLRequest` when the SQL + carries no LIMIT.""" return self._limit @property @@ -412,15 +419,53 @@ def _extract_sql_references(self) -> DatastoreSearchSQLRequest: Runs after `_check_sql_is_select`, so we know we have a single SELECT / WITH. CTE aliases are excluded from `_resource_ids` (they're defined inline, not external tables). LIMIT is - required — the service uses it to build pagination links and - to cap the streaming response; missing LIMIT raises a clean - ValidationError up front. + required (`_REQUIRE_LIMIT`) — the service uses it to build + pagination links and to cap the streaming response; missing + LIMIT raises a clean ValidationError up front. The dump-SQL + subclass relaxes the rule: exports cover the full result set, + so `_limit` stays `None` when absent. """ self._resource_ids, self._function_names = parse_sql_references(self.sql) - self._limit, self._offset = parse_sql_pagination(self.sql) + self._limit, self._offset = parse_sql_pagination( + self.sql, require_limit=self._REQUIRE_LIMIT, + ) return self +class DatastoreDumpSQLRequest(DatastoreSearchSQLRequest): + """Query parameters for `GET /datastore/dump/query`. + + Same vetted-SELECT contract as `DatastoreSearchSQLRequest` (single + SELECT / WITH, table + function extraction, `extra="forbid"`), but + the result is exported as a **file**: `LIMIT` is optional — absent + means the full result set; present it is honored as written, and + the JSON path's row cap does not apply. `format` picks the export + format. + """ + + _REQUIRE_LIMIT: ClassVar[bool] = False + + sql: str = Field( + description=( + "A single read-only `SELECT` / `WITH` statement. Reference " + "resources by their id. `LIMIT` is optional — absent exports " + "the full result set; present it is honored as written." + ), + examples=[ + 'SELECT * FROM "balancing_auction_results_2025" ' + "WHERE accepted = true" + ], + ) + + format: DumpFormat = Field( + default="csv", + description=( + "Export format: `csv` | `gzip` (gzipped CSV) | `ndjson` | " + "`parquet`." + ), + ) + + class DatastoreInfoRequest(BaseModel): """Query parameters for `GET /api/3/datastore_info`. diff --git a/datastore/schemas/validators.py b/datastore/schemas/validators.py index 6d44cd8..0799589 100644 --- a/datastore/schemas/validators.py +++ b/datastore/schemas/validators.py @@ -326,9 +326,9 @@ def parse_sql_references(sql: str, *, dialect: str = "postgres") -> tuple[list[s def parse_sql_pagination( - sql: str, *, dialect: str = "postgres", -) -> tuple[int, int]: - """Extract `(limit, offset)` from a SELECT. LIMIT is required. + sql: str, *, dialect: str = "postgres", require_limit: bool = True, +) -> tuple[int | None, int]: + """Extract `(limit, offset)` from a SELECT. `datastore_search_sql` lets callers ship raw SQL but the API still wants page metadata + links — so we parse the LIMIT/OFFSET out of @@ -338,6 +338,12 @@ def parse_sql_pagination( paginate properly and so an unbounded SELECT can't lock streaming open. + `require_limit=False` (download mode — the result is exported as a + file, not streamed as a page) allows the LIMIT to be absent, in + which case `limit` comes back as `None` and the full result set is + meant. OFFSET without LIMIT is rejected either way — BigQuery + can't execute a bare OFFSET, so fail fast with a clean message. + OFFSET defaults to 0 when absent. Non-integer LIMIT/OFFSET expressions (e.g. `LIMIT @x`) raise too — pagination needs a constant. @@ -351,20 +357,27 @@ def parse_sql_pagination( raise ValueError(f"could not parse SQL: {e}") from e limit_node = tree.args.get("limit") + limit: int | None = None if limit_node is None: - raise ValueError( - "SQL must include a LIMIT clause (e.g. 'LIMIT 100'); " - "an explicit page size is required for pagination links " - "and to prevent unbounded SELECTs" - ) - limit_expr = limit_node.expression if isinstance(limit_node, exp.Limit) else None - if not isinstance(limit_expr, exp.Literal) or not limit_expr.is_int: - raise ValueError( - "LIMIT must be a constant integer literal" + if require_limit: + raise ValueError( + "SQL must include a LIMIT clause (e.g. 'LIMIT 100'); " + "an explicit page size is required for pagination links " + "and to prevent unbounded SELECTs" + ) + if tree.args.get("offset") is not None: + raise ValueError("OFFSET without LIMIT is not supported") + else: + limit_expr = ( + limit_node.expression if isinstance(limit_node, exp.Limit) else None ) - limit = int(limit_expr.this) - if limit < 0: - raise ValueError("LIMIT must be >= 0") + if not isinstance(limit_expr, exp.Literal) or not limit_expr.is_int: + raise ValueError( + "LIMIT must be a constant integer literal" + ) + limit = int(limit_expr.this) + if limit < 0: + raise ValueError("LIMIT must be >= 0") offset = 0 offset_node = tree.args.get("offset") diff --git a/datastore/services/dump.py b/datastore/services/dump.py deleted file mode 100644 index 9607ba6..0000000 --- a/datastore/services/dump.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Service for multi-shard streaming of `/datastore/dump/`. - -Single-shard exports are served as 302 redirects (bytes flow GCS → -client, never through us). When BigQuery shards an export (>1 GB -CSV/NDJSON), the endpoint falls back to **server-side stream-concat -through this module**: we fetch each shard from GCS via async httpx -and forward bytes to the client as a single download. - -Resource profile per active stream-concat dump: - - Memory: one ~64 KiB chunk in flight per active shard (we walk - shards serially, so peak ≈ one chunk). - - CPU: byte forwarding plus a single newline scan per CSV shard to - strip its header. Essentially zero. - - Threads: none — everything runs on the asyncio loop via httpx. - - Network: full dump size through our server (the unavoidable cost - of the "one URL → one file" contract). - -Parquet shards aren't supported here because their footers can't be -byte-concat'd; the engine refuses multi-shard Parquet at 1 GB. -""" - -from __future__ import annotations - -import zlib -from collections.abc import AsyncIterator - -import httpx - -# Read-side network knobs. No total timeout (a multi-GB stream legitimately -# takes minutes); just a generous per-chunk read timeout so a dead GCS -# connection doesn't hang us forever. -_TIMEOUT = httpx.Timeout(connect=10.0, read=60.0, write=10.0, pool=10.0) -# Chunk size we pull from GCS / forward to the client. -_CHUNK_BYTES = 64 * 1024 - - -async def stream_csv_shards(urls: list[str]) -> AsyncIterator[bytes]: - """Stream-concat CSV shards. Header from the first shard only; the - first newline of each subsequent shard is dropped (BigQuery emits - a header row per shard when `header=true`).""" - 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() - if i == 0: - async for chunk in resp.aiter_bytes(_CHUNK_BYTES): - yield chunk - else: - async for chunk in _skip_first_line( - resp.aiter_bytes(_CHUNK_BYTES) - ): - yield chunk - - -async def stream_ndjson_shards(urls: list[str]) -> AsyncIterator[bytes]: - """Stream-concat NDJSON shards. Each shard is independent - newline-delimited JSON, so pure byte concatenation produces a - valid combined stream.""" - async with httpx.AsyncClient(timeout=_TIMEOUT) as client: - for url in urls: - async with client.stream("GET", url) as resp: - resp.raise_for_status() - async for chunk in resp.aiter_bytes(_CHUNK_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]: - """Drop bytes up to and including the first `\\n`, then forward - the rest unchanged. Used to strip the duplicate CSV header on - non-first shards. Memory bound: bytes of the header line plus - one chunk.""" - pending = bytearray() - async for chunk in chunks: - pending.extend(chunk) - idx = pending.find(b"\n") - if idx >= 0: - yield bytes(pending[idx + 1:]) - pending.clear() - break - async for chunk in chunks: - yield chunk diff --git a/datastore/services/read.py b/datastore/services/read.py index 1eb1c50..6b12b32 100644 --- a/datastore/services/read.py +++ b/datastore/services/read.py @@ -131,16 +131,7 @@ async def search_sql_datastore( OFFSET so callers can follow `_links.next` / `prev` without re-editing their SQL. """ - allowed = get_allowed_sql_functions( - context.config.DATASTORE_ENGINE, - override_path=context.config.SQL_FUNCTIONS_ALLOW_FILE, - ) - disallowed = sorted(set(data_dict.get("function_names", [])) - allowed) - if disallowed: - raise ValidationError( - f"sql uses disallowed function(s): {', '.join(disallowed)}", - fields={"sql": [f"disallowed: {', '.join(disallowed)}"]}, - ) + _ensure_allowed_sql_functions(context, data_dict.get("function_names", [])) limit = data_dict["limit"] offset = data_dict["offset"] @@ -184,6 +175,52 @@ async def search_sql_datastore( ) +async def dump_sql_datastore( + context: RequestContext, data_dict: dict[str, Any] +) -> list[str]: + """Export a vetted SELECT's result via the engine; return signed URLs. + + The download-mode sibling of `search_sql_datastore`: same function + allow-list gate, but **no LIMIT clamp** — the result lands in GCS as + a file, not in a streamed JSON envelope, so a LIMIT in the SQL is + honored as written and its absence means the full result set. + + Dispatches on the read-only engine; the EXPORT-side elevation to rw + credentials is engine-internal (BigQuery writes shards under the rw + identity). + """ + _ensure_allowed_sql_functions(context, data_dict.get("function_names", [])) + engine = get_datastore_engine(context, mode="ro") + # `dump_sql` is a native coroutine (its blocking calls are offloaded + # internally) — await it directly, no `to_thread`. + return await engine.dump_sql( + data_dict["sql"], + data_dict["fmt"], + resource_ids=data_dict["resource_ids"], + function_names=data_dict["function_names"], + ) + + +def _ensure_allowed_sql_functions( + context: RequestContext, function_names: list[str], +) -> None: + """Reject SQL that calls functions outside the engine's allow-list. + + Shared by `search_sql_datastore` and `dump_sql_datastore` so both + entry points enforce the identical gate. + """ + allowed = get_allowed_sql_functions( + context.config.DATASTORE_ENGINE, + override_path=context.config.SQL_FUNCTIONS_ALLOW_FILE, + ) + disallowed = sorted(set(function_names) - allowed) + if disallowed: + raise ValidationError( + f"sql uses disallowed function(s): {', '.join(disallowed)}", + fields={"sql": [f"disallowed: {', '.join(disallowed)}"]}, + ) + + async def info_datastore( context: RequestContext, data_dict: dict[str, Any] ) -> DatastoreInfoResponse.Result: diff --git a/datastore/services/streaming.py b/datastore/services/streaming.py index 4c572af..20f993c 100644 --- a/datastore/services/streaming.py +++ b/datastore/services/streaming.py @@ -25,6 +25,11 @@ CSV / TSV rows are embedded inside a JSON string value, so each row's text is JSON-escaped via `orjson.dumps(s)[1:-1]` before being yielded between the records field's opening / closing `"` quotes. + +`zip_archive_writer` at the bottom is the odd one out: it serves the +download endpoints rather than `datastore_search`, and packs several +already-exported files into one archive. Same discipline — bytes are +yielded a chunk at a time and never accumulated. """ from __future__ import annotations @@ -32,7 +37,8 @@ import base64 import csv import io -from collections.abc import Iterator +import zipfile +from collections.abc import AsyncIterator, Iterator from decimal import Decimal from typing import Any @@ -315,3 +321,85 @@ def _json_string_inner(s: str) -> bytes: value chunk-by-chunk across many rows without materialising it. """ return orjson.dumps(s)[1:-1] + + +# --------------------------------------------------------------------------- +# Zip archive over already-exported files (multi-file download) +# --------------------------------------------------------------------------- + +# Read size from storage, and roughly how often bytes reach the client. +_ZIP_CHUNK_BYTES = 1 << 20 + +# Per-operation timeout while pulling a part. The client's default is +# tuned for CKAN's small JSON calls; a multi-GB shard needs far more +# headroom, and a timeout mid-archive truncates the download. +_ZIP_FETCH_TIMEOUT = 300.0 + + +class _ZipSink: + """The write-only file object `zipfile` builds the archive into. + + `ZipFile` probes for `tell()`; finding none it switches to its + non-seekable mode, which writes each member's sizes and CRC in a + trailing data descriptor instead of seeking back to patch the local + header. That is precisely what makes the archive streamable. + """ + + def __init__(self) -> None: + self._buf = bytearray() + + def write(self, data: bytes) -> int: + self._buf += data + return len(data) + + def flush(self) -> None: + return None + + def __len__(self) -> int: + return len(self._buf) + + def drain(self) -> bytes: + """Hand over everything buffered so far and reset.""" + out = bytes(self._buf) + del self._buf[:] + return out + + +async def zip_archive_writer( + http: Any, + members: list[tuple[str, str]], +) -> AsyncIterator[bytes]: + """Stream a zip of `members` — `(filename, url)` pairs — as it builds. + + Each URL is fetched a chunk at a time and framed straight into the + archive, so peak memory is one chunk rather than one file. `http` is + an `httpx.AsyncClient` supplied by the caller (services never build + their own). + + Entries are **stored, never deflated**: the payload is parquet, + already compressed, so deflating would burn 30–100 MB/s of CPU for + no size win. The archive is pure framing, and `force_zip64` keeps a + member over 4 GiB from blowing up at close. + + This is the one download path where the server carries the bytes, so + it inherits the costs: no `Content-Length` (hence no progress bar), + and no range support, so a dropped connection restarts the transfer. + """ + sink = _ZipSink() + with zipfile.ZipFile( + sink, mode="w", compression=zipfile.ZIP_STORED, allowZip64=True + ) as archive: + for filename, url in members: + with archive.open(filename, mode="w", force_zip64=True) as entry: + async with http.stream( + "GET", url, timeout=_ZIP_FETCH_TIMEOUT + ) as response: + response.raise_for_status() + async for chunk in response.aiter_bytes(_ZIP_CHUNK_BYTES): + entry.write(chunk) + if len(sink) >= _ZIP_CHUNK_BYTES: + yield sink.drain() + # Member closed → its data descriptor is in the buffer. + yield sink.drain() + # Archive closed → central directory + EOCD. + yield sink.drain() diff --git a/tests/test_cors.py b/tests/test_cors.py index 72fc125..2815cb1 100644 --- a/tests/test_cors.py +++ b/tests/test_cors.py @@ -12,10 +12,9 @@ from contextlib import contextmanager import pytest -from fastapi.testclient import TestClient - from datastore.core.config import get_config from datastore.main import create_app +from fastapi.testclient import TestClient @contextmanager diff --git a/tests/test_datastore_dump.py b/tests/test_datastore_dump.py index 47756ca..46b5096 100644 --- a/tests/test_datastore_dump.py +++ b/tests/test_datastore_dump.py @@ -1,29 +1,28 @@ """Tests for `GET /datastore/dump/{resource_id}`. -Engine returns a list of signed-URL shards: - - len == 1 → endpoint 302s to the URL (no server bandwidth). - - len > 1 → endpoint stream-concats shards from GCS via async httpx. +The engine returns one signed URL (csv / gzip / ndjson shards are +composed into a single object), so every format 302s. Only a sharded +parquet export comes back as several URLs, which the endpoint fetches +and streams as one zip. -We patch `BigQueryBackend.dump` to control how many "shards" the -engine reports, and patch `httpx.AsyncClient` for the stream-concat -tests so they don't try to fetch real URLs. +We patch `BigQueryBackend.dump` to control what the engine reports, and +`app.state.http` to serve the signed URLs from memory. """ from __future__ import annotations -import gzip -from collections.abc import AsyncIterator +import io +import zipfile from typing import Any from unittest.mock import MagicMock, patch +import httpx import pytest -from datastore.core.exceptions import PayloadTooLargeError, ServerError from datastore.infrastructure.engines.bigquery import BigQueryBackend -from datastore.infrastructure.engines.bigquery.backend import ( - _build_export_select, +from datastore.infrastructure.engines.bigquery.export import ( + _export_select_list, _is_export_too_large, ) -from datastore.services.dump import _skip_first_line from fastapi.testclient import TestClient from tests.conftest import FakeCKAN @@ -40,6 +39,17 @@ async def fake(self: BigQueryBackend, resource_id: str, fmt: str) -> list[str]: return patch.object(BigQueryBackend, "dump", fake) +def stub_signed_urls(client: TestClient, parts: dict[str, bytes]) -> None: + """Serve `parts` (url → body) from `app.state.http`, the client the + zip writer fetches the signed URLs with.""" + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=parts[str(request.url)]) + + client.app.state.http = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + ) + + # --- single shard: 302 redirect ------------------------------------------- @@ -63,94 +73,58 @@ def test_each_format_supports_single_shard_redirect( assert response.status_code == 302 -# --- multi-shard: stream-concat ------------------------------------------- +# --- several shards: one streamed zip ------------------------------------- -def test_multi_shard_csv_stream_concat_dedups_header( - client: TestClient, -) -> None: - """Header from shard 1 only; shards 2..N have their first line - dropped before bytes hit the client.""" - shards = { - "url-1": b"col1,col2\na,1\nb,2\n", - "url-2": b"col1,col2\nc,3\nd,4\n", - "url-3": b"col1,col2\ne,5\n", +def test_multi_file_parquet_streams_one_zip(client: TestClient) -> None: + """A sharded parquet export is fetched back and framed into a single + zip, so the caller still gets one file from one URL.""" + parts = { + "https://signed/p0.parquet": b"PAR1-first-shard-bytes", + "https://signed/p1.parquet": b"PAR1-second-shard-bytes", } + stub_signed_urls(client, parts) - with _patch_dump(list(shards.keys())), _patch_httpx_stream(shards): - response = client.get(DUMP_URL, follow_redirects=False) + with _patch_dump(list(parts)): + response = client.get(DUMP_URL, params={"format": "parquet"}) assert response.status_code == 200 - assert response.headers["content-type"].startswith("text/csv") - assert response.headers["content-disposition"] == ( - 'attachment; filename="balancing_auction_results_2025.csv"' + assert response.headers["content-type"] == "application/zip" + assert ( + response.headers["content-disposition"] + == 'attachment; filename="balancing_auction_results_2025.zip"' ) - # Header once, then all rows in order. - assert response.text.splitlines() == [ - "col1,col2", - "a,1", "b,2", - "c,3", "d,4", - "e,5", - ] + archive = zipfile.ZipFile(io.BytesIO(response.content)) + # `testzip` walks every member and checks its CRC — the framing has + # to be right, not merely parseable. + assert archive.testzip() is None + assert archive.namelist() == [ + "balancing_auction_results_2025_01.parquet", + "balancing_auction_results_2025_02.parquet", + ] + assert [archive.read(n) for n in archive.namelist()] == list(parts.values()) -def test_multi_shard_ndjson_pure_byte_concat(client: TestClient) -> None: - """Each NDJSON shard is self-contained; bytes concatenate cleanly.""" - shards = { - "url-1": b'{"id":1}\n{"id":2}\n', - "url-2": b'{"id":3}\n', - } - with _patch_dump(list(shards.keys())), _patch_httpx_stream(shards): - response = client.get( - DUMP_URL, params={"format": "ndjson"}, follow_redirects=False, - ) - assert response.status_code == 200 - assert response.headers["content-type"].startswith("application/x-ndjson") - assert response.text == ( - '{"id":1}\n{"id":2}\n' - '{"id":3}\n' - ) +def test_zip_members_are_stored_not_deflated(client: TestClient) -> None: + """Parquet is already compressed — deflating it would cost CPU for + no size win, so members go in uncompressed.""" + body = b"PAR1" + b"x" * 4096 + parts = {"https://signed/a.parquet": body, "https://signed/b.parquet": body} + stub_signed_urls(client, parts) -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, - ) + with _patch_dump(list(parts)): + response = client.get(DUMP_URL, params={"format": "parquet"}) - 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", - ] + archive = zipfile.ZipFile(io.BytesIO(response.content)) + for info in archive.infolist(): + assert info.compress_type == zipfile.ZIP_STORED + assert info.compress_size == len(body) # --- error paths ---------------------------------------------------------- -def test_too_large_parquet_returns_413(client: TestClient) -> None: - with _patch_dump(PayloadTooLargeError("exceeds 1 GB after parquet export")): - response = client.get(DUMP_URL, params={"format": "parquet"}) - assert response.status_code == 413 - assert response.json()["error"]["__type"] == "Payload Too Large" - - def test_unknown_format_returns_validation_error(client: TestClient) -> None: response = client.get(DUMP_URL, params={"format": "xml"}) assert response.status_code == 400 @@ -197,7 +171,7 @@ def test_build_export_select_iso_casts_timestamp_and_datetime() -> None: _bq_field("delivery_local", "DATETIME"), _bq_field("delivery_day", "DATE"), ] - select = _build_export_select(schema, fmt="csv") + select = _export_select_list(schema, fmt="csv") assert ( "FORMAT_TIMESTAMP('%Y-%m-%dT%H:%M:%S', `delivery_start`, 'UTC')" in select @@ -215,7 +189,7 @@ def test_build_export_select_iso_casts_timestamp_and_datetime() -> None: def test_build_export_select_parquet_returns_star() -> None: schema = [_bq_field("delivery_start", "TIMESTAMP")] - assert _build_export_select(schema, fmt="parquet") == "*" + assert _export_select_list(schema, fmt="parquet") == "*" def test_build_export_select_parquet_casts_json_columns() -> None: @@ -224,26 +198,12 @@ def test_build_export_select_parquet_casts_json_columns() -> None: _bq_field("bidder_metadata", "JSON"), _bq_field("delivery_start", "TIMESTAMP"), ] - assert _build_export_select(schema, fmt="parquet") == ( + assert _export_select_list(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( - schema, fmt="csv", - ) - - -def _bq_field(name: str, field_type: str) -> Any: - f = MagicMock() - f.name = name - f.field_type = field_type - return f - - # --- helpers: too-large heuristic ----------------------------------------- @@ -259,125 +219,6 @@ def test_unrelated_error_is_not_classified_as_too_large() -> None: assert _is_export_too_large(RuntimeError("auth failed")) is False -# --- helpers: CSV header-skip ------------------------------------------- - - -def test_skip_first_line_drops_header_and_forwards_rest() -> None: - """`_skip_first_line` strips up to and including the first `\\n`, - then byte-forwards everything else unchanged.""" - import asyncio - - async def chunks() -> AsyncIterator[bytes]: - yield b"col1,col2\n" - yield b"a,1\n" - yield b"b,2\n" - - async def run() -> bytes: - out = bytearray() - async for chunk in _skip_first_line(chunks()): - out.extend(chunk) - return bytes(out) - - assert asyncio.run(run()) == b"a,1\nb,2\n" - - -def test_skip_first_line_handles_header_split_across_chunks() -> None: - """The newline may not arrive in the first chunk — verify the - buffer accumulates until the newline is found.""" - import asyncio - - async def chunks() -> AsyncIterator[bytes]: - yield b"col1," # no newline yet - yield b"col2\n" - yield b"row,1\n" - - async def run() -> bytes: - out = bytearray() - async for chunk in _skip_first_line(chunks()): - out.extend(chunk) - return bytes(out) - - assert asyncio.run(run()) == b"row,1\n" - - -# --- engine: placeholder + bucket-missing guards ------------------------- - - -def test_dump_polling_releases_event_loop_between_reloads() -> None: - """Polling loop should `asyncio.sleep` between `job.reload` calls - so other coroutines on the same loop keep running. Verified by - interleaving a ticker — if the dump call hogged the loop, the - ticker would barely advance during the polls. - - The job is flagged with `error_result` once it transitions to DONE - so `dump()` raises immediately after the polling loop, without - reaching the post-extract GCS read.""" - import asyncio - - # `_engine_with_storage` stubs `google.cloud.storage` and gives the - # backend a real `table.modified` so the pre-extract cache lookup - # (empty here → cache miss → polling branch) doesn't blow up. - backend, storage_client = _engine_with_storage([]) - bucket_obj = storage_client.bucket.return_value - # Cache lookup returns no shards → fall through into the extract / - # poll branch. The post-extract retry is never reached because the - # job ends with an error. - bucket_obj.list_blobs.return_value = [] - - job = MagicMock() - job.state = "PENDING" - reload_calls = 0 - - def fake_reload() -> None: - nonlocal reload_calls - reload_calls += 1 - if reload_calls >= 3: - job.state = "DONE" - # Flag an error so dump raises right after the loop and - # doesn't try to read the GCS shard list. - job.error_result = {"message": "test-only error"} - - job.error_result = None - job.reload = fake_reload - backend.client.query.return_value = job - - # Speed up the test: 50 ms poll interval. - with patch( - "datastore.infrastructure.engines.bigquery.backend" - "._DUMP_POLL_INTERVAL_SECONDS", - 0.05, - ): - async def run() -> int: - ticks = 0 - - async def tick() -> None: - nonlocal ticks - while True: - await asyncio.sleep(0) - ticks += 1 - - ticker = asyncio.create_task(tick()) - try: - with pytest.raises(ServerError, match="test-only error"): - await backend.dump("res-1", "csv") - finally: - ticker.cancel() - try: - await ticker - except asyncio.CancelledError: - pass - return ticks - - ticks = asyncio.run(run()) - # 2 sleeps × 50 ms = 100 ms minimum of loop-yielding time; - # the ticker should rack up far more iterations than the - # number of reload calls if the loop is genuinely free. - assert ticks > reload_calls * 10, ( - f"event loop appears blocked during polling: " - f"only {ticks} ticker passes for {reload_calls} reloads" - ) - - # --- engine: GCS-backed cache by table.modified -------------------------- @@ -434,34 +275,43 @@ def _engine_with_storage(storage_blobs: list[Any]) -> tuple[BigQueryBackend, Any return backend, storage_client +def _bq_field(name: str, field_type: str) -> Any: + f = MagicMock() + f.name = name + f.field_type = field_type + return f + + def _fake_blob(name: str, signed_url: str = "https://signed/x") -> Any: + import datetime as _dt + blob = MagicMock() blob.name = name blob.generate_signed_url.return_value = signed_url + # A real datetime so the GC age gate can compare it. Fresh by + # default → spared; tests wanting it swept set an older value. + blob.time_created = _dt.datetime.now(_dt.timezone.utc) return blob -def test_dump_cache_hit_skips_extract_job() -> None: - """When GCS already has shards for this `(rid, fmt, table.modified)`, - `dump()` returns signed URLs straight from the cache — no - `client.query` call to BigQuery.""" +def _run_dump(backend: BigQueryBackend, resource_id: str, fmt: str) -> list[str]: + """Run `dump()` and flush its background cleanup (compose part + deletion, revision GC) so delete assertions are deterministic.""" import asyncio - blob = _fake_blob("dumps/res-1/csv/.csv", "https://cached") - backend, _ = _engine_with_storage([blob]) + async def go() -> list[str]: + urls = await backend.dump(resource_id, fmt) + if backend._cleanup_tasks: + await asyncio.gather(*list(backend._cleanup_tasks)) + return urls - urls = asyncio.run(backend.dump("res-1", "csv")) - - assert urls == ["https://cached"] - # No extract job submitted — that's the whole point of caching. - assert backend.client.query.call_count == 0 + return asyncio.run(go()) def test_dump_cache_miss_submits_extract_then_returns_urls() -> None: - """First call to `list_blobs` returns empty (cache miss); - `dump()` submits the extract, then `list_blobs` returns the - written shards on the post-extract retry.""" - import asyncio + """First call to `list_blobs` returns empty (cache miss); `dump()` + submits the extract, then the csv shards are composed into ONE object + and its single signed URL is returned.""" new_blob = _fake_blob("dumps/res-1/csv/_000.csv", "https://fresh") backend, storage_client = _engine_with_storage([]) @@ -469,62 +319,62 @@ def test_dump_cache_miss_submits_extract_then_returns_urls() -> None: # Pre-extract: empty. Post-extract refresh: one shard. GC sweep: # same one shard (nothing stale to delete on first dump ever). bucket_obj.list_blobs.side_effect = [[], [new_blob], [new_blob]] + # The composite object's signed URL (csv always composes). + bucket_obj.blob.return_value.generate_signed_url.return_value = ( + "https://composed" + ) # Job goes straight to DONE without errors. - job = MagicMock() - job.state = "DONE" - job.error_result = None + job = MagicMock() # `job.result()` returns without raising backend.client.query.return_value = job - urls = asyncio.run(backend.dump("res-1", "csv")) + urls = _run_dump(backend, "res-1", "csv") - assert urls == ["https://fresh"] + # csv shards are composed into one object → a single 302 URL. + assert urls == ["https://composed"] + bucket_obj.blob.return_value.compose.assert_called() # Exactly one extract submitted on cache miss. 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 +def test_dump_gzip_exports_headerless_and_composes() -> None: + """`format=gzip` exports header-less GZIP shards under its own cache + prefix; the composed object gets one gzip header member in front, so + the download is a plain 302 with no compression work in the pod.""" 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 + job = MagicMock() # `job.result()` returns without raising backend.client.query.return_value = job - urls = asyncio.run(backend.dump("res-1", "gzip")) + _run_dump(backend, "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 "header=false" in sql # header comes from the composed member assert "_*.csv.gz" in sql + prefixes = {c.kwargs["prefix"] for c in bucket_obj.list_blobs.call_args_list} + assert all(p.startswith("dumps/res-1/gzip/") for p in prefixes) 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 + job = MagicMock() # `job.result()` returns without raising backend.client.query.return_value = job - urls = asyncio.run(backend.dump("res-1", "parquet")) + urls = _run_dump(backend, "res-1", "parquet") assert urls == ["https://fresh"] sql = backend.client.query.call_args.args[0] @@ -532,85 +382,45 @@ def test_dump_parquet_export_uses_wildcard_uri() -> None: 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"), - ] +def test_dump_exports_into_an_attempt_and_publishes_success() -> None: + """The whole-table dump uses the same layout as the SQL dump: shards + under `//`, `_SUCCESS` written last.""" + shard = _fake_blob("dumps/res-1/parquet//att1/part_000.parquet", "https://p0") 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 - storage from growing unbounded across table updates. The current - revision's blobs stay.""" - import asyncio - import datetime as dt - + bucket_obj.list_blobs.side_effect = [[], [shard], [shard]] + backend.client.query.return_value = MagicMock() + + urls = _run_dump(backend, "res-1", "parquet") + + assert urls == ["https://p0"] + # The export URI is scoped to a private attempt dir under the shared + # revision prefix: `//part_*.ext`. + uri = backend.client.query.call_args.args[0].split("uri='")[1].split("'")[0] + path, _, filename = uri.removeprefix("gs://bkt/").rpartition("/") + assert filename == "part_*.parquet" + rev, attempt = path.removeprefix("dumps/res-1/parquet/").split("/") + assert rev and attempt and rev != attempt + # `_SUCCESS` published, and it is the only object created. + created = [c.args[0] for c in bucket_obj.blob.call_args_list] + assert [n.rsplit("/", 1)[1] for n in created] == ["_SUCCESS"] + + +def test_dump_ignores_an_attempt_without_success() -> None: + """A half-written attempt from a concurrent request is invisible to + the whole-table dump too — it re-exports rather than serving it.""" + inflight = _fake_blob("dumps/res-1/parquet//att1/part_000.parquet") + fresh = _fake_blob("dumps/res-1/parquet//att2/part_000.parquet", "https://ok") backend, storage_client = _engine_with_storage([]) bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[inflight], [fresh], []] + backend.client.query.return_value = MagicMock() - # Match the rev that backend.dump() computes from table.modified - # (`_engine_with_storage` sets it to 2026-01-01 UTC). - table_modified = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) - rev = f"{int(table_modified.timestamp() * 1_000_000):x}" - new_blob = _fake_blob( - f"dumps/res-1/csv/{rev}_000.csv", "https://fresh" - ) - old_blob_a = _fake_blob("dumps/res-1/csv/oldrev1_000.csv") - old_blob_b = _fake_blob("dumps/res-1/csv/oldrev2_000.csv") - - # Calls in order: - # 1) pre-extract cache lookup (prefix=dumps/.../) → empty - # 2) post-extract refresh (prefix=dumps/.../) → [new] - # 3) GC sweep (prefix=dumps/res-1/csv/) → [new, old_a, old_b] - bucket_obj.list_blobs.side_effect = [ - [], - [new_blob], - [new_blob, old_blob_a, old_blob_b], - ] - - job = MagicMock() - job.state = "DONE" - job.error_result = None - backend.client.query.return_value = job + urls = _run_dump(backend, "res-1", "parquet") - urls = asyncio.run(backend.dump("res-1", "csv")) - - assert urls == ["https://fresh"] - # The current revision must not be deleted. - assert new_blob.delete.call_count == 0 - # Both older revisions get cleaned up. - assert old_blob_a.delete.call_count == 1 - assert old_blob_b.delete.call_count == 1 - - -def test_dump_cache_hit_does_not_delete_anything() -> None: - """A cache hit must not trigger GC — there's no new revision to - supersede the existing one.""" - import asyncio - - cached = _fake_blob("dumps/res-1/csv/_000.csv", "https://cached") - backend, _ = _engine_with_storage([cached]) - - urls = asyncio.run(backend.dump("res-1", "csv")) - - assert urls == ["https://cached"] - # No extract → no GC. - assert cached.delete.call_count == 0 + assert urls == ["https://ok"] + assert backend.client.query.call_count == 1 # exported its own attempt + assert inflight.delete.call_count == 0 # never touched def test_dump_cache_key_changes_when_table_modified_advances() -> None: @@ -627,7 +437,7 @@ def test_dump_cache_key_changes_when_table_modified_advances() -> None: # Each dump-on-cache-hit lists twice: ro lookup + rw re-fetch for # signing. Both calls must return the same blob. table.modified = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) - first_old = _fake_blob("dumps/res-1/csv/old.csv", "https://old") + first_old = _fake_blob("dumps/res-1/csv/old.composed.csv", "https://old") bucket_obj.list_blobs.side_effect = [[first_old], [first_old]] asyncio.run(backend.dump("res-1", "csv")) # Both calls used the same prefix (the cache-hit rev) — take the @@ -636,7 +446,7 @@ def test_dump_cache_key_changes_when_table_modified_advances() -> None: # Bump the table; new call hits a different prefix. table.modified = dt.datetime(2026, 2, 1, tzinfo=dt.timezone.utc) - second_new = _fake_blob("dumps/res-1/csv/new.csv", "https://new") + second_new = _fake_blob("dumps/res-1/csv/new.composed.csv", "https://new") bucket_obj.list_blobs.side_effect = [[second_new], [second_new]] asyncio.run(backend.dump("res-1", "csv")) second_prefix = bucket_obj.list_blobs.call_args_list[-2].kwargs["prefix"] @@ -646,71 +456,4 @@ def test_dump_cache_key_changes_when_table_modified_advances() -> None: ) -def test_dump_returns_empty_list_in_placeholder_mode() -> None: - import asyncio - - backend = BigQueryBackend(mode="ro") - assert asyncio.run(backend.dump("res-1", "csv")) == [] - - -def test_dump_raises_when_export_bucket_unset() -> None: - import asyncio - - backend = BigQueryBackend(mode="ro") - backend.client = MagicMock() - backend.config = MagicMock() - backend.config.BIGQUERY_PROJECT = "proj-1" - backend.config.BIGQUERY_DATASET = "ds-1" - backend.config.BIGQUERY_EXPORT_BUCKET = "" - - with pytest.raises(ServerError, match="BIGQUERY_EXPORT_BUCKET"): - asyncio.run(backend.dump("res-1", "csv")) - - # --- test infrastructure -------------------------------------------------- - - -def _patch_httpx_stream(shards: dict[str, bytes]): - """Patch `httpx.AsyncClient.stream` so its async context-manager - returns a fake `Response` whose `aiter_bytes` walks the bytes for - that URL. Lets us drive the stream-concat helpers without hitting - the network.""" - - def make_resp(data: bytes) -> Any: - resp = MagicMock() - resp.raise_for_status = MagicMock(return_value=None) - - async def aiter_bytes(chunk_size: int = 64 * 1024): - # Walk the fixture bytes in `chunk_size` slices so callers - # see real chunk boundaries (matters for the header-skip - # path which may span chunks). - for i in range(0, len(data), chunk_size): - yield data[i:i + chunk_size] - - resp.aiter_bytes = aiter_bytes - return resp - - class FakeStreamCtx: - def __init__(self, url: str) -> None: - self._url = url - - async def __aenter__(self) -> Any: - return make_resp(shards[self._url]) - - async def __aexit__(self, *a: Any) -> None: - pass - - class FakeClient: - async def __aenter__(self) -> "FakeClient": - return self - - async def __aexit__(self, *a: Any) -> None: - pass - - def stream(self, method: str, url: str) -> FakeStreamCtx: - return FakeStreamCtx(url) - - return patch( - "datastore.services.dump.httpx.AsyncClient", - MagicMock(return_value=FakeClient()), - ) diff --git a/tests/test_datastore_dump_sql.py b/tests/test_datastore_dump_sql.py new file mode 100644 index 0000000..1a231a7 --- /dev/null +++ b/tests/test_datastore_dump_sql.py @@ -0,0 +1,1126 @@ +"""Engine-level tests for `BigQueryBackend.dump_sql` (search_sql downloads). + +Reuses `_engine_with_storage` / `_fake_blob` / `_bq_field` from +`tests.test_datastore_dump` — same stubbing pattern: `backend.client` is +a MagicMock, `_build_bq_client` / `_build_storage_client` are +monkeypatched on the instance, and `list_blobs` `side_effect` lists +drive cache hit / miss. + +Call order on a cache miss: + + client.query #1 — dry run (kwargs `job_config.dry_run is True`) + client.query #2 — EXPORT DATA (positional SQL) + list_blobs — pre-check (cacheable only) → post-export refresh + → GC sweep +""" + +from __future__ import annotations + +import asyncio +import datetime as dt +import hashlib +import io +import threading +import zipfile +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from datastore.core.exceptions import ( + NotFoundError, + PayloadTooLargeError, + ServerError, + ValidationError, +) +from datastore.infrastructure.engines.bigquery import BigQueryBackend +from datastore.infrastructure.engines.bigquery.export import _delete_old_cache +from datastore.infrastructure.engines.bigquery.lib import qualify_table_refs +from fastapi.testclient import TestClient + +from tests.conftest import FakeCKAN +from tests.test_datastore_dump import ( + _bq_field, + _engine_with_storage, + _fake_blob, + stub_signed_urls, +) + +DUMP_SQL_URL = "/datastore/dump/query" + +_NOW = dt.datetime.now(dt.timezone.utc) + + +def _dry_job(schema: list[Any] | None = None) -> Any: + """A dry-run query job: only `.schema` matters.""" + job = MagicMock() + job.schema = schema if schema is not None else [_bq_field("a", "INT64")] + return job + + +def _export_job() -> Any: + """An EXPORT DATA job whose `result()` returns without raising.""" + return MagicMock() + + +def _blob(name: str, url: str = "https://signed/x", age_hours: float = 0.0) -> Any: + """`_fake_blob` with a real `time_created` so the GC age gate can + compare it against a datetime cutoff.""" + blob = _fake_blob(name, url) + blob.time_created = _NOW - dt.timedelta(hours=age_hours) + return blob + + +def _run( + backend: BigQueryBackend, + sql: str = "SELECT * FROM res1", + fmt: str = "csv", + resource_ids: list[str] | None = None, + function_names: list[str] | None = None, +) -> list[str]: + async def go() -> list[str]: + urls = await backend.dump_sql( + sql, + fmt, + resource_ids=["res1"] if resource_ids is None else resource_ids, + function_names=function_names or [], + ) + # Part-deletion + revision GC run off the response path; flush + # them so delete assertions are deterministic. + if backend._cleanup_tasks: + await asyncio.gather(*list(backend._cleanup_tasks)) + return urls + + return asyncio.run(go()) + + +def _expected_prefix(sql: str, fmt: str = "csv") -> str: + """Recompute the cache prefix the engine derives for the fixture + config (project `proj-1`, dataset `ds-1`, single table `res1` + modified 2026-01-01 UTC).""" + qualified = qualify_table_refs(sql, project="proj-1", dataset="ds-1") + qhash = hashlib.sha256(qualified.encode()).hexdigest()[:16] + us = int( + dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc).timestamp() * 1_000_000 + ) + rev = hashlib.sha256(f"res1:{us}".encode()).hexdigest()[:16] + return f"dumps/{qhash}/{fmt}/{rev}/" + + +def _attempt(base: str, *files: str, attempt: str = "att1") -> list[Any]: + """One **published** attempt: the named files plus `_SUCCESS`.""" + return [ + *(_blob(f"{base}{attempt}/{f}", f"https://signed/{f}") for f in files), + _blob(f"{base}{attempt}/_SUCCESS"), + ] + + +# --- cache behaviour -------------------------------------------------------- + + +def test_cache_hit_skips_dry_run_and_export() -> None: + """Blobs already under the (qhash, fmt, rev) prefix → signed URLs + straight from GCS; zero BigQuery jobs.""" + backend, _ = _engine_with_storage(_attempt("dumps/h/csv/rev/", "data.csv")) + + urls = _run(backend) + + assert urls == ["https://signed/data.csv"] + assert backend.client.query.call_count == 0 + + +def test_cache_miss_dry_runs_then_exports() -> None: + new_blob = _blob("dumps/h/csv/rev_000.csv", "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]] + bucket_obj.blob.return_value.generate_signed_url.return_value = ( + "https://composed" + ) + backend.client.query.side_effect = [_dry_job(), _export_job()] + + urls = _run(backend) + + # csv shards compose into one object → single signed URL. + assert urls == ["https://composed"] + assert backend.client.query.call_count == 2 + # First call is the RO dry run; second is the EXPORT DATA statement. + dry_call = backend.client.query.call_args_list[0] + assert dry_call.kwargs["job_config"].dry_run is True + export_sql = backend.client.query.call_args_list[1].args[0] + assert "EXPORT DATA OPTIONS(" in export_sql + assert "dumps/" in export_sql + assert "overwrite=true" in export_sql + # The user SQL rides in subquery position; projection comes from the + # dry run's output schema; no `_id` ordering is injected. + assert "FROM (" in export_sql + assert "`a`" in export_sql + assert "ORDER BY" not in export_sql + + +def test_cache_prefix_matches_qhash_and_table_rev_scheme() -> None: + """Pre-check prefix is `dumps// + //`.""" + backend, storage_client = _engine_with_storage( + _attempt("dumps/h/csv/rev/", "data.csv") + ) + bucket_obj = storage_client.bucket.return_value + + sql = "SELECT * FROM res1" + _run(backend, sql=sql) + + pre = bucket_obj.list_blobs.call_args_list[0].kwargs["prefix"] + assert pre == _expected_prefix(sql) + + +def test_cache_prefix_stable_across_identical_calls() -> None: + backend, storage_client = _engine_with_storage( + _attempt("dumps/h/csv/rev/", "data.csv") + ) + bucket_obj = storage_client.bucket.return_value + + _run(backend) + _run(backend) + + prefixes = { + c.kwargs["prefix"] for c in bucket_obj.list_blobs.call_args_list + } + assert len(prefixes) == 1 + + +def test_rev_changes_when_any_referenced_table_changes() -> None: + """Multi-table SQL: bumping either table's `modified` produces a new + revision prefix; the other table alone can't satisfy the cache.""" + backend, storage_client = _engine_with_storage( + _attempt("dumps/h/csv/rev/", "data.csv") + ) + bucket_obj = storage_client.bucket.return_value + + t1, t2 = MagicMock(), MagicMock() + t1.modified = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) + t2.modified = dt.datetime(2026, 1, 2, tzinfo=dt.timezone.utc) + sql = "SELECT * FROM r1 JOIN r2 ON true" + + backend.client.get_table.side_effect = [t1, t2] + _run(backend, sql=sql, resource_ids=["r1", "r2"]) + first = bucket_obj.list_blobs.call_args_list[0].kwargs["prefix"] + + t2_bumped = MagicMock() + t2_bumped.modified = dt.datetime(2026, 3, 1, tzinfo=dt.timezone.utc) + backend.client.get_table.side_effect = [t1, t2_bumped] + _run(backend, sql=sql, resource_ids=["r1", "r2"]) + second = bucket_obj.list_blobs.call_args_list[-1].kwargs["prefix"] + + assert first != second + + +def test_non_deterministic_sql_bypasses_cache() -> None: + """`now()`-style SQL never reads the cache and writes under a fresh + uuid rev each run — two calls, two exports, two prefixes.""" + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + b1 = _blob("dumps/h/csv/u1_000.csv", "https://one") + b2 = _blob("dumps/h/csv/u2_000.csv", "https://two") + # No pre-check list per run — only post-export refresh + GC sweep. + bucket_obj.list_blobs.side_effect = [[b1], [b1], [b2], [b2]] + backend.client.query.side_effect = [ + _dry_job(), _export_job(), _dry_job(), _export_job(), + ] + + _run(backend, function_names=["now"]) + _run(backend, function_names=["now"]) + + assert backend.client.query.call_count == 4 + refresh_1 = bucket_obj.list_blobs.call_args_list[0].kwargs["prefix"] + refresh_2 = bucket_obj.list_blobs.call_args_list[2].kwargs["prefix"] + assert refresh_1 != refresh_2 + + +def test_table_modified_none_is_non_cacheable() -> None: + """A table without `modified` metadata can't participate in a stable + rev — treated like non-deterministic SQL (uuid rev per run).""" + backend, storage_client = _engine_with_storage([]) + backend.client.get_table.return_value.modified = None + bucket_obj = storage_client.bucket.return_value + b1 = _blob("y1", "https://one") + b2 = _blob("y2", "https://two") + bucket_obj.list_blobs.side_effect = [[b1], [b1], [b2], [b2]] + backend.client.query.side_effect = [ + _dry_job(), _export_job(), _dry_job(), _export_job(), + ] + + _run(backend) + _run(backend) + + assert backend.client.query.call_count == 4 + refresh_1 = bucket_obj.list_blobs.call_args_list[0].kwargs["prefix"] + refresh_2 = bucket_obj.list_blobs.call_args_list[2].kwargs["prefix"] + assert refresh_1 != refresh_2 + + +def test_gc_spares_young_attempts_and_reaps_old_revisions() -> None: + """The sweep is age-gated by the signed-URL lifetime: a young sibling + attempt may still be exporting or hold live URLs, so it survives; a + superseded revision past the expiry is reclaimed. This request's own + attempt always stays. + + Uses `parquet` so the compose path doesn't also delete shards.""" + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + + sql = "SELECT * FROM res1" + prefix = _expected_prefix(sql, fmt="parquet") + base = prefix.rstrip("/").rsplit("/", 1)[0] + current = _blob(f"{prefix}att1/part_000.parquet", "https://fresh") + old_rev = _blob(f"{base}/oldrev/att9/part_000.parquet", age_hours=3.0) + young_sibling = _blob(f"{prefix}att2/part_000.parquet", age_hours=0.0) + + bucket_obj.list_blobs.side_effect = [ + [], # pre-check (cache miss) + [current], # post-export refresh + [current, old_rev, young_sibling], # GC sweep + ] + backend.client.query.side_effect = [_dry_job(), _export_job()] + + urls = _run(backend, sql=sql, fmt="parquet") + + assert urls == ["https://fresh"] + assert current.delete.call_count == 0 # our own attempt + assert young_sibling.delete.call_count == 0 # may still be exporting + assert old_rev.delete.call_count == 1 # past the URL expiry + + +def test_gc_stale_blobs_no_age_gate_deletes_all_non_current() -> None: + """`_delete_old_cache(min_age=None)` (the cacheable path) deletes every + blob outside the current revision, ignoring age.""" + keep = "dumps/h/csv/current" + current = _blob(f"{keep}_000.csv") + old = _blob("dumps/h/csv/old_000.csv", age_hours=5.0) + young = _blob("dumps/h/csv/young_000.csv", age_hours=0.0) + rw_gcs = MagicMock() + rw_gcs.list_blobs.return_value = [current, old, young] + + deleted = _delete_old_cache( + rw_gcs, sweep_prefix="dumps/h/csv/", keep_prefix=keep, + min_age=None, + ) + + assert deleted == 2 + assert current.delete.call_count == 0 + assert old.delete.call_count == 1 + assert young.delete.call_count == 1 + + +def test_gc_stale_blobs_age_gate_keeps_young() -> None: + """`_delete_old_cache(min_age=…)` (the non-cacheable path) deletes only + superseded blobs older than the cutoff; younger ones (whose signed + URLs may still be live) survive.""" + keep = "dumps/h/csv/current" + current = _blob(f"{keep}_000.csv") + old = _blob("dumps/h/csv/old_000.csv", age_hours=5.0) + young = _blob("dumps/h/csv/young_000.csv", age_hours=0.0) + rw_gcs = MagicMock() + rw_gcs.list_blobs.return_value = [current, old, young] + + deleted = _delete_old_cache( + rw_gcs, sweep_prefix="dumps/h/csv/", keep_prefix=keep, + min_age=dt.timedelta(hours=1), + ) + + assert deleted == 1 + assert current.delete.call_count == 0 + assert old.delete.call_count == 1 + assert young.delete.call_count == 0 + + +# --- output ordering --------------------------------------------------------- + + +def test_export_orders_by_id_when_output_carries_it() -> None: + """No user ORDER BY + `_id` in the output → outer `ORDER BY _id`, + matching the JSON path's default ordering (a subquery's ORDER BY + without LIMIT is ignored by BigQuery, so it must sit outside).""" + new_blob = _blob("o_000.csv", "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]] + backend.client.query.side_effect = [ + _dry_job([_bq_field("_id", "INT64"), _bq_field("a", "INT64")]), + _export_job(), + ] + + _run(backend, sql="SELECT * FROM res1") + + export_sql = backend.client.query.call_args_list[1].args[0] + assert export_sql.endswith(") ORDER BY `_id`") + + +def test_export_hoists_user_order_by_to_outer_query() -> None: + """A user ORDER BY on output columns is copied to the outer query so + the file (and shard concat) actually follows it.""" + new_blob = _blob("o_000.csv", "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]] + backend.client.query.side_effect = [ + _dry_job([_bq_field("a", "INT64")]), + _export_job(), + ] + + _run(backend, sql="SELECT a FROM res1 ORDER BY a DESC") + + export_sql = backend.client.query.call_args_list[1].args[0] + # The outer (post-subquery) ORDER BY carries the user's key; sqlglot + # renders the postgres NULLS ordering explicitly for BigQuery. + outer = export_sql.rsplit(")", 1)[1] + assert "ORDER BY `a` DESC" in outer + + +def test_export_skips_outer_order_for_non_output_sort_keys() -> None: + """ORDER BY on a column that isn't in the output can't be hoisted + past the subquery boundary — the export stays unordered rather than + failing.""" + new_blob = _blob("o_000.csv", "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]] + backend.client.query.side_effect = [ + _dry_job([_bq_field("a", "INT64")]), + _export_job(), + ] + + _run(backend, sql="SELECT a FROM res1 ORDER BY b") + + export_sql = backend.client.query.call_args_list[1].args[0] + # The inner ORDER BY b remains inside the parens; no outer ORDER BY. + assert export_sql.rstrip().endswith(")") + + +# --- per-format export SQL -------------------------------------------------- + + +def test_parquet_export_casts_json_columns_from_dry_run_schema() -> None: + """BigQuery can't export native JSON to parquet — the dry run's + output schema drives a TO_JSON_STRING cast.""" + new_blob = _blob("z_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]] + backend.client.query.side_effect = [ + _dry_job([_bq_field("id", "INT64"), _bq_field("meta", "JSON")]), + _export_job(), + ] + + _run(backend, fmt="parquet") + + export_sql = backend.client.query.call_args_list[1].args[0] + assert "format='PARQUET'" in export_sql + assert "TO_JSON_STRING(`meta`)" in export_sql + + +def test_csv_export_iso_casts_timestamps_from_dry_run_schema() -> None: + """CSV downloads render TIMESTAMP identically to `datastore_search` + and `/datastore/dump` (shared `format_select_column`).""" + new_blob = _blob("z_000.csv", "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]] + backend.client.query.side_effect = [ + _dry_job([_bq_field("ts", "TIMESTAMP")]), + _export_job(), + ] + + _run(backend, fmt="csv") + + export_sql = backend.client.query.call_args_list[1].args[0] + assert "format='CSV'" in export_sql + # header=false: the single header is composed in as a separate member. + assert "header=false" in export_sql + assert ( + "FORMAT_TIMESTAMP('%Y-%m-%dT%H:%M:%S', `ts`, 'UTC')" in export_sql + ) + + +def test_multi_shard_parquet_returns_every_url() -> None: + """Parquet isn't composable, so all shard URLs come back.""" + shards = [ + _blob("p_000.parquet", "https://p0"), + _blob("p_001.parquet", "https://p1"), + ] + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], shards, shards] + backend.client.query.side_effect = [_dry_job(), _export_job()] + + assert _run(backend, fmt="parquet") == ["https://p0", "https://p1"] + + +# --- compose (single-URL) ---------------------------------------------------- + + +def _distinct_blob_factory() -> tuple[Any, dict[str, Any]]: + """`bucket.blob(name)` side_effect returning a distinct mock per name + (each with a `url:` signed URL). Real GCS gives distinct objects; + the default MagicMock returns one shared blob, which hides compose.""" + made: dict[str, Any] = {} + + def factory(name: str) -> Any: + m = made.get(name) + if m is None: + m = MagicMock() + m.name = name + m.generate_signed_url.return_value = f"url:{name}" + made[name] = m + return m + + return factory, made + + +def test_csv_multishard_composes_header_and_shards_into_one_url() -> None: + """csv: a header member + the header-less shards are composed into ONE + object; the parts are deleted; a single composite URL comes back.""" + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + shard0 = _blob("dumps/h/csv/rev/att1/part_000.csv") + shard1 = _blob("dumps/h/csv/rev/att1/part_001.csv") + bucket_obj.list_blobs.side_effect = [[], [shard0, shard1], [shard0, shard1]] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="csv") + + composite = next(n for n in made if n.endswith("data.csv")) + header = next(n for n in made if n.endswith("header.csv")) + assert urls == [f"url:{composite}"] # single URL + made[header].upload_from_string.assert_called_once() # header written + made[composite].compose.assert_called_once() + sources = made[composite].compose.call_args.args[0] + assert made[header] is sources[0] # header sorts first + assert shard0 in sources and shard1 in sources + assert shard0.delete.called and shard1.delete.called # parts removed + + +def test_ndjson_multishard_composes_without_header() -> None: + """ndjson: shards compose directly (no header member) into one URL.""" + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + shard0 = _blob("dumps/h/ndjson/rev/att1/part_000.json") + shard1 = _blob("dumps/h/ndjson/rev/att1/part_001.json") + bucket_obj.list_blobs.side_effect = [[], [shard0, shard1], [shard0, shard1]] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="ndjson") + + assert not any(n.endswith("header.json") for n in made) # no header + composite = next(n for n in made if n.endswith("data.json")) + assert urls == [f"url:{composite}"] + sources = made[composite].compose.call_args.args[0] + assert shard0 in sources and shard1 in sources + + +def test_ndjson_single_shard_is_not_composed() -> None: + """A lone ndjson shard is already one URL — no compose, 302 to it.""" + new_blob = _blob("dumps/h/ndjson/rev/att1/part_000.json", "https://one") + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], [new_blob], [new_blob]] + backend.client.query.side_effect = [_dry_job(), _export_job()] + + urls = _run(backend, fmt="ndjson") + + assert urls == ["https://one"] + # Only the `_SUCCESS` marker is created — no composite, no header. + created = [c.args[0] for c in bucket_obj.blob.call_args_list] + assert [n.rsplit("/", 1)[1] for n in created] == ["_SUCCESS"] + + +def test_compose_chains_beyond_32_sources() -> None: + """>32 shards → compose is chained in batches of ≤32 (the composite + itself counts as one source in later calls).""" + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + shards = [_blob(f"dumps/h/ndjson/rev/att1/part_{i:03d}.json") for i in range(40)] + bucket_obj.list_blobs.side_effect = [[], shards, shards] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + _run(backend, fmt="ndjson") + + composite = made[next(n for n in made if n.endswith("data.json"))] + # 40 sources → compose(32), then compose(composite + 8) = 2 calls. + assert composite.compose.call_count == 2 + first = composite.compose.call_args_list[0].args[0] + second = composite.compose.call_args_list[1].args[0] + assert len(first) == 32 + assert second[0] is composite and len(second) == 1 + 8 + + +def test_attempt_without_success_is_invisible() -> None: + """An attempt mid-export (or one that died) has no `_SUCCESS`, so it + is never served — this request exports its own attempt instead, and + never deletes the other one.""" + inflight = [ + _blob("dumps/h/csv/rev/att1/part_000.csv"), + _blob("dumps/h/csv/rev/att1/part_001.csv"), + ] + fresh = _blob("dumps/h/csv/rev/att2/part_000.csv") + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [inflight, [fresh], []] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="csv") + + assert backend.client.query.call_count == 2 # exported its own + composite = next(n for n in made if n.endswith("data.csv")) + assert urls == [f"url:{composite}"] + for b in inflight: + assert b.delete.call_count == 0 # left alone + + +def test_published_attempt_wins_over_an_inflight_one() -> None: + """A finished attempt is served even while a sibling is still + exporting — the in-flight objects are never mixed in.""" + done = _attempt("dumps/h/csv/rev/", "data.csv") + inflight = _blob("dumps/h/csv/rev/att2/part_000.csv", "https://partial") + backend, _ = _engine_with_storage([*done, inflight]) + + urls = _run(backend) + + assert urls == ["https://signed/data.csv"] + assert backend.client.query.call_count == 0 + + +def test_success_marker_is_never_composed_or_signed() -> None: + """`_SUCCESS` is written after the compose returns, and compose takes + this attempt's explicit shard list — never a fresh listing. It also + never reaches the served set. A stray member would break gzip + outright, not merely garble it.""" + shards = [ + _blob("dumps/h/csv/rev/att1/part_000.csv"), + _blob("dumps/h/csv/rev/att1/part_001.csv"), + ] + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], shards, shards] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="csv") + + composite = made[next(n for n in made if n.endswith("data.csv"))] + sources = composite.compose.call_args.args[0] + assert not any("_SUCCESS" in getattr(s, "name", "") for s in sources) + assert not any("_SUCCESS" in u for u in urls) + made[next(n for n in made if n.endswith("_SUCCESS"))]\ + .upload_from_string.assert_called_once_with(b"") + + +def test_compose_deletes_shards_and_header_but_not_the_marker() -> None: + """After compose the shards and the synthesized header are garbage + and go in the background. `_SUCCESS` must survive — deleting it would + silently un-publish the attempt, and the next request would re-export.""" + shards = [ + _blob("dumps/h/csv/rev/att1/part_000.csv"), + _blob("dumps/h/csv/rev/att1/part_001.csv"), + ] + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], shards, shards] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + _run(backend, fmt="csv") + + for shard in shards: + assert shard.delete.call_count == 1 + assert made[next(n for n in made if n.endswith("header.csv"))].delete.called + marker = made[next(n for n in made if n.endswith("_SUCCESS"))] + assert marker.delete.call_count == 0 + assert not made[next(n for n in made if n.endswith("data.csv"))].delete.called + + +def test_response_does_not_wait_for_shard_deletion() -> None: + """Deleting the compose sources is fire-and-forget: the caller gets + its URLs while the deletes are still in flight. + + The gate keeps every delete blocked until after the engine has + returned — so if the deletion were ever awaited on the response path, + this would block rather than assert. + """ + gate = threading.Event() + shards = [ + _blob("dumps/h/csv/rev/att1/part_000.csv"), + _blob("dumps/h/csv/rev/att1/part_001.csv"), + ] + for shard in shards: + shard.delete.side_effect = lambda: gate.wait(5) + + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[], shards, shards] + backend.client.query.side_effect = [_dry_job(), _export_job()] + + async def go() -> list: + urls = await backend.dump_sql( + "SELECT * FROM res1", "csv", resource_ids=["res1"], function_names=[] + ) + pending = [t for t in backend._cleanup_tasks if not t.done()] + gate.set() + await asyncio.gather(*backend._cleanup_tasks) + return [urls, pending] + + urls, pending = asyncio.run(go()) + + assert urls # served without waiting + assert pending # …while cleanup was still running + for shard in shards: + assert shard.delete.call_count == 1 + + +def test_leftover_shards_are_not_served_beside_the_composite() -> None: + """A published attempt still holds its pre-compose shards until the + background delete runs; only the composed object is handed out.""" + base = "dumps/h/csv/rev/att1/" + backend, _ = _engine_with_storage([ + _blob(f"{base}data.csv", "https://composed"), + _blob(f"{base}part_000.csv", "https://part0"), + _blob(f"{base}header.csv", "https://header"), + _blob(f"{base}_SUCCESS"), + ]) + + urls = _run(backend) + + assert urls == ["https://composed"] + + +# --- cache-hit validation (self-healing) -------------------------------------- + + +def test_legacy_multishard_csv_hit_is_rebuilt_to_composite() -> None: + """A pre-compose cache (multiple raw csv shards, no composite) must + not be served as a stream: the hit is invalidated, the stale blobs + are deleted, and the export+compose reruns → single 302 URL.""" + legacy = [ + _blob("dumps/h/csv/rev/att1/part_000.csv"), + _blob("dumps/h/csv/rev/att1/part_001.csv"), + ] + fresh = _blob("dumps/h/csv/rev_000.csv", "https://fresh-shard") + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [ + legacy, # pre-check (ro): invalid hit (no composite) + legacy, # re-list (rw) for the clear + [fresh], # post-export refresh + [], # GC sweep + ] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="csv") + + for b in legacy: + assert b.delete.called # stale cache cleared + assert backend.client.query.call_count == 2 # re-exported + composite = next(n for n in made if n.endswith("data.csv")) + assert urls == [f"url:{composite}"] # single URL again + + +def test_lone_raw_csv_shard_hit_is_rebuilt() -> None: + """One raw (non-composed) csv shard = a run that died between export + and compose; it is header-less, so serving it would drop the header. + The hit is invalidated and rebuilt.""" + raw = _blob("dumps/h/csv/rev/att1/part_000.csv") + fresh = _blob("dumps/h/csv/rev_000.csv", "https://fresh-shard") + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [[raw], [raw], [fresh], []] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="csv") + + assert raw.delete.called + composite = next(n for n in made if n.endswith("data.csv")) + assert urls == [f"url:{composite}"] + + +def test_ndjson_single_shard_hit_is_valid() -> None: + """ndjson needs no header member, so a published single shard is a + perfectly good hit — served directly, no `data.json`.""" + backend, _ = _engine_with_storage( + _attempt("dumps/h/ndjson/rev/", "part_000.json") + ) + + urls = _run(backend, fmt="ndjson") + + assert urls == ["https://signed/part_000.json"] + assert backend.client.query.call_count == 0 + + +def test_ndjson_multishard_hit_is_rebuilt() -> None: + """Multiple ndjson blobs on a hit (legacy cache) → invalidated and + recomposed into one URL.""" + legacy = [ + _blob("dumps/h/ndjson/rev/att1/part_000.json"), + _blob("dumps/h/ndjson/rev/att1/part_001.json"), + ] + fresh = [ + _blob("dumps/h/ndjson/rev/att1/part_000.json"), + _blob("dumps/h/ndjson/rev/att1/part_001.json"), + ] + backend, storage_client = _engine_with_storage([]) + bucket_obj = storage_client.bucket.return_value + bucket_obj.list_blobs.side_effect = [legacy, legacy, fresh, []] + backend.client.query.side_effect = [_dry_job(), _export_job()] + factory, made = _distinct_blob_factory() + bucket_obj.blob.side_effect = factory + + urls = _run(backend, fmt="ndjson") + + for b in legacy: + assert b.delete.called + composite = next(n for n in made if n.endswith("data.json")) + assert urls == [f"url:{composite}"] + + +# --- guards & error mapping -------------------------------------------------- + + +def test_placeholder_mode_returns_empty_list() -> None: + backend = BigQueryBackend(mode="ro") + urls = asyncio.run(backend.dump_sql( + "SELECT 1", "csv", resource_ids=[], function_names=[], + )) + assert urls == [] + + +def test_non_ro_mode_rejected() -> None: + backend = BigQueryBackend(mode="rw") + backend.client = MagicMock() + with pytest.raises(ServerError, match="read-only"): + asyncio.run(backend.dump_sql( + "SELECT 1", "csv", resource_ids=[], function_names=[], + )) + + +def test_bucket_unset_raises_server_error() -> None: + backend = BigQueryBackend(mode="ro") + backend.client = MagicMock() + backend.config = MagicMock() + backend.config.BIGQUERY_EXPORT_BUCKET = "" + with pytest.raises(ServerError, match="BIGQUERY_EXPORT_BUCKET"): + asyncio.run(backend.dump_sql( + "SELECT 1", "csv", resource_ids=[], function_names=[], + )) + + +def test_missing_table_raises_not_found() -> None: + from google.api_core.exceptions import NotFound + + backend, _ = _engine_with_storage([]) + backend.client.get_table.side_effect = NotFound("gone") + with pytest.raises(NotFoundError, match="res1"): + _run(backend) + + +def test_export_job_too_large_maps_to_413() -> None: + """A ">1 GB single URI" failure raised by the export job surfaces as + 413, not a generic 500.""" + backend, storage_client = _engine_with_storage([]) + storage_client.bucket.return_value.list_blobs.side_effect = [[]] + failing = MagicMock() + failing.result.side_effect = RuntimeError( + "Cannot export more than 1 GB to a single URI; use the wildcard" + ) + backend.client.query.side_effect = [_dry_job(), failing] + + with pytest.raises(PayloadTooLargeError, match="exceeds 1 GB"): + _run(backend, fmt="parquet") + + +def test_export_job_failure_maps_to_server_error() -> None: + """Any other export failure is a 500 carrying the upstream message.""" + backend, storage_client = _engine_with_storage([]) + storage_client.bucket.return_value.list_blobs.side_effect = [[]] + failing = MagicMock() + failing.result.side_effect = RuntimeError("quota exceeded") + backend.client.query.side_effect = [_dry_job(), failing] + + with pytest.raises(ServerError, match="quota exceeded"): + _run(backend) + + +def test_dry_run_failure_maps_to_validation_error() -> None: + """SQL that doesn't compile against the real schema fails the RO dry + run → clean 400, before any RW/export work.""" + backend, storage_client = _engine_with_storage([]) + storage_client.bucket.return_value.list_blobs.side_effect = [[]] + backend.client.query.side_effect = RuntimeError("Unrecognized name: nope") + with pytest.raises(ValidationError, match="BigQuery validation"): + _run(backend) + + +def test_zero_table_sql_exports_without_get_table() -> None: + """`SELECT 1`-style queries reference no tables: no `get_table` + calls, stable empty-pairs rev, export still runs. (parquet → no + compose, so the single shard's URL comes straight back.)""" + new_blob = _blob("s_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]] + backend.client.query.side_effect = [_dry_job(), _export_job()] + + urls = _run(backend, sql="SELECT 1", resource_ids=[], fmt="parquet") + + assert urls == ["https://fresh"] + assert backend.client.get_table.call_count == 0 + + +# ============================================================================= +# Endpoint: GET /datastore/dump/query +# ============================================================================= + + +def _patch_dump_sql(urls_or_exc: list[str] | Exception): + """Patch `BigQueryBackend.dump_sql` to return URLs or raise. + + Returns `(patcher, calls)` — `calls` collects the kwargs of every + invocation so tests can assert the endpoint forwards the parsed SQL + facts verbatim.""" + calls: list[dict[str, Any]] = [] + + async def fake( + self: BigQueryBackend, + sql: str, + fmt: str, + *, + resource_ids: list[str], + function_names: list[str], + ) -> list[str]: + calls.append({ + "sql": sql, + "fmt": fmt, + "resource_ids": resource_ids, + "function_names": function_names, + }) + if isinstance(urls_or_exc, Exception): + raise urls_or_exc + return urls_or_exc + + return patch.object(BigQueryBackend, "dump_sql", fake), calls + + +def test_forwards_sql_and_names_verbatim( + client: TestClient, fake_ckan: FakeCKAN, +) -> None: + """The engine receives the SQL exactly as sent (LIMIT intact) plus + the schema-parsed table / function names.""" + patcher, calls = _patch_dump_sql(["https://x/1.json"]) + sql = 'SELECT count(*) FROM "balancing_auction_results_2025" LIMIT 7' + with patcher: + response = client.get( + DUMP_SQL_URL, + params={"sql": sql, "format": "ndjson"}, + follow_redirects=False, + ) + assert response.status_code == 302 + assert calls == [{ + "sql": sql, + "fmt": "ndjson", + "resource_ids": ["balancing_auction_results_2025"], + "function_names": ["count"], + }] + + +def test_format_defaults_to_csv(client: TestClient) -> None: + patcher, calls = _patch_dump_sql(["https://x/1.csv"]) + with patcher: + response = client.get( + DUMP_SQL_URL, + params={"sql": "SELECT 1 LIMIT 10"}, + follow_redirects=False, + ) + assert response.status_code == 302 + assert calls[0]["fmt"] == "csv" + + +def test_limit_is_optional(client: TestClient) -> None: + """No LIMIT → the full result set is exported.""" + patcher, _ = _patch_dump_sql(["https://x/1.csv"]) + with patcher: + response = client.get( + DUMP_SQL_URL, + params={"sql": "SELECT 1"}, + follow_redirects=False, + ) + assert response.status_code == 302 + + +def test_limit_above_search_cap_allowed(client: TestClient) -> None: + """The SEARCH_RESULT_ROWS_MAX clamp is a JSON-envelope rule; a dump + honors the SQL's LIMIT as written.""" + patcher, calls = _patch_dump_sql(["https://x/1.csv"]) + with patcher: + response = client.get( + DUMP_SQL_URL, + params={"sql": "SELECT 1 LIMIT 50000"}, + follow_redirects=False, + ) + assert response.status_code == 302 + assert "LIMIT 50000" in calls[0]["sql"] + + +def test_offset_without_limit_rejected(client: TestClient) -> None: + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT 1 OFFSET 10", + }) + assert response.status_code == 400 + body = response.json() + assert body["error"]["__type"] == "Validation Error" + assert "OFFSET" in body["error"]["message"] + + +def test_bogus_format_rejected(client: TestClient) -> None: + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT 1 LIMIT 5", "format": "xml", + }) + assert response.status_code == 400 + assert response.json()["error"]["__type"] == "Validation Error" + + +def test_missing_sql_names_the_field(client: TestClient) -> None: + """`/datastore/dump/query` resolves to the SQL route (declared before + `/{resource_id}`, so `query` is a reserved resource name) — a missing + `sql` param is a validation error on this endpoint, not a 404 dump + of a table called 'query'.""" + response = client.get(DUMP_SQL_URL) + assert response.status_code == 400 + body = response.json() + assert body["error"]["__type"] == "Validation Error" + assert "sql" in body["error"]["fields"] + + +def test_every_single_file_format_redirects(client: TestClient) -> None: + """Composed to one object by the engine → always a 302; the pod + never touches the bytes for any format.""" + for fmt in ("csv", "gzip", "ndjson", "parquet"): + patcher, _ = _patch_dump_sql([f"https://signed/one.{fmt}"]) + with patcher: + response = client.get( + DUMP_SQL_URL, + params={"sql": "SELECT 1 LIMIT 5", "format": fmt}, + follow_redirects=False, + ) + assert response.status_code == 302, fmt + + +def test_multi_file_parquet_streams_one_zip(client: TestClient) -> None: + """A sharded parquet export is streamed back as a single zip — not a + URL list, and not a 413. Members are named after the `query` base.""" + parts = { + "https://signed/p0.parquet": b"PAR1-shard-zero", + "https://signed/p1.parquet": b"PAR1-shard-one", + } + stub_signed_urls(client, parts) + + patcher, _ = _patch_dump_sql(list(parts)) + with patcher: + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT 1 LIMIT 5", "format": "parquet", + }) + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/zip" + assert ( + response.headers["content-disposition"] + == 'attachment; filename="query.zip"' + ) + + archive = zipfile.ZipFile(io.BytesIO(response.content)) + assert archive.testzip() is None + assert archive.namelist() == ["query_01.parquet", "query_02.parquet"] + assert [archive.read(n) for n in archive.namelist()] == list(parts.values()) + + +def test_payload_too_large_error_maps_to_413(client: TestClient) -> None: + patcher, _ = _patch_dump_sql( + PayloadTooLargeError("exported as multiple parquet shards"), + ) + with patcher: + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT 1 LIMIT 5", "format": "parquet", + }) + assert response.status_code == 413 + assert response.json()["error"]["__type"] == "Payload Too Large" + + +def test_disallowed_function_rejected_before_engine( + client: TestClient, +) -> None: + """The function allow-list applies here too — the service raises + before any engine/export work.""" + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT pg_read_file('/etc/passwd') LIMIT 1", + }) + assert response.status_code == 400 + assert "pg_read_file" in response.json()["error"]["message"].lower() + + +def test_unknown_table_returns_404( + client: TestClient, fake_ckan: FakeCKAN, +) -> None: + response = client.get(DUMP_SQL_URL, params={ + "sql": 'SELECT * FROM "does-not-exist" LIMIT 10', + }) + assert response.status_code == 404 + + +def test_denied_key_returns_403( + client: TestClient, fake_ckan: FakeCKAN, +) -> None: + fake_ckan.deny("test-token") + response = client.get(DUMP_SQL_URL, params={ + "sql": 'SELECT * FROM "balancing_auction_results_2025" LIMIT 10', + }) + assert response.status_code == 403 + + +def test_join_authorizes_each_table( + client: TestClient, fake_ckan: FakeCKAN, +) -> None: + fake_ckan.add_resource("other_table", package_id="pkg-balancing-2025") + before = fake_ckan.authorize_calls + patcher, _ = _patch_dump_sql(["https://x/1.csv"]) + with patcher: + response = client.get( + DUMP_SQL_URL, + params={ + "sql": ( + 'SELECT a.id FROM "balancing_auction_results_2025" a ' + 'JOIN "other_table" b ON a.id = b.id LIMIT 10' + ), + }, + follow_redirects=False, + ) + assert response.status_code == 302 + assert fake_ckan.authorize_calls - before == 2 + + +def test_unconfigured_engine_returns_500(client: TestClient) -> None: + """Placeholder engine (no BQ creds in the test env) exports nothing; + the endpoint refuses to serve an empty file and 500s explicitly.""" + response = client.get(DUMP_SQL_URL, params={ + "sql": "SELECT 1 LIMIT 5", + }) + assert response.status_code == 500 + assert "not configured" in response.json()["error"]["message"] diff --git a/tests/test_datastore_search_sql.py b/tests/test_datastore_search_sql.py index d21f115..f9a0202 100644 --- a/tests/test_datastore_search_sql.py +++ b/tests/test_datastore_search_sql.py @@ -21,7 +21,10 @@ from __future__ import annotations import pytest -from datastore.schemas.validators import parse_sql_references +from datastore.schemas.validators import ( + parse_sql_pagination, + parse_sql_references, +) from fastapi.testclient import TestClient from tests.conftest import FakeCKAN @@ -326,3 +329,40 @@ def test_each_table_authorized_once_for_joins( }) assert response.status_code == 200 assert fake_ckan.authorize_calls - before == 2 + + +# 8. Download param moved to /datastore/dump/query -------------------------- + +def test_download_param_no_longer_accepted(client: TestClient) -> None: + """SQL downloads live at `GET /datastore/dump/query?sql=&format=` now; + `extra="forbid"` rejects the retired `download` param here.""" + response = client.get(SQL_URL, params={ + "sql": "SELECT 1 LIMIT 10", "download": "csv", + }) + assert response.status_code == 400 + assert response.json()["error"]["__type"] == "Validation Error" + + +# 9. parse_sql_pagination (unit) --------------------------------------------- + + +def test_parse_sql_pagination_requires_limit_by_default() -> None: + with pytest.raises(ValueError, match="LIMIT"): + parse_sql_pagination("SELECT 1") + + +def test_parse_sql_pagination_optional_limit_absent() -> None: + assert parse_sql_pagination("SELECT 1", require_limit=False) == (None, 0) + + +def test_parse_sql_pagination_optional_limit_present() -> None: + """A LIMIT/OFFSET present in the SQL is honored as written even when + not required.""" + assert parse_sql_pagination( + "SELECT 1 LIMIT 5 OFFSET 2", require_limit=False, + ) == (5, 2) + + +def test_parse_sql_pagination_offset_without_limit_rejected() -> None: + with pytest.raises(ValueError, match="OFFSET without LIMIT"): + parse_sql_pagination("SELECT 1 OFFSET 10", require_limit=False) diff --git a/tests/test_read_service.py b/tests/test_read_service.py index 105111e..5d70e15 100644 --- a/tests/test_read_service.py +++ b/tests/test_read_service.py @@ -18,9 +18,17 @@ from collections.abc import Iterator from types import SimpleNamespace from typing import Any +from unittest.mock import patch +import pytest from datastore.core.config import Config -from datastore.services.read import _build_pagination_links, search_datastore +from datastore.core.exceptions import ValidationError +from datastore.infrastructure.engines.bigquery import BigQueryBackend +from datastore.services.read import ( + _build_pagination_links, + dump_sql_datastore, + search_datastore, +) def _ctx() -> SimpleNamespace: @@ -302,3 +310,56 @@ def test_links_prev_clamps_to_zero_on_partial_first_page() -> None: "/path", limit=50, offset=20, total=100, ) assert "offset=0" in links["prev"] + + +# --- dump_sql_datastore ---------------------------------------------------- + + +def test_dump_sql_service_rejects_disallowed_functions() -> None: + """The same allow-list gate as `search_sql_datastore` runs before + any engine work in download mode.""" + with pytest.raises(ValidationError, match="pg_read_file"): + asyncio.run(dump_sql_datastore(_ctx(), { + "sql": "SELECT pg_read_file('/x') LIMIT 1", + "fmt": "csv", + "resource_ids": [], + "function_names": ["pg_read_file"], + })) + + +def test_dump_sql_service_forwards_args_and_returns_urls() -> None: + """No LIMIT clamp, no reshaping — the service forwards the parsed + SQL facts to the engine and returns its signed URLs as-is.""" + captured: dict[str, Any] = {} + + async def fake( + self: BigQueryBackend, + sql: str, + fmt: str, + *, + resource_ids: list[str], + function_names: list[str], + ) -> list[str]: + captured.update( + sql=sql, + fmt=fmt, + resource_ids=resource_ids, + function_names=function_names, + ) + return ["https://signed/1"] + + with patch.object(BigQueryBackend, "dump_sql", fake): + urls = asyncio.run(dump_sql_datastore(_ctx(), { + "sql": "SELECT * FROM r1 LIMIT 50000", + "fmt": "ndjson", + "resource_ids": ["r1"], + "function_names": ["count"], + })) + + assert urls == ["https://signed/1"] + assert captured == { + "sql": "SELECT * FROM r1 LIMIT 50000", + "fmt": "ndjson", + "resource_ids": ["r1"], + "function_names": ["count"], + }