-
Notifications
You must be signed in to change notification settings - Fork 0
1506: feat: enable scheduler rest api by default #7
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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; | ||
|
|
@@ -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() | ||
|
|
@@ -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)) | ||
| })? | ||
|
Comment on lines
174
to
177
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Don't serialize raw backend errors into the public REST body. These branches return 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
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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) = | ||
|
|
@@ -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 | ||
|
|
@@ -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(_)) => { | ||
|
|
@@ -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, | ||
|
|
@@ -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() | ||
|
|
@@ -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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The DOT endpoints shouldn't succeed on a missing job. Line 432 and Line 453 return plain 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
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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. |
||
|
|
@@ -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()) | ||
| } | ||
|
|
@@ -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) => { | ||
|
|
@@ -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() | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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>)) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The 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 #[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.
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Severity: medium 🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Version endpoint incorrectly uses scheduler state handlerMedium Severity The new
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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( | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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")] | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| pub disable_rest: bool, | ||
| } | ||
|
|
||
| /// Configurations for the ballista scheduler of scheduling jobs and tasks | ||
|
|
@@ -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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧩 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
Either keep this field out of the public struct surface, use a builder, or defer to the next breaking release. 🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| impl Default for SchedulerConfig { | ||
|
|
@@ -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, | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -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) | ||
|
|
||


There was a problem hiding this comment.
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: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.
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.