diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index 477af876b0..e2eebbb80c 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -37,7 +37,7 @@ use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, error, tracing, warn}; use common_time::timezone::parse_timezone; use futures_util::StreamExt; -use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; +use session::context::{Channel, QueryContext, QueryContextBuilder, QueryContextRef}; use session::hints::READ_PREFERENCE_HINT; use snafu::{OptionExt, ResultExt}; use table::TableRef; @@ -89,7 +89,7 @@ impl GreptimeRequestHandler { })?; let header = request.header.as_ref(); - let query_ctx = create_query_context(header, hints)?; + let query_ctx = create_query_context(Channel::Grpc, header, hints)?; let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?; query_ctx.set_current_user(user_info); @@ -286,7 +286,10 @@ pub(crate) async fn auth( }) } +/// Creates a new `QueryContext` from the provided request header and extensions. +/// Strongly recommend setting an appropriate channel, as this is very helpful for statistics. pub(crate) fn create_query_context( + channel: Channel, header: Option<&RequestHeader>, mut extensions: Vec<(String, String)>, ) -> Result { @@ -321,7 +324,8 @@ pub(crate) fn create_query_context( let mut ctx_builder = QueryContextBuilder::default() .current_catalog(catalog) .current_schema(schema) - .timezone(timezone); + .timezone(timezone) + .channel(channel); if let Some(x) = extensions .iter() @@ -397,6 +401,7 @@ mod tests { ..Default::default() }; let query_context = create_query_context( + Channel::Unknown, Some(&header), vec![ ("auto_create_table".to_string(), "true".to_string()), diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 8658be1b88..fb9d605e70 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -27,7 +27,7 @@ use common_error::status_code::StatusCode; use common_time::util::current_time_rfc3339; use promql_parser::parser::value::ValueType; use query::parser::PromQuery; -use session::context::QueryContext; +use session::context::{Channel, QueryContext}; use snafu::OptionExt; use tonic::{Request, Response}; @@ -77,7 +77,8 @@ impl PrometheusGateway for PrometheusGatewayService { }; let header = inner.header.as_ref(); - let query_ctx = create_query_context(header, Default::default())?; + let query_ctx = create_query_context(Channel::Promql, header, Default::default())?; + let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?; query_ctx.set_current_user(user_info);