fix: prometheus api only returns 200 (#4471)

fix: prometheus api returns http status other than 200
GreptimeTeam · Jul 31, 2024 · c66d309 · c66d309
1 parent 6560507
commit c66d309
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 44 deletions.
diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs
@@ -752,7 +752,7 @@ fn log_error_if_necessary(error: &Error) {
 impl IntoResponse for Error {
     fn into_response(self) -> Response {
         let error_msg = self.output_msg();
-        let status = status_code_to_http_status(self.status_code());
+        let status = status_code_to_http_status(&self.status_code());
 
         log_error_if_necessary(&self);
 

diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs
@@ -124,10 +124,7 @@ impl PrometheusGatewayService {
             match retrieve_metric_name_and_result_type(&query.query) {
                 Ok((metric_name, result_type)) => (metric_name.unwrap_or_default(), result_type),
                 Err(err) => {
-                    return PrometheusJsonResponse::error(
-                        err.status_code().to_string(),
-                        err.output_msg(),
-                    )
+                    return PrometheusJsonResponse::error(err.status_code(), err.output_msg())
                 }
             };
         // range query only returns matrix

diff --git a/src/servers/src/http/error_result.rs b/src/servers/src/http/error_result.rs
@@ -82,13 +82,13 @@ impl IntoResponse for ErrorResponse {
             HeaderValue::from(execution_time),
         );
         let status = StatusCode::from_u32(code).unwrap_or(StatusCode::Unknown);
-        let status_code = status_code_to_http_status(status);
+        let status_code = status_code_to_http_status(&status);
 
         (status_code, resp).into_response()
     }
 }
 
-pub fn status_code_to_http_status(status_code: StatusCode) -> HttpStatusCode {
+pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode {
     match status_code {
         StatusCode::Success | StatusCode::Cancelled => HttpStatusCode::OK,
 

diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs
@@ -133,7 +133,7 @@ pub async fn format_query(
         }
         Err(reason) => {
             let err = InvalidQuerySnafu { reason }.build();
-            PrometheusJsonResponse::error(err.status_code().to_string(), err.output_msg())
+            PrometheusJsonResponse::error(err.status_code(), err.output_msg())
         }
     }
 }
@@ -201,9 +201,7 @@ pub async fn instant_query(
     let result = handler.do_query(&prom_query, query_ctx).await;
     let (metric_name, result_type) = match retrieve_metric_name_and_result_type(&prom_query.query) {
         Ok((metric_name, result_type)) => (metric_name.unwrap_or_default(), result_type),
-        Err(err) => {
-            return PrometheusJsonResponse::error(err.status_code().to_string(), err.output_msg())
-        }
+        Err(err) => return PrometheusJsonResponse::error(err.status_code(), err.output_msg()),
     };
     PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await
 }
@@ -254,9 +252,7 @@ pub async fn range_query(
 
     let result = handler.do_query(&prom_query, query_ctx).await;
     let metric_name = match retrieve_metric_name_and_result_type(&prom_query.query) {
-        Err(err) => {
-            return PrometheusJsonResponse::error(err.status_code().to_string(), err.output_msg())
-        }
+        Err(err) => return PrometheusJsonResponse::error(err.status_code(), err.output_msg()),
         Ok((metric_name, _)) => metric_name.unwrap_or_default(),
     };
     PrometheusJsonResponse::from_query_result(result, metric_name, ValueType::Matrix).await
@@ -335,9 +331,7 @@ pub async fn labels_query(
     let mut labels = match get_all_column_names(&catalog, &schema, &handler.catalog_manager()).await
     {
         Ok(labels) => labels,
-        Err(e) => {
-            return PrometheusJsonResponse::error(e.status_code().to_string(), e.output_msg())
-        }
+        Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()),
     };
     // insert the special metric name label
     let _ = labels.insert(METRIC_NAME.to_string());
@@ -385,10 +379,7 @@ pub async fn labels_query(
             if err.status_code() != StatusCode::TableNotFound
                 && err.status_code() != StatusCode::TableColumnNotFound
             {
-                return PrometheusJsonResponse::error(
-                    err.status_code().to_string(),
-                    err.output_msg(),
-                );
+                return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
             }
         }
     }
@@ -705,7 +696,7 @@ pub async fn label_values_query(
         {
             Ok(table_names) => table_names,
             Err(e) => {
-                return PrometheusJsonResponse::error(e.status_code().to_string(), e.output_msg());
+                return PrometheusJsonResponse::error(e.status_code(), e.output_msg());
             }
         };
         table_names.sort_unstable();
@@ -717,10 +708,7 @@ pub async fn label_values_query(
             {
                 Ok(table_names) => table_names,
                 Err(e) => {
-                    return PrometheusJsonResponse::error(
-                        e.status_code().to_string(),
-                        e.output_msg(),
-                    );
+                    return PrometheusJsonResponse::error(e.status_code(), e.output_msg());
                 }
             };
         let mut field_columns = field_columns.into_iter().collect::<Vec<_>>();
@@ -730,7 +718,10 @@ pub async fn label_values_query(
 
     let queries = params.matches.0;
     if queries.is_empty() {
-        return PrometheusJsonResponse::error("Invalid argument", "match[] parameter is required");
+        return PrometheusJsonResponse::error(
+            StatusCode::InvalidArguments,
+            "match[] parameter is required",
+        );
     }
 
     let start = params.start.unwrap_or_else(yesterday_rfc3339);
@@ -758,10 +749,7 @@ pub async fn label_values_query(
             if err.status_code() != StatusCode::TableNotFound
                 && err.status_code() != StatusCode::TableColumnNotFound
             {
-                return PrometheusJsonResponse::error(
-                    err.status_code().to_string(),
-                    err.output_msg(),
-                );
+                return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
             }
         }
     }
@@ -957,7 +945,10 @@ pub async fn series_query(
         queries = form_params.matches.0;
     }
     if queries.is_empty() {
-        return PrometheusJsonResponse::error("Unsupported", "match[] parameter is required");
+        return PrometheusJsonResponse::error(
+            StatusCode::Unsupported,
+            "match[] parameter is required",
+        );
     }
     let start = params
         .start
@@ -1007,7 +998,7 @@ pub async fn series_query(
         )
         .await
         {
-            return PrometheusJsonResponse::error(err.status_code().to_string(), err.output_msg());
+            return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
         }
     }
     let merge_map = merge_map

diff --git a/src/servers/src/http/prometheus_resp.rs b/src/servers/src/http/prometheus_resp.rs
@@ -37,6 +37,7 @@ use super::prometheus::{
     PromData, PromQueryResult, PromSeriesMatrix, PromSeriesVector, PrometheusResponse,
 };
 use crate::error::{CollectRecordbatchSnafu, Result, UnexpectedResultSnafu};
+use crate::http::error_result::status_code_to_http_status;
 
 #[derive(Debug, Default, Serialize, Deserialize, JsonSchema, PartialEq)]
 pub struct PrometheusJsonResponse {
@@ -50,6 +51,8 @@ pub struct PrometheusJsonResponse {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub warnings: Option<Vec<String>>,
 
+    #[serde(skip)]
+    pub status_code: Option<StatusCode>,
     // placeholder for header value
     #[serde(skip)]
     #[serde(default)]
@@ -64,8 +67,14 @@ impl IntoResponse for PrometheusJsonResponse {
             serde_json::to_string(&self.resp_metrics).ok()
         };
 
+        let http_code = self.status_code.map(|c| status_code_to_http_status(&c));
+
         let mut resp = Json(self).into_response();
 
+        if let Some(http_code) = http_code {
+            *resp.status_mut() = http_code;
+        }
+
         if let Some(m) = metrics.and_then(|m| HeaderValue::from_str(&m).ok()) {
             resp.headers_mut().insert(&GREPTIME_DB_HEADER_METRICS, m);
         }
@@ -75,18 +84,18 @@ impl IntoResponse for PrometheusJsonResponse {
 }
 
 impl PrometheusJsonResponse {
-    pub fn error<S1, S2>(error_type: S1, reason: S2) -> Self
+    pub fn error<S1>(error_type: StatusCode, reason: S1) -> Self
     where
         S1: Into<String>,
-        S2: Into<String>,
     {
         PrometheusJsonResponse {
             status: "error".to_string(),
             data: PrometheusResponse::default(),
             error: Some(reason.into()),
-            error_type: Some(error_type.into()),
+            error_type: Some(error_type.to_string()),
             warnings: None,
             resp_metrics: Default::default(),
+            status_code: Some(error_type),
         }
     }
 
@@ -98,6 +107,7 @@ impl PrometheusJsonResponse {
             error_type: None,
             warnings: None,
             resp_metrics: Default::default(),
+            status_code: None,
         }
     }
 
@@ -124,9 +134,10 @@ impl PrometheusJsonResponse {
                             result_type,
                         )?)
                     }
-                    OutputData::AffectedRows(_) => {
-                        Self::error("Unexpected", "expected data result, but got affected rows")
-                    }
+                    OutputData::AffectedRows(_) => Self::error(
+                        StatusCode::Unexpected,
+                        "expected data result, but got affected rows",
+                    ),
                 };
 
             if let Some(physical_plan) = result.meta.plan {
@@ -158,7 +169,7 @@ impl PrometheusJsonResponse {
                         ..Default::default()
                     }))
                 } else {
-                    Self::error(err.status_code().to_string(), err.output_msg())
+                    Self::error(err.status_code(), err.output_msg())
                 }
             }
         }

diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs
@@ -577,6 +577,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) {
         error_type: None,
         warnings: None,
         resp_metrics: Default::default(),
+        status_code: None,
     };
     assert_eq!(instant_query_result, expected);
 
@@ -628,6 +629,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) {
         error_type: None,
         warnings: None,
         resp_metrics: Default::default(),
+        status_code: None,
     };
     assert_eq!(range_query_result, expected);
 
@@ -660,6 +662,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) {
         error_type: None,
         warnings: None,
         resp_metrics: Default::default(),
+        status_code: None,
     };
     assert_eq!(range_query_result, expected);
 

diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs
@@ -520,19 +520,31 @@ pub async fn test_prom_http_api(store_type: StorageType) {
         .header("Content-Type", "application/x-www-form-urlencoded")
         .send()
         .await;
-    assert_eq!(res.status(), StatusCode::OK);
+    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
+    let prom_resp = res.json::<PrometheusJsonResponse>().await;
+    assert_eq!(prom_resp.status, "error");
+    assert!(prom_resp
+        .error_type
+        .is_some_and(|err| err.eq_ignore_ascii_case("TableNotFound")));
+    assert!(prom_resp
+        .error
+        .is_some_and(|err| err.eq_ignore_ascii_case("Table not found: greptime.public.up")));
 
     // label values
     // should return error if there is no match[]
     let res = client
         .get("/v1/prometheus/api/v1/label/instance/values")
         .send()
         .await;
-    assert_eq!(res.status(), StatusCode::OK);
+    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
     let prom_resp = res.json::<PrometheusJsonResponse>().await;
     assert_eq!(prom_resp.status, "error");
-    assert!(prom_resp.error.is_some_and(|err| !err.is_empty()));
-    assert!(prom_resp.error_type.is_some_and(|err| !err.is_empty()));
+    assert!(prom_resp
+        .error
+        .is_some_and(|err| err.eq_ignore_ascii_case("match[] parameter is required")));
+    assert!(prom_resp
+        .error_type
+        .is_some_and(|err| err.eq_ignore_ascii_case("InvalidArguments")));
 
     // single match[]
     let res = client