Add SQL download endpoint; compose downloads to a single file, zip sharded parquet - #14
Conversation
GET /datastore/dump/sql?sql=<SELECT>&format=csv|gzip|ndjson|parquet exports
the result of a vetted read-only SELECT as a file through the BigQuery
EXPORT DATA -> GCS pipeline, reusing the /datastore/dump/{resource_id}
response shaping (302 single shard / stream-concat multi-shard / 413
multi-shard parquet). The CKAN action API (datastore_search_sql) stays
pure JSON; sql is a reserved resource name on the dump route.
- DatastoreDumpSQLRequest (subclasses the search-SQL request): LIMIT is
optional and uncapped in download mode; OFFSET-without-LIMIT rejected.
parse_sql_pagination gains a require_limit flag.
- Engine dump_sql: GCS cache keyed on sql_dumps/<qhash>/<fmt>/<rev>
(qhash = sha256 of qualified SQL, rev = sha256 over every referenced
table's modified). Non-deterministic SQL bypasses the cache (fresh uuid
rev). RO dry-run preflight yields the output schema + clean 400s; the
user ORDER BY is preserved across shards; superseded revisions are GC'd
immediately (age-gated only for the non-cacheable path).
- Refactor dump() + dump_sql() over one shared _cached_export workflow
plus helpers (_get_table, _run_export_job, _gc_stale_blobs, _sign_blobs).
- dump_sql_datastore service shares the SQL function allow-list gate with
search_sql_datastore.
- Shared DumpFormat in core/constants.
- Engine/endpoint/service tests and docs (API.md, CLAUDE.md, README).
…dump/query Download logic (cache -> export -> compose -> sign) moves out of backend.py and services/dump.py into a dedicated engine module.
EXPORT DATA requires a wildcard URI and shards by write parallelism, so parquet can come back as several files that can't be composed. The endpoint now fetches the parts and streams one zip (stored, not deflated) instead of returning a URL list; single-file exports still 302 to the signed URL.
📝 WalkthroughWalkthroughAdds ChangesDatastore export flow
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant DumpEndpoint
participant ReadService
participant BigQueryExport
participant GCS
Client->>DumpEndpoint: GET /datastore/dump/query
DumpEndpoint->>ReadService: validate and authorize SQL references
ReadService->>BigQueryExport: export vetted SQL
BigQueryExport->>GCS: read cache or export and compose objects
GCS-->>BigQueryExport: return signed URLs
BigQueryExport-->>DumpEndpoint: return URLs
DumpEndpoint-->>Client: redirect or streamed ZIP response
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (3)
tests/test_datastore_dump_sql.py (1)
1-15: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffEngine-level
dump_sqltests could live undertests/engines/bigquery/.Lines 108-730 are BigQuery-engine tests (mocked
client.query,list_blobs,_delete_old_cache); lines 733-970 are endpoint tests. CLAUDE.md's tree reservestests/engines/bigquery/test_*.pyfor the former. Splitting is optional, but worth considering before this file grows further.As per coding guidelines: "auth and engine tests should mirror the corresponding production directory structure".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_datastore_dump_sql.py` around lines 1 - 15, Split tests/test_datastore_dump_sql.py according to test scope: move the BigQuery engine tests, including mocked client.query, list_blobs, and _delete_old_cache cases, into tests/engines/bigquery/test_*.py, while keeping the endpoint tests in their current test module. Preserve shared fixtures and all existing test behavior.Source: Coding guidelines
datastore/infrastructure/engines/bigquery/export.py (2)
265-267: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBuild the rw BigQuery client lazily.
_build_bq_client("rw")loads credentials and constructs abigquery.Clientsynchronously on the request path for every call, including pure cache hits where it's never used — the path the docstring advertises as "no BigQuery". Only_run_export_jobneeds it.♻️ Suggested change
- rw_bq = backend._build_bq_client("rw") ro_gcs = backend._build_storage_client("ro").bucket(bucket) rw_gcs = backend._build_storage_client("rw").bucket(bucket)and inside the miss branch:
export_sql, header_bytes = await build_export_sql() - await _run_export_job(rw_bq, export_sql, what=what, fmt=fmt) + rw_bq = backend._build_bq_client("rw") + await _run_export_job(rw_bq, export_sql, what=what, fmt=fmt)Note
tests/test_datastore_dump.py::_engine_with_storagestubs_build_bq_clientunconditionally, so this stays test-compatible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@datastore/infrastructure/engines/bigquery/export.py` around lines 265 - 267, Move the _build_bq_client("rw") call out of the request setup and into the cache-miss branch immediately before the _run_export_job invocation. Keep storage clients initialized as currently, and ensure cache-hit paths never construct the BigQuery client while preserving the existing test stub compatibility.
375-390: 🩺 Stability & Availability | 🔵 Trivial
job.result()on the default executor can saturate it under concurrent dumps.
asyncio.to_threaduses the interpreter's defaultThreadPoolExecutor(min(32, cpu+4)workers), shared with every otherto_threadcall in the app (signing, listing, deletes). AnEXPORT DATAcan hold a worker for minutes, so a handful of concurrent large dumps can stall unrelated offloaded work. Consider a dedicated bounded executor for export waits, or pollingjob.done()withasyncio.sleepinstead of blocking a worker.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@datastore/infrastructure/engines/bigquery/export.py` around lines 375 - 390, The _run_export_job function blocks the shared default executor while waiting on job.result(), allowing concurrent exports to starve unrelated to_thread work. Move the blocking export wait to a dedicated bounded executor, or replace it with nonblocking job.done() polling and asyncio.sleep; keep the existing query submission and completion/error behavior intact.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@CLAUDE.md`:
- Line 489: Update the SQL download description in CLAUDE.md to match
download_response: state that composed outputs, including gzip, are returned via
a 302 redirect, and multi-file parquet is streamed as a zip archive. Remove the
incorrect claim about streamed gzip and JSON URL lists, aligning the wording
with §5.3 step 7 and the table entry.
In `@datastore/api/endpoints/dump.py`:
- Around line 124-135: Update the dump endpoint to delegate resource dumping
through a new dump_datastore service function, mirroring dump_sql_datastore in
the read service. Move get_datastore_engine and engine.dump handling into
dump_datastore, while keeping authorization, request parsing, and
download_response construction in dump.
In `@datastore/infrastructure/engines/bigquery/export.py`:
- Around line 538-556: The /datastore/dump/{resource_id} flow must validate
resource_id before _signed_urls uses it in the Content-Disposition filename.
Reject non-ASCII characters and quotes, carriage returns, or line feeds with the
route’s existing invalid-input response, while preserving signing for valid
identifiers.
- Around line 99-137: Update the dump flow around the _prepare_download call to
set cacheable based on whether table.modified is present, matching dump_sql’s
handling: disable caching when table.modified is None and retain caching for
modified tables. Keep the existing revision and export behavior unchanged.
- Around line 291-293: Remove the `or rw_blobs` fallback in the re-list flow
around `_serve_with_cached`: only assign `blobs` to the validated servable
result. If validation returns no blobs, route through the existing
export/rebuild branch rather than signing raw `rw_blobs`, restructuring the
surrounding `if blobs` flow as needed.
---
Nitpick comments:
In `@datastore/infrastructure/engines/bigquery/export.py`:
- Around line 265-267: Move the _build_bq_client("rw") call out of the request
setup and into the cache-miss branch immediately before the _run_export_job
invocation. Keep storage clients initialized as currently, and ensure cache-hit
paths never construct the BigQuery client while preserving the existing test
stub compatibility.
- Around line 375-390: The _run_export_job function blocks the shared default
executor while waiting on job.result(), allowing concurrent exports to starve
unrelated to_thread work. Move the blocking export wait to a dedicated bounded
executor, or replace it with nonblocking job.done() polling and asyncio.sleep;
keep the existing query submission and completion/error behavior intact.
In `@tests/test_datastore_dump_sql.py`:
- Around line 1-15: Split tests/test_datastore_dump_sql.py according to test
scope: move the BigQuery engine tests, including mocked client.query,
list_blobs, and _delete_old_cache cases, into tests/engines/bigquery/test_*.py,
while keeping the endpoint tests in their current test module. Preserve shared
fixtures and all existing test behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: cf5e61c4-b1c3-434c-bc34-2cf9c8fdb453
📒 Files selected for processing (18)
API.mdCLAUDE.mdREADME.mddatastore/api/endpoints/datastore.pydatastore/api/endpoints/dump.pydatastore/core/constants.pydatastore/infrastructure/engines/base.pydatastore/infrastructure/engines/bigquery/backend.pydatastore/infrastructure/engines/bigquery/export.pydatastore/schemas/request.pydatastore/schemas/validators.pydatastore/services/dump.pydatastore/services/read.pydatastore/services/streaming.pytests/test_datastore_dump.pytests/test_datastore_dump_sql.pytests/test_datastore_search_sql.pytests/test_read_service.py
💤 Files with no reviewable changes (1)
- datastore/services/dump.py
d11c71d to
ae56c33
Compare
| rw_gcs = backend._build_storage_client("rw").bucket(bucket) | ||
|
|
||
| # 1. cache lookup — and self-heal an unusable entry. | ||
| blobs = await _list_files_sorted(ro_gcs, prefix) if cacheable else [] |
There was a problem hiding this comment.
If individual A calls
select * from x
And then individual B calls
select * from x
Because of this line, the request made by individual B, will clear the cache, so by the time that individual A receives the GCS url, the files will be either incomplete, or invalid (Might end up mixed with individual B request)
There was a problem hiding this comment.
@luccasmmg this is fixed now so nothing deletes another request's work any more.
A and B each export into their own attempt directory, so their EXPORT DATA
jobs can't touch each other's files, and an entry that doesn't validate is
just skipped instead of wiped.
The cleanup sweep is also age-gated now — it used to delete superseded
revisions immediately, which is the same bug from the other side: a client
who started downloading a minute ago still holds a signed URL pointing at
those objects. Nothing is deleted until it's older than the URL expiry
(BIGQUERY_EXPORT_URL_EXPIRY_HOURS, default 1h).
| if fmt == "ndjson" and len(blobs) == 1: | ||
| return blobs | ||
| return None | ||
| return blobs |
There was a problem hiding this comment.
If individual A calls
select * from x
And then individual B calls
select * from x
I think if Individual B requests with format=parquet, due to the fact that we dont check the composed section, he will receive incomplete data?
There was a problem hiding this comment.
An attempt is now only servable once it has a _SUCCESS marker, and that marker is written after the export and the compose. So a listing taken mid-export finds nothing to serve and the reader exports its own attempt
instead of reading half a dataset.
| for i, url in enumerate(urls) | ||
| ] | ||
| return StreamingResponse( | ||
| zip_archive_writer(request.app.state.http, members), |
There was a problem hiding this comment.
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.
@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?
Concurrent exports for the same cache key no longer clear each other's files or serve a half-written one. Also folds sql_dumps/ into dumps/.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/test_datastore_dump_sql.py (1)
292-332: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider driving both
min_agebranches throughdump_sqlinstead of calling_delete_old_cachedirectly.These two tests reach into a private engine helper, while
test_gc_spares_young_attempts_and_reaps_old_revisionsalready covers the age-gated branch end-to-end. Keeping themin_age=Nonecase at the same layer (cacheable dump → sweep) would keep the suite on the public surface and stop these from pinning the helper signature.As per coding guidelines: "Tests must use public APIs and layer-appropriate interfaces; fixtures must not reach into production internals through back doors."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_datastore_dump_sql.py` around lines 292 - 332, The tests test_gc_stale_blobs_no_age_gate_deletes_all_non_current and test_gc_stale_blobs_age_gate_keeps_young call the private _delete_old_cache helper directly. Rewrite them to exercise the public dump_sql flow, configuring cacheable and non-cacheable dumps to cover min_age=None and the age-gated behavior, while preserving the existing assertions about which blobs are deleted and the current revision being retained.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/test_datastore_dump_sql.py`:
- Around line 292-332: The tests
test_gc_stale_blobs_no_age_gate_deletes_all_non_current and
test_gc_stale_blobs_age_gate_keeps_young call the private _delete_old_cache
helper directly. Rewrite them to exercise the public dump_sql flow, configuring
cacheable and non-cacheable dumps to cover min_age=None and the age-gated
behavior, while preserving the existing assertions about which blobs are deleted
and the current revision being retained.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d44daac8-d912-4b3b-81ea-bc79ec36e368
📒 Files selected for processing (6)
CLAUDE.mdREADME.mddatastore/infrastructure/engines/bigquery/export.pytests/test_cors.pytests/test_datastore_dump.pytests/test_datastore_dump_sql.py
🚧 Files skipped from review as they are similar to previous changes (3)
- README.md
- datastore/infrastructure/engines/bigquery/export.py
- tests/test_datastore_dump.py
Issue: #15
What changed
Added
GET /datastore/dump/query— download a filtered result set by SQL. Same validation and per-table auth asdatastore_search_sql, butLIMITis optional and uncapped.Cached until a referenced table changes — first request exports + composes in single file then later requests are a signed redirect with no BigQuery job. Cache key includes every referenced table's
modified, so any table change invalidates it. Non-deterministic SQL (now(), …) always re-exports.Downloads are a GCS redirect (302) —
EXPORT DATAshards are composed server-side into one object for ndjson / csv / gzip csv, so bytes go storage → client and the API stays out of the byte path.csv / gzip get one header row — shards export header-less and a synthesized header is composed in front.
Parquet: single file when the export doesn't shard; zip when it does. Parquet shards can't be composed (footer + magic bytes), and
EXPORT DATAgives no control over file count — it shards by write parallelism, not size, so a ~40 MB result can land as 7 files. When sharded, the API fetches the parts and streams one zip (stored, not deflated — parquet is already compressed). No way found to merge parquet shards without reading and rewriting the data.Moved the download pipeline out of
backend.py/services/dump.pyintoinfrastructure/engines/bigquery/export.py.Background cleanup — compose sources and stale revisions are deleted off the response path; unusable cache entries self-heal by rebuilding.
Overview of download flow
flowchart TD Req["GET /datastore/dump/{rid}<br/>GET /datastore/dump/query?sql=…"] --> Cache{"cached for this<br/>table version?"} Cache -- hit --> Sign["sign V4 URL(s)"] Cache -- miss --> Export["EXPORT DATA → shards in GCS<br/>(uri must be a wildcard)"] Export --> Fmt{"format"} Fmt -- "csv · gzip · ndjson" --> Compose["GCS compose → ONE object<br/>(+ synthesized header for csv/gzip)"] Fmt -- "parquet" --> Keep["shards stay as-is<br/>(can't be composed)"] Compose --> Sign Keep --> Sign Compose -.-> GC["background: delete sources<br/>+ sweep old revisions"] Sign --> Count{"how many files?"} Count -- "1" --> Redirect["302 → signed GCS URL<br/>bytes: GCS → client (resumable)"] Count -- "N" --> Zip["200 streamed zip of the parts<br/>bytes: GCS → API → client"] classDef out fill:#ecfdf5,stroke:#059669,color:#064e3b classDef bg fill:#f5f5f5,stroke:#9ca3af,color:#374151 class Redirect,Zip out class GC bgSummary by CodeRabbit
New Features
Documentation
LIMITrules and file-export usage.