-
Notifications
You must be signed in to change notification settings - Fork 0
Add SQL download endpoint; compose downloads to a single file, zip sharded parquet #14
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
f8acdf9
Add SQL query result download endpoint (/datastore/dump/sql)
sagargg c6e925b
Move the dump pipeline into bigquery/export.py; rename /dump/sql to /…
sagargg 641c543
Stream a zip when a parquet export shards into multiple files
sagargg ae56c33
Fix ruff import order in test_cors.py
sagargg f0e98e5
Isolate each export in its own attempt directory, published by _SUCCESS
sagargg File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should probably remove this for the dumps query, they are either redirects or in the one case they are not (parquet) they are already zipped, could make things faster
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@luccasmmg I didn't quite understand you comment. this streaming response only needed for Parquet files when the export is split into more than one sharded Parquet file?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nevermind