From e6f8a4923d51098e127542e1e4e0b8f36ae3b1d2 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 24 Jun 2026 15:43:02 +0200 Subject: [PATCH] Add async query jobs with incremental polling 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. --- beacon-api/src/axum/client/mod.rs | 34 +- beacon-api/src/axum/client/query.rs | 502 -------------- beacon-api/src/axum/client/query/execute.rs | 194 ++++++ beacon-api/src/axum/client/query/explain.rs | 126 ++++ beacon-api/src/axum/client/query/jobs.rs | 344 ++++++++++ beacon-api/src/axum/client/query/metrics.rs | 47 ++ beacon-api/src/axum/client/query/mod.rs | 176 +++++ beacon-config/src/lib.rs | 42 ++ beacon-core/src/api.rs | 64 +- beacon-core/src/lib.rs | 1 + beacon-core/src/query_job.rs | 452 +++++++++++++ beacon-core/src/query_result.rs | 26 + beacon-core/src/runtime.rs | 695 +++++++++++++++++++- 13 files changed, 2156 insertions(+), 547 deletions(-) delete mode 100644 beacon-api/src/axum/client/query.rs create mode 100644 beacon-api/src/axum/client/query/execute.rs create mode 100644 beacon-api/src/axum/client/query/explain.rs create mode 100644 beacon-api/src/axum/client/query/jobs.rs create mode 100644 beacon-api/src/axum/client/query/metrics.rs create mode 100644 beacon-api/src/axum/client/query/mod.rs create mode 100644 beacon-core/src/query_job.rs diff --git a/beacon-api/src/axum/client/mod.rs b/beacon-api/src/axum/client/mod.rs index 2996adc3..b4508c28 100644 --- a/beacon-api/src/axum/client/mod.rs +++ b/beacon-api/src/axum/client/mod.rs @@ -28,12 +28,16 @@ pub struct ClientApiDoc; #[allow(deprecated)] pub(crate) fn setup_client_router() -> (Router>, utoipa::openapi::OpenApi) { OpenApiRouter::with_openapi(ClientApiDoc::openapi()) - .routes(routes!(query::query)) - .routes(routes!(query::parse_query)) - .routes(routes!(query::query_metrics)) - .routes(routes!(query::explain_query)) - .routes(routes!(query::explain_analyze_query)) - .routes(routes!(query::available_columns)) + .routes(routes!(query::execute::query)) + .routes(routes!(query::jobs::submit_query_job)) + .routes(routes!(query::jobs::query_job_status, query::jobs::cancel_query_job)) + .routes(routes!(query::jobs::query_job_stream)) + .routes(routes!(query::jobs::query_job_result)) + .routes(routes!(query::execute::parse_query)) + .routes(routes!(query::metrics::query_metrics)) + .routes(routes!(query::explain::explain_query)) + .routes(routes!(query::explain::explain_analyze_query)) + .routes(routes!(query::execute::available_columns)) .routes(routes!(datasets::datasets)) .routes(routes!(datasets::list_datasets)) .routes(routes!(datasets::list_dataset_schema)) @@ -87,5 +91,23 @@ mod tests { Some("#/components/schemas/Query"), "expected /api/query requestBody to $ref the Query schema" ); + + // The async query-job endpoints are registered. + let paths = spec.pointer("/paths").and_then(|p| p.as_object()).expect("paths"); + for (path, method) in [ + ("/api/query/jobs", "post"), + ("/api/query/jobs/{query_id}", "get"), + ("/api/query/jobs/{query_id}", "delete"), + ("/api/query/jobs/{query_id}/stream", "get"), + ("/api/query/jobs/{query_id}/result", "get"), + ] { + let entry = paths + .get(path) + .unwrap_or_else(|| panic!("expected path `{path}` in OpenAPI doc")); + assert!( + entry.get(method).is_some(), + "expected `{method}` operation on `{path}`" + ); + } } } diff --git a/beacon-api/src/axum/client/query.rs b/beacon-api/src/axum/client/query.rs deleted file mode 100644 index 4c8dc0bc..00000000 --- a/beacon-api/src/axum/client/query.rs +++ /dev/null @@ -1,502 +0,0 @@ -//! Query execution endpoints for the client HTTP API. - -use ::axum::{ - body::Body, - extract::{Path, State}, - http::{header, HeaderMap, HeaderName, HeaderValue, StatusCode}, - response::{IntoResponse, Response}, - Json, -}; -use beacon_core::runtime::Runtime; -use beacon_core::{ - api::{QueryMetricsView, QueryRequest}, - query::Query, - query_result::QueryOutputFile, -}; -use futures::TryStreamExt; -use std::sync::Arc; - -/// Resolves the HTTP super-user flag from the request's `Authorization` header. -/// -/// Only HTTP basic auth elevates a request: this transport has no bearer concept -/// (bearer tokens are Flight-SQL-only), so non-`Basic` schemes are left untouched -/// rather than rejected. -/// -/// - No `Authorization` header, or a non-`Basic` scheme (e.g. `Bearer …`) → -/// anonymous, read-only (`Ok(false)`). -/// - Valid admin basic credentials → super-user, DDL/DML allowed (`Ok(true)`). -/// - A `Basic` header that fails validation → `Err(UNAUTHORIZED)` so bad -/// credentials surface as an error instead of silently degrading to read-only. -fn resolve_super_user( - headers: &HeaderMap, - admin: &beacon_config::AdminConfig, -) -> Result { - let is_basic = headers - .get(header::AUTHORIZATION) - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value.starts_with("Basic ")); - - if is_basic { - crate::axum::auth::verify_basic_auth_header(headers, admin)?; - Ok(true) - } else { - Ok(false) - } -} - -/// Executes a query against the runtime and streams the result to the client. -/// -/// The response is either an Arrow IPC stream (zstd-compressed) for in-memory -/// results or a file download when the query produced a materialized output -/// file (CSV, Parquet, Arrow, JSON, ODV, NetCDF, GeoParquet). -#[tracing::instrument(level = "info", skip(state, headers))] -#[utoipa::path( - tag = "query", - post, - path = "/api/query", - request_body = Query, - responses( - ( - status = 200, - description = "Query results in the format requested by the query. The \ - default is a zstd-compressed Arrow IPC stream; file output formats \ - (CSV, Parquet, Arrow, ODV, NetCDF, GeoParquet) are returned as \ - a file download. The result's query id is returned in the \ - `x-beacon-query-id` response header.", - content_type = "application/vnd.apache.arrow.stream" - ), - (status = 400, description = "Invalid, unsupported, or disabled query"), - (status = 401, description = "Basic credentials were supplied but are invalid"), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -pub(crate) async fn query( - State(state): State>, - headers: HeaderMap, - Json(query_obj): Json, -) -> Result, (StatusCode, Json)> { - let query = query_obj.into_query().map_err(|err| { - tracing::error!("Error parsing beacon query: {}", err); - (StatusCode::BAD_REQUEST, Json(err.to_string())) - })?; - - // SQL over the HTTP client API is gated by `sql.enable` (JSON is always - // allowed); the Flight SQL transport has its own `flight_sql.enable`. - if matches!(query.inner, beacon_core::query::InnerQuery::Sql(_)) - && !state.config().sql.enable - { - return Err(( - StatusCode::BAD_REQUEST, - Json("SQL queries are not enabled".to_string()), - )); - } - - // HTTP client queries are read-only by default; supplying valid admin basic - // credentials elevates the request to super-user, allowing DDL/DML (e.g. - // CREATE EXTERNAL TABLE, CREATE/INSERT on managed tables) over HTTP. This - // mirrors the Flight SQL transport, where basic auth resolves to an admin - // `AuthContext`. - let is_super_user = resolve_super_user(&headers, &state.config().admin).map_err(|status| { - (status, Json("invalid admin credentials".to_string())) - })?; - let query_result = state.run_query(query, is_super_user).await.map_err(|err| { - tracing::error!("Error running beacon query: {}", err); - (StatusCode::BAD_REQUEST, Json(err.to_string())) - })?; - - match query_result.query_output { - beacon_core::query_result::QueryOutput::File(query_output_file) => { - handle_query_output_file(query_output_file, query_result.query_id).await - } - beacon_core::query_result::QueryOutput::Stream(arrow_output_stream) => { - let ipc_options = arrow::ipc::writer::IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .map_err(|err| { - tracing::error!("failed to configure Arrow IPC zstd compression: {err}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json("failed to configure Arrow IPC compression".to_string()), - ) - })?; - - let query_id_header = HeaderValue::from_str(&query_result.query_id.to_string()) - .map_err(|err| { - tracing::error!("failed to encode query id header: {err}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json("failed to encode response headers".to_string()), - ) - })?; - // 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"); - - let schema = arrow_output_stream.schema(); - - // Stream Arrow IPC directly so large result sets do not need to be buffered in memory. - let axum_stream = axum_streams::StreamBodyAs::arrow_ipc_with_options_errors( - schema, - arrow_output_stream.map_err(|e| ::axum::Error::new(Box::new(e))), - ipc_options, - ) - .header("x-beacon-query-id", query_id_header) - .header(header::CONTENT_DISPOSITION, content_disposition); - - Ok(axum_stream.into_response()) - } - } -} - -/// Converts a completed file-backed query result into a streamed HTTP response. -async fn handle_query_output_file( - output_file: QueryOutputFile, - query_id: uuid::Uuid, -) -> Result, (StatusCode, Json)> { - match output_file { - QueryOutputFile::Csv(named_temp_file) => { - file_stream_response(named_temp_file.path(), "text/csv", "csv", query_id).await - } - QueryOutputFile::Parquet(named_temp_file) => { - file_stream_response( - named_temp_file.path(), - "application/vnd.apache.parquet", - "parquet", - query_id, - ) - .await - } - QueryOutputFile::Ipc(named_temp_file) => { - file_stream_response( - named_temp_file.path(), - "application/vnd.apache.arrow.file", - "arrow", - query_id, - ) - .await - } - QueryOutputFile::Json(named_temp_file) => { - file_stream_response(named_temp_file.path(), "application/json", "json", query_id).await - } - QueryOutputFile::Odv(named_temp_file) => { - file_stream_response(named_temp_file.path(), "application/zip", "zip", query_id).await - } - QueryOutputFile::NetCDF(named_temp_file) => { - file_stream_response(named_temp_file.path(), "application/netcdf", "nc", query_id).await - } - QueryOutputFile::GeoParquet(named_temp_file) => { - file_stream_response( - named_temp_file.path(), - "application/vnd.apache.arrow.geo+parquet", - "geoparquet", - query_id, - ) - .await - } - } -} - -/// Streams a temporary result file to the client with the appropriate content headers. -async fn file_stream_response( - file_path: &std::path::Path, - content_type: &str, - file_ext: &str, - query_id: uuid::Uuid, -) -> Result, (StatusCode, Json)> { - let file = tokio::fs::File::open(file_path).await.map_err(|err| { - tracing::error!("failed to open query result file {file_path:?}: {err}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json("failed to open query result".to_string()), - ) - })?; - let stream = tokio_util::io::ReaderStream::new(file); - let inner_stream = Body::from_stream(stream); - Ok(( - [ - (header::CONTENT_TYPE, content_type), - ( - header::CONTENT_DISPOSITION, - format!("attachment; filename=\"output.{}\"", file_ext).as_str(), - ), - ( - HeaderName::from_static("x-beacon-query-id"), - query_id.to_string().as_str(), - ), - ], - inner_stream, - ) - .into_response()) -} - -/// Validates a query body by parsing it. Returns 200 if the payload deserializes -/// into a [`QueryRequest`] — the query is not executed. -#[tracing::instrument(level = "info")] -#[utoipa::path( - tag = "query", - post, - path = "/api/parse-query", - request_body = Query, - responses( - (status = 200, description = "The query body is valid (it was not executed)"), - (status = 400, description = "Malformed query body"), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -pub(crate) async fn parse_query(Json(_query_obj): Json) -> StatusCode { - StatusCode::OK -} - -/// Returns recorded planner metrics for a previously executed query. -#[tracing::instrument(level = "info", skip(state))] -#[utoipa::path( - tag = "query", - get, - path = "/api/query/metrics/{query_id}", - params( - ("query_id" = String, Path, description = "UUID of a previously executed query") - ), - responses( - (status = 200, description = "Recorded metrics for the query", body = QueryMetricsView), - (status = 400, description = "The query id is not a valid UUID"), - (status = 404, description = "No metrics recorded for the given query id"), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -pub(crate) async fn query_metrics( - State(state): State>, - Path(query_id): Path, -) -> Result, (StatusCode, Json)> { - let query_id = uuid::Uuid::parse_str(&query_id).map_err(|_| { - ( - StatusCode::BAD_REQUEST, - Json("Invalid UUID format".to_string()), - ) - })?; - - let metrics = state.get_query_metrics(query_id).ok_or_else(|| { - ( - StatusCode::NOT_FOUND, - Json("Query ID not found".to_string()), - ) - })?; - - Ok(Json(metrics)) -} - -/// Returns a JSON-encoded explanation of the plan the runtime would produce for -/// the supplied query without executing it. -#[tracing::instrument(level = "info", skip(state))] -#[utoipa::path( - tag = "query", - post, - path = "/api/explain-query", - request_body = Query, - responses( - ( - status = 200, - description = "JSON explanation of the plan the runtime would produce \ - for the query (the query is not executed)", - content_type = "application/json" - ), - (status = 400, description = "Invalid or unsupported query"), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -pub(crate) async fn explain_query( - State(state): State>, - Json(query_obj): Json, -) -> Result, (StatusCode, Json)> { - let result = state.explain_client_query(query_obj).await; - match result { - Ok(explanation) => Ok(( - [ - (header::CONTENT_TYPE, "application/json"), - (header::CONTENT_DISPOSITION, "attachment"), - ], - Body::from(explanation), - ) - .into_response()), - Err(err) => { - tracing::error!("Error explaining beacon query: {}", err); - Err((StatusCode::BAD_REQUEST, Json(err.to_string()))) - } - } -} - -/// Runs the supplied query and returns its physical plan annotated with per-node -/// runtime metrics as PostgreSQL-style JSON (pgjson) — the `EXPLAIN ANALYZE` -/// analog of `/api/explain-query`. Unlike that endpoint, the query is executed. -#[tracing::instrument(level = "info", skip(state, headers))] -#[utoipa::path( - tag = "query", - post, - path = "/api/explain-analyze-query", - request_body = Query, - responses( - ( - status = 200, - description = "pgjson explanation of the query's physical plan annotated \ - with per-node runtime metrics (the query IS executed to collect them)", - content_type = "application/json" - ), - (status = 400, description = "Invalid or unsupported query"), - (status = 401, description = "Basic credentials were supplied but are invalid"), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -pub(crate) async fn explain_analyze_query( - State(state): State>, - headers: HeaderMap, - Json(query_obj): Json, -) -> Result, (StatusCode, Json)> { - let query = query_obj.into_query().map_err(|err| { - tracing::error!("Error parsing beacon query: {}", err); - (StatusCode::BAD_REQUEST, Json(err.to_string())) - })?; - - // EXPLAIN ANALYZE executes the query, so it is gated by `sql.enable` exactly - // like `/api/query` (JSON is always allowed). Without this, SQL could be run - // through this endpoint while SQL is disabled, bypassing the restriction. - if matches!(query.inner, beacon_core::query::InnerQuery::Sql(_)) - && !state.config().sql.enable - { - return Err(( - StatusCode::BAD_REQUEST, - Json("SQL queries are not enabled".to_string()), - )); - } - - // It also resolves admin vs anonymous the same way `/api/query` does: - // anonymous is read-only, valid admin basic auth elevates to super-user - // (allowing DDL/DML, with the same side effects). - let is_super_user = resolve_super_user(&headers, &state.config().admin) - .map_err(|status| (status, Json("invalid admin credentials".to_string())))?; - let result = state - .explain_analyze_client_query(query, is_super_user) - .await; - match result { - Ok(explanation) => Ok(( - [ - (header::CONTENT_TYPE, "application/json"), - (header::CONTENT_DISPOSITION, "attachment"), - ], - Body::from(explanation), - ) - .into_response()), - Err(err) => { - tracing::error!("Error explain-analyzing beacon query: {}", err); - Err((StatusCode::BAD_REQUEST, Json(err.to_string()))) - } - } -} - -/// Backward-compatible endpoint for clients that still request the default schema as column names. -#[tracing::instrument(level = "info", skip(state))] -#[utoipa::path( - tag = "query", - get, - path = "/api/query/available-columns", - responses( - (status = 200, description = "Column names of the default table schema", body = Vec), - ), - security( - (), - ("basic-auth" = []), - ("bearer" = []) - ) -)] -#[deprecated = "Use /api/default-table-schema instead"] -pub(crate) async fn available_columns(State(state): State>) -> Json> { - Json( - state - .list_default_table_schema_view() - .await - .fields - .iter() - .map(|f| f.name.clone()) - .collect(), - ) -} - -#[cfg(test)] -mod tests { - use ::axum::http::{HeaderMap, StatusCode}; - use base64::{engine::general_purpose, Engine as _}; - - use super::resolve_super_user; - - /// Builds an `Authorization: Basic ...` header map for the given credentials. - fn basic_auth_headers(username: &str, password: &str) -> HeaderMap { - let value = format!( - "Basic {}", - general_purpose::STANDARD.encode(format!("{username}:{password}")) - ); - let mut headers = HeaderMap::new(); - headers.insert("Authorization", value.parse().unwrap()); - headers - } - - fn test_admin() -> beacon_config::AdminConfig { - beacon_config::AdminConfig { - username: "beacon-admin".to_string(), - password: "beacon-password".to_string(), - } - } - - /// Valid admin credentials elevate the HTTP request to super-user so DDL/DML - /// is allowed over HTTP. - #[test] - fn valid_admin_credentials_resolve_to_super_user() { - let admin = test_admin(); - let headers = basic_auth_headers(&admin.username, &admin.password); - assert_eq!(resolve_super_user(&headers, &admin), Ok(true)); - } - - /// No credentials at all stays anonymous (read-only), not an error. - #[test] - fn missing_authorization_header_is_anonymous() { - let headers = HeaderMap::new(); - assert_eq!(resolve_super_user(&headers, &test_admin()), Ok(false)); - } - - /// Credentials that are present but wrong are rejected rather than silently - /// degrading to read-only. - #[test] - fn wrong_credentials_are_rejected() { - let headers = basic_auth_headers("not-the-admin", "wrong-password"); - assert_eq!( - resolve_super_user(&headers, &test_admin()), - Err(StatusCode::UNAUTHORIZED) - ); - } - - /// A non-Basic scheme (e.g. a bearer token, which the endpoint's OpenAPI - /// advertises) must not be treated as failed basic auth: it falls through to - /// anonymous/read-only rather than being rejected with 401. - #[test] - fn bearer_token_is_anonymous_not_rejected() { - let mut headers = HeaderMap::new(); - headers.insert("Authorization", "Bearer some-token".parse().unwrap()); - assert_eq!(resolve_super_user(&headers, &test_admin()), Ok(false)); - } -} diff --git a/beacon-api/src/axum/client/query/execute.rs b/beacon-api/src/axum/client/query/execute.rs new file mode 100644 index 00000000..8b910c80 --- /dev/null +++ b/beacon-api/src/axum/client/query/execute.rs @@ -0,0 +1,194 @@ +//! One-shot query execution and validation. +//! +//! `POST /api/query` runs a query and returns its result inline (an Arrow IPC +//! stream for in-memory results, or a file download for materialized output +//! formats). `POST /api/parse-query` validates a body without executing it, and +//! the deprecated `GET /api/query/available-columns` lists the default table's +//! columns. + +use ::axum::{ + body::Body, + extract::State, + http::{header, HeaderMap, HeaderValue, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use beacon_core::runtime::Runtime; +use beacon_core::{ + api::QueryRequest, + query::Query, + query_result::{QueryOutput, QueryOutputFile}, +}; +use futures::TryStreamExt; +use std::sync::Arc; + +use super::{ensure_sql_allowed, file_stream_response, invalid_credentials, resolve_super_user}; + +/// Executes a query against the runtime and streams the result to the client. +/// +/// The response is either an Arrow IPC stream (zstd-compressed) for in-memory +/// results or a file download when the query produced a materialized output +/// file (CSV, Parquet, Arrow, JSON, ODV, NetCDF, GeoParquet). +#[tracing::instrument(level = "info", skip(state, headers))] +#[utoipa::path( + tag = "query", + post, + path = "/api/query", + request_body = Query, + responses( + ( + status = 200, + description = "Query results in the format requested by the query. The \ + default is a zstd-compressed Arrow IPC stream; file output formats \ + (CSV, Parquet, Arrow, ODV, NetCDF, GeoParquet) are returned as \ + a file download. The result's query id is returned in the \ + `x-beacon-query-id` response header.", + content_type = "application/vnd.apache.arrow.stream" + ), + (status = 400, description = "Invalid, unsupported, or disabled query"), + (status = 401, description = "Basic credentials were supplied but are invalid"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn query( + State(state): State>, + headers: HeaderMap, + Json(query_obj): Json, +) -> Result, (StatusCode, Json)> { + let query = query_obj.into_query().map_err(|err| { + tracing::error!("Error parsing beacon query: {}", err); + (StatusCode::BAD_REQUEST, Json(err.to_string())) + })?; + + ensure_sql_allowed(&query, &state)?; + + // HTTP client queries are read-only by default; supplying valid admin basic + // credentials elevates the request to super-user, allowing DDL/DML (e.g. + // CREATE EXTERNAL TABLE, CREATE/INSERT on managed tables) over HTTP. This + // mirrors the Flight SQL transport, where basic auth resolves to an admin + // `AuthContext`. + let is_super_user = + resolve_super_user(&headers, &state.config().admin).map_err(invalid_credentials)?; + let query_result = state.run_query(query, is_super_user).await.map_err(|err| { + tracing::error!("Error running beacon query: {}", err); + (StatusCode::BAD_REQUEST, Json(err.to_string())) + })?; + + match query_result.query_output { + QueryOutput::File(query_output_file) => { + handle_query_output_file(query_output_file, query_result.query_id).await + } + QueryOutput::Stream(arrow_output_stream) => { + let ipc_options = arrow::ipc::writer::IpcWriteOptions::default() + .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) + .map_err(|err| { + tracing::error!("failed to configure Arrow IPC zstd compression: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to configure Arrow IPC compression".to_string()), + ) + })?; + + let query_id_header = HeaderValue::from_str(&query_result.query_id.to_string()) + .map_err(|err| { + tracing::error!("failed to encode query id header: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to encode response headers".to_string()), + ) + })?; + // 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"); + + let schema = arrow_output_stream.schema(); + + // Stream Arrow IPC directly so large result sets do not need to be buffered in memory. + let axum_stream = axum_streams::StreamBodyAs::arrow_ipc_with_options_errors( + schema, + arrow_output_stream.map_err(|e| { + // The 200 OK status and headers are already on the wire by the + // time batches stream, so a mid-stream error cannot surface as an + // HTTP error code: the client just sees a truncated Arrow IPC + // stream. Log it here so the failure is at least observable + // server-side instead of being silently swallowed. + tracing::error!("error producing Arrow stream batch: {e}"); + ::axum::Error::new(Box::new(e)) + }), + ipc_options, + ) + .header("x-beacon-query-id", query_id_header) + .header(header::CONTENT_DISPOSITION, content_disposition); + + Ok(axum_stream.into_response()) + } + } +} + +/// Converts a completed file-backed query result into a streamed HTTP response. +async fn handle_query_output_file( + output_file: QueryOutputFile, + query_id: uuid::Uuid, +) -> Result, (StatusCode, Json)> { + file_stream_response( + output_file.path(), + output_file.content_type(), + output_file.extension(), + query_id, + ) + .await +} + +/// Validates a query body by parsing it. Returns 200 if the payload deserializes +/// into a [`QueryRequest`] — the query is not executed. +#[tracing::instrument(level = "info")] +#[utoipa::path( + tag = "query", + post, + path = "/api/parse-query", + request_body = Query, + responses( + (status = 200, description = "The query body is valid (it was not executed)"), + (status = 400, description = "Malformed query body"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn parse_query(Json(_query_obj): Json) -> StatusCode { + StatusCode::OK +} + +/// Backward-compatible endpoint for clients that still request the default schema as column names. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + get, + path = "/api/query/available-columns", + responses( + (status = 200, description = "Column names of the default table schema", body = Vec), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +#[deprecated = "Use /api/default-table-schema instead"] +pub(crate) async fn available_columns(State(state): State>) -> Json> { + Json( + state + .list_default_table_schema_view() + .await + .fields + .iter() + .map(|f| f.name.clone()) + .collect(), + ) +} diff --git a/beacon-api/src/axum/client/query/explain.rs b/beacon-api/src/axum/client/query/explain.rs new file mode 100644 index 00000000..74ab5143 --- /dev/null +++ b/beacon-api/src/axum/client/query/explain.rs @@ -0,0 +1,126 @@ +//! Plan explanation endpoints. +//! +//! `POST /api/explain-query` returns the logical plan as JSON without executing +//! the query; `POST /api/explain-analyze-query` executes it and returns the +//! physical plan annotated with per-node runtime metrics (pgjson). + +use ::axum::{ + body::Body, + extract::State, + http::{header, HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use beacon_core::runtime::Runtime; +use beacon_core::{api::QueryRequest, query::Query}; +use std::sync::Arc; + +use super::{ensure_sql_allowed, invalid_credentials, resolve_super_user}; + +/// Returns a JSON-encoded explanation of the plan the runtime would produce for +/// the supplied query without executing it. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + post, + path = "/api/explain-query", + request_body = Query, + responses( + ( + status = 200, + description = "JSON explanation of the plan the runtime would produce \ + for the query (the query is not executed)", + content_type = "application/json" + ), + (status = 400, description = "Invalid or unsupported query"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn explain_query( + State(state): State>, + Json(query_obj): Json, +) -> Result, (StatusCode, Json)> { + let result = state.explain_client_query(query_obj).await; + match result { + Ok(explanation) => Ok(( + [ + (header::CONTENT_TYPE, "application/json"), + (header::CONTENT_DISPOSITION, "attachment"), + ], + Body::from(explanation), + ) + .into_response()), + Err(err) => { + tracing::error!("Error explaining beacon query: {}", err); + Err((StatusCode::BAD_REQUEST, Json(err.to_string()))) + } + } +} + +/// Runs the supplied query and returns its physical plan annotated with per-node +/// runtime metrics as PostgreSQL-style JSON (pgjson) — the `EXPLAIN ANALYZE` +/// analog of `/api/explain-query`. Unlike that endpoint, the query is executed. +#[tracing::instrument(level = "info", skip(state, headers))] +#[utoipa::path( + tag = "query", + post, + path = "/api/explain-analyze-query", + request_body = Query, + responses( + ( + status = 200, + description = "pgjson explanation of the query's physical plan annotated \ + with per-node runtime metrics (the query IS executed to collect them)", + content_type = "application/json" + ), + (status = 400, description = "Invalid or unsupported query"), + (status = 401, description = "Basic credentials were supplied but are invalid"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn explain_analyze_query( + State(state): State>, + headers: HeaderMap, + Json(query_obj): Json, +) -> Result, (StatusCode, Json)> { + let query = query_obj.into_query().map_err(|err| { + tracing::error!("Error parsing beacon query: {}", err); + (StatusCode::BAD_REQUEST, Json(err.to_string())) + })?; + + // EXPLAIN ANALYZE executes the query, so it is gated by `sql.enable` exactly + // like `/api/query`. Without this, SQL could be run through this endpoint + // while SQL is disabled, bypassing the restriction. + ensure_sql_allowed(&query, &state)?; + + // It also resolves admin vs anonymous the same way `/api/query` does: + // anonymous is read-only, valid admin basic auth elevates to super-user + // (allowing DDL/DML, with the same side effects). + let is_super_user = + resolve_super_user(&headers, &state.config().admin).map_err(invalid_credentials)?; + let result = state + .explain_analyze_client_query(query, is_super_user) + .await; + match result { + Ok(explanation) => Ok(( + [ + (header::CONTENT_TYPE, "application/json"), + (header::CONTENT_DISPOSITION, "attachment"), + ], + Body::from(explanation), + ) + .into_response()), + Err(err) => { + tracing::error!("Error explain-analyzing beacon query: {}", err); + Err((StatusCode::BAD_REQUEST, Json(err.to_string()))) + } + } +} diff --git a/beacon-api/src/axum/client/query/jobs.rs b/beacon-api/src/axum/client/query/jobs.rs new file mode 100644 index 00000000..31d5549b --- /dev/null +++ b/beacon-api/src/axum/client/query/jobs.rs @@ -0,0 +1,344 @@ +//! Asynchronous query jobs: submit a query, then poll for status and results. +//! +//! Unlike `POST /api/query`, which commits `200 OK` before the first batch (so a +//! mid-stream failure can only truncate the response), a job decouples execution +//! from delivery. Streamable jobs hand back Arrow batches incrementally via +//! `GET /api/query/jobs/{id}/stream`; file jobs materialize fully and are +//! downloaded via `GET /api/query/jobs/{id}/result`. Either way a mid-execution +//! error surfaces cleanly as a terminal `failed` status. + +use ::axum::{ + body::Body, + extract::{Path, State}, + http::{header, HeaderMap, HeaderName, HeaderValue, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use beacon_core::runtime::Runtime; +use beacon_core::{ + api::{QueryJobStatusView, QueryJobSubmitView, QueryRequest}, + query::Query, + query_job::{CancelOutcome, JobKind, PollOutcome, QueryJobState}, +}; +use std::sync::Arc; + +use super::{ensure_sql_allowed, file_stream_response, invalid_credentials, parse_query_id, + resolve_super_user}; + +/// Serializes a chunk of batches as a self-contained, zstd-compressed Arrow IPC +/// stream (schema + batches), matching the `/api/query` streaming content type. +fn encode_arrow_ipc( + schema: arrow::datatypes::SchemaRef, + batches: &[arrow::array::RecordBatch], +) -> Result, (StatusCode, Json)> { + let options = arrow::ipc::writer::IpcWriteOptions::default() + .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) + .map_err(|err| { + tracing::error!("failed to configure Arrow IPC zstd compression: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to configure Arrow IPC compression".to_string()), + ) + })?; + let mut buf = Vec::new(); + { + let mut writer = + arrow::ipc::writer::StreamWriter::try_new_with_options(&mut buf, &schema, options) + .map_err(|err| { + tracing::error!("failed to create Arrow IPC writer: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to serialize query result".to_string()), + ) + })?; + for batch in batches { + writer.write(batch).map_err(|err| { + tracing::error!("failed to write Arrow IPC batch: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to serialize query result".to_string()), + ) + })?; + } + writer.finish().map_err(|err| { + tracing::error!("failed to finish Arrow IPC stream: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to serialize query result".to_string()), + ) + })?; + } + Ok(buf) +} + +/// Submits a query for asynchronous execution and returns its job id. +/// +/// Planning and validation happen synchronously, so an invalid query is rejected +/// here with `400`; execution then runs in the background. Results are retrieved +/// by polling: streamable jobs via `GET /api/query/jobs/{id}/stream`, file jobs +/// via `GET /api/query/jobs/{id}` then `GET /api/query/jobs/{id}/result`. Auth +/// mirrors `/api/query` (valid admin basic credentials elevate to super-user). +#[tracing::instrument(level = "info", skip(state, headers))] +#[utoipa::path( + tag = "query", + post, + path = "/api/query/jobs", + request_body = Query, + responses( + (status = 202, description = "Job accepted; poll for status and results", body = QueryJobSubmitView), + (status = 400, description = "Invalid, unsupported, or disabled query"), + (status = 401, description = "Basic credentials were supplied but are invalid"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn submit_query_job( + State(state): State>, + headers: HeaderMap, + Json(query_obj): Json, +) -> Result, (StatusCode, Json)> { + let query = query_obj.into_query().map_err(|err| { + tracing::error!("Error parsing beacon query: {}", err); + (StatusCode::BAD_REQUEST, Json(err.to_string())) + })?; + + ensure_sql_allowed(&query, &state)?; + + let is_super_user = + resolve_super_user(&headers, &state.config().admin).map_err(invalid_credentials)?; + + let (query_id, kind) = state + .submit_query_job(query, is_super_user) + .await + .map_err(|err| { + tracing::error!("Error submitting beacon query job: {}", err); + (StatusCode::BAD_REQUEST, Json(err.to_string())) + })?; + + let query_id_header = HeaderValue::from_str(&query_id.to_string()).map_err(|err| { + tracing::error!("failed to encode query id header: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to encode response headers".to_string()), + ) + })?; + + let body = QueryJobSubmitView { + query_id: query_id.to_string(), + kind, + }; + Ok(( + StatusCode::ACCEPTED, + [(HeaderName::from_static("x-beacon-query-id"), query_id_header)], + Json(body), + ) + .into_response()) +} + +/// Returns the current status of an async query job. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + get, + path = "/api/query/jobs/{query_id}", + params( + ("query_id" = String, Path, description = "UUID returned by POST /api/query/jobs") + ), + responses( + (status = 200, description = "Current job status", body = QueryJobStatusView), + (status = 400, description = "The query id is not a valid UUID"), + (status = 404, description = "No job with the given id"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn query_job_status( + State(state): State>, + Path(query_id): Path, +) -> Result, (StatusCode, Json)> { + let query_id = parse_query_id(&query_id)?; + let snapshot = state.query_job_snapshot(query_id).ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json("Query job not found".to_string()), + ) + })?; + Ok(Json(QueryJobStatusView::from(snapshot))) +} + +/// Long-polls a streamable job for its next batches. +/// +/// Returns a zstd-compressed Arrow IPC stream (`200`) when batches are ready, a +/// `204` when none became ready within the server-side wait (still running), or +/// a JSON terminal body (`{"state":"completed"}` / `{"state":"failed","error":…}`) +/// once the job finishes. A mid-execution error therefore surfaces cleanly here +/// after any already-produced batches were delivered — unlike `/api/query`, where +/// it can only truncate the stream. Batches are delivered at-most-once: a poll +/// response that does not reach the client loses those rows. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + get, + path = "/api/query/jobs/{query_id}/stream", + params( + ("query_id" = String, Path, description = "UUID of a streamable job") + ), + 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"), + (status = 400, description = "The query id is not a valid UUID"), + (status = 404, description = "No job with the given id"), + (status = 409, description = "The job is a file job, or was cancelled"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn query_job_stream( + State(state): State>, + Path(query_id): Path, +) -> Result, (StatusCode, Json)> { + let id = parse_query_id(&query_id)?; + let wait = state.query_job_poll_wait(); + match state.poll_query_job_stream(id, wait).await { + PollOutcome::Batches { schema, batches } => { + let body = encode_arrow_ipc(schema, &batches)?; + Ok(( + [( + header::CONTENT_TYPE, + HeaderValue::from_static("application/vnd.apache.arrow.stream"), + )], + body, + ) + .into_response()) + } + PollOutcome::Pending => Ok(StatusCode::NO_CONTENT.into_response()), + PollOutcome::Completed => { + Ok(Json(serde_json::json!({ "state": "completed" })).into_response()) + } + PollOutcome::Failed(error) => { + Ok(Json(serde_json::json!({ "state": "failed", "error": error })).into_response()) + } + PollOutcome::NotStreamable => Err(( + StatusCode::CONFLICT, + Json("job produces a file result; poll status and download via /result".to_string()), + )), + PollOutcome::Cancelled => Err(( + StatusCode::CONFLICT, + Json("query job was cancelled".to_string()), + )), + PollOutcome::NotFound => Err(( + StatusCode::NOT_FOUND, + Json("Query job not found".to_string()), + )), + } +} + +/// Downloads the materialized result of a completed file job. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + get, + path = "/api/query/jobs/{query_id}/result", + params( + ("query_id" = String, Path, description = "UUID of a file job") + ), + responses( + (status = 200, description = "The materialized result file"), + (status = 400, description = "Bad UUID, or the job failed (body carries the error)"), + (status = 404, description = "No job with the given id"), + (status = 409, description = "Streamable job, still running, or cancelled"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn query_job_result( + State(state): State>, + Path(query_id): Path, +) -> Result, (StatusCode, Json)> { + let id = parse_query_id(&query_id)?; + let snapshot = state.query_job_snapshot(id).ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json("Query job not found".to_string()), + ) + })?; + + if matches!(snapshot.kind, JobKind::Streamable) { + return Err(( + StatusCode::CONFLICT, + Json("streamable job; consume results via /stream".to_string()), + )); + } + + match snapshot.state { + QueryJobState::Running => Err(( + StatusCode::CONFLICT, + Json("query job is still running; result requires completion".to_string()), + )), + QueryJobState::Failed { error } => Err((StatusCode::BAD_REQUEST, Json(error))), + QueryJobState::Cancelled => Err(( + StatusCode::CONFLICT, + Json("query job was cancelled".to_string()), + )), + QueryJobState::Succeeded => { + let file = snapshot.file.ok_or_else(|| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("succeeded job has no result file".to_string()), + ) + })?; + file_stream_response(&file.path, file.content_type, file.file_ext, id).await + } + } +} + +/// Cancels a running job, aborting its background execution. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + delete, + path = "/api/query/jobs/{query_id}", + params( + ("query_id" = String, Path, description = "UUID of the job to cancel") + ), + responses( + (status = 200, description = "The job was cancelled"), + (status = 400, description = "The query id is not a valid UUID"), + (status = 404, description = "No job with the given id"), + (status = 409, description = "The job had already finished"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn cancel_query_job( + State(state): State>, + Path(query_id): Path, +) -> Result<(StatusCode, Json), (StatusCode, Json)> { + let id = parse_query_id(&query_id)?; + match state.cancel_query_job(id) { + CancelOutcome::Cancelled => Ok((StatusCode::OK, Json("cancelled".to_string()))), + CancelOutcome::AlreadyFinished => Err(( + StatusCode::CONFLICT, + Json("query job already finished".to_string()), + )), + CancelOutcome::NotFound => Err(( + StatusCode::NOT_FOUND, + Json("Query job not found".to_string()), + )), + } +} diff --git a/beacon-api/src/axum/client/query/metrics.rs b/beacon-api/src/axum/client/query/metrics.rs new file mode 100644 index 00000000..6cd3bf04 --- /dev/null +++ b/beacon-api/src/axum/client/query/metrics.rs @@ -0,0 +1,47 @@ +//! Recorded query metrics lookup. + +use ::axum::{ + extract::{Path, State}, + http::StatusCode, + Json, +}; +use beacon_core::{api::QueryMetricsView, runtime::Runtime}; +use std::sync::Arc; + +use super::parse_query_id; + +/// Returns recorded planner metrics for a previously executed query. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "query", + get, + path = "/api/query/metrics/{query_id}", + params( + ("query_id" = String, Path, description = "UUID of a previously executed query") + ), + responses( + (status = 200, description = "Recorded metrics for the query", body = QueryMetricsView), + (status = 400, description = "The query id is not a valid UUID"), + (status = 404, description = "No metrics recorded for the given query id"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn query_metrics( + State(state): State>, + Path(query_id): Path, +) -> Result, (StatusCode, Json)> { + let query_id = parse_query_id(&query_id)?; + + let metrics = state.get_query_metrics(query_id).ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + Json("Query ID not found".to_string()), + ) + })?; + + Ok(Json(metrics)) +} diff --git a/beacon-api/src/axum/client/query/mod.rs b/beacon-api/src/axum/client/query/mod.rs new file mode 100644 index 00000000..7dee7243 --- /dev/null +++ b/beacon-api/src/axum/client/query/mod.rs @@ -0,0 +1,176 @@ +//! Query endpoints for the client HTTP API, organized by concern: +//! +//! - [`execute`] — one-shot query execution, body validation, and column listing. +//! - [`jobs`] — asynchronous query jobs (submit / status / stream / result / cancel). +//! - [`explain`] — plan explanation (`EXPLAIN` and `EXPLAIN ANALYZE`). +//! - [`metrics`] — recorded planner/runtime metrics for a previously run query. +//! +//! Request helpers shared across those submodules (auth resolution, the +//! `sql.enable` gate, path-UUID parsing, and file streaming) live in this module. + +use ::axum::{ + body::Body, + http::{header, HeaderMap, HeaderName, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use beacon_core::query::Query; +use beacon_core::runtime::Runtime; + +pub(crate) mod execute; +pub(crate) mod explain; +pub(crate) mod jobs; +pub(crate) mod metrics; + +/// Resolves the HTTP super-user flag from the request's `Authorization` header. +/// +/// Only HTTP basic auth elevates a request: this transport has no bearer concept +/// (bearer tokens are Flight-SQL-only), so non-`Basic` schemes are left untouched +/// rather than rejected. +/// +/// - No `Authorization` header, or a non-`Basic` scheme (e.g. `Bearer …`) → +/// anonymous, read-only (`Ok(false)`). +/// - Valid admin basic credentials → super-user, DDL/DML allowed (`Ok(true)`). +/// - A `Basic` header that fails validation → `Err(UNAUTHORIZED)` so bad +/// credentials surface as an error instead of silently degrading to read-only. +fn resolve_super_user( + headers: &HeaderMap, + admin: &beacon_config::AdminConfig, +) -> Result { + let is_basic = headers + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("Basic ")); + + if is_basic { + crate::axum::auth::verify_basic_auth_header(headers, admin)?; + Ok(true) + } else { + Ok(false) + } +} + +/// Maps a `resolve_super_user` failure to the standard 401 JSON response. +fn invalid_credentials(status: StatusCode) -> (StatusCode, Json) { + (status, Json("invalid admin credentials".to_string())) +} + +/// Rejects SQL queries when `sql.enable` is off (JSON queries are always allowed). +/// +/// SQL over the HTTP client API is gated by `sql.enable`; the Flight SQL transport +/// has its own `flight_sql.enable`. Endpoints that execute a query (run, submit a +/// job, or `EXPLAIN ANALYZE`) call this so SQL cannot slip through while disabled. +fn ensure_sql_allowed( + query: &Query, + state: &Runtime, +) -> Result<(), (StatusCode, Json)> { + if matches!(query.inner, beacon_core::query::InnerQuery::Sql(_)) && !state.config().sql.enable { + return Err(( + StatusCode::BAD_REQUEST, + Json("SQL queries are not enabled".to_string()), + )); + } + Ok(()) +} + +/// Parses a path UUID, mapping a malformed value to `400`. +fn parse_query_id(raw: &str) -> Result)> { + uuid::Uuid::parse_str(raw) + .map_err(|_| (StatusCode::BAD_REQUEST, Json("Invalid UUID format".to_string()))) +} + +/// Streams a temporary result file to the client with the appropriate content headers. +async fn file_stream_response( + file_path: &std::path::Path, + content_type: &str, + file_ext: &str, + query_id: uuid::Uuid, +) -> Result, (StatusCode, Json)> { + let file = tokio::fs::File::open(file_path).await.map_err(|err| { + tracing::error!("failed to open query result file {file_path:?}: {err}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json("failed to open query result".to_string()), + ) + })?; + let stream = tokio_util::io::ReaderStream::new(file); + let inner_stream = Body::from_stream(stream); + Ok(( + [ + (header::CONTENT_TYPE, content_type), + ( + header::CONTENT_DISPOSITION, + format!("attachment; filename=\"output.{}\"", file_ext).as_str(), + ), + ( + HeaderName::from_static("x-beacon-query-id"), + query_id.to_string().as_str(), + ), + ], + inner_stream, + ) + .into_response()) +} + +#[cfg(test)] +mod tests { + use ::axum::http::{HeaderMap, StatusCode}; + use base64::{engine::general_purpose, Engine as _}; + + use super::resolve_super_user; + + /// Builds an `Authorization: Basic ...` header map for the given credentials. + fn basic_auth_headers(username: &str, password: &str) -> HeaderMap { + let value = format!( + "Basic {}", + general_purpose::STANDARD.encode(format!("{username}:{password}")) + ); + let mut headers = HeaderMap::new(); + headers.insert("Authorization", value.parse().unwrap()); + headers + } + + fn test_admin() -> beacon_config::AdminConfig { + beacon_config::AdminConfig { + username: "beacon-admin".to_string(), + password: "beacon-password".to_string(), + } + } + + /// Valid admin credentials elevate the HTTP request to super-user so DDL/DML + /// is allowed over HTTP. + #[test] + fn valid_admin_credentials_resolve_to_super_user() { + let admin = test_admin(); + let headers = basic_auth_headers(&admin.username, &admin.password); + assert_eq!(resolve_super_user(&headers, &admin), Ok(true)); + } + + /// No credentials at all stays anonymous (read-only), not an error. + #[test] + fn missing_authorization_header_is_anonymous() { + let headers = HeaderMap::new(); + assert_eq!(resolve_super_user(&headers, &test_admin()), Ok(false)); + } + + /// Credentials that are present but wrong are rejected rather than silently + /// degrading to read-only. + #[test] + fn wrong_credentials_are_rejected() { + let headers = basic_auth_headers("not-the-admin", "wrong-password"); + assert_eq!( + resolve_super_user(&headers, &test_admin()), + Err(StatusCode::UNAUTHORIZED) + ); + } + + /// A non-Basic scheme (e.g. a bearer token, which the endpoint's OpenAPI + /// advertises) must not be treated as failed basic auth: it falls through to + /// anonymous/read-only rather than being rejected with 401. + #[test] + fn bearer_token_is_anonymous_not_rejected() { + let mut headers = HeaderMap::new(); + headers.insert("Authorization", "Bearer some-token".parse().unwrap()); + assert_eq!(resolve_super_user(&headers, &test_admin()), Ok(false)); + } +} diff --git a/beacon-config/src/lib.rs b/beacon-config/src/lib.rs index afbdd10c..6276b7c8 100644 --- a/beacon-config/src/lib.rs +++ b/beacon-config/src/lib.rs @@ -61,6 +61,26 @@ pub struct RuntimeConfig { pub st_within_point_cache_size: usize, pub enable_sys_info: bool, pub batch_size: usize, + /// Async query-job buffering and lifecycle settings. + pub query_jobs: QueryJobConfig, +} + +/// Settings for the async query-job (submit/poll) subsystem. +#[derive(Debug, Clone)] +pub struct QueryJobConfig { + /// Process-wide in-memory budget (bytes) for buffered streamable batches + /// across all jobs. Once exhausted, further batches spill to disk. + pub buffer_memory_bytes: u64, + /// Per-job spill cap (bytes); `0` means unbounded. Exceeding it fails the job. + pub max_spill_bytes: u64, + /// Maximum number of query jobs executing in the background concurrently. + pub max_concurrent: usize, + /// How long a terminal (succeeded/failed/cancelled) job is retained before eviction. + pub ttl_secs: u64, + /// Abort a streamable job not polled within this idle window. + pub idle_secs: u64, + /// How long a stream poll waits for the next batch before returning empty. + pub poll_wait_ms: u64, } #[derive(Debug, Clone)] @@ -295,6 +315,20 @@ struct RawConfig { #[envconfig(from = "BEACON_BATCH_SIZE", default = "64000")] beacon_batch_size: usize, + // Async query-job (submit/poll) subsystem + #[envconfig(from = "BEACON_QUERY_JOB_BUFFER_MEMORY_BYTES", default = "268435456")] + query_job_buffer_memory_bytes: u64, + #[envconfig(from = "BEACON_QUERY_JOB_MAX_SPILL_BYTES", default = "4294967296")] + query_job_max_spill_bytes: u64, + #[envconfig(from = "BEACON_QUERY_JOB_MAX_CONCURRENT", default = "8")] + query_job_max_concurrent: usize, + #[envconfig(from = "BEACON_QUERY_JOB_TTL_SECS", default = "600")] + query_job_ttl_secs: u64, + #[envconfig(from = "BEACON_QUERY_JOB_IDLE_SECS", default = "120")] + query_job_idle_secs: u64, + #[envconfig(from = "BEACON_QUERY_JOB_POLL_WAIT_MS", default = "5000")] + query_job_poll_wait_ms: u64, + /// Whether to split streams into 16k row slices for better memory management and parallelism. #[envconfig(from = "BEACON_ENABLE_BBF_SPLIT_STREAMS_SLICE", default = "false")] bbf_split_streams_slice: bool, @@ -355,6 +389,14 @@ impl From for Config { st_within_point_cache_size: raw.st_within_point_cache_size, enable_sys_info: raw.enable_sys_info, batch_size: raw.beacon_batch_size, + query_jobs: QueryJobConfig { + buffer_memory_bytes: raw.query_job_buffer_memory_bytes, + max_spill_bytes: raw.query_job_max_spill_bytes, + max_concurrent: raw.query_job_max_concurrent, + ttl_secs: raw.query_job_ttl_secs, + idle_secs: raw.query_job_idle_secs, + poll_wait_ms: raw.query_job_poll_wait_ms, + }, }, sql: SqlConfig { enable: raw.enable_sql, diff --git a/beacon-core/src/api.rs b/beacon-core/src/api.rs index a22ce966..06b46551 100644 --- a/beacon-core/src/api.rs +++ b/beacon-core/src/api.rs @@ -3,12 +3,12 @@ use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use crate::metrics::ConsolidatedMetrics; use arrow::datatypes::{Field, Schema}; use beacon_data_lake::crawler::{CrawlReport, CrawlerDefinition, TableNaming}; use beacon_datafusion_ext::format_ext::DatasetMetadata; use beacon_datafusion_ext::table_ext::TableDefinition; use beacon_functions::function_doc::FunctionDoc; -use crate::metrics::ConsolidatedMetrics; use serde_json::{Map, Value}; use utoipa::ToSchema; @@ -196,6 +196,59 @@ impl TryFrom for QueryMetricsView { } } +/// Submission acknowledgement for an async query job, returned by +/// `POST /api/query/jobs`. The `query_id` is used to poll status, stream +/// results, download a file result, or cancel the job. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, ToSchema)] +pub struct QueryJobSubmitView { + /// The job's unique identifier (UUID), also returned in `x-beacon-query-id`. + pub query_id: String, + /// Whether results are delivered incrementally (`streamable`) or as a file. + pub kind: crate::query_job::JobKind, +} + +/// Current status of an async query job, returned by `GET /api/query/jobs/{id}`. +#[derive(Debug, Clone, serde::Serialize, ToSchema)] +#[serde(tag = "state", rename_all = "lowercase")] +pub enum QueryJobStatusView { + /// The job is still executing. + Running { + /// The job's delivery model. + kind: crate::query_job::JobKind, + }, + /// The job finished successfully. For File jobs the result is downloadable. + Succeeded { + /// The job's delivery model. + kind: crate::query_job::JobKind, + }, + /// The job failed; `error` describes why. + Failed { + /// The job's delivery model. + kind: crate::query_job::JobKind, + /// Human-readable failure reason. + error: String, + }, + /// The job was cancelled by the client or the idle sweeper. + Cancelled { + /// The job's delivery model. + kind: crate::query_job::JobKind, + }, +} + +impl From for QueryJobStatusView { + fn from(snapshot: crate::query_job::QueryJobSnapshot) -> Self { + let kind = snapshot.kind; + match snapshot.state { + crate::query_job::QueryJobState::Running => QueryJobStatusView::Running { kind }, + crate::query_job::QueryJobState::Succeeded => QueryJobStatusView::Succeeded { kind }, + crate::query_job::QueryJobState::Failed { error } => { + QueryJobStatusView::Failed { kind, error } + } + crate::query_job::QueryJobState::Cancelled => QueryJobStatusView::Cancelled { kind }, + } + } +} + /// The storage format and options of a registered table, as a flattened /// configuration object. Internal (double-underscore) option keys are stripped. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, ToSchema)] @@ -237,7 +290,9 @@ impl TryFrom> for TableConfigView { /// How a crawler turns a discovered group of files into a table name. Mirrors the /// data-lake [`TableNaming`] so the API surface need not depend on its internals. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize, ToSchema)] +#[derive( + Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize, ToSchema, +)] #[serde(rename_all = "snake_case")] pub enum TableNamingView { /// Use the leaf component of the group's base prefix (`argo/floats` -> `floats`). @@ -462,7 +517,10 @@ mod table_config_redaction_tests { assert!(!json.contains("super-secret-password")); // The encrypted material (ciphertext/nonce) must not leak either. assert!(!json.contains("ciphertext")); - assert_eq!(view.config.get("secret"), Some(&Value::String("***".to_string()))); + assert_eq!( + view.config.get("secret"), + Some(&Value::String("***".to_string())) + ); // Non-secret connection options remain visible. assert!(json.contains("db.internal")); } diff --git a/beacon-core/src/lib.rs b/beacon-core/src/lib.rs index d85b8bcc..2afd7be1 100644 --- a/beacon-core/src/lib.rs +++ b/beacon-core/src/lib.rs @@ -2,6 +2,7 @@ pub mod api; pub mod metrics; pub mod parser; pub mod query; +pub mod query_job; pub mod query_result; pub mod runtime; mod statement_plan; diff --git a/beacon-core/src/query_job.rs b/beacon-core/src/query_job.rs new file mode 100644 index 00000000..a4191deb --- /dev/null +++ b/beacon-core/src/query_job.rs @@ -0,0 +1,452 @@ +//! Async query jobs: submit a query, then poll for incremental results. +//! +//! The synchronous streaming endpoint (`/api/query`) commits `200 OK` before the +//! first batch, so a mid-stream execution error can only truncate the response — +//! it can never surface as a clean error. Query jobs decouple execution from +//! delivery: a job runs in the background, and the client polls for status and +//! results. A query that fails mid-execution becomes an observable `Failed` +//! state instead of a silently truncated stream. +//! +//! Two job kinds: +//! - [`JobKind::Streamable`] — Arrow IPC / the default in-memory result. Batches +//! are delivered incrementally as they are produced, via a +//! [`SpillableBatchBuffer`]. +//! - [`JobKind::File`] — any explicit file `output` (Parquet, NetCDF, …). The +//! full file is materialized first, then downloaded. +//! +//! Buffered batches are bounded by a process-wide [`BufferBudget`] counted in +//! bytes across all jobs; once the budget is exhausted, further batches spill to +//! a per-job temp file rather than growing memory. Delivery is drain-on-read +//! (at-most-once): a poll response that never reaches the client loses those rows. + +use std::collections::VecDeque; +use std::io::{Seek, SeekFrom, Write}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Instant; + +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use arrow::ipc::reader::StreamReader; +use arrow::ipc::writer::StreamWriter; +use tempfile::NamedTempFile; +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use utoipa::ToSchema; + +use crate::query_result::QueryOutputFile; + +/// Which delivery model a job uses, decided at submit time from the query's +/// `output` format. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, ToSchema)] +#[serde(rename_all = "lowercase")] +pub enum JobKind { + /// Incremental Arrow batch delivery via the stream poll endpoint. + Streamable, + /// A fully-materialized file, downloaded once complete. + File, +} + +/// Terminal-or-running state of a job. +#[derive(Debug, Clone)] +pub enum QueryJobState { + /// Still executing. + Running, + /// Finished successfully (file ready, or stream fully produced). + Succeeded, + /// Execution failed; carries the error message. + Failed { error: String }, + /// Cancelled by the client or the idle sweeper. + Cancelled, +} + +impl QueryJobState { + fn is_terminal(&self) -> bool { + !matches!(self, QueryJobState::Running) + } +} + +/// Process-wide budget for in-memory buffered batch bytes, shared by every +/// streamable job's buffer. +#[derive(Debug)] +pub struct BufferBudget { + used: AtomicU64, + limit: u64, +} + +impl BufferBudget { + pub fn new(limit: u64) -> Arc { + Arc::new(Self { + used: AtomicU64::new(0), + limit, + }) + } + + /// Try to reserve `bytes` of the in-memory budget. Returns `true` on success. + fn try_reserve(&self, bytes: u64) -> bool { + let mut current = self.used.load(Ordering::Relaxed); + loop { + if current.saturating_add(bytes) > self.limit { + return false; + } + match self.used.compare_exchange_weak( + current, + current + bytes, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return true, + Err(observed) => current = observed, + } + } + } + + fn release(&self, bytes: u64) { + self.used.fetch_sub(bytes, Ordering::Relaxed); + } + + /// Current in-memory bytes reserved across all jobs (for tests/metrics). + pub fn used(&self) -> u64 { + self.used.load(Ordering::Relaxed) + } +} + +/// A buffered batch, either resident in memory or spilled to the job's temp file. +enum Slot { + InMemory { batch: RecordBatch, bytes: u64 }, + Spilled { offset: u64, len: u64 }, +} + +enum Terminal { + Completed, + Failed(String), +} + +struct BufferInner { + queue: VecDeque, + spill_bytes: u64, + terminal: Option, +} + +/// A FIFO buffer of Arrow batches that keeps recent batches in memory (bounded by +/// the shared [`BufferBudget`]) and spills the overflow to a temp file. Producers +/// call [`push`](Self::push) / [`finish`](Self::finish); the consumer calls +/// [`drain`](Self::drain). +pub struct SpillableBatchBuffer { + schema: SchemaRef, + budget: Arc, + tmp_dir: PathBuf, + max_spill_bytes: u64, + inner: parking_lot::Mutex, + /// Lazily-created spill file; the `NamedTempFile` deletes on drop. + spill: parking_lot::Mutex>, + notify: Notify, +} + +/// What a [`SpillableBatchBuffer::drain`] / poll yielded. +pub enum DrainOutcome { + /// One or more batches were delivered (more may still follow). + Batches(Vec), + /// The buffer is empty and the producer finished successfully. + Completed, + /// The buffer is empty and the producer failed. + Failed(String), + /// Nothing became available within the wait; the job is still running. + Pending, +} + +impl SpillableBatchBuffer { + pub fn new( + schema: SchemaRef, + budget: Arc, + tmp_dir: PathBuf, + max_spill_bytes: u64, + ) -> Arc { + Arc::new(Self { + schema, + budget, + tmp_dir, + max_spill_bytes, + inner: parking_lot::Mutex::new(BufferInner { + queue: VecDeque::new(), + spill_bytes: 0, + terminal: None, + }), + spill: parking_lot::Mutex::new(None), + notify: Notify::new(), + }) + } + + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// Append a produced batch. Keeps it in memory if the global budget allows, + /// otherwise spills it to disk. Returns `Err` if the per-job spill cap would + /// be exceeded — the caller should then `finish(Failed(..))`. + pub fn push(&self, batch: RecordBatch) -> Result<(), String> { + let bytes = batch.get_array_memory_size() as u64; + if self.budget.try_reserve(bytes) { + self.inner + .lock() + .queue + .push_back(Slot::InMemory { batch, bytes }); + self.notify.notify_one(); + return Ok(()); + } + + // Budget exhausted: spill the batch to the job's temp file. + let blob = serialize_batch(&self.schema, &batch).map_err(|e| e.to_string())?; + let len = blob.len() as u64; + + let offset = { + let mut inner = self.inner.lock(); + if self.max_spill_bytes != 0 + && inner.spill_bytes.saturating_add(len) > self.max_spill_bytes + { + return Err(format!( + "query job spill limit exceeded ({} bytes)", + self.max_spill_bytes + )); + } + let offset = inner.spill_bytes; + inner.spill_bytes += len; + offset + }; + + self.write_spill(offset, &blob).map_err(|e| e.to_string())?; + self.inner + .lock() + .queue + .push_back(Slot::Spilled { offset, len }); + self.notify.notify_one(); + Ok(()) + } + + /// Mark the producer as finished. Wakes any pending poll. + pub fn finish(&self, success: Result<(), String>) { + let terminal = match success { + Ok(()) => Terminal::Completed, + Err(error) => Terminal::Failed(error), + }; + self.inner.lock().terminal = Some(terminal); + self.notify.notify_one(); + } + + /// Drain up to `max` batches in FIFO order, waiting up to `wait` for the first + /// one. Releases in-memory budget for delivered batches and reads spilled ones + /// back from disk. + pub async fn drain(&self, max: usize, wait: std::time::Duration) -> DrainOutcome { + loop { + // Register for notification *before* checking, so a push between the + // check and the await is not lost. + let notified = self.notify.notified(); + + let popped = { + let mut inner = self.inner.lock(); + if inner.queue.is_empty() { + match &inner.terminal { + Some(Terminal::Completed) => return DrainOutcome::Completed, + Some(Terminal::Failed(e)) => return DrainOutcome::Failed(e.clone()), + None => None, + } + } else { + let take = max.min(inner.queue.len()); + Some(inner.queue.drain(..take).collect::>()) + } + }; + + if let Some(slots) = popped { + return DrainOutcome::Batches(self.materialize(slots)); + } + + // Queue empty and not terminal: wait for a push/finish or time out. + match tokio::time::timeout(wait, notified).await { + Ok(()) => continue, + Err(_) => return DrainOutcome::Pending, + } + } + } + + /// Decode popped slots into batches, releasing budget for in-memory ones. + fn materialize(&self, slots: Vec) -> Vec { + let mut out = Vec::with_capacity(slots.len()); + for slot in slots { + match slot { + Slot::InMemory { batch, bytes } => { + self.budget.release(bytes); + out.push(batch); + } + 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}"); + } + }, + } + } + out + } + + fn write_spill(&self, offset: u64, blob: &[u8]) -> std::io::Result<()> { + let mut guard = self.spill.lock(); + if guard.is_none() { + *guard = Some( + tempfile::Builder::new() + .prefix("beacon_qjob_spill_") + .suffix(".arrows") + .tempfile_in(&self.tmp_dir)?, + ); + } + let file = guard.as_mut().expect("spill file present").as_file_mut(); + file.seek(SeekFrom::Start(offset))?; + file.write_all(blob)?; + file.flush()?; + Ok(()) + } + + fn read_spill(&self, offset: u64, len: u64) -> std::io::Result { + let path = { + let guard = self.spill.lock(); + guard + .as_ref() + .map(|f| f.path().to_path_buf()) + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::NotFound, "spill file missing") + })? + }; + use std::io::Read; + let mut file = std::fs::File::open(path)?; + file.seek(SeekFrom::Start(offset))?; + let mut buf = vec![0u8; len as usize]; + file.read_exact(&mut buf)?; + deserialize_batch(&buf) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())) + } +} + +/// Drop releases any still-buffered in-memory budget so a cancelled/evicted job +/// does not leak the global allowance. +impl Drop for SpillableBatchBuffer { + fn drop(&mut self) { + let inner = self.inner.get_mut(); + let mut to_release = 0u64; + for slot in inner.queue.drain(..) { + if let Slot::InMemory { bytes, .. } = slot { + to_release += bytes; + } + } + if to_release > 0 { + self.budget.release(to_release); + } + } +} + +/// Serialize a single batch as a self-contained Arrow IPC stream (schema + batch). +fn serialize_batch( + schema: &SchemaRef, + batch: &RecordBatch, +) -> Result, arrow::error::ArrowError> { + let mut buf = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema)?; + writer.write(batch)?; + writer.finish()?; + } + Ok(buf) +} + +/// Decode the first batch from a self-contained Arrow IPC stream blob. +fn deserialize_batch(bytes: &[u8]) -> Result { + let mut reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)?; + reader.next().transpose()?.ok_or_else(|| { + arrow::error::ArrowError::IpcError("spilled batch blob contained no batch".to_string()) + }) +} + +/// Metadata needed to stream a File-kind job's result back to the client. +#[derive(Debug, Clone)] +pub struct FileResultMeta { + pub path: PathBuf, + pub content_type: &'static str, + pub file_ext: &'static str, +} + +/// A registered job: its kind, state, lifecycle timestamps, and either a stream +/// buffer (streamable) or a materialized output file (file). +pub struct QueryJob { + pub kind: JobKind, + pub state: QueryJobState, + /// When the job reached a terminal state (drives TTL eviction). + pub finished_at: Option, + /// Last time a stream poll touched this job (drives idle-abort for streamable jobs). + pub last_poll_at: Instant, + /// Background execution task handle, for cancellation. Cleared on terminal. + pub handle: Option>, + /// Streamable jobs only: the incremental batch buffer. + pub buffer: Option>, + /// File jobs only: the materialized result (owns the `NamedTempFile`). + pub output: Option, +} + +impl QueryJob { + pub fn new_running(kind: JobKind, buffer: Option>) -> Self { + Self { + kind, + state: QueryJobState::Running, + finished_at: None, + last_poll_at: Instant::now(), + handle: None, + buffer, + output: None, + } + } + + pub fn is_terminal(&self) -> bool { + self.state.is_terminal() + } +} + +/// Read-only view of a job for the status/result handlers, captured without +/// holding the registry lock during async IO. +#[derive(Debug, Clone)] +pub struct QueryJobSnapshot { + pub kind: JobKind, + pub state: QueryJobState, + /// Present for succeeded File jobs. + pub file: Option, +} + +/// Result of a stream poll, mapped by the API layer to an HTTP response. +pub enum PollOutcome { + /// Batches drained this round (more may follow). Carries the schema so the + /// API can serialize a self-contained Arrow IPC stream. + Batches { + schema: SchemaRef, + batches: Vec, + }, + /// Still running, no batch became ready within the wait. + Pending, + /// Stream fully delivered and the query succeeded. + Completed, + /// The query failed; carries the error message. + Failed(String), + /// The job is a File job — use the status/result endpoints instead. + NotStreamable, + /// The job was cancelled. + Cancelled, + /// No job with that id. + NotFound, +} + +/// Outcome of a cancellation request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CancelOutcome { + /// A running job was cancelled. + Cancelled, + /// The job had already finished; nothing to cancel. + AlreadyFinished, + /// No job with that id. + NotFound, +} diff --git a/beacon-core/src/query_result.rs b/beacon-core/src/query_result.rs index da0e2c8a..ccfab80a 100644 --- a/beacon-core/src/query_result.rs +++ b/beacon-core/src/query_result.rs @@ -65,6 +65,32 @@ impl QueryOutputFile { QueryOutputFile::GeoParquet(file) => file.path(), } } + + /// The HTTP `Content-Type` for downloading this output file. + pub fn content_type(&self) -> &'static str { + match self { + QueryOutputFile::Csv(_) => "text/csv", + QueryOutputFile::Ipc(_) => "application/vnd.apache.arrow.file", + QueryOutputFile::Json(_) => "application/json", + QueryOutputFile::Parquet(_) => "application/vnd.apache.parquet", + QueryOutputFile::NetCDF(_) => "application/netcdf", + QueryOutputFile::Odv(_) => "application/zip", + QueryOutputFile::GeoParquet(_) => "application/vnd.apache.arrow.geo+parquet", + } + } + + /// The file extension (without the dot) for the downloaded file name. + pub fn extension(&self) -> &'static str { + match self { + QueryOutputFile::Csv(_) => "csv", + QueryOutputFile::Ipc(_) => "arrow", + QueryOutputFile::Json(_) => "json", + QueryOutputFile::Parquet(_) => "parquet", + QueryOutputFile::NetCDF(_) => "nc", + QueryOutputFile::Odv(_) => "zip", + QueryOutputFile::GeoParquet(_) => "geoparquet", + } + } } impl From for QueryOutputFile { diff --git a/beacon-core/src/runtime.rs b/beacon-core/src/runtime.rs index 9fafd0ed..85282f35 100644 --- a/beacon-core/src/runtime.rs +++ b/beacon-core/src/runtime.rs @@ -1,7 +1,12 @@ //! High-level Beacon runtime shared by the API transports. -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{Arc, Weak}, + time::{Duration, Instant}, +}; +use crate::metrics::{ConsolidatedMetrics, MetricsTracker}; use arrow::{ array::AsArray, datatypes::{SchemaRef, UInt64Type}, @@ -15,7 +20,6 @@ use beacon_datafusion_ext::{ stats_cache::beacon_file_statistics_cache, }; use beacon_functions::function_doc::FunctionDoc; -use crate::metrics::{ConsolidatedMetrics, MetricsTracker}; use datafusion::{ catalog::TableFunctionImpl, execution::{ @@ -29,14 +33,22 @@ use parking_lot::Mutex; use crate::{ api::{ - CrawlReportView, CrawlerView, CreateCrawlerRequest, CreateExternalTableRequest, DatasetInfo, - FunctionInfo, QueryMetricsView, QueryRequest, SchemaView, TableConfigView, + CrawlReportView, CrawlerView, CreateCrawlerRequest, CreateExternalTableRequest, + DatasetInfo, FunctionInfo, QueryMetricsView, QueryRequest, SchemaView, TableConfigView, }, parser::{beacon_parser::BeaconParser, statement::BeaconStatement}, + query_job::{ + BufferBudget, CancelOutcome, FileResultMeta, JobKind, PollOutcome, QueryJob, + QueryJobSnapshot, QueryJobState, SpillableBatchBuffer, + }, query_result::{ArrowOutputStream, QueryOutput, QueryOutputFile, QueryResult}, sys::{self, SystemInfo}, }; +/// Maximum number of batches a single stream poll drains in one response, to +/// bound the response size regardless of how far the producer has run ahead. +const MAX_BATCHES_PER_POLL: usize = 32; + /// Beacon's single execution layer: startup, catalog access, queries, SQL, and files. pub struct Runtime { session_ctx: Arc, @@ -44,6 +56,12 @@ pub struct Runtime { listing_table_factory: Arc, crawler_manager: Arc, query_metrics: Arc>>, + /// Registry of async query jobs (submit/poll), keyed by query id. + query_jobs: Arc>>, + /// Process-wide in-memory budget shared by all streamable job buffers. + query_job_budget: Arc, + /// Caps how many query jobs execute in the background concurrently. + query_job_semaphore: Arc, /// The configuration this runtime was built with. Owned, not process-global. config: Arc, } @@ -144,8 +162,7 @@ impl Runtime { // Build the crawler manager once the data lake and tables exist, then publish // it through the handle so `CREATE/RUN/DROP CRAWLER` actions can reach it. - let events_available = - config.storage.enable_fs_events || config.storage.enable_s3_events; + let events_available = config.storage.enable_fs_events || config.storage.enable_s3_events; let crawler_manager = beacon_data_lake::crawler::CrawlerManager::new( session_ctx.clone(), file_formats.clone(), @@ -156,12 +173,25 @@ impl Runtime { crawler_manager.init().await?; let _ = crawler_handle.set(crawler_manager.clone()); + let job_cfg = &config.runtime.query_jobs; + let query_jobs = Arc::new(Mutex::new(HashMap::new())); + let query_job_budget = BufferBudget::new(job_cfg.buffer_memory_bytes); + let query_job_semaphore = Arc::new(tokio::sync::Semaphore::new(job_cfg.max_concurrent.max(1))); + spawn_query_job_sweeper( + Arc::downgrade(&query_jobs), + Duration::from_secs(job_cfg.ttl_secs), + Duration::from_secs(job_cfg.idle_secs), + ); + Ok(Self { session_ctx, file_formats, listing_table_factory, crawler_manager, query_metrics: Arc::new(Mutex::new(HashMap::new())), + query_jobs, + query_job_budget, + query_job_semaphore, config, }) } @@ -251,6 +281,31 @@ impl Runtime { query: crate::query::Query, is_super_user: bool, ) -> anyhow::Result { + let (plan, query_id, query_json, output) = self.prepare_plan(query, is_super_user).await?; + + match output { + Some(output) => { + self.run_query_to_file(plan, output, query_id, query_json) + .await + } + None => self.run_query_to_stream(plan, query_id, query_json).await, + } + } + + /// Lower, validate, and assign an id to a query without executing it. Shared + /// by the synchronous `run_query` and the async `submit_query_job` paths, so + /// parse/validation errors surface to the caller (HTTP 400) before any work + /// is scheduled. + async fn prepare_plan( + &self, + query: crate::query::Query, + is_super_user: bool, + ) -> anyhow::Result<( + datafusion::logical_expr::LogicalPlan, + uuid::Uuid, + serde_json::Value, + Option, + )> { let query_id = uuid::Uuid::new_v4(); let query_json = serde_json::to_value(&query)?; let crate::query::Query { inner, output } = query; @@ -258,10 +313,7 @@ impl Runtime { let plan = self.lower_query(inner).await?; crate::statement_plan::validate_query_plan(&plan, is_super_user)?; - match output { - Some(output) => self.run_query_to_file(plan, output, query_id, query_json).await, - None => self.run_query_to_stream(plan, query_id, query_json).await, - } + Ok((plan, query_id, query_json, output)) } /// Stream the plan's results, wrapping the stream so output rows/bytes are @@ -298,7 +350,11 @@ impl Runtime { // `Output::parse` wraps the (already validated) plan in a `COPY TO` the // temp file; this COPY is beacon-generated, so it is not re-validated. let (copy_plan, output_file) = output - .parse(self.session_ctx.as_ref(), &self.config.storage.tmp_dir, plan) + .parse( + self.session_ctx.as_ref(), + &self.config.storage.tmp_dir, + plan, + ) .await?; let output_file = QueryOutputFile::from(output_file); @@ -323,6 +379,199 @@ impl Runtime { }) } + /// Submit a query for asynchronous execution and return its id plus delivery + /// kind. Planning and validation happen synchronously (so bad queries error + /// here, surfaced as HTTP 400); execution runs on a background task whose + /// results are retrieved via [`Self::poll_query_job_stream`] (streamable) or + /// [`Self::query_job_snapshot`] + the file result download (file). + /// + /// Classification: no `output` or `output.format = Ipc` → [`JobKind::Streamable`]; + /// any other output format → [`JobKind::File`]. + pub async fn submit_query_job( + self: &Arc, + query: crate::query::Query, + is_super_user: bool, + ) -> anyhow::Result<(uuid::Uuid, JobKind)> { + 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)); + + if is_file { + let output = output.expect("file job classified from Some(output)"); + self.query_jobs + .lock() + .insert(query_id, QueryJob::new_running(JobKind::File, None)); + + let this = self.clone(); + let handle = tokio::spawn(async move { + let _permit = this.query_job_semaphore.clone().acquire_owned().await; + let result = this + .run_query_to_file(plan, output, query_id, query_json) + .await; + let mut jobs = this.query_jobs.lock(); + if let Some(job) = jobs.get_mut(&query_id) { + if matches!(job.state, QueryJobState::Running) { + match result { + Ok(QueryResult { + query_output: QueryOutput::File(file), + .. + }) => { + job.output = Some(file); + job.state = QueryJobState::Succeeded; + } + Ok(_) => { + job.state = QueryJobState::Failed { + error: "file job did not produce a file output".to_string(), + }; + } + Err(error) => { + job.state = QueryJobState::Failed { + error: error.to_string(), + }; + } + } + job.finished_at = Some(Instant::now()); + job.handle = None; + } + } + }); + self.attach_job_handle(query_id, handle); + return Ok((query_id, JobKind::File)); + } + + // Streamable: build the physical stream now (so plan errors surface at + // submit), then spawn a producer that drains it into the spill buffer. + let stream = + crate::statement_plan::execute_statement_plan(&self.session_ctx, plan).await?; + let schema = stream.schema(); + let buffer = SpillableBatchBuffer::new( + schema, + self.query_job_budget.clone(), + self.config.storage.tmp_dir.clone(), + self.config.runtime.query_jobs.max_spill_bytes, + ); + self.query_jobs.lock().insert( + query_id, + QueryJob::new_running(JobKind::Streamable, Some(buffer.clone())), + ); + + 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; + let outcome = produce_stream_into_buffer(stream, &buffer, &metrics).await; + buffer.finish(outcome.clone()); + if outcome.is_ok() { + this.query_metrics + .lock() + .insert(query_id, metrics.get_consolidated_metrics()); + } + let mut jobs = this.query_jobs.lock(); + if let Some(job) = jobs.get_mut(&query_id) { + if matches!(job.state, QueryJobState::Running) { + job.state = match outcome { + Ok(()) => QueryJobState::Succeeded, + Err(error) => QueryJobState::Failed { error }, + }; + job.finished_at = Some(Instant::now()); + job.handle = None; + } + } + }); + self.attach_job_handle(query_id, handle); + Ok((query_id, JobKind::Streamable)) + } + + /// Store a job's background task handle, unless the task already finished (or + /// the job was cancelled) before we got here — in which case the handle is + /// simply dropped. + fn attach_job_handle(&self, query_id: uuid::Uuid, handle: tokio::task::JoinHandle<()>) { + let mut jobs = self.query_jobs.lock(); + match jobs.get_mut(&query_id) { + Some(job) if matches!(job.state, QueryJobState::Running) => { + job.handle = Some(handle); + } + _ => {} + } + } + + /// Read-only snapshot of a job for the status/result handlers. + pub fn query_job_snapshot(&self, query_id: uuid::Uuid) -> Option { + let jobs = self.query_jobs.lock(); + let job = jobs.get(&query_id)?; + let file = job.output.as_ref().map(|output| FileResultMeta { + path: output.path().to_path_buf(), + content_type: output.content_type(), + file_ext: output.extension(), + }); + Some(QueryJobSnapshot { + kind: job.kind, + state: job.state.clone(), + file, + }) + } + + /// Long-poll a streamable job for its next batches, waiting up to `wait`. + pub async fn poll_query_job_stream( + &self, + query_id: uuid::Uuid, + wait: Duration, + ) -> PollOutcome { + let buffer = { + let mut jobs = self.query_jobs.lock(); + let Some(job) = jobs.get_mut(&query_id) else { + return PollOutcome::NotFound; + }; + if matches!(job.kind, JobKind::File) { + return PollOutcome::NotStreamable; + } + if matches!(job.state, QueryJobState::Cancelled) { + return PollOutcome::Cancelled; + } + job.last_poll_at = Instant::now(); + job.buffer.clone() + }; + let Some(buffer) = buffer else { + return PollOutcome::NotFound; + }; + + use crate::query_job::DrainOutcome; + match buffer.drain(MAX_BATCHES_PER_POLL, wait).await { + DrainOutcome::Batches(batches) => PollOutcome::Batches { + schema: buffer.schema(), + batches, + }, + DrainOutcome::Completed => PollOutcome::Completed, + DrainOutcome::Failed(error) => PollOutcome::Failed(error), + DrainOutcome::Pending => PollOutcome::Pending, + } + } + + /// Cancel a running job, aborting its background task. + pub fn cancel_query_job(&self, query_id: uuid::Uuid) -> CancelOutcome { + let mut jobs = self.query_jobs.lock(); + let Some(job) = jobs.get_mut(&query_id) else { + return CancelOutcome::NotFound; + }; + if job.is_terminal() { + return CancelOutcome::AlreadyFinished; + } + if let Some(handle) = job.handle.take() { + handle.abort(); + } + job.state = QueryJobState::Cancelled; + job.finished_at = Some(Instant::now()); + CancelOutcome::Cancelled + } + + /// The default stream-poll wait, from config. + pub fn query_job_poll_wait(&self) -> Duration { + Duration::from_millis(self.config.runtime.query_jobs.poll_wait_ms) + } + /// Lower the body of a client query (JSON or SQL) to a `LogicalPlan` without /// executing or validating it (permission checks and output formatting happen /// in `run_query` on the lowered plan). @@ -340,14 +589,11 @@ impl Runtime { /// Lower a SQL statement (SELECT, DDL/DML, or a beacon custom statement) to a /// `LogicalPlan`. The result is validated in `run_query` before execution. - async fn lower_sql( - &self, - sql: &str, - ) -> anyhow::Result { + async fn lower_sql(&self, sql: &str) -> anyhow::Result { match Self::parse_beacon_statement(sql)? { - BeaconStatement::CreateMaterializedView(statement) => { - Ok(crate::statement_plan::create_materialized_view_plan(statement)) - } + BeaconStatement::CreateMaterializedView(statement) => Ok( + crate::statement_plan::create_materialized_view_plan(statement), + ), BeaconStatement::Refresh(statement) => { Ok(crate::statement_plan::refresh_plan(statement)) } @@ -553,8 +799,13 @@ impl Runtime { } pub async fn list_table_config(&self, table_name: String) -> Option { - let provider = self.session_ctx.table_provider(table_name.as_str()).await.ok()?; - let config = beacon_data_lake::definition_from_provider(&table_name, provider.as_ref()).ok()?; + let provider = self + .session_ctx + .table_provider(table_name.as_str()) + .await + .ok()?; + let config = + beacon_data_lake::definition_from_provider(&table_name, provider.as_ref()).ok()?; match TableConfigView::try_from(config) { Ok(config) => Some(config), Err(error) => { @@ -612,10 +863,14 @@ impl Runtime { offset: Option, limit: Option, ) -> anyhow::Result> { - Ok( - beacon_data_lake::list_datasets(&self.session_ctx, &self.file_formats, offset, limit, pattern) - .await?, + Ok(beacon_data_lake::list_datasets( + &self.session_ctx, + &self.file_formats, + offset, + limit, + pattern, ) + .await?) } pub async fn total_datasets(&self) -> anyhow::Result { @@ -693,6 +948,84 @@ impl Runtime { } } +/// Drive a DataFusion result stream into a streamable job's spill buffer, +/// recording output metrics per batch. Returns `Ok` on clean completion or an +/// error string (execution failure, or a spill-cap breach reported by `push`). +async fn produce_stream_into_buffer( + mut stream: datafusion::execution::SendableRecordBatchStream, + buffer: &SpillableBatchBuffer, + metrics: &MetricsTracker, +) -> Result<(), String> { + loop { + match stream.try_next().await { + Ok(Some(batch)) => { + metrics.add_output_rows(batch.num_rows() as u64); + metrics.add_output_bytes(batch.get_array_memory_size() as u64); + buffer.push(batch)?; + } + Ok(None) => return Ok(()), + Err(error) => return Err(error.to_string()), + } + } +} + +/// Background sweeper that evicts terminal query jobs past their TTL and aborts +/// streamable jobs that have not been polled within the idle window. Holds only a +/// `Weak` reference to the job registry, so it stops once the runtime is dropped. +fn spawn_query_job_sweeper( + jobs: Weak>>, + ttl: Duration, + idle: Duration, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(Duration::from_secs(5)); + loop { + ticker.tick().await; + let Some(jobs) = jobs.upgrade() else { + break; // runtime dropped + }; + let mut map = jobs.lock(); + sweep_jobs(&mut map, Instant::now(), ttl, idle); + } + }); +} + +/// Abort idle streamable jobs and evict terminal jobs past their TTL. Extracted +/// from the sweeper loop so the lifecycle policy can be unit-tested deterministically. +fn sweep_jobs( + map: &mut HashMap, + now: Instant, + ttl: Duration, + idle: Duration, +) { + // Abort streamable jobs that no client is polling anymore. + let idle_ids: Vec = map + .iter() + .filter(|(_, job)| { + matches!(job.kind, JobKind::Streamable) + && matches!(job.state, QueryJobState::Running) + && now.duration_since(job.last_poll_at) > idle + }) + .map(|(id, _)| *id) + .collect(); + for id in idle_ids { + if let Some(job) = map.get_mut(&id) { + if let Some(handle) = job.handle.take() { + handle.abort(); + } + job.state = QueryJobState::Cancelled; + job.finished_at = Some(now); + } + } + + // Evict terminal jobs whose retention window has elapsed (dropping their + // buffers/output files, which deletes the temp files). + map.retain(|_, job| match job.finished_at { + Some(finished) => now.duration_since(finished) <= ttl, + None => true, + }); +} + /// Assemble an injection-safe `CREATE EXTERNAL TABLE` statement from a structured /// request. Identifiers (table name, `STORED AS` type, partition columns) are /// validated as bare words; `LOCATION` and `OPTIONS` values are emitted as escaped @@ -779,7 +1112,9 @@ mod materialized_view_tests { #[tokio::test(flavor = "multi_thread")] async fn materialized_view_create_query_refresh_and_drop() { - let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())).await.expect("runtime should start"); + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); let suffix = uuid::Uuid::new_v4().simple(); let mv = format!("mv_test_{suffix}"); @@ -851,7 +1186,9 @@ mod materialized_view_tests { #[tokio::test(flavor = "multi_thread")] async fn materialized_view_handles_zero_row_result() { - let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())).await.expect("runtime should start"); + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); let mv = format!("mv_empty_{}", uuid::Uuid::new_v4().simple()); // A query with no rows must still create a queryable, Parquet-backed view. @@ -915,12 +1252,22 @@ mod client_query_tests { /// same pipeline as `run_sql`, returning a streamed result. #[tokio::test(flavor = "multi_thread")] async fn json_query_runs_through_unified_path() { - let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())).await.expect("runtime should start"); + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); let suffix = uuid::Uuid::new_v4().simple(); let table = format!("json_q_{suffix}"); - run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT, b BIGINT)")).await; - run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1, 2), (3, 4)")).await; + run_sql( + &runtime, + &format!("CREATE TABLE {table} (a BIGINT, b BIGINT)"), + ) + .await; + run_sql( + &runtime, + &format!("INSERT INTO {table} VALUES (1, 2), (3, 4)"), + ) + .await; let batches = runtime .run_query( @@ -944,7 +1291,9 @@ mod client_query_tests { /// file download. #[tokio::test(flavor = "multi_thread")] async fn query_with_output_format_produces_file() { - let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())).await.expect("runtime should start"); + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); let suffix = uuid::Uuid::new_v4().simple(); let table = format!("out_{suffix}"); @@ -1027,7 +1376,9 @@ mod client_query_tests { /// operation (super-user-only). #[tokio::test(flavor = "multi_thread")] async fn non_super_user_is_gated_by_validation() { - let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())).await.expect("runtime should start"); + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); let suffix = uuid::Uuid::new_v4().simple(); let table = format!("val_{suffix}"); @@ -1036,7 +1387,10 @@ mod client_query_tests { // Non-super-user: read-only SELECT is allowed. runtime - .run_query(crate::query::Query::sql(format!("SELECT * FROM {table}")), false) + .run_query( + crate::query::Query::sql(format!("SELECT * FROM {table}")), + false, + ) .await .expect("non-super SELECT should be allowed") .into_record_stream() @@ -1121,10 +1475,16 @@ mod restart_tests { let config = std::sync::Arc::new(beacon_config::Config::load().unwrap()); // First runtime: create a base table with data and a view over it. - let runtime = Runtime::new(config.clone()).await.expect("runtime should start"); + let runtime = Runtime::new(config.clone()) + .await + .expect("runtime should start"); run_sql(&runtime, &format!("CREATE TABLE {base} (a BIGINT)")).await; run_sql(&runtime, &format!("INSERT INTO {base} VALUES (1), (2)")).await; - run_sql(&runtime, &format!("CREATE VIEW {view} AS SELECT a FROM {base}")).await; + run_sql( + &runtime, + &format!("CREATE VIEW {view} AS SELECT a FROM {base}"), + ) + .await; drop(runtime); // A fresh runtime rebuilds the catalog purely from the persisted @@ -1211,9 +1571,7 @@ mod external_table_sql_tests { #[test] fn rejects_unsafe_identifiers() { assert!(build_create_external_table_sql(&req("bad name", "PARQUET", "x/")).is_err()); - assert!( - build_create_external_table_sql(&req("t; DROP TABLE u", "PARQUET", "x/")).is_err() - ); + assert!(build_create_external_table_sql(&req("t; DROP TABLE u", "PARQUET", "x/")).is_err()); assert!(build_create_external_table_sql(&req("t", "PARQUET'", "x/")).is_err()); let mut r = req("t", "PARQUET", "x/"); @@ -1278,3 +1636,268 @@ mod crawler_admin_tests { assert!(runtime.drop_crawler(&name).await.is_err()); } } + +#[cfg(test)] +mod query_job_tests { + use super::{sweep_jobs, Runtime}; + use crate::api::QueryRequest; + use crate::query::Query; + use crate::query_job::{CancelOutcome, JobKind, PollOutcome, QueryJob, QueryJobState}; + use arrow::record_batch::RecordBatch; + use futures::TryStreamExt; + use std::collections::HashMap; + use std::sync::Arc; + use std::time::{Duration, Instant}; + + /// Boot a runtime (wrapped in `Arc`, as the job API requires), optionally + /// tweaking the config first. + async fn runtime_with(tweak: impl FnOnce(&mut beacon_config::Config)) -> Arc { + let mut config = beacon_config::Config::load().unwrap(); + tweak(&mut config); + Arc::new( + Runtime::new(Arc::new(config)) + .await + .expect("runtime should start"), + ) + } + + /// Seed tables via the synchronous super-user path. + async fn run_sql(runtime: &Runtime, sql: &str) { + runtime + .run_query(Query::sql(sql.to_string()), true) + .await + .expect("sql should run") + .into_record_stream() + .expect("streamed result") + .try_collect::>() + .await + .expect("sql should drain"); + } + + fn json_query(value: serde_json::Value) -> Query { + serde_json::from_value::(value) + .expect("query request should deserialize") + .into_query() + .expect("query request should convert") + } + + /// Poll a streamable job to a terminal state, collecting every batch. + async fn drain_job( + runtime: &Arc, + id: uuid::Uuid, + ) -> Result, String> { + let mut all = Vec::new(); + loop { + match runtime + .poll_query_job_stream(id, Duration::from_millis(500)) + .await + { + PollOutcome::Batches { batches, .. } => all.extend(batches), + PollOutcome::Pending => continue, + PollOutcome::Completed => return Ok(all), + PollOutcome::Failed(error) => return Err(error), + PollOutcome::NotStreamable + | PollOutcome::Cancelled + | PollOutcome::NotFound => panic!("unexpected poll outcome for {id}"), + } + } + } + + /// A streamable job delivers all produced batches and ends `Completed`. + #[tokio::test(flavor = "multi_thread")] + async fn streamable_job_delivers_all_batches() { + let runtime = runtime_with(|_| {}).await; + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("job_ok_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1), (2), (3)")).await; + + let (id, kind) = runtime + .submit_query_job(Query::sql(format!("SELECT a FROM {table}")), false) + .await + .expect("submit should succeed"); + assert_eq!(kind, JobKind::Streamable); + + let batches = drain_job(&runtime, id).await.expect("job should succeed"); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 3, "all inserted rows should be delivered"); + } + + /// A query that fails mid-execution surfaces as a `Failed` terminal poll — + /// the case the synchronous streaming endpoint cannot report. + #[tokio::test(flavor = "multi_thread")] + async fn streamable_job_reports_midstream_failure() { + let runtime = runtime_with(|_| {}).await; + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("job_err_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1), (2)")).await; + + // `a - 1` is zero for a = 1, so the integer division errors at execution + // (and cannot be constant-folded away since `a` is a column). + let (id, _) = runtime + .submit_query_job( + Query::sql(format!("SELECT 1 / (a - 1) AS x FROM {table}")), + false, + ) + .await + .expect("submit should succeed (planning is valid)"); + + let error = drain_job(&runtime, id) + .await + .expect_err("job should fail at execution"); + assert!(!error.is_empty(), "failure should carry an error message"); + + // Status reflects the failure too. + let snapshot = runtime.query_job_snapshot(id).expect("job present"); + assert!(matches!(snapshot.state, QueryJobState::Failed { .. })); + } + + /// With a tiny memory budget every batch spills to disk; results must still + /// come back correctly and the global budget returns to zero. + #[tokio::test(flavor = "multi_thread")] + async fn streamable_job_spills_under_memory_pressure() { + let runtime = runtime_with(|c| c.runtime.query_jobs.buffer_memory_bytes = 1).await; + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("job_spill_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (10), (20), (30)")).await; + + let (id, _) = runtime + .submit_query_job(Query::sql(format!("SELECT a FROM {table}")), false) + .await + .expect("submit should succeed"); + + let batches = drain_job(&runtime, id).await.expect("job should succeed"); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 3, "rows must survive the spill round-trip"); + assert_eq!( + runtime.query_job_budget.used(), + 0, + "all in-memory budget should be released after draining" + ); + } + + /// A file job materializes a result that the snapshot exposes for download. + #[tokio::test(flavor = "multi_thread")] + async fn file_job_materializes_result() { + let runtime = runtime_with(|_| {}).await; + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("job_file_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1), (2)")).await; + + let (id, kind) = runtime + .submit_query_job( + json_query(serde_json::json!({ + "from": table, + "select": ["a"], + "output": { "format": "parquet" }, + })), + false, + ) + .await + .expect("submit should succeed"); + assert_eq!(kind, JobKind::File); + + // Poll status until terminal. + let snapshot = loop { + let snapshot = runtime.query_job_snapshot(id).expect("job present"); + match snapshot.state { + QueryJobState::Running => { + tokio::time::sleep(Duration::from_millis(20)).await; + } + _ => break snapshot, + } + }; + assert!(matches!(snapshot.state, QueryJobState::Succeeded)); + let file = snapshot.file.expect("succeeded file job has a result file"); + let len = std::fs::metadata(&file.path).expect("result file exists").len(); + assert!(len > 0, "result file should be non-empty"); + } + + /// Cancelling a job drives it terminal; a second cancel is a no-op. + #[tokio::test(flavor = "multi_thread")] + async fn cancel_drives_job_terminal() { + let runtime = runtime_with(|_| {}).await; + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("job_cancel_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1)")).await; + + let (id, _) = runtime + .submit_query_job(Query::sql(format!("SELECT a FROM {table}")), false) + .await + .expect("submit should succeed"); + + // The job may finish before we cancel (tiny query): both are valid first + // outcomes. Either way the job is terminal afterward. + let first = runtime.cancel_query_job(id); + assert!(matches!( + first, + CancelOutcome::Cancelled | CancelOutcome::AlreadyFinished + )); + assert!( + !matches!( + runtime.query_job_snapshot(id).expect("job present").state, + QueryJobState::Running + ), + "job should be terminal after cancel" + ); + assert_eq!(runtime.cancel_query_job(id), CancelOutcome::AlreadyFinished); + } + + /// The sweeper aborts idle running streamable jobs and evicts terminal jobs + /// past their TTL. + #[test] + fn sweep_aborts_idle_and_evicts_expired() { + let ttl = Duration::from_secs(600); + let idle = Duration::from_secs(120); + let now = Instant::now(); + let old = now + .checked_sub(Duration::from_secs(10_000)) + .unwrap_or(now); + + let mut map: HashMap = HashMap::new(); + + // Idle running streamable job → should be cancelled (now terminal). + let idle_id = uuid::Uuid::new_v4(); + let mut idle_job = QueryJob::new_running(JobKind::Streamable, None); + idle_job.last_poll_at = old; + map.insert(idle_id, idle_job); + + // Terminal job finished long ago → should be evicted. + let expired_id = uuid::Uuid::new_v4(); + let mut expired = QueryJob::new_running(JobKind::File, None); + expired.state = QueryJobState::Succeeded; + expired.finished_at = Some(old); + map.insert(expired_id, expired); + + // Fresh running job → untouched. + let fresh_id = uuid::Uuid::new_v4(); + map.insert(fresh_id, QueryJob::new_running(JobKind::Streamable, None)); + + sweep_jobs(&mut map, now, ttl, idle); + + assert!( + matches!( + map.get(&idle_id).map(|j| &j.state), + Some(QueryJobState::Cancelled) + ), + "idle streamable job should be cancelled" + ); + assert!(!map.contains_key(&expired_id), "expired job should be evicted"); + assert!( + matches!( + map.get(&fresh_id).map(|j| &j.state), + Some(QueryJobState::Running) + ), + "fresh job should be left running" + ); + } +}