1506: feat: enable scheduler rest api by default - #7
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request primarily focuses on enhancing the scheduler's REST API. It enables the REST API by default, streamlines error handling across various API endpoints using a new custom error response structure, and provides a configuration option to disable the API if needed. These changes aim to improve the usability and robustness of the scheduler's external interface. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
WalkthroughThe changes introduce structured error responses for the REST API layer, replacing StatusCode error returns with a new SchedulerErrorResponse type across multiple handler methods. Compile-time and runtime toggles are added for REST API control through feature flags and configuration settings. A new "/api/version" endpoint is added to the routing, and the default feature set is extended to include rest-api. Error handling is unified with JSON payload responses for consistency. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Code ReviewOverall this PR makes a reasonable improvement by enabling the REST API by default and introducing a structured JSON error response type. Here are my findings: Bugs / Correctness Issues1. // api/mod.rs
.route("/api/version", get(handlers::get_scheduler_state::<T, U>))
2. } else {
Ok("Not Found".to_string()) // handlers.rs:426 and ~448
}When the job doesn't exist these handlers silently return HTTP 200 with body Design / Architecture3. 4. Should be: Security / Information Disclosure5. Internal error messages exposed in HTTP responses SchedulerErrorResponse::with_error(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Error occurred while getting the execution graph for job '{job_id}' reason: {}", err),
)The Minor / Style6. Error context is dropped silently in several places // Several handlers still use:
.map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))?The 7. Test CoverageNo tests were added for:
At minimum a unit test for Summary
The core direction (structured JSON errors, default-on REST API, runtime disable flag) is solid. Addressing the two bugs and the security concern would be the priority items before merging. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables the scheduler's REST API by default, which is a great enhancement for usability. The implementation includes a new configuration flag --disable-rest to opt-out, and improves API error handling by introducing a structured SchedulerErrorResponse.
My review includes a couple of minor suggestions to improve the new API endpoint for versioning and to fix a typo in a command-line help message. Overall, the changes are well-implemented and align with the goal of the pull request.
| ) -> 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>)) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
|
|
||
| #[cfg(feature = "rest-api")] | ||
| /// Should the rest api be disabled | ||
| #[arg(long, default_value_t = false, help = "Should the rest api be disable")] |
There was a problem hiding this comment.
There's a small typo in the help message. It should be 'disabled' instead of 'disable'.
| #[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")] |
🤖 Augment PR SummarySummary: This PR enables the scheduler REST API by default and standardizes REST-facing error handling. Changes:
Technical Notes: gRPC and REST continue to share the same Axum server/router; error responses are now consistently JSON-encoded with status codes. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| ) -> 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>)) |
There was a problem hiding this comment.
/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
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
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.
| 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)) |
There was a problem hiding this comment.
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:218ballista/scheduler/src/api/handlers.rs:288
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
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.
|
|
||
| #[cfg(feature = "rest-api")] | ||
| /// Should the rest api be disabled | ||
| #[arg(long, default_value_t = false, help = "Should the rest api be disable")] |
There was a problem hiding this comment.
The --disable-rest help text is user-facing and currently says "Should the rest api be disable" (grammar/typo).
Severity: low
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
| ) -> 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>)) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ballista/scheduler/src/api/handlers.rs (1)
281-345:⚠️ Potential issue | 🟠 MajorReturn 404 when the job is missing.
The
elsebranch currently returns200 {"stages":[]}for an unknownjob_id. That makes a missing job indistinguishable from a job with no stages and is inconsistent withget_jobandcancel_job. ReturnErr(SchedulerErrorResponse::new(StatusCode::NOT_FOUND))instead.Suggested fix
- } else { - Ok(Json(QueryStagesResponse { stages: vec![] })) - } + } else { + Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) + }🤖 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 281 - 345, The handler currently treats a missing job (None from data_server.state.task_manager.get_job_execution_graph) as Ok(Json(QueryStagesResponse { stages: vec![] })), but should return a 404 like get_job and cancel_job; replace the else branch to return Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) when graph is None (ensure you reference get_job_execution_graph, QueryStagesResponse and SchedulerErrorResponse), and add any needed imports so the NOT_FOUND status compiles.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ballista/scheduler/src/api/handlers.rs`:
- Around line 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).
- Around line 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.
In `@ballista/scheduler/src/config.rs`:
- Around line 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.
---
Outside diff comments:
In `@ballista/scheduler/src/api/handlers.rs`:
- Around line 281-345: The handler currently treats a missing job (None from
data_server.state.task_manager.get_job_execution_graph) as
Ok(Json(QueryStagesResponse { stages: vec![] })), but should return a 404 like
get_job and cancel_job; replace the else branch to return
Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) when graph is None
(ensure you reference get_job_execution_graph, QueryStagesResponse and
SchedulerErrorResponse), and add any needed imports so the NOT_FOUND status
compiles.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: bd94c56f-6cb8-4e7b-838b-5f603a6e1cd2
📒 Files selected for processing (5)
ballista/scheduler/Cargo.tomlballista/scheduler/src/api/handlers.rsballista/scheduler/src/api/mod.rsballista/scheduler/src/config.rsballista/scheduler/src/scheduler_process.rs
| .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)) | ||
| })? |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| ) -> 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()) | ||
| } |
There was a problem hiding this comment.
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).
There was a problem hiding this comment.
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.
| #[cfg(feature = "rest-api")] | ||
| /// Should the rest api be disabled | ||
| pub disable_rest: bool, |
There was a problem hiding this comment.
🧩 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.
value:good-but-wont-fix; category:bug; feedback: The Claude 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. |
value:good-to-have; category:bug; feedback: This code is not changed in the reviewed Pull Request but the Claude 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. |
value:useful; category:bug; feedback: The Claude AI reviewer is correct! The idea of the Pull Request author is to provide a minimal REST API that catches all requests to /api/* and returns an error explaining that the REST API cargo feature is not enabled. But to implement it he uses the always enabled gRPC-related modules. Instead he should use |
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. |
2 participants
1506: To review by AI