servers/grpc/
prom_query_gateway.rs1use std::sync::Arc;
19
20use api::v1::prometheus_gateway_server::PrometheusGateway;
21use api::v1::promql_request::Promql;
22use api::v1::{PromqlRequest, PromqlResponse, ResponseHeader};
23use async_trait::async_trait;
24use auth::UserProviderRef;
25use common_error::ext::ErrorExt;
26use common_error::status_code::StatusCode;
27use common_time::util::current_time_rfc3339;
28use promql_parser::parser::value::ValueType;
29use query::parser::PromQuery;
30use session::context::{Channel, QueryContext};
31use snafu::OptionExt;
32use tonic::{Request, Response};
33
34use crate::error::InvalidQuerySnafu;
35use crate::grpc::TonicResult;
36use crate::grpc::context_auth::auth;
37use crate::grpc::greptime_handler::create_query_context;
38use crate::http::prometheus::{PrometheusJsonResponse, retrieve_metric_name_and_result_type};
39use crate::prometheus_handler::PrometheusHandlerRef;
40
41pub struct PrometheusGatewayService {
42 handler: PrometheusHandlerRef,
43 user_provider: Option<UserProviderRef>,
44}
45
46#[async_trait]
47impl PrometheusGateway for PrometheusGatewayService {
48 async fn handle(&self, req: Request<PromqlRequest>) -> TonicResult<Response<PromqlResponse>> {
49 let mut is_range_query = false;
50 let inner = req.into_inner();
51 let prom_query = match inner.promql.context(InvalidQuerySnafu {
52 reason: "Expecting non-empty PromqlRequest.",
53 })? {
54 Promql::RangeQuery(range_query) => {
55 is_range_query = true;
56 PromQuery {
57 query: range_query.query,
58 start: range_query.start,
59 end: range_query.end,
60 step: range_query.step,
61 lookback: range_query.lookback,
62 }
63 }
64 Promql::InstantQuery(instant_query) => {
65 let time = if instant_query.time.is_empty() {
66 current_time_rfc3339()
67 } else {
68 instant_query.time
69 };
70 PromQuery {
71 query: instant_query.query,
72 start: time.clone(),
73 end: time,
74 step: String::from("1s"),
75 lookback: instant_query.lookback,
76 }
77 }
78 };
79
80 let header = inner.header.as_ref();
81 let query_ctx = create_query_context(Channel::Promql, header, Default::default())?;
82
83 let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?;
84 query_ctx.set_current_user(user_info);
85
86 let json_response = self
87 .handle_inner(prom_query, query_ctx, is_range_query)
88 .await;
89 let json_bytes = serde_json::to_string(&json_response).unwrap().into_bytes();
90
91 let response = Response::new(PromqlResponse {
92 header: Some(ResponseHeader {
93 status: Some(api::v1::Status {
94 status_code: StatusCode::Success as _,
95 ..Default::default()
96 }),
97 }),
98 body: json_bytes,
99 });
100 Ok(response)
101 }
102}
103
104impl PrometheusGatewayService {
105 pub fn new(handler: PrometheusHandlerRef, user_provider: Option<UserProviderRef>) -> Self {
106 Self {
107 handler,
108 user_provider,
109 }
110 }
111
112 async fn handle_inner(
113 &self,
114 query: PromQuery,
115 ctx: Arc<QueryContext>,
116 is_range_query: bool,
117 ) -> PrometheusJsonResponse {
118 let db = ctx.get_db_string();
119 let _timer = crate::metrics::METRIC_SERVER_GRPC_PROM_REQUEST_TIMER
120 .with_label_values(&[db.as_str()])
121 .start_timer();
122
123 let result = self.handler.do_query(&query, ctx).await;
124 let (metric_name, mut result_type) =
125 match retrieve_metric_name_and_result_type(&query.query) {
126 Ok((metric_name, result_type)) => (metric_name, result_type),
127 Err(err) => {
128 return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
129 }
130 };
131 if is_range_query {
133 result_type = ValueType::Matrix;
134 };
135
136 PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await
137 }
138}