Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ballista/scheduler/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ required-features = ["build-binary"]

[features]
build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"]
default = ["build-binary", "substrait"]
default = ["build-binary", "substrait", "rest-api"]
# job info can cache stage plans, in some cases where
# task plans can be re-computed, cache behavior may need to be disabled.
disable-stage-plan-cache = []
Expand Down
63 changes: 37 additions & 26 deletions ballista/scheduler/src/api/handlers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::scheduler_process::SchedulerErrorResponse;
use crate::scheduler_server::SchedulerServer;
use crate::scheduler_server::event::QueryStageSchedulerEvent;
use crate::state::execution_graph::ExecutionStage;
Expand Down Expand Up @@ -122,15 +123,13 @@ pub async fn get_jobs<
U: AsExecutionPlan + Send + Sync + 'static,
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
) -> Result<impl IntoResponse, StatusCode> {
// TODO: Display last seen information in UI
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
let state = &data_server.state;

let jobs = state
.task_manager
.get_jobs()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let jobs =
state.task_manager.get_jobs().await.map_err(|_| {
SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)
})?;

let jobs: Vec<JobResponse> = jobs
.iter()
Expand Down Expand Up @@ -166,17 +165,17 @@ pub async fn get_job<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path(job_id): Path<String>,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
let graph = data_server
.state
.task_manager
.get_job_execution_graph(&job_id)
.await
.map_err(|err| {
tracing::error!("Error occurred while getting the execution graph for job '{job_id}': {err:?}");
StatusCode::INTERNAL_SERVER_ERROR
tracing::error!("Error occurred while getting the execution graph for job '{job_id}' reason: {err:?}");
SchedulerErrorResponse::with_error(StatusCode::INTERNAL_SERVER_ERROR, format!("Error occurred while getting the execution graph for job '{job_id}' reason: {}", err))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Several 500 responses now embed the underlying error string in the HTTP JSON via SchedulerErrorResponse::with_error(...). Can you confirm this level of detail is OK to expose to REST clients (it can leak internal information) vs logging it and returning a generic message?

Severity: medium

Other Locations
  • ballista/scheduler/src/api/handlers.rs:218
  • ballista/scheduler/src/api/handlers.rs:288

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The Augment AI reviewer is correct! The internal errors should not be added to the JSON responses as is, because they may contain sensitive information. The errors should be just logged and the JSON response should contain some generic information about the problem. Only users with access to the logs should see the full details.

})?
Comment on lines 174 to 177

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Don't serialize raw backend errors into the public REST body.

These branches return err/e.to_string() directly to clients. With REST now default-on, that leaks internal details and makes the HTTP contract depend on unstable backend error text. Log the full error, but return a fixed client-safe message here.

Example hardening pattern
-            SchedulerErrorResponse::with_error(
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Error getting job status: {}", err),
-            )
+            SchedulerErrorResponse::with_error(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to get job status".to_string(),
+            )

Apply the same pattern to the other handlers in this block.

Also applies to: 218-221, 287-291

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@ballista/scheduler/src/api/handlers.rs` around lines 174 - 177, The map_err
closures in handlers.rs that currently log the full backend error but return
err.to_string() to clients (e.g., the closure building the "Error occurred while
getting the execution graph for job '{job_id}' reason: {err:?}" message) must be
changed to avoid leaking internals: keep tracing::error!(.., err = ?err) to log
the detailed error, but replace the returned SchedulerErrorResponse body with a
fixed, client-safe message like "Internal server error while retrieving
execution graph" (do similarly for the other map_err branches referenced around
the other handlers at the indicated ranges). Update the closures that construct
SchedulerErrorResponse (the map_err for the execution graph and the similar
closures at the other ranges) so the status stays INTERNAL_SERVER_ERROR but the
formatted message is a constant, non-sensitive string while logging includes the
original err for diagnostics.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The internal errors should not be added to the JSON responses as is, because they may contain sensitive information. The errors should be just logged and the JSON response should contain some generic information about the problem. Only users with access to the logs should see the full details.

.ok_or(StatusCode::NOT_FOUND)?;
.ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?;
let stage_plan = format!("{:?}", graph);
let job = graph.as_ref();
let (plain_status, job_status) =
Expand Down Expand Up @@ -207,7 +206,7 @@ pub async fn cancel_job<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path(job_id): Path<String>,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
// 404 if the job doesn't exist
let job_status = data_server
.state
Expand All @@ -216,9 +215,12 @@ pub async fn cancel_job<
.await
.map_err(|err| {
tracing::error!("Error getting job status: {err:?}");
StatusCode::INTERNAL_SERVER_ERROR
SchedulerErrorResponse::with_error(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Error getting job status: {}", err),
)
})?
.ok_or(StatusCode::NOT_FOUND)?;
.ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?;

match &job_status.status {
None | Some(Status::Queued(_)) | Some(Status::Running(_)) => {
Expand All @@ -229,11 +231,13 @@ pub async fn cancel_job<
tracing::error!(
"Error getting query stage event loop sender: {err:?}"
);
StatusCode::INTERNAL_SERVER_ERROR
SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)
})?
.post_event(QueryStageSchedulerEvent::JobCancel(job_id))
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
.map_err(|_| {
SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)
})?;

Ok((
StatusCode::OK,
Expand Down Expand Up @@ -274,13 +278,18 @@ pub async fn get_query_stages<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path(job_id): Path<String>,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
if let Some(graph) = data_server
.state
.task_manager
.get_job_execution_graph(&job_id)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
.map_err(|e| {
SchedulerErrorResponse::with_error(
StatusCode::INTERNAL_SERVER_ERROR,
e.to_string(),
)
})?
{
let stages = graph
.as_ref()
Expand Down Expand Up @@ -409,16 +418,16 @@ pub async fn get_job_dot_graph<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path(job_id): Path<String>,
) -> Result<String, StatusCode> {
) -> Result<String, SchedulerErrorResponse> {
if let Some(graph) = data_server
.state
.task_manager
.get_job_execution_graph(&job_id)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
.map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))?
{
ExecutionGraphDot::generate(graph.as_ref())
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
.map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))
} else {
Ok("Not Found".to_string())
}
Comment on lines +421 to 433

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

The DOT endpoints shouldn't succeed on a missing job.

Line 432 and Line 453 return plain "Not Found" with HTTP 200. That breaks REST semantics, and get_job_svg_graph downstream turns the same condition into 400 Cannot parse graph. Return a structured 404 here instead.

Suggested fix
-    } else {
-        Ok("Not Found".to_string())
-    }
+    } else {
+        Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND))
+    }

Apply the same change in both endpoints.

Also applies to: 442-454

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@ballista/scheduler/src/api/handlers.rs` around lines 421 - 433, The DOT
endpoints currently return Ok("Not Found") when get_job_execution_graph returns
None; change those branches to return a structured 404 error instead of a 200
string. Replace the else return in the functions that call
data_server.state.task_manager.get_job_execution_graph(...) (the handlers that
then call ExecutionGraphDot::generate) so they return
Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) when the job graph is
missing; apply the same change to both DOT-related endpoints (the functions that
invoke ExecutionGraphDot::generate and the similar SVG endpoint).

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:good-to-have; category:bug; feedback: This code is not changed in the reviewed Pull Request but the CodeRabbit AI reviewer is correct! Instead of returning status=200 with text body "Not found" these endpoints should return 404 with an optional JSON body with more details.

Expand All @@ -430,16 +439,16 @@ pub async fn get_query_stage_dot_graph<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path((job_id, stage_id)): Path<(String, usize)>,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
if let Some(graph) = data_server
.state
.task_manager
.get_job_execution_graph(&job_id)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
.map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))?
{
ExecutionGraphDot::generate_for_query_stage(graph.as_ref(), stage_id)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
.map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))
} else {
Ok("Not Found".to_string())
}
Expand All @@ -451,7 +460,7 @@ pub async fn get_job_svg_graph<
>(
State(data_server): State<Arc<SchedulerServer<T, U>>>,
Path(job_id): Path<String>,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, SchedulerErrorResponse> {
let dot = get_job_dot_graph(State(data_server.clone()), Path(job_id)).await?;
match graphviz_rust::parse(&dot) {
Ok(graph) => {
Expand All @@ -460,7 +469,9 @@ pub async fn get_job_svg_graph<
&mut PrinterContext::default(),
vec![CommandArg::Format(Format::Svg)],
)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
.map_err(|_| {
SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR)
})?;

let svg = String::from_utf8_lossy(&result).to_string();
Ok(Response::builder()
Expand Down
1 change: 1 addition & 0 deletions ballista/scheduler/src/api/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ pub fn get_routes<
) -> Router {
let router = Router::new()
.route("/api/state", get(handlers::get_scheduler_state::<T, U>))
.route("/api/version", get(handlers::get_scheduler_state::<T, U>))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The /api/version endpoint reuses the get_scheduler_state handler, which returns both the scheduler start time and the version. For an endpoint named /api/version, it would be more idiomatic and less surprising for clients if it returned only the version information.

Consider creating a new, simpler handler for this route that doesn't require state and returns only the version.

For example, you could add a new handler in ballista/scheduler/src/api/handlers.rs:

#[derive(Debug, serde::Serialize)]
pub struct VersionResponse {
    version: &'static str,
}

pub async fn get_version() -> impl IntoResponse {
    Json(VersionResponse {
        version: BALLISTA_VERSION,
    })
}

And then update this line to use it:

.route("/api/version", get(handlers::get_version))

This would make the API clearer and more maintainable.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:good-but-wont-fix; category:bug; feedback: The Gemini AI reviewer is correct! The new REST endpoint /api/version is exactly the same as /api/state. But the goal is to extend the /api/state endpoint with more information in a follow-up Pull Request. Any user who don't need the extra information should use the new /api/version endpoint.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/api/version is currently routed to get_scheduler_state, so it returns the full state payload (including started) and duplicates /api/state; is that intentional for clients expecting a version-only endpoint? If not, consider wiring this path to a dedicated version handler.

Severity: medium

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:good-but-wont-fix; category:bug; feedback: The Augment AI reviewer is correct! The new REST endpoint /api/version is exactly the same as /api/state. But the goal is to extend the /api/state endpoint with more information in a follow-up Pull Request. Any user who don't need the extra information should use the new /api/version endpoint.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Version endpoint incorrectly uses scheduler state handler

Medium Severity

The new /api/version route is wired to handlers::get_scheduler_state, the same handler used by /api/state. This handler returns a SchedulerStateResponse containing both started (the scheduler start timestamp) and version, making it a duplicate of /api/state rather than a dedicated version endpoint. Consumers hitting /api/version would receive unexpected extra fields (like started), and the two routes are fully redundant.

Fix in Cursor Fix in Web

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:good-but-wont-fix; category:bug; feedback: The Bugbot AI reviewer is correct! The new REST endpoint /api/version is exactly the same as /api/state. But the goal is to extend the /api/state endpoint with more information in a follow-up Pull Request. Any user who don't need the extra information should use the new /api/version endpoint.

.route("/api/executors", get(handlers::get_executors::<T, U>))
.route("/api/jobs", get(handlers::get_jobs::<T, U>))
.route(
Expand Down
12 changes: 12 additions & 0 deletions ballista/scheduler/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,11 @@ pub struct Config {
help = "The interval to check expired or dead executors"
)]
pub expire_dead_executor_interval_seconds: u64,

#[cfg(feature = "rest-api")]
/// Should the rest api be disabled
#[arg(long, default_value_t = false, help = "Should the rest api be disable")]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

There's a small typo in the help message. It should be 'disabled' instead of 'disable'.

Suggested change
#[arg(long, default_value_t = false, help = "Should the rest api be disable")]
#[arg(long, default_value_t = false, help = "Should the rest api be disabled")]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The --disable-rest help text is user-facing and currently says "Should the rest api be disable" (grammar/typo).

Severity: low

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

pub disable_rest: bool,
}

/// Configurations for the ballista scheduler of scheduling jobs and tasks
Expand Down Expand Up @@ -245,6 +250,9 @@ pub struct SchedulerConfig {
pub override_create_grpc_client_endpoint: Option<EndpointOverrideFn>,
/// Whether to use TLS when connecting to executors (for flight proxy)
pub use_tls: bool,
#[cfg(feature = "rest-api")]
/// Should the rest api be disabled
pub disable_rest: bool,
Comment on lines +253 to +255

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== rest-api default feature ==\n'
rg -n -C2 '^default = .*rest-api' ballista/scheduler/Cargo.toml

printf '\n== SchedulerConfig definition ==\n'
rg -n -C2 'pub struct SchedulerConfig|#\[non_exhaustive\]|disable_rest' ballista/scheduler/src/config.rs

printf '\n== SchedulerConfig struct literals in the repo ==\n'
rg -n -C2 --type=rust '\bSchedulerConfig\s*\{'

Repository: martin-augment/datafusion-ballista

Length of output: 3848


Avoid adding a public field to SchedulerConfig without semver protection.

SchedulerConfig is public and exhaustively constructible (no #[non_exhaustive]). Adding disable_rest breaks downstream crates that construct it via struct literals, and because rest-api is a default feature (Line 37, Cargo.toml), this affects default builds. The field addition under #[cfg(feature = "rest-api")] also causes the struct's shape to vary by feature flag, compounding the breaking change.

Either keep this field out of the public struct surface, use a builder, or defer to the next breaking release.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@ballista/scheduler/src/config.rs` around lines 253 - 255, The new public
field disable_rest on SchedulerConfig is a breaking change; instead make the
field private and feature-gated or move it out of the public struct surface and
expose feature-gated accessors or a builder; specifically, remove or change the
pub disable_rest to a private field inside SchedulerConfig (or to a separate
internal config struct), add a getter method (e.g.,
SchedulerConfig::disable_rest() behind the rest-api feature) or introduce a
builder for constructing SchedulerConfig with that option, or mark
SchedulerConfig #[non_exhaustive] if you intend to accept the breaking change in
a major release—pick one approach and apply it consistently so downstream crates
aren't broken by feature-flag variations.

}

impl Default for SchedulerConfig {
Expand Down Expand Up @@ -273,6 +281,8 @@ impl Default for SchedulerConfig {
override_physical_codec: None,
override_create_grpc_client_endpoint: None,
use_tls: false,
#[cfg(feature = "rest-api")]
disable_rest: false,
}
}
}
Expand Down Expand Up @@ -520,6 +530,8 @@ impl TryFrom<Config> for SchedulerConfig {
override_session_builder: None,
override_create_grpc_client_endpoint: None,
use_tls: false,
#[cfg(feature = "rest-api")]
disable_rest: opt.disable_rest,
};

Ok(config)
Expand Down
77 changes: 70 additions & 7 deletions ballista/scheduler/src/scheduler_process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
use crate::flight_proxy_service::BallistaFlightProxyService;

use arrow_flight::flight_service_server::FlightServiceServer;
use axum::Json;
use axum::response::{IntoResponse, Response};
use ballista_core::BALLISTA_VERSION;
use ballista_core::error::BallistaError;
use ballista_core::extension::BallistaConfigGrpcEndpoint;
Expand Down Expand Up @@ -130,17 +132,42 @@ pub async fn start_grpc_service<
tonic_builder.add_service(ExternalScalerServer::new(scheduler.clone()));

let tonic = tonic_builder.routes().into_axum_router();
let tonic = tonic.fallback(|| async { (StatusCode::NOT_FOUND, "404 - Not Found") });

// registering default handler for unmatched requests
let tonic =
tonic.fallback(|| async { SchedulerErrorResponse::new(StatusCode::NOT_FOUND) });

#[cfg(feature = "rest-api")]
let axum = get_routes(Arc::new(scheduler));
#[cfg(feature = "rest-api")]
let final_route = axum
.merge(tonic)
.into_make_service_with_connect_info::<SocketAddr>();
let final_route = if config.disable_rest {
tonic
.route(
"/api/{*path}",
axum::routing::any(|| async {
SchedulerErrorResponse::with_error(
StatusCode::NOT_FOUND,
"Rest api has been disabled at startup".to_string(),
)
}),
)
.into_make_service_with_connect_info::<SocketAddr>()
} else {
let axum = get_routes(Arc::new(scheduler));
axum.merge(tonic)
.into_make_service_with_connect_info::<SocketAddr>()
};

#[cfg(not(feature = "rest-api"))]
let final_route = tonic.into_make_service_with_connect_info::<SocketAddr>();
let final_route = tonic
.route(
"/api/{*path}",
axum::routing::any(|| async {
SchedulerErrorResponse::with_error(
StatusCode::NOT_FOUND,
"Rest api has been disabled at compile time".to_string(),
)
}),
)
.into_make_service_with_connect_info::<SocketAddr>();

let listener = tokio::net::TcpListener::bind(&address)
.await
Expand All @@ -165,3 +192,39 @@ pub async fn start_server(

start_grpc_service(address, scheduler).await
}

#[derive(Debug, serde::Serialize)]
pub(crate) struct SchedulerErrorResponse {
#[serde(skip)]
status_code: StatusCode,
http_code: u16,
reason: Option<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}

impl SchedulerErrorResponse {
pub(crate) fn new(status_code: StatusCode) -> Self {
Self {
status_code,
reason: status_code.canonical_reason(),
http_code: status_code.as_u16(),
error: None,
}
}
pub(crate) fn with_error(status_code: StatusCode, error: String) -> Self {
Self {
status_code,
reason: status_code.canonical_reason(),
http_code: status_code.as_u16(),
error: Some(error),
}
}
}

impl IntoResponse for SchedulerErrorResponse {
fn into_response(self) -> Response {
let status = self.status_code;
(status, Json(self)).into_response()
}
}
Loading