feat: add limit to label value api (GreptimeTeam#6795)

* feat: add limit to label value api

Signed-off-by: Ning Sun <sunning@greptime.com>

* feat: limit for special labels

Signed-off-by: Ning Sun <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>

Author: sunng87

src/servers/src/http/prometheus.rs

@@ -978,6 +978,7 @@ pub struct LabelValueQuery {
     #[serde(flatten)]
     matches: Matches,
     db: Option<String>,
+    limit: Option<usize>,
 }
 
 #[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::<Vec<_>>();
         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<String>, limit: Option<usize>) {
+    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<String> {

tests-integration/tests/http.rs

@@ -978,6 +978,62 @@ pub async fn test_prom_http_api(store_type: StorageType) {
         serde_json::from_value::<PrometheusResponse>(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::<PrometheusJsonResponse>(&res.text().await).unwrap();
+    assert_eq!(body.status, "success");
+    assert_eq!(
+        body.data,
+        serde_json::from_value::<PrometheusResponse>(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::<PrometheusJsonResponse>(&res.text().await).unwrap();
+    assert_eq!(body.status, "success");
+    assert_eq!(
+        body.data,
+        serde_json::from_value::<PrometheusResponse>(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::<PrometheusJsonResponse>(&res.text().await).unwrap();
+    assert_eq!(body.status, "success");
+    assert_eq!(
+        body.data,
+        serde_json::from_value::<PrometheusResponse>(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::<PrometheusJsonResponse>(&res.text().await).unwrap();
+    assert_eq!(body.status, "success");
+    assert_eq!(
+        body.data,
+        serde_json::from_value::<PrometheusResponse>(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")