You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mi...@apache.org on 2026/09/12 21:11:14 UTC
(datafusion-ballista) branch main updated: feat(scheduler): document REST API with OpenAPI and serve /api/openapi.json (#2397)
This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new ae9bff026 feat(scheduler): document REST API with OpenAPI and serve /api/openapi.json (#2397)
ae9bff026 is described below
commit ae9bff026ffd18b256d39b10890a2ca96e1fc836
Author: debaditya <82...@users.noreply.github.com>
AuthorDate: Sun Sep 13 02:41:07 2026 +0530
feat(scheduler): document REST API with OpenAPI and serve /api/openapi.json (#2397)
* feat(scheduler): document REST API with OpenAPI and serve /api/openapi.json (#2351)
* fix(scheduler): address review feedback on OpenAPI spec, core schema derives, and health probes
* test(api-types): enable utoipa feature for tests via self dev-dependency
---
Cargo.lock | 28 +++
Cargo.toml | 1 +
ballista/api-types/Cargo.toml | 9 +
ballista/api-types/src/dto.rs | 79 ++++++++
ballista/core/Cargo.toml | 2 +
ballista/core/src/serde/scheduler/mod.rs | 2 +
ballista/scheduler/Cargo.toml | 10 +-
ballista/scheduler/src/api/handlers.rs | 197 ++++++++++++++++---
ballista/scheduler/src/api/health.rs | 27 ++-
ballista/scheduler/src/api/mod.rs | 5 +
ballista/scheduler/src/api/openapi.rs | 320 +++++++++++++++++++++++++++++++
ballista/scheduler/src/api/routes.rs | 3 +-
docs/source/user-guide/scheduler.md | 1 +
13 files changed, 657 insertions(+), 27 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 3f4e07f05..52a29e98c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1028,7 +1028,10 @@ dependencies = [
name = "ballista-api-types"
version = "54.0.0"
dependencies = [
+ "ballista-api-types",
"serde",
+ "serde_json",
+ "utoipa",
]
[[package]]
@@ -1141,6 +1144,7 @@ dependencies = [
"tonic-prost",
"tonic-prost-build",
"url",
+ "utoipa",
"uuid",
]
@@ -1259,6 +1263,7 @@ dependencies = [
"tracing-appender",
"tracing-subscriber",
"url-escape",
+ "utoipa",
"uuid",
]
@@ -7781,6 +7786,29 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
+[[package]]
+name = "utoipa"
+version = "5.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8bde15df68e80b16c7d16b9616e80770ad158988daa56a27dccd1e55558b0160"
+dependencies = [
+ "indexmap 2.14.0",
+ "serde",
+ "serde_json",
+ "utoipa-gen",
+]
+
+[[package]]
+name = "utoipa-gen"
+version = "5.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6ba0b99ee52df3028635d93840c797102da61f8a7bb3cf751032455895b52ef8"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
[[package]]
name = "uuid"
version = "1.26.0"
diff --git a/Cargo.toml b/Cargo.toml
index 0b4f34224..66ab42dab 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -98,6 +98,7 @@ tokio = { version = "1" }
tokio-stream = { version = "0.1" }
url = { version = "2.5" }
url-escape = { version = "0.1.2" }
+utoipa = { version = "5.5", default-features = false, features = ["macros"] }
uuid = { version = "1.23", features = ["v4", "v7"] }
# cargo build --profile release-lto
diff --git a/ballista/api-types/Cargo.toml b/ballista/api-types/Cargo.toml
index 1715b1843..12b818ad9 100644
--- a/ballista/api-types/Cargo.toml
+++ b/ballista/api-types/Cargo.toml
@@ -26,5 +26,14 @@ authors = { workspace = true }
edition = { workspace = true }
rust-version = { workspace = true }
+[features]
+default = []
+utoipa = ["dep:utoipa"]
+
[dependencies]
serde = { workspace = true, features = ["derive"] }
+utoipa = { workspace = true, optional = true }
+
+[dev-dependencies]
+ballista-api-types = { path = ".", features = ["utoipa"] }
+serde_json = { workspace = true }
diff --git a/ballista/api-types/src/dto.rs b/ballista/api-types/src/dto.rs
index 732df63dc..f91d37dda 100644
--- a/ballista/api-types/src/dto.rs
+++ b/ballista/api-types/src/dto.rs
@@ -23,6 +23,7 @@ use std::collections::BTreeMap;
/// Summary of one job, served by `GET /api/jobs` and `GET /api/job/{job_id}`.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct JobResponse {
/// A `String` rather than `ballista_core::JobId` so this crate stays a
/// serde-only leaf. `JobId` is `#[serde(transparent)]` over `String`, so
@@ -57,6 +58,7 @@ pub struct JobResponse {
/// Terminal or in-flight state of a single task.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub enum TaskStatus {
/// Task is currently executing.
Running,
@@ -73,6 +75,7 @@ pub enum TaskStatus {
/// Per-task detail within a stage.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct TaskSummary {
/// task id
pub id: usize,
@@ -105,6 +108,7 @@ pub struct TaskSummary {
/// Five-number summary over a stage's tasks, used to spot skew.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct Percentiles {
/// Smallest observed value.
pub min: u64,
@@ -120,6 +124,7 @@ pub struct Percentiles {
/// Summary of one query stage, served by `GET /api/job/{job_id}/stages`.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct QueryStageSummary {
/// Stage id, as a string.
pub stage_id: String,
@@ -150,6 +155,7 @@ pub struct QueryStageSummary {
/// Response body for `GET /api/job/{job_id}/stages`.
#[derive(Debug, Clone, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct QueryStagesResponse {
/// One summary per stage, in stage order.
pub stages: Vec<QueryStageSummary>,
@@ -159,6 +165,7 @@ pub struct QueryStagesResponse {
/// parameter, and part of the REST contract, so it lives alongside the
/// response types rather than with the HTTP handlers.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum PlanFormat {
/// `?plan_format=default` => plain indent, no metrics
@@ -176,3 +183,75 @@ pub enum PlanFormat {
/// A `BTreeMap` so key order is deterministic: the same job must serialize
/// identically whether it is served live or replayed from a stored log.
pub type JobConfig = BTreeMap<String, String>;
+
+#[cfg(all(test, feature = "utoipa"))]
+mod tests {
+ use super::*;
+ use std::collections::BTreeSet;
+ use utoipa::ToSchema;
+
+ fn get_schema_property_names<T: ToSchema>() -> BTreeSet<String> {
+ let schema_ref = T::schema();
+ match schema_ref {
+ utoipa::openapi::RefOr::T(utoipa::openapi::Schema::Object(obj)) => {
+ obj.properties.keys().cloned().collect()
+ }
+ _ => panic!("Expected Object schema"),
+ }
+ }
+
+ fn get_json_property_names<T: Serialize>(value: &T) -> BTreeSet<String> {
+ let json = serde_json::to_value(value).expect("serialize to value");
+ match json {
+ serde_json::Value::Object(map) => map.keys().cloned().collect(),
+ _ => panic!("Expected JSON Object"),
+ }
+ }
+
+ #[test]
+ fn test_dto_schemas() {
+ assert_eq!(JobResponse::name(), "JobResponse");
+ assert_eq!(TaskStatus::name(), "TaskStatus");
+ assert_eq!(TaskSummary::name(), "TaskSummary");
+ assert_eq!(Percentiles::name(), "Percentiles");
+ assert_eq!(QueryStageSummary::name(), "QueryStageSummary");
+ assert_eq!(QueryStagesResponse::name(), "QueryStagesResponse");
+ assert_eq!(PlanFormat::name(), "PlanFormat");
+
+ let task_summary = TaskSummary {
+ id: 1,
+ status: TaskStatus::Running,
+ partition_id: vec![0],
+ scheduled_time: 100,
+ launch_time: 100,
+ start_exec_time: 100,
+ end_exec_time: 200,
+ exec_duration: 100,
+ finish_time: 200,
+ input_rows: 10,
+ output_rows: 10,
+ };
+ assert_eq!(
+ get_schema_property_names::<TaskSummary>(),
+ get_json_property_names(&task_summary)
+ );
+
+ let percentiles = Percentiles {
+ min: 1,
+ p25: 2,
+ median: 3,
+ p75: 4,
+ max: 5,
+ };
+ assert_eq!(
+ get_schema_property_names::<Percentiles>(),
+ get_json_property_names(&percentiles)
+ );
+
+ let query_stages = QueryStagesResponse { stages: vec![] };
+ assert_eq!(
+ get_schema_property_names::<QueryStagesResponse>(),
+ get_json_property_names(&query_stages)
+ );
+ }
+}
diff --git a/ballista/core/Cargo.toml b/ballista/core/Cargo.toml
index 35c47082e..b9da2bac0 100644
--- a/ballista/core/Cargo.toml
+++ b/ballista/core/Cargo.toml
@@ -42,6 +42,7 @@ docsrs = []
# Used for testing ONLY: causes all values to hash to the same value (test for collisions)
force_hash_collisions = ["datafusion/force_hash_collisions"]
spark-compat = ["dep:datafusion-spark"]
+utoipa = ["dep:utoipa"]
[dependencies]
arrow-flight = { workspace = true }
@@ -65,6 +66,7 @@ tokio-stream = { workspace = true, features = ["net"] }
tonic = { workspace = true }
tonic-prost = { workspace = true }
url = { workspace = true }
+utoipa = { workspace = true, optional = true }
uuid = { workspace = true }
[dev-dependencies]
diff --git a/ballista/core/src/serde/scheduler/mod.rs b/ballista/core/src/serde/scheduler/mod.rs
index f5a9c65ca..e025864a7 100644
--- a/ballista/core/src/serde/scheduler/mod.rs
+++ b/ballista/core/src/serde/scheduler/mod.rs
@@ -204,6 +204,7 @@ pub struct ExecutorMetadata {
/// Specification of an executor, indicating its runtime-assigned vcore count.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct ExecutorSpecification {
/// Virtual cores assigned to this executor at runtime — analogous to YARN
/// vcores (`yarn.nodemanager.resource.cpu-vcores`) or Spark's
@@ -241,6 +242,7 @@ impl ExecutorSpecification {
/// Operating system level specification of an executor
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct ExecutorOperatingSystemSpecification {
/// System name
pub system_name: String,
diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml
index 0f3708c6e..a9f49b79f 100644
--- a/ballista/scheduler/Cargo.toml
+++ b/ballista/scheduler/Cargo.toml
@@ -54,7 +54,14 @@ disable-stage-plan-cache = []
graphviz-support = ["dep:graphviz-rust"]
keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"]
prometheus-metrics = ["prometheus", "once_cell"]
-rest-api = ["dep:ballista-api-types", "dep:ballista-history", "dep:serde_json"]
+rest-api = [
+ "ballista-api-types?/utoipa",
+ "ballista-core/utoipa",
+ "dep:ballista-api-types",
+ "dep:ballista-history",
+ "dep:serde_json",
+ "dep:utoipa",
+]
spark-compat = ["ballista-core/spark-compat"]
substrait = ["dep:datafusion-substrait"]
@@ -93,6 +100,7 @@ tracing = { workspace = true, optional = true }
tracing-appender = { workspace = true, optional = true }
tracing-subscriber = { workspace = true, optional = true }
url-escape = { workspace = true }
+utoipa = { workspace = true, optional = true }
uuid = { workspace = true }
[[test]]
diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs
index 345168c09..1b195a92d 100644
--- a/ballista/scheduler/src/api/handlers.rs
+++ b/ballista/scheduler/src/api/handlers.rs
@@ -21,7 +21,7 @@ use axum::{
extract::{Path, State},
response::{IntoResponse, Response},
};
-use ballista_api_types::dto::{JobResponse, PlanFormat};
+use ballista_api_types::dto::{JobResponse, PlanFormat, QueryStagesResponse};
use ballista_core::BALLISTA_VERSION;
use ballista_core::serde::protobuf::job_status::Status;
use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric};
@@ -42,29 +42,29 @@ use http::{HeaderMap, StatusCode, header::CONTENT_TYPE};
use std::collections::HashMap;
use std::sync::Arc;
-#[derive(Debug, serde::Serialize)]
-struct SchedulerStateResponse {
- started: u128,
- version: &'static str,
- datafusion_version: &'static str,
- substrait_support: bool,
- keda_support: bool,
- prometheus_support: bool,
- graphviz_support: bool,
- spark_support: bool,
- scheduling_policy: String,
+#[derive(Debug, serde::Serialize, utoipa::ToSchema)]
+pub struct SchedulerStateResponse {
+ pub started: u128,
+ pub version: &'static str,
+ pub datafusion_version: &'static str,
+ pub substrait_support: bool,
+ pub keda_support: bool,
+ pub prometheus_support: bool,
+ pub graphviz_support: bool,
+ pub spark_support: bool,
+ pub scheduling_policy: String,
#[serde(skip_serializing_if = "Option::is_none")]
- advertise_flight_endpoint: Option<String>,
- enable_embedded_flight_proxy: bool,
+ pub advertise_flight_endpoint: Option<String>,
+ pub enable_embedded_flight_proxy: bool,
}
-#[derive(Debug, serde::Serialize)]
-struct SchedulerVersionResponse {
- version: &'static str,
- datafusion_version: &'static str,
+#[derive(Debug, serde::Serialize, utoipa::ToSchema)]
+pub struct SchedulerVersionResponse {
+ pub version: &'static str,
+ pub datafusion_version: &'static str,
}
-#[derive(Debug, serde::Serialize)]
+#[derive(Debug, serde::Serialize, utoipa::ToSchema)]
pub struct ExecutorResponse {
pub id: String,
pub host: String,
@@ -75,7 +75,7 @@ pub struct ExecutorResponse {
pub os_info: ExecutorOperatingSystemSpecification,
}
-#[derive(Debug, serde::Serialize)]
+#[derive(Debug, serde::Serialize, utoipa::ToSchema)]
#[serde(tag = "type", content = "value", rename_all = "snake_case")]
#[allow(clippy::enum_variant_names)]
pub enum ExecutorMetricResponse {
@@ -102,14 +102,15 @@ impl ExecutorMetricResponse {
}
}
-#[derive(Debug, serde::Serialize)]
-struct CancelJobResponse {
+#[derive(Debug, serde::Serialize, utoipa::ToSchema)]
+pub struct CancelJobResponse {
pub cancelled: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub reason: Option<String>,
}
-#[derive(Debug, serde::Deserialize, Default)]
+#[derive(Debug, serde::Deserialize, Default, utoipa::IntoParams, utoipa::ToSchema)]
+#[into_params(parameter_in = Query)]
pub struct JobQueryParams {
/// Controls plan format
pub plan_format: Option<PlanFormat>,
@@ -167,6 +168,14 @@ pub async fn get_webtui<
Ok(Redirect::to(&target))
}
+#[utoipa::path(
+ get,
+ path = "/api/state",
+ tag = "state",
+ responses(
+ (status = 200, description = "Scheduler state and feature configuration", body = SchedulerStateResponse)
+ )
+)]
pub async fn get_scheduler_state<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -196,6 +205,14 @@ pub async fn get_scheduler_state<
Json(response)
}
+#[utoipa::path(
+ get,
+ path = "/api/version",
+ tag = "version",
+ responses(
+ (status = 200, description = "Scheduler and DataFusion version information", body = SchedulerVersionResponse)
+ )
+)]
pub async fn get_scheduler_version() -> impl IntoResponse {
let response = SchedulerVersionResponse {
version: BALLISTA_VERSION,
@@ -204,6 +221,14 @@ pub async fn get_scheduler_version() -> impl IntoResponse {
Json(response)
}
+#[utoipa::path(
+ get,
+ path = "/api/executors",
+ tag = "executors",
+ responses(
+ (status = 200, description = "List of registered executors", body = Vec<ExecutorResponse>)
+ )
+)]
pub async fn get_executors<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -234,6 +259,18 @@ pub async fn get_executors<
Json(executors)
}
+#[utoipa::path(
+ get,
+ path = "/api/executor/{executor_id}",
+ tag = "executors",
+ params(
+ ("executor_id" = String, Path, description = "Unique executor identifier")
+ ),
+ responses(
+ (status = 200, description = "Executor details", body = ExecutorResponse),
+ (status = 404, description = "Executor not found", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_executor_info<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -267,6 +304,15 @@ pub async fn get_executor_info<
.ok_or(SchedulerErrorResponse::new(StatusCode::NOT_FOUND))
}
+#[utoipa::path(
+ get,
+ path = "/api/jobs",
+ tag = "jobs",
+ responses(
+ (status = 200, description = "List of all jobs", body = Vec<JobResponse>),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_jobs<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -288,6 +334,20 @@ pub async fn get_jobs<
Ok(Json(jobs))
}
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}",
+ tag = "jobs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier"),
+ JobQueryParams
+ ),
+ responses(
+ (status = 200, description = "Job details", body = JobResponse),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_job<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -313,6 +373,20 @@ pub async fn get_job<
)))
}
+#[utoipa::path(
+ patch,
+ path = "/api/job/{job_id}",
+ tag = "jobs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier")
+ ),
+ responses(
+ (status = 200, description = "Job cancellation accepted", body = CancelJobResponse),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse),
+ (status = 409, description = "Job is in terminal state", body = CancelJobResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn cancel_job<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -380,6 +454,20 @@ pub async fn cancel_job<
}
}
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}/stages",
+ tag = "jobs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier"),
+ JobQueryParams
+ ),
+ responses(
+ (status = 200, description = "Query stages summary", body = QueryStagesResponse),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_query_stages<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -411,6 +499,19 @@ pub async fn get_query_stages<
}
}
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}/dot",
+ tag = "graphs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier")
+ ),
+ responses(
+ (status = 200, description = "Job execution graph in Graphviz DOT format", content_type = "text/plain", body = String),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_job_dot_graph<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -438,6 +539,20 @@ pub async fn get_job_dot_graph<
}
}
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}/stage/{stage_id}/dot",
+ tag = "graphs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier"),
+ ("stage_id" = usize, Path, description = "Stage identifier")
+ ),
+ responses(
+ (status = 200, description = "Query stage graph in Graphviz DOT format", content_type = "text/plain", body = String),
+ (status = 404, description = "Job or stage not found", body = SchedulerErrorResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_query_stage_dot_graph<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -459,6 +574,20 @@ pub async fn get_query_stage_dot_graph<
}
}
#[cfg(feature = "graphviz-support")]
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}/dot_svg",
+ tag = "graphs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier")
+ ),
+ responses(
+ (status = 200, description = "Job execution graph rendered as SVG", content_type = "image/svg+xml", body = String),
+ (status = 400, description = "DOT parsing or graphviz execution error", body = SchedulerErrorResponse),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse),
+ (status = 500, description = "Internal server error", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_job_svg_graph<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -492,6 +621,16 @@ pub async fn get_job_svg_graph<
}
}
+#[utoipa::path(
+ get,
+ path = "/api/metrics",
+ tag = "metrics",
+ responses(
+ (status = 200, description = "Prometheus text metrics", content_type = "text/plain; version=0.0.4; charset=utf-8", body = String),
+ (status = 204, description = "No metrics collected"),
+ (status = 500, description = "Internal server error")
+ )
+)]
pub async fn get_scheduler_metrics<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
@@ -514,6 +653,18 @@ pub async fn get_scheduler_metrics<
}
}
+#[utoipa::path(
+ get,
+ path = "/api/job/{job_id}/config",
+ tag = "jobs",
+ params(
+ ("job_id" = String, Path, description = "Unique job identifier")
+ ),
+ responses(
+ (status = 200, description = "Job session configuration properties", body = std::collections::BTreeMap<String, String>),
+ (status = 404, description = "Job not found", body = SchedulerErrorResponse)
+ )
+)]
pub async fn get_job_config<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
diff --git a/ballista/scheduler/src/api/health.rs b/ballista/scheduler/src/api/health.rs
index d0ee4a852..293617b3f 100644
--- a/ballista/scheduler/src/api/health.rs
+++ b/ballista/scheduler/src/api/health.rs
@@ -55,11 +55,34 @@ pub(crate) fn health_routes<
.with_state(scheduler_server)
}
-async fn healthz() -> Response {
+#[cfg_attr(
+ feature = "rest-api",
+ utoipa::path(
+ get,
+ path = "/healthz",
+ tag = "health",
+ responses(
+ (status = 200, description = "Scheduler liveness probe", content_type = "text/plain", body = String)
+ )
+ )
+)]
+pub(crate) async fn healthz() -> Response {
(StatusCode::OK, "ok").into_response()
}
-async fn readyz<
+#[cfg_attr(
+ feature = "rest-api",
+ utoipa::path(
+ get,
+ path = "/readyz",
+ tag = "health",
+ responses(
+ (status = 200, description = "Scheduler is ready to accept queries", content_type = "text/plain", body = String),
+ (status = 503, description = "Scheduler is not ready (insufficient executors)", content_type = "text/plain", body = String)
+ )
+ )
+)]
+pub(crate) async fn readyz<
T: AsLogicalPlan + Clone + Send + Sync + 'static,
U: AsExecutionPlan + Send + Sync + 'static,
>(
diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs
index 1e865d84b..e9a2dba6f 100644
--- a/ballista/scheduler/src/api/mod.rs
+++ b/ballista/scheduler/src/api/mod.rs
@@ -18,9 +18,13 @@ pub(crate) mod dto_build;
mod handlers;
mod health;
#[cfg(feature = "rest-api")]
+pub mod openapi;
+#[cfg(feature = "rest-api")]
mod routes;
pub(super) use health::health_routes;
#[cfg(feature = "rest-api")]
+pub use openapi::{ApiDoc, openapi_spec};
+#[cfg(feature = "rest-api")]
pub use routes::get_routes;
use axum::response::{IntoResponse, Response};
@@ -37,6 +41,7 @@ pub(super) fn route_disabled(reason: String) -> Router {
}
#[derive(Debug, serde::Serialize)]
+#[cfg_attr(feature = "rest-api", derive(utoipa::ToSchema))]
pub(crate) struct SchedulerErrorResponse {
#[serde(skip)]
status_code: StatusCode,
diff --git a/ballista/scheduler/src/api/openapi.rs b/ballista/scheduler/src/api/openapi.rs
new file mode 100644
index 000000000..8265ce0e2
--- /dev/null
+++ b/ballista/scheduler/src/api/openapi.rs
@@ -0,0 +1,320 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! OpenAPI specification for the Ballista scheduler REST API.
+
+use crate::api::SchedulerErrorResponse;
+use crate::api::handlers::{
+ CancelJobResponse, ExecutorMetricResponse, ExecutorResponse, JobQueryParams,
+ SchedulerStateResponse, SchedulerVersionResponse,
+};
+use axum::Json;
+use axum::response::IntoResponse;
+use ballista_api_types::dto::{
+ JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse,
+ TaskStatus, TaskSummary,
+};
+use ballista_core::serde::scheduler::{
+ ExecutorOperatingSystemSpecification, ExecutorSpecification,
+};
+use std::sync::LazyLock;
+use utoipa::OpenApi;
+
+/// OpenAPI documentation structure for the Ballista scheduler REST API.
+#[derive(OpenApi)]
+#[openapi(
+ info(
+ title = "Apache DataFusion Ballista Scheduler REST API",
+ version = env!("CARGO_PKG_VERSION"),
+ description = "REST API for Apache DataFusion Ballista Scheduler",
+ license(
+ name = "Apache-2.0",
+ url = "https://www.apache.org/licenses/LICENSE-2.0"
+ )
+ ),
+ paths(
+ crate::api::health::healthz,
+ crate::api::health::readyz,
+ crate::api::handlers::get_scheduler_state,
+ crate::api::handlers::get_scheduler_version,
+ crate::api::handlers::get_executors,
+ crate::api::handlers::get_executor_info,
+ crate::api::handlers::get_jobs,
+ crate::api::handlers::get_job,
+ crate::api::handlers::cancel_job,
+ crate::api::handlers::get_job_config,
+ crate::api::handlers::get_query_stages,
+ crate::api::handlers::get_job_dot_graph,
+ crate::api::handlers::get_query_stage_dot_graph,
+ crate::api::handlers::get_scheduler_metrics,
+ crate::api::openapi::get_openapi_spec,
+ ),
+ components(
+ schemas(
+ SchedulerStateResponse,
+ SchedulerVersionResponse,
+ ExecutorResponse,
+ ExecutorMetricResponse,
+ ExecutorSpecification,
+ ExecutorOperatingSystemSpecification,
+ CancelJobResponse,
+ JobQueryParams,
+ SchedulerErrorResponse,
+ JobResponse,
+ TaskStatus,
+ TaskSummary,
+ Percentiles,
+ QueryStageSummary,
+ QueryStagesResponse,
+ PlanFormat,
+ )
+ ),
+ tags(
+ (name = "health", description = "Kubernetes liveness and readiness probes"),
+ (name = "state", description = "Scheduler state and feature configuration"),
+ (name = "version", description = "Version information"),
+ (name = "executors", description = "Executor management and metrics"),
+ (name = "jobs", description = "Job execution, monitoring, and cancellation"),
+ (name = "graphs", description = "Execution graph visualization"),
+ (name = "metrics", description = "Prometheus cluster metrics"),
+ (name = "openapi", description = "OpenAPI specification"),
+ )
+)]
+pub struct ApiDoc;
+
+#[cfg(feature = "graphviz-support")]
+#[derive(OpenApi)]
+#[openapi(paths(crate::api::handlers::get_job_svg_graph))]
+struct GraphvizApiDoc;
+
+static OPENAPI_SPEC: LazyLock<utoipa::openapi::OpenApi> = LazyLock::new(|| {
+ #[allow(unused_mut)]
+ let mut spec = ApiDoc::openapi();
+ #[cfg(feature = "graphviz-support")]
+ spec.merge(GraphvizApiDoc::openapi());
+ spec
+});
+
+/// Generate or retrieve the cached OpenAPI specification for the scheduler REST API.
+pub fn openapi_spec() -> utoipa::openapi::OpenApi {
+ OPENAPI_SPEC.clone()
+}
+
+/// Handler for `GET /api/openapi.json`.
+#[utoipa::path(
+ get,
+ path = "/api/openapi.json",
+ tag = "openapi",
+ responses(
+ (status = 200, description = "OpenAPI specification document in JSON format", body = Object)
+ )
+)]
+pub async fn get_openapi_spec() -> impl IntoResponse {
+ Json(openapi_spec())
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::config::SchedulerConfig;
+ use crate::metrics::default_metrics_collector;
+ use crate::scheduler_server::SchedulerServer;
+ use crate::test_utils::test_cluster_context;
+ use axum::body::Body;
+ use axum::http::{Request, StatusCode};
+ use ballista_core::serde::BallistaCodec;
+ use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+ use std::collections::BTreeSet;
+ use std::sync::Arc;
+ use tower::ServiceExt;
+ use utoipa::ToSchema;
+
+ fn get_schema_property_names<T: ToSchema>() -> BTreeSet<String> {
+ let schema_ref = T::schema();
+ match schema_ref {
+ utoipa::openapi::RefOr::T(utoipa::openapi::Schema::Object(obj)) => {
+ obj.properties.keys().cloned().collect()
+ }
+ _ => panic!("Expected Object schema"),
+ }
+ }
+
+ fn get_json_property_names<T: serde::Serialize>(value: &T) -> BTreeSet<String> {
+ let json = serde_json::to_value(value).expect("serialize to value");
+ match json {
+ serde_json::Value::Object(map) => map.keys().cloned().collect(),
+ _ => panic!("Expected JSON Object"),
+ }
+ }
+
+ #[test]
+ fn test_openapi_spec_serialization() {
+ let spec = openapi_spec();
+ let json = serde_json::to_string_pretty(&spec)
+ .expect("OpenAPI spec must serialize to JSON");
+ assert!(!json.is_empty());
+
+ let value: serde_json::Value = serde_json::from_str(&json)
+ .expect("Serialized OpenAPI spec must be valid JSON");
+ assert_eq!(
+ value["info"]["title"],
+ "Apache DataFusion Ballista Scheduler REST API"
+ );
+ }
+
+ #[test]
+ fn test_openapi_spec_contains_all_registered_paths() {
+ let spec = openapi_spec();
+ let paths: Vec<&str> = spec.paths.paths.keys().map(|k| k.as_str()).collect();
+
+ let expected_paths = [
+ "/healthz",
+ "/readyz",
+ "/api/openapi.json",
+ "/api/state",
+ "/api/version",
+ "/api/executors",
+ "/api/executor/{executor_id}",
+ "/api/jobs",
+ "/api/job/{job_id}",
+ "/api/job/{job_id}/config",
+ "/api/job/{job_id}/stages",
+ "/api/job/{job_id}/dot",
+ "/api/job/{job_id}/stage/{stage_id}/dot",
+ "/api/metrics",
+ ];
+
+ for path in expected_paths {
+ assert!(
+ paths.contains(&path),
+ "OpenAPI spec is missing registered path: {path}. Found paths: {paths:?}"
+ );
+ }
+
+ #[cfg(feature = "graphviz-support")]
+ assert!(
+ paths.contains(&"/api/job/{job_id}/dot_svg"),
+ "OpenAPI spec is missing graphviz path: /api/job/{{job_id}}/dot_svg"
+ );
+ }
+
+ #[test]
+ fn test_openapi_spec_contains_all_schemas() {
+ let spec = openapi_spec();
+ let components = spec.components.expect("OpenAPI components must be present");
+ let schemas = components.schemas;
+
+ let expected_schemas = [
+ "SchedulerStateResponse",
+ "SchedulerVersionResponse",
+ "ExecutorResponse",
+ "ExecutorMetricResponse",
+ "ExecutorSpecification",
+ "ExecutorOperatingSystemSpecification",
+ "CancelJobResponse",
+ "JobQueryParams",
+ "SchedulerErrorResponse",
+ "JobResponse",
+ "TaskStatus",
+ "TaskSummary",
+ "Percentiles",
+ "QueryStageSummary",
+ "QueryStagesResponse",
+ "PlanFormat",
+ ];
+
+ for schema_name in expected_schemas {
+ assert!(
+ schemas.contains_key(schema_name),
+ "OpenAPI spec is missing schema: {schema_name}. Found schemas: {:?}",
+ schemas.keys().collect::<Vec<_>>()
+ );
+ }
+ }
+
+ #[test]
+ fn test_schema_property_alignment_with_serde() {
+ assert_eq!(
+ get_schema_property_names::<ExecutorOperatingSystemSpecification>(),
+ get_json_property_names(&ExecutorOperatingSystemSpecification::default())
+ );
+
+ assert_eq!(
+ get_schema_property_names::<ExecutorSpecification>(),
+ get_json_property_names(&ExecutorSpecification::default())
+ );
+
+ let version_response = SchedulerVersionResponse {
+ version: "1.0",
+ datafusion_version: "1.0",
+ };
+ assert_eq!(
+ get_schema_property_names::<SchedulerVersionResponse>(),
+ get_json_property_names(&version_response)
+ );
+
+ let cancel_response = CancelJobResponse {
+ cancelled: true,
+ reason: Some("cancelled by user".to_string()),
+ };
+ assert_eq!(
+ get_schema_property_names::<CancelJobResponse>(),
+ get_json_property_names(&cancel_response)
+ );
+ }
+
+ #[tokio::test]
+ async fn test_openapi_json_endpoint() {
+ let config = SchedulerConfig::default();
+ let server = SchedulerServer::new(
+ "localhost:50050".to_owned(),
+ test_cluster_context(),
+ BallistaCodec::default(),
+ Arc::new(config),
+ default_metrics_collector().unwrap(),
+ );
+ let router =
+ crate::api::get_routes::<LogicalPlanNode, PhysicalPlanNode>(Arc::new(server));
+
+ let response = router
+ .oneshot(
+ Request::builder()
+ .uri("/api/openapi.json")
+ .body(Body::empty())
+ .unwrap(),
+ )
+ .await
+ .expect("router oneshot");
+
+ assert_eq!(response.status(), StatusCode::OK);
+
+ let body_bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
+ .await
+ .expect("read body");
+ let spec_value: serde_json::Value = serde_json::from_slice(&body_bytes)
+ .expect("Response body must be valid JSON");
+
+ assert_eq!(
+ spec_value["info"]["title"],
+ "Apache DataFusion Ballista Scheduler REST API"
+ );
+ assert!(spec_value["paths"]["/api/state"].is_object());
+ assert!(spec_value["paths"]["/api/jobs"].is_object());
+ assert!(spec_value["paths"]["/healthz"].is_object());
+ assert!(spec_value["paths"]["/readyz"].is_object());
+ }
+}
diff --git a/ballista/scheduler/src/api/routes.rs b/ballista/scheduler/src/api/routes.rs
index 6a3341934..8a292c020 100644
--- a/ballista/scheduler/src/api/routes.rs
+++ b/ballista/scheduler/src/api/routes.rs
@@ -10,7 +10,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-use crate::api::handlers;
+use crate::api::{handlers, openapi};
use crate::config::SchedulerConfig;
use crate::scheduler_server::SchedulerServer;
use axum::{Router, routing::get};
@@ -31,6 +31,7 @@ pub fn get_routes<
let router = Router::new()
.route(&web_tui_route, get(handlers::get_webtui))
+ .route("/api/openapi.json", get(openapi::get_openapi_spec))
.route("/api/state", get(handlers::get_scheduler_state::<T, U>))
.route("/api/version", get(handlers::get_scheduler_version))
.route("/api/executors", get(handlers::get_executors::<T, U>))
diff --git a/docs/source/user-guide/scheduler.md b/docs/source/user-guide/scheduler.md
index 2311ea42e..03aa92fc3 100644
--- a/docs/source/user-guide/scheduler.md
+++ b/docs/source/user-guide/scheduler.md
@@ -75,6 +75,7 @@ The scheduler also provides a REST API that allows jobs to be monitored.
| API | Method | Description |
| -------------------------------------- | ------ | ----------------------------------------------------------------- |
+| /api/openapi.json | GET | Return OpenAPI v3 specification document for the REST API. |
| /api/jobs | GET | Get a list of jobs that have been submitted to the cluster. |
| /api/job/{job_id} | GET | Get a summary of a submitted job. |
| /api/job/{job_id}/dot | GET | Produce a query plan in DOT (graphviz) format. |
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org