Conversation
Adds a submit/poll query API alongside the existing synchronous
/api/query endpoint, so query errors can surface cleanly instead of
truncating an already-committed Arrow stream.
- POST /api/query/jobs submits a query; planning/validation run
synchronously (400 on bad query), execution runs in the background.
- Streamable jobs (Arrow IPC / default) deliver batches incrementally
via GET /api/query/jobs/{id}/stream (long-poll: Arrow chunk, 204, or a
terminal JSON completed/failed). A mid-execution error becomes an
observable `failed` status after the good batches were delivered.
- File jobs (Parquet/NetCDF/ODV/GeoParquet) materialize fully, then
download via GET /api/query/jobs/{id}/result.
- GET /api/query/jobs/{id} reports status; DELETE cancels a running job.
Produced batches are bounded by a process-wide in-memory budget counted
in bytes across all jobs; overflow spills to a per-job temp file rather
than growing memory. A sweeper evicts terminal jobs past their TTL and
aborts idle streamable jobs; a semaphore caps background concurrency.
The client query handlers are reorganized into a query/ module
(execute, jobs, explain, metrics) with shared request helpers.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces an asynchronous query job model (submit + poll) to make mid-stream execution failures observable to clients, while keeping the existing synchronous POST /api/query endpoint intact. It also refactors the Beacon client query handlers into a query/ module layout and adds configurable buffering/spilling + lifecycle management for async jobs.
Changes:
- Added async query job execution to
beacon-core(job registry, background execution, polling, cancellation, sweeper). - Added spillable, globally budgeted buffering for streamable job batches, plus new runtime config knobs (
BEACON_QUERY_JOB_*). - Refactored and extended the axum client query API to register new
/api/query/jobs*endpoints and split query handlers by concern.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| beacon-core/src/runtime.rs | Adds async job submission/poll/cancel APIs, job registry, and sweeper integration into the runtime. |
| beacon-core/src/query_result.rs | Adds content_type() / extension() helpers for output-file downloads. |
| beacon-core/src/query_job.rs | Introduces spillable batch buffering + job state types used by the async job subsystem. |
| beacon-core/src/lib.rs | Exposes the new query_job module. |
| beacon-core/src/api.rs | Adds API view types for async job submit + status responses. |
| beacon-config/src/lib.rs | Adds RuntimeConfig.query_jobs and env-configured defaults for job buffering/lifecycle. |
| beacon-api/src/axum/client/query/mod.rs | Creates shared query helpers (auth resolution, SQL gate, UUID parsing, file streaming). |
| beacon-api/src/axum/client/query/execute.rs | Moves one-shot query execution/validation endpoints into query/execute. |
| beacon-api/src/axum/client/query/explain.rs | Moves plan explanation endpoints into query/explain. |
| beacon-api/src/axum/client/query/metrics.rs | Moves metrics lookup into query/metrics. |
| beacon-api/src/axum/client/query/jobs.rs | Adds submit/status/stream/result/cancel endpoints for async query jobs. |
| beacon-api/src/axum/client/query.rs | Removes the legacy monolithic query handler module (replaced by query/ submodules). |
| beacon-api/src/axum/client/mod.rs | Registers the new split query routes and asserts job routes exist in the OpenAPI test. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+395
to
+400
| use crate::query::output::OutputFormat; | ||
|
|
||
| let (plan, query_id, query_json, output) = | ||
| self.prepare_plan(query, is_super_user).await?; | ||
|
|
||
| let is_file = matches!(&output, Some(o) if !matches!(o.format, OutputFormat::Ipc)); |
Comment on lines
+388
to
+390
| /// Classification: no `output` or `output.format = Ipc` → [`JobKind::Streamable`]; | ||
| /// any other output format → [`JobKind::File`]. | ||
| pub async fn submit_query_job( |
|
|
||
| let this = self.clone(); | ||
| let handle = tokio::spawn(async move { | ||
| let _permit = this.query_job_semaphore.clone().acquire_owned().await; |
| let this = self.clone(); | ||
| let metrics = MetricsTracker::new(query_json, query_id); | ||
| let handle = tokio::spawn(async move { | ||
| let _permit = this.query_job_semaphore.clone().acquire_owned().await; |
Comment on lines
+281
to
+286
| Slot::Spilled { offset, len } => match self.read_spill(offset, len) { | ||
| Ok(batch) => out.push(batch), | ||
| Err(e) => { | ||
| tracing::error!("failed to read spilled query-job batch: {e}"); | ||
| } | ||
| }, |
Comment on lines
+104
to
+106
| // Static content-disposition value is ASCII and cannot fail at runtime; wrapped for parity. | ||
| let content_disposition = | ||
| HeaderValue::from_static("application/vnd.apache.arrow.stream"); |
Comment on lines
+192
to
+194
| responses( | ||
| (status = 200, description = "Arrow IPC batches, or a JSON terminal status", content_type = "application/vnd.apache.arrow.stream"), | ||
| (status = 204, description = "No batch ready yet; the job is still running"), |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
POST /api/querystreams Arrow IPC directly from a live DataFusion stream. The200 OKand headers are committed before the first batch, so a mid-stream execution error cannot surface as an HTTP error — the client just receives a truncated Arrow stream. This adds an async job model so query results can be polled and failures reported cleanly.Created
New submit/poll API alongside the untouched synchronous
/api/query:POST /api/query/jobs— plans + validates synchronously (invalid query →400), then runs execution in the background. Returns202+{ query_id, kind }.GET /api/query/jobs/{id}— status:running/succeeded/failed { error }/cancelled.GET /api/query/jobs/{id}/stream— long-poll for streamable jobs: an Arrow IPC chunk (200),204when nothing's ready yet, or a terminal JSON{"state":"completed"}/{"state":"failed","error":…}. A mid-execution error becomes an observablefailedstatus after the good batches were delivered.GET /api/query/jobs/{id}/result— download the materialized file for file-format jobs (Parquet/NetCDF/ODV/GeoParquet).DELETE /api/query/jobs/{id}— cancel a running job.Bounded, spillable buffering
Produced batches are held in a process-wide in-memory budget counted in bytes across all jobs; once exhausted, further batches spill to a per-job temp file instead of growing memory. Large results stay bounded. Delivery is drain-on-read (at-most-once) — documented on the endpoint.
Lifecycle
A background sweeper evicts terminal jobs past their TTL (deleting temp/spill files) and aborts idle streamable jobs; a semaphore caps background execution concurrency. All settings are configurable (
BEACON_QUERY_JOB_*).Refactor
The client query handlers were reorganized from one ~800-line
query.rsinto aquery/module —execute,jobs,explain,metrics— with shared request helpers (resolve_super_user,ensure_sql_allowed,parse_query_id,file_stream_response) deduped intomod.rs.Tests
client_query_testsand all beacon-api tests pass; the OpenAPI test now asserts the five/api/query/jobs*routes are registered.Notes
/api/query/metrics/{id}); submit still gates execution via basic auth.