feat: impl instant query and add tests (#1452)

* feat: impl instant query and add tests

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* clean up

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/api/Cargo.toml

@@ -10,7 +10,7 @@ common-base = { path = "../common/base" }
 common-error = { path = "../common/error" }
 common-time = { path = "../common/time" }
 datatypes = { path = "../datatypes" }
-greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "57e6a032db70264ec394cbb8a2f327ad33705d76" }
+greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "a26c40c004f998180b8acd853b22f083773f36b9" }
 prost.workspace = true
 snafu = { version = "0.7", features = ["backtraces"] }
 tonic.workspace = true

src/client/src/client.rs

@@ -16,6 +16,7 @@ use std::sync::Arc;
 
 use api::v1::greptime_database_client::GreptimeDatabaseClient;
 use api::v1::health_check_client::HealthCheckClient;
+use api::v1::prometheus_gateway_client::PrometheusGatewayClient;
 use api::v1::HealthCheckRequest;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use common_grpc::channel_manager::ChannelManager;
@@ -156,6 +157,11 @@ impl Client {
         })
     }
 
+    pub fn make_prometheus_gateway_client(&self) -> Result<PrometheusGatewayClient<Channel>> {
+        let (_, channel) = self.find_channel()?;
+        Ok(PrometheusGatewayClient::new(channel))
+    }
+
     pub async fn health_check(&self) -> Result<()> {
         let (_, channel) = self.find_channel()?;
         let mut client = HealthCheckClient::new(channel);

src/servers/src/grpc/prom_query_gateway.rs

@@ -20,8 +20,10 @@ use api::v1::promql_request::Promql;
 use api::v1::{PromqlRequest, PromqlResponse, ResponseHeader};
 use async_trait::async_trait;
 use query::parser::PromQuery;
+use snafu::OptionExt;
 use tonic::{Request, Response};
 
+use crate::error::InvalidQuerySnafu;
 use crate::grpc::handler::create_query_context;
 use crate::grpc::TonicResult;
 use crate::prom::{retrieve_metric_name, PromHandlerRef, PromJsonResponse};
@@ -34,14 +36,21 @@ pub struct PrometheusGatewayService {
 impl PrometheusGateway for PrometheusGatewayService {
     async fn handle(&self, req: Request<PromqlRequest>) -> TonicResult<Response<PromqlResponse>> {
         let inner = req.into_inner();
-        let prom_query = match inner.promql {
-            Some(Promql::RangeQuery(range_query)) => PromQuery {
+        let prom_query = match inner.promql.context(InvalidQuerySnafu {
+            reason: "Expecting non-empty PromqlRequest.",
+        })? {
+            Promql::RangeQuery(range_query) => PromQuery {
                 query: range_query.query,
                 start: range_query.start,
                 end: range_query.end,
                 step: range_query.step,
             },
-            _ => unimplemented!(),
+            Promql::InstantQuery(instant_query) => PromQuery {
+                query: instant_query.query,
+                start: instant_query.time.clone(),
+                end: instant_query.time,
+                step: String::from("1s"),
+            },
         };
 
         let query_context = create_query_context(inner.header.as_ref());

src/servers/src/prom.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 //! prom supply the prometheus HTTP API Server compliance
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::net::SocketAddr;
 use std::sync::Arc;
 
@@ -156,30 +156,30 @@ impl Server for PromServer {
     }
 }
 
-#[derive(Debug, Default, Serialize, Deserialize, JsonSchema)]
+#[derive(Debug, Default, Serialize, Deserialize, JsonSchema, PartialEq)]
 pub struct PromSeries {
-    metric: HashMap<String, String>,
-    values: Vec<(f64, String)>,
+    pub metric: HashMap<String, String>,
+    pub values: Vec<(f64, String)>,
 }
 
-#[derive(Debug, Default, Serialize, Deserialize, JsonSchema)]
+#[derive(Debug, Default, Serialize, Deserialize, JsonSchema, PartialEq)]
 pub struct PromData {
     #[serde(rename = "resultType")]
-    result_type: String,
-    result: Vec<PromSeries>,
+    pub result_type: String,
+    pub result: Vec<PromSeries>,
 }
 
-#[derive(Debug, Default, Serialize, Deserialize, JsonSchema)]
+#[derive(Debug, Default, Serialize, Deserialize, JsonSchema, PartialEq)]
 pub struct PromJsonResponse {
-    status: String,
-    data: PromData,
+    pub status: String,
+    pub data: PromData,
     #[serde(skip_serializing_if = "Option::is_none")]
-    error: Option<String>,
+    pub error: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "errorType")]
-    error_type: Option<String>,
+    pub error_type: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    warnings: Option<Vec<String>>,
+    pub warnings: Option<Vec<String>>,
 }
 
 impl PromJsonResponse {
@@ -281,7 +281,7 @@ impl PromJsonResponse {
         })?;
 
         let metric_name = (METRIC_NAME.to_string(), metric_name);
-        let mut buffer = HashMap::<Vec<(String, String)>, Vec<(f64, String)>>::new();
+        let mut buffer = BTreeMap::<Vec<(String, String)>, Vec<(f64, String)>>::new();
 
         for batch in batches.iter() {
             // prepare things...
@@ -338,7 +338,7 @@ impl PromJsonResponse {
                 metric: tags.into_iter().collect(),
                 values,
             })
-            .collect();
+            .collect::<Vec<_>>();
 
         let data = PromData {
             result_type: "matrix".to_string(),