feat: support prometheus format_query endpoint (#2731)

* feat: support prometheus format_query endpoint Signed-off-by: tison <[email protected]> * add test Signed-off-by: tison <[email protected]> --------- Signed-off-by: tison <[email protected]>
GreptimeTeam · Nov 10, 2023 · f5eede4 · f5eede4
1 parent 22ee45f
commit f5eede4
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 1 deletion.
diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs
@@ -66,7 +66,7 @@ use self::influxdb::{influxdb_health, influxdb_ping, influxdb_write_v1, influxdb
 use crate::configurator::ConfiguratorRef;
 use crate::error::{AlreadyStartedSnafu, Result, StartHttpSnafu};
 use crate::http::prometheus::{
-    instant_query, label_values_query, labels_query, range_query, series_query,
+    format_query, instant_query, label_values_query, labels_query, range_query, series_query,
 };
 use crate::metrics::{
     HTTP_TRACK_METRICS, METRIC_HTTP_REQUESTS_ELAPSED, METRIC_HTTP_REQUESTS_TOTAL,
@@ -623,6 +623,10 @@ impl HttpServer {
 
     fn route_prometheus<S>(&self, prometheus_handler: PrometheusHandlerRef) -> Router<S> {
         Router::new()
+            .route(
+                "/format_query",
+                routing::post(format_query).get(format_query),
+            )
             .route("/query", routing::post(instant_query).get(instant_query))
             .route("/query_range", routing::post(range_query).get(range_query))
             .route("/labels", routing::post(labels_query).get(labels_query))

diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs
@@ -70,6 +70,7 @@ pub enum PrometheusResponse {
     Labels(Vec<String>),
     Series(Vec<HashMap<String, String>>),
     LabelValues(Vec<String>),
+    FormatQuery(String),
 }
 
 impl Default for PrometheusResponse {
@@ -290,6 +291,33 @@ impl PrometheusJsonResponse {
     }
 }
 
+#[derive(Debug, Default, Serialize, Deserialize, JsonSchema)]
+pub struct FormatQuery {
+    query: Option<String>,
+}
+
+#[axum_macros::debug_handler]
+pub async fn format_query(
+    State(_handler): State<PrometheusHandlerRef>,
+    Query(params): Query<InstantQuery>,
+    Extension(_query_ctx): Extension<QueryContextRef>,
+    Form(form_params): Form<InstantQuery>,
+) -> Json<PrometheusJsonResponse> {
+    let _timer = crate::metrics::METRIC_HTTP_PROMQL_FORMAT_QUERY_ELAPSED.start_timer();
+
+    let query = params.query.or(form_params.query).unwrap_or_default();
+    match promql_parser::parser::parse(&query) {
+        Ok(expr) => {
+            let pretty = expr.prettify();
+            PrometheusJsonResponse::success(PrometheusResponse::FormatQuery(pretty))
+        }
+        Err(reason) => {
+            let err = InvalidQuerySnafu { reason }.build();
+            PrometheusJsonResponse::error(err.status_code().to_string(), err.output_msg())
+        }
+    }
+}
+
 #[derive(Debug, Default, Serialize, Deserialize, JsonSchema)]
 pub struct InstantQuery {
     query: Option<String>,

diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs
@@ -100,6 +100,11 @@ lazy_static! {
         "servers opentsdb line write elapsed"
     )
     .unwrap();
+    pub static ref METRIC_HTTP_PROMQL_FORMAT_QUERY_ELAPSED: Histogram = register_histogram!(
+        "servers_http_promql_format_query_elapsed",
+        "servers http promql format query elapsed"
+    )
+    .unwrap();
     pub static ref METRIC_HTTP_PROMQL_INSTANT_QUERY_ELAPSED: Histogram = register_histogram!(
         "servers_http_promql_instant_query_elapsed",
         "servers http promql instant query elapsed"

diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs
@@ -341,6 +341,17 @@ pub async fn test_prom_http_api(store_type: StorageType) {
     let (app, mut guard) = setup_test_prom_app_with_frontend(store_type, "promql_api").await;
     let client = TestClient::new(app);
 
+    // format_query
+    let res = client
+        .get("/v1/prometheus/api/v1/format_query?query=foo/bar")
+        .send()
+        .await;
+    assert_eq!(res.status(), StatusCode::OK);
+    assert_eq!(
+        res.text().await.as_str(),
+        r#"{"status":"success","data":"foo / bar"}"#
+    );
+
     // instant query
     let res = client
         .get("/v1/prometheus/api/v1/query?query=up&time=1")