diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 8509fa053f..529f4e952a 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -978,6 +978,7 @@ pub struct LabelValueQuery { #[serde(flatten)] matches: Matches, db: Option, + limit: Option, } #[axum_macros::debug_handler] @@ -1025,6 +1026,7 @@ pub async fn label_values_query( } } table_names.sort_unstable(); + truncate_results(&mut table_names, params.limit); return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(table_names)); } else if label_name == FIELD_NAME_LABEL { let field_columns = handle_schema_err!( @@ -1033,12 +1035,14 @@ pub async fn label_values_query( .unwrap_or_default(); let mut field_columns = field_columns.into_iter().collect::>(); field_columns.sort_unstable(); + truncate_results(&mut field_columns, params.limit); return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(field_columns)); } else if label_name == SCHEMA_LABEL || label_name == DATABASE_LABEL { let catalog_manager = handler.catalog_manager(); match retrieve_schema_names(&query_ctx, catalog_manager, params.matches.0).await { - Ok(schema_names) => { + Ok(mut schema_names) => { + truncate_results(&mut schema_names, params.limit); return PrometheusJsonResponse::success(PrometheusResponse::LabelValues( schema_names, )); @@ -1100,9 +1104,19 @@ pub async fn label_values_query( let mut label_values: Vec<_> = label_values.into_iter().collect(); label_values.sort_unstable(); + truncate_results(&mut label_values, params.limit); + PrometheusJsonResponse::success(PrometheusResponse::LabelValues(label_values)) } +fn truncate_results(label_values: &mut Vec, limit: Option) { + if let Some(limit) = limit { + if limit > 0 && label_values.len() >= limit { + label_values.truncate(limit); + } + } +} + /// Take metric name from the [VectorSelector]. /// It takes the name in the selector or removes the name matcher. fn take_metric_name(selector: &mut VectorSelector) -> Option { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 5262e05d4f..dde7ca23e3 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -978,6 +978,62 @@ pub async fn test_prom_http_api(store_type: StorageType) { serde_json::from_value::(json!(["val"])).unwrap() ); + // limit + let res = client + .get("/v1/prometheus/api/v1/label/host/values?match[]=demo&start=0&end=600&limit=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!(["host1"])).unwrap() + ); + + // limit 0 + let res = client + .get("/v1/prometheus/api/v1/label/host/values?match[]=demo&start=0&end=600&limit=0") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!(["host1", "host2"])).unwrap() + ); + + // limit 10 + let res = client + .get("/v1/prometheus/api/v1/label/host/values?match[]=demo&start=0&end=600&limit=10") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!(["host1", "host2"])).unwrap() + ); + + // special labels limit + let res = client + .get("/v1/prometheus/api/v1/label/__schema__/values?start=0&end=600&limit=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = serde_json::from_str::(&res.text().await).unwrap(); + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!([ + "greptime_private", + "information_schema", + ])) + .unwrap() + ); + // query an empty database should return nothing let res = client .get("/v1/prometheus/api/v1/label/host/values?match[]=demo&start=0&end=600")