1use std::borrow::Borrow;
18use std::collections::{BTreeMap, HashMap, HashSet};
19use std::hash::{Hash, Hasher};
20use std::sync::Arc;
21
22use arrow::array::{Array, AsArray};
23use arrow::datatypes::{
24 Date32Type, Date64Type, Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type,
25 Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType,
26 UInt8Type, UInt16Type, UInt32Type, UInt64Type,
27};
28use arrow_schema::{DataType, IntervalUnit};
29use auth::{PermissionTableTarget, PermissionTableTargets, UserProviderRef};
30use axum::extract::{Path, Query, State};
31use axum::{Extension, Form};
32use catalog::CatalogManagerRef;
33use common_decimal::Decimal128;
34use common_error::ext::ErrorExt;
35use common_error::status_code::StatusCode;
36use common_query::native_histogram::is_native_histogram_value_type;
37use common_query::{Output, OutputData};
38use common_recordbatch::{RecordBatch, RecordBatches};
39use common_telemetry::{debug, tracing};
40use common_time::util::{current_time_rfc3339, yesterday_rfc3339};
41use common_time::{Date, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
42use common_version::OwnedBuildInfo;
43use datafusion_common::ScalarValue;
44use datatypes::prelude::ConcreteDataType;
45use datatypes::schema::{ColumnSchema, SchemaRef};
46use datatypes::types::jsonb_to_string;
47use futures::future::join_all;
48use futures::{StreamExt, TryStreamExt};
49use itertools::Itertools;
50use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
51use promql_parser::parser::token::TokenType;
52use promql_parser::parser::value::ValueType;
53use promql_parser::parser::{
54 AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier,
55 MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector,
56};
57use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement};
58use serde::de::{self, MapAccess, Visitor};
59use serde::{Deserialize, Serialize};
60use serde_json::Value;
61use session::context::{Channel, QueryContext, QueryContextRef};
62use snafu::{Location, OptionExt, ResultExt, ensure};
63use store_api::metric_engine_consts::{
64 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
65};
66use table::TableRef;
67use table::metadata::TableInfo;
68use table::requests::{
69 METRIC_TEMPORALITY_DELTA, SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE,
70 SEMANTIC_METRIC_UNIT, SEMANTIC_VALUE_MIXED,
71};
72
73pub use super::result::prometheus_resp::{PromSampleValue, PrometheusJsonResponse};
74use crate::error::{
75 CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu,
76 NotSupportedSnafu, ParseTimestampSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu,
77};
78use crate::http::authorize::select_database;
79use crate::http::header::collect_plan_metrics;
80use crate::otlp::metrics::ucum_to_openmetrics_unit;
81use crate::prom_store::{FIELD_NAME_LABEL, METRIC_NAME_LABEL, is_database_selection_label};
82use crate::prometheus_handler::{
83 ParsedPromQuery, PrometheusHandlerRef, resolve_schema_from_matchers,
84};
85
86#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
88#[serde(deny_unknown_fields)]
89pub struct PromSeriesVector {
90 pub metric: BTreeMap<String, String>,
91 #[serde(skip_serializing_if = "Option::is_none")]
92 pub value: Option<(f64, String)>,
93 #[serde(skip_serializing_if = "Option::is_none")]
94 pub histogram: Option<(f64, PromNativeHistogram)>,
95}
96
97#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
99#[serde(deny_unknown_fields)]
100pub struct PromSeriesMatrix {
101 pub metric: BTreeMap<String, String>,
102 #[serde(skip_serializing_if = "Vec::is_empty", default)]
103 pub values: Vec<(f64, PromSampleValue)>,
104 #[serde(skip_serializing_if = "Vec::is_empty", default)]
105 pub histograms: Vec<(f64, PromNativeHistogram)>,
106}
107
108#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
110pub struct PromNativeHistogram {
111 pub count: String,
113 pub sum: String,
115 pub buckets: Vec<(u8, String, String, String)>,
121}
122
123#[derive(Debug, Serialize, Deserialize, PartialEq)]
125#[serde(untagged)]
126pub enum PromQueryResult {
127 Matrix(Vec<PromSeriesMatrix>),
128 Vector(Vec<PromSeriesVector>),
129 Scalar(#[serde(skip_serializing_if = "Option::is_none")] Option<(f64, String)>),
130 String(#[serde(skip_serializing_if = "Option::is_none")] Option<(f64, String)>),
131}
132
133impl Default for PromQueryResult {
134 fn default() -> Self {
135 PromQueryResult::Matrix(Default::default())
136 }
137}
138
139#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
140pub struct PromData {
141 #[serde(rename = "resultType")]
142 pub result_type: String,
143 pub result: PromQueryResult,
144}
145
146#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
148pub struct PromMetadata {
149 #[serde(rename = "type")]
150 pub metric_type: String,
151 pub unit: String,
152 pub help: String,
153}
154
155#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
158pub struct Column(Arc<String>);
159
160impl From<&str> for Column {
161 fn from(s: &str) -> Self {
162 Self(Arc::new(s.to_string()))
163 }
164}
165
166#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
167#[serde(untagged)]
168pub enum PrometheusResponse {
169 PromData(PromData),
170 Labels(Vec<String>),
171 Series(Vec<HashMap<Column, String>>),
172 Metadata(BTreeMap<String, Vec<PromMetadata>>),
173 LabelValues(Vec<String>),
174 FormatQuery(String),
175 BuildInfo(OwnedBuildInfo),
176 #[serde(skip_deserializing)]
177 ParseResult(promql_parser::parser::Expr),
178 Exemplars(Vec<crate::exemplar::query::Series>),
179 #[default]
180 None,
181}
182
183impl PrometheusResponse {
184 pub(super) fn append(&mut self, other: PrometheusResponse) {
188 match (self, other) {
189 (
190 PrometheusResponse::PromData(PromData {
191 result: PromQueryResult::Matrix(lhs),
192 ..
193 }),
194 PrometheusResponse::PromData(PromData {
195 result: PromQueryResult::Matrix(rhs),
196 ..
197 }),
198 ) => {
199 lhs.extend(rhs);
200 }
201
202 (
203 PrometheusResponse::PromData(PromData {
204 result: PromQueryResult::Vector(lhs),
205 ..
206 }),
207 PrometheusResponse::PromData(PromData {
208 result: PromQueryResult::Vector(rhs),
209 ..
210 }),
211 ) => {
212 lhs.extend(rhs);
213 }
214 _ => {
215 }
217 }
218 }
219
220 pub fn is_none(&self) -> bool {
221 matches!(self, PrometheusResponse::None)
222 }
223}
224
225#[derive(Debug, Default, Deserialize)]
226pub struct ExemplarQuery {
227 query: Option<String>,
228 start: Option<String>,
229 end: Option<String>,
230 db: Option<String>,
231 timeout: Option<String>,
232}
233
234#[axum_macros::debug_handler]
235#[tracing::instrument(
236 skip_all,
237 fields(protocol = "prometheus", request_type = "query_exemplars")
238)]
239pub async fn query_exemplars(
240 user_provider: Option<Extension<UserProviderRef>>,
241 State(handler): State<PrometheusHandlerRef>,
242 Query(params): Query<ExemplarQuery>,
243 Extension(mut ctx): Extension<QueryContext>,
244 Form(form): Form<ExemplarQuery>,
245) -> PrometheusJsonResponse {
246 if let Err(error) = select_database(
247 &mut ctx,
248 params.db.or(form.db).as_deref(),
249 user_provider.as_ref().map(|Extension(provider)| provider),
250 )
251 .await
252 {
253 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
254 }
255 ctx.set_channel(Channel::Prometheus);
256 let ctx = Arc::new(ctx);
257 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
258 .with_label_values(&[ctx.get_db_string().as_str(), "query_exemplars"])
259 .start_timer();
260 let query = match crate::exemplar::query::Query::parse(
261 params.query.or(form.query).unwrap_or_default(),
262 ¶ms.start.or(form.start).unwrap_or_default(),
263 ¶ms.end.or(form.end).unwrap_or_default(),
264 &ctx.current_schema(),
265 ) {
266 Ok(query) => query,
267 Err(error) => {
268 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
269 }
270 };
271 let timeout = if let Some(timeout) = params.timeout.or(form.timeout) {
272 match humantime::parse_duration(&timeout) {
273 Ok(timeout) => Some(timeout),
274 Err(error) => {
275 return PrometheusJsonResponse::error(
276 StatusCode::InvalidArguments,
277 error.to_string(),
278 );
279 }
280 }
281 } else {
282 ctx.query_timeout()
283 };
284 let future = handler.query_exemplars(query, ctx.clone());
285 let result = if let Some(timeout) = timeout.filter(|timeout| !timeout.is_zero()) {
286 match tokio::time::timeout(timeout, future).await {
287 Ok(result) => result,
288 Err(_) => {
289 return PrometheusJsonResponse::error(
290 StatusCode::Cancelled,
291 "exemplar query timed out".to_string(),
292 );
293 }
294 }
295 } else {
296 future.await
297 };
298 match result {
299 Ok(series) => {
300 let mut response =
301 PrometheusJsonResponse::success(PrometheusResponse::Exemplars(series));
302 response.warnings = ctx.warning().map(|warning| vec![warning]);
303 response
304 }
305 Err(error) => PrometheusJsonResponse::error(error.status_code(), error.output_msg()),
306 }
307}
308
309#[derive(Debug, Default, Serialize, Deserialize)]
310pub struct FormatQuery {
311 query: Option<String>,
312}
313
314#[derive(Debug, Default, Serialize, Deserialize)]
316pub struct MetadataQuery {
317 db: Option<String>,
318 limit: Option<usize>,
319 metric: Option<String>,
320}
321
322#[axum_macros::debug_handler]
323#[tracing::instrument(
324 skip_all,
325 fields(protocol = "prometheus", request_type = "format_query")
326)]
327pub async fn format_query(
328 State(_handler): State<PrometheusHandlerRef>,
329 Query(params): Query<InstantQuery>,
330 Extension(_query_ctx): Extension<QueryContext>,
331 Form(form_params): Form<InstantQuery>,
332) -> PrometheusJsonResponse {
333 let query = params.query.or(form_params.query).unwrap_or_default();
334 match promql_parser::parser::parse(&query) {
335 Ok(expr) => {
336 let pretty = expr.prettify();
337 PrometheusJsonResponse::success(PrometheusResponse::FormatQuery(pretty))
338 }
339 Err(reason) => {
340 let err = InvalidQuerySnafu { reason }.build();
341 PrometheusJsonResponse::error(err.status_code(), err.output_msg())
342 }
343 }
344}
345
346#[derive(Debug, Default, Serialize, Deserialize)]
347pub struct BuildInfoQuery {}
348
349#[axum_macros::debug_handler]
350#[tracing::instrument(
351 skip_all,
352 fields(protocol = "prometheus", request_type = "build_info_query")
353)]
354pub async fn build_info_query() -> PrometheusJsonResponse {
355 let build_info = common_version::build_info().clone();
356 PrometheusJsonResponse::success(PrometheusResponse::BuildInfo(build_info.into()))
357}
358
359#[axum_macros::debug_handler]
360#[tracing::instrument(
361 skip_all,
362 fields(protocol = "prometheus", request_type = "metadata_query")
363)]
364pub async fn metadata_query(
365 user_provider: Option<Extension<UserProviderRef>>,
366 State(handler): State<PrometheusHandlerRef>,
367 Query(params): Query<MetadataQuery>,
368 Extension(mut query_ctx): Extension<QueryContext>,
369) -> PrometheusJsonResponse {
370 if let Err(error) = select_database(
371 &mut query_ctx,
372 params.db.as_deref(),
373 user_provider.as_ref().map(|Extension(provider)| provider),
374 )
375 .await
376 {
377 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
378 }
379 let query_ctx = Arc::new(query_ctx);
380
381 if let Err(err) = handler.check_query_permission(&[], &query_ctx).await {
382 return PrometheusJsonResponse::error(err.status_code(), err.to_string());
383 }
384 if let Some(metric) = ¶ms.metric
385 && let Err(err) = handler
386 .check_query_target_permission(
387 current_schema_metric_targets(&query_ctx, std::slice::from_ref(metric)),
388 &query_ctx,
389 )
390 .await
391 {
392 return PrometheusJsonResponse::error(err.status_code(), err.to_string());
393 }
394
395 let mut metadata =
396 match retrieve_metric_metadata(&query_ctx, handler.catalog_manager(), ¶ms).await {
397 Ok(metadata) => metadata,
398 Err(err) => {
399 return PrometheusJsonResponse::error(
400 StatusCode::InvalidArguments,
401 err.to_string(),
402 );
403 }
404 };
405 let allowed_metric_names = match handler
406 .filter_metadata_metric_names(
407 metadata.keys().cloned().collect(),
408 query_ctx.current_schema().as_str(),
409 &query_ctx,
410 )
411 .await
412 {
413 Ok(metric_names) => metric_names,
414 Err(err) => {
415 return PrometheusJsonResponse::error(err.status_code(), err.to_string());
416 }
417 }
418 .into_iter()
419 .collect::<HashSet<_>>();
420 metadata.retain(|metric, _| allowed_metric_names.contains(metric));
421 if let Some(limit) = params.limit {
422 metadata = metadata.into_iter().take(limit).collect();
423 }
424
425 PrometheusJsonResponse::success(PrometheusResponse::Metadata(metadata))
426}
427
428#[derive(Debug, Default, Serialize, Deserialize)]
429pub struct InstantQuery {
430 query: Option<String>,
431 lookback: Option<String>,
432 time: Option<String>,
433 timeout: Option<String>,
434 db: Option<String>,
435}
436
437macro_rules! try_call_return_response {
440 (@output_msg $handle: expr, $status_code: expr) => {
441 match $handle {
442 Ok(res) => res,
443 Err(err) => {
444 let msg = err.output_msg();
445 return PrometheusJsonResponse::error($status_code, msg);
446 }
447 }
448 };
449 ($handle: expr, $status_code: expr) => {
450 match $handle {
451 Ok(res) => res,
452 Err(err) => {
453 let msg = err.to_string();
454 return PrometheusJsonResponse::error($status_code, msg);
455 }
456 }
457 };
458 ($handle: expr) => {
459 match $handle {
460 Ok(res) => res,
461 Err(err) => {
462 let status_code = err.status_code();
463 let msg = err.to_string();
464 return PrometheusJsonResponse::error(status_code, msg);
465 }
466 }
467 };
468}
469
470#[axum_macros::debug_handler]
471#[tracing::instrument(
472 skip_all,
473 fields(protocol = "prometheus", request_type = "instant_query")
474)]
475pub async fn instant_query(
476 user_provider: Option<Extension<UserProviderRef>>,
477 State(handler): State<PrometheusHandlerRef>,
478 Query(params): Query<InstantQuery>,
479 Extension(mut query_ctx): Extension<QueryContext>,
480 Form(form_params): Form<InstantQuery>,
481) -> PrometheusJsonResponse {
482 let time = params
484 .time
485 .or(form_params.time)
486 .unwrap_or_else(current_time_rfc3339);
487 let prom_query = PromQuery {
488 query: params.query.or(form_params.query).unwrap_or_default(),
489 start: time.clone(),
490 end: time,
491 step: "1s".to_string(),
492 lookback: params
493 .lookback
494 .or(form_params.lookback)
495 .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
496 alias: None,
497 };
498
499 if let Err(error) = select_database(
501 &mut query_ctx,
502 params.db.as_deref(),
503 user_provider.as_ref().map(|Extension(provider)| provider),
504 )
505 .await
506 {
507 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
508 }
509 let query_ctx = Arc::new(query_ctx);
510 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
511 .with_label_values(&[query_ctx.get_db_string().as_str(), "instant_query"])
512 .start_timer();
513 let prom_query = try_call_return_response!(
514 @output_msg ParsedPromQuery::parse(prom_query, &query_ctx),
515 StatusCode::InvalidArguments
516 );
517 let promql_expr = prom_query.expr();
518
519 let metric_name_discovery =
520 try_call_return_response!(find_metric_name_not_equal_matchers(promql_expr));
521 if let Some(discovery) = metric_name_discovery {
522 debug!("Find metric name matchers: {:?}", discovery.name_matchers);
523
524 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
525 let (static_targets, unresolved_selectors) =
526 try_call_return_response!(static_promql_targets(promql_expr, &query_ctx));
527 try_call_return_response!(
528 handler
529 .check_query_target_permission(static_targets, &query_ctx,)
530 .await
531 );
532 if unresolved_selectors != 1 {
533 try_call_return_response!(
534 handler
535 .check_query_target_permission(PermissionTableTargets::Unresolved, &query_ctx,)
536 .await
537 );
538 return do_instant_query(&handler, prom_query, query_ctx).await;
539 }
540
541 let schema = discovery
542 .schema
543 .unwrap_or_else(|| query_ctx.current_schema());
544 let metric_names = try_call_return_response!(
545 handler
546 .query_metric_names(discovery.name_matchers, &schema, &query_ctx)
547 .await
548 );
549
550 debug!("Find metric names: {:?}", metric_names);
551
552 let prom_queries = expand_metric_name_queries(&prom_query, metric_names);
553 try_call_return_response!(
554 handler
555 .check_query_permission_parsed(&prom_queries, &query_ctx)
556 .await
557 );
558
559 if prom_queries.is_empty() {
560 let result_type = promql_expr.value_type();
561
562 return PrometheusJsonResponse::success(PrometheusResponse::PromData(PromData {
563 result_type: result_type.to_string(),
564 ..Default::default()
565 }));
566 }
567
568 let responses = join_all(prom_queries.into_iter().map(|prom_query| {
569 let query_ctx = query_ctx.clone();
570 let handler = handler.clone();
571
572 async move { do_instant_query(&handler, prom_query, query_ctx).await }
573 }))
574 .await;
575
576 responses
577 .into_iter()
578 .reduce(|mut acc, resp| {
579 acc.append_query_response(resp);
580 acc
581 })
582 .unwrap()
583 } else {
584 do_instant_query(&handler, prom_query, query_ctx).await
585 }
586}
587
588async fn do_instant_query(
590 handler: &PrometheusHandlerRef,
591 prom_query: ParsedPromQuery,
592 query_ctx: QueryContextRef,
593) -> PrometheusJsonResponse {
594 let (metric_name, result_type) = retrieve_metric_name_and_result_type(prom_query.expr());
595 let query_id = query_ctx.remote_query_id().map(str::to_string);
596 let result = handler.do_query_parsed(prom_query, query_ctx).await;
597 PrometheusJsonResponse::from_query_result(result, metric_name, result_type, query_id.as_deref())
598 .await
599}
600
601#[derive(Debug, Default, Serialize, Deserialize)]
602pub struct RangeQuery {
603 query: Option<String>,
604 start: Option<String>,
605 end: Option<String>,
606 step: Option<String>,
607 lookback: Option<String>,
608 timeout: Option<String>,
609 db: Option<String>,
610}
611
612#[axum_macros::debug_handler]
613#[tracing::instrument(
614 skip_all,
615 fields(protocol = "prometheus", request_type = "range_query")
616)]
617pub async fn range_query(
618 user_provider: Option<Extension<UserProviderRef>>,
619 State(handler): State<PrometheusHandlerRef>,
620 Query(params): Query<RangeQuery>,
621 Extension(mut query_ctx): Extension<QueryContext>,
622 Form(form_params): Form<RangeQuery>,
623) -> PrometheusJsonResponse {
624 let prom_query = PromQuery {
625 query: params.query.or(form_params.query).unwrap_or_default(),
626 start: params.start.or(form_params.start).unwrap_or_default(),
627 end: params.end.or(form_params.end).unwrap_or_default(),
628 step: params.step.or(form_params.step).unwrap_or_default(),
629 lookback: params
630 .lookback
631 .or(form_params.lookback)
632 .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
633 alias: None,
634 };
635
636 if let Err(error) = select_database(
638 &mut query_ctx,
639 params.db.as_deref(),
640 user_provider.as_ref().map(|Extension(provider)| provider),
641 )
642 .await
643 {
644 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
645 }
646 let query_ctx = Arc::new(query_ctx);
647 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
648 .with_label_values(&[query_ctx.get_db_string().as_str(), "range_query"])
649 .start_timer();
650 let prom_query = try_call_return_response!(
651 @output_msg ParsedPromQuery::parse(prom_query, &query_ctx),
652 StatusCode::InvalidArguments
653 );
654 let promql_expr = prom_query.expr();
655 if promql_expr.value_type() == ValueType::Matrix {
658 let err = query::promql::error::RangeVectorInRangeQuerySnafu.build();
659 return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
660 }
661
662 let metric_name_discovery =
663 try_call_return_response!(find_metric_name_not_equal_matchers(promql_expr));
664 if let Some(discovery) = metric_name_discovery {
665 debug!("Find metric name matchers: {:?}", discovery.name_matchers);
666
667 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
668 let (static_targets, unresolved_selectors) =
669 try_call_return_response!(static_promql_targets(promql_expr, &query_ctx));
670 try_call_return_response!(
671 handler
672 .check_query_target_permission(static_targets, &query_ctx,)
673 .await
674 );
675 if unresolved_selectors != 1 {
676 try_call_return_response!(
677 handler
678 .check_query_target_permission(PermissionTableTargets::Unresolved, &query_ctx,)
679 .await
680 );
681 return do_range_query(&handler, prom_query, query_ctx).await;
682 }
683
684 let schema = discovery
685 .schema
686 .unwrap_or_else(|| query_ctx.current_schema());
687 let metric_names = try_call_return_response!(
688 handler
689 .query_metric_names(discovery.name_matchers, &schema, &query_ctx)
690 .await
691 );
692
693 debug!("Find metric names: {:?}", metric_names);
694
695 let prom_queries = expand_metric_name_queries(&prom_query, metric_names);
696 try_call_return_response!(
697 handler
698 .check_query_permission_parsed(&prom_queries, &query_ctx)
699 .await
700 );
701
702 if prom_queries.is_empty() {
703 return PrometheusJsonResponse::success(PrometheusResponse::PromData(PromData {
704 result_type: ValueType::Matrix.to_string(),
705 ..Default::default()
706 }));
707 }
708
709 let responses = join_all(prom_queries.into_iter().map(|prom_query| {
710 let query_ctx = query_ctx.clone();
711 let handler = handler.clone();
712
713 async move { do_range_query(&handler, prom_query, query_ctx).await }
714 }))
715 .await;
716
717 responses
719 .into_iter()
720 .reduce(|mut acc, resp| {
721 acc.append_query_response(resp);
722 acc
723 })
724 .unwrap()
725 } else {
726 do_range_query(&handler, prom_query, query_ctx).await
727 }
728}
729
730async fn do_range_query(
732 handler: &PrometheusHandlerRef,
733 prom_query: ParsedPromQuery,
734 query_ctx: QueryContextRef,
735) -> PrometheusJsonResponse {
736 let (metric_name, _) = retrieve_metric_name_and_result_type(prom_query.expr());
737 let query_id = query_ctx.remote_query_id().map(str::to_string);
738 let result = handler
741 .do_query_parsed(prom_query.with_unordered_output(), query_ctx)
742 .await;
743 PrometheusJsonResponse::from_query_result(
744 result,
745 metric_name,
746 ValueType::Matrix,
747 query_id.as_deref(),
748 )
749 .await
750}
751
752#[derive(Debug, Default, Serialize)]
753struct Matches(Vec<String>);
754
755#[derive(Debug, Default, Serialize, Deserialize)]
756pub struct LabelsQuery {
757 start: Option<String>,
758 end: Option<String>,
759 lookback: Option<String>,
760 #[serde(flatten)]
761 matches: Matches,
762 db: Option<String>,
763}
764
765impl<'de> Deserialize<'de> for Matches {
767 fn deserialize<D>(deserializer: D) -> std::result::Result<Matches, D::Error>
768 where
769 D: de::Deserializer<'de>,
770 {
771 struct MatchesVisitor;
772
773 impl<'d> Visitor<'d> for MatchesVisitor {
774 type Value = Vec<String>;
775
776 fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
777 formatter.write_str("a string")
778 }
779
780 fn visit_map<M>(self, mut access: M) -> std::result::Result<Self::Value, M::Error>
781 where
782 M: MapAccess<'d>,
783 {
784 let mut matches = Vec::new();
785 while let Some((key, value)) = access.next_entry::<String, String>()? {
786 if key == "match[]" {
787 matches.push(value);
788 }
789 }
790 Ok(matches)
791 }
792 }
793 Ok(Matches(deserializer.deserialize_map(MatchesVisitor)?))
794 }
795}
796
797macro_rules! handle_schema_err {
804 ($result:expr) => {
805 match $result {
806 Ok(v) => Some(v),
807 Err(err) => {
808 if err.status_code() == StatusCode::TableNotFound
809 || err.status_code() == StatusCode::TableColumnNotFound
810 {
811 None
813 } else {
814 return PrometheusJsonResponse::error(err.status_code(), err.output_msg());
815 }
816 }
817 }
818 };
819}
820
821#[axum_macros::debug_handler]
822#[tracing::instrument(
823 skip_all,
824 fields(protocol = "prometheus", request_type = "labels_query")
825)]
826pub async fn labels_query(
827 user_provider: Option<Extension<UserProviderRef>>,
828 State(handler): State<PrometheusHandlerRef>,
829 Query(params): Query<LabelsQuery>,
830 Extension(mut query_ctx): Extension<QueryContext>,
831 Form(form_params): Form<LabelsQuery>,
832) -> PrometheusJsonResponse {
833 if let Err(error) = select_database(
834 &mut query_ctx,
835 params.db.as_deref(),
836 user_provider.as_ref().map(|Extension(provider)| provider),
837 )
838 .await
839 {
840 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
841 }
842 let query_ctx = Arc::new(query_ctx);
843
844 let mut queries = params.matches.0;
845 if queries.is_empty() {
846 queries = form_params.matches.0;
847 }
848
849 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
850 .with_label_values(&[query_ctx.get_db_string().as_str(), "labels_query"])
851 .start_timer();
852
853 let prom_queries = if queries.is_empty() {
854 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
855 Vec::new()
856 } else {
857 let start = params
858 .start
859 .or(form_params.start)
860 .unwrap_or_else(yesterday_rfc3339);
861 let end = params
862 .end
863 .or(form_params.end)
864 .unwrap_or_else(current_time_rfc3339);
865 let lookback = params
866 .lookback
867 .or(form_params.lookback)
868 .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string());
869 let prom_queries = queries
870 .into_iter()
871 .map(|query| PromQuery {
872 query,
873 start: start.clone(),
874 end: end.clone(),
875 step: DEFAULT_LOOKBACK_STRING.to_string(),
876 lookback: lookback.clone(),
877 alias: None,
878 })
879 .collect::<Vec<_>>();
880 let prom_queries = try_call_return_response!(
881 prom_queries
882 .into_iter()
883 .map(|query| ParsedPromQuery::parse(query, &query_ctx))
884 .collect::<Result<Vec<_>>>()
885 );
886 try_call_return_response!(
887 handler
888 .check_query_permission_parsed(&prom_queries, &query_ctx)
889 .await
890 );
891 prom_queries
892 };
893
894 let catalog = query_ctx.current_catalog();
895 let schema = query_ctx.current_schema();
896
897 if prom_queries.is_empty() {
899 let (mut labels, table_names) =
900 match get_all_column_names(catalog, &schema, &handler.catalog_manager()).await {
901 Ok(result) => result,
902 Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()),
903 };
904 try_call_return_response!(
905 handler
906 .check_query_target_permission(
907 current_schema_metric_targets(&query_ctx, &table_names),
908 &query_ctx,
909 )
910 .await
911 );
912 let _ = labels.insert(METRIC_NAME.to_string());
913 let mut labels_vec = labels.into_iter().collect::<Vec<_>>();
914 labels_vec.sort_unstable();
915 return PrometheusJsonResponse::success(PrometheusResponse::Labels(labels_vec));
916 }
917
918 let mut labels = match resolved_promql_targets(&prom_queries, &query_ctx) {
920 Some(targets) => {
921 match get_target_column_names(&targets, &handler.catalog_manager(), &query_ctx).await {
922 Ok(labels) => labels,
923 Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()),
924 }
925 }
926 None => match get_all_column_names(catalog, &schema, &handler.catalog_manager()).await {
927 Ok((labels, _)) => labels,
928 Err(e) => return PrometheusJsonResponse::error(e.status_code(), e.output_msg()),
929 },
930 };
931 let _ = labels.insert(METRIC_NAME.to_string());
932
933 let mut fetched_labels = HashSet::new();
934 let _ = fetched_labels.insert(METRIC_NAME.to_string());
935
936 let mut merge_map = HashMap::new();
937 for prom_query in prom_queries {
938 let result = handler.do_query_parsed(prom_query, query_ctx.clone()).await;
939 handle_schema_err!(
940 retrieve_labels_name_from_query_result(result, &mut fetched_labels, &mut merge_map)
941 .await
942 );
943 }
944
945 fetched_labels.retain(|l| labels.contains(l));
947 let mut sorted_labels: Vec<String> = fetched_labels.into_iter().collect();
948 sorted_labels.sort();
949 let merge_map = merge_map
950 .into_iter()
951 .map(|(k, v)| (k, Value::from(v)))
952 .collect();
953 let mut resp = PrometheusJsonResponse::success(PrometheusResponse::Labels(sorted_labels));
954 resp.resp_metrics = merge_map;
955 resp
956}
957
958async fn get_all_column_names(
960 catalog: &str,
961 schema: &str,
962 manager: &CatalogManagerRef,
963) -> std::result::Result<(HashSet<String>, Vec<String>), catalog::error::Error> {
964 let mut labels = HashSet::new();
965 let mut table_names = Vec::new();
966 let mut tables = manager.tables(catalog, schema, None);
967 while let Some(table) = tables.try_next().await? {
968 if is_internal_physical_metric_table(&table) {
969 continue;
970 }
971 table_names.push(table.table_info().name.clone());
972 extend_tag_column_names(&mut labels, &table);
973 }
974
975 Ok((labels, table_names))
976}
977
978async fn get_target_column_names(
979 targets: &[PermissionTableTarget],
980 manager: &CatalogManagerRef,
981 query_ctx: &QueryContext,
982) -> std::result::Result<HashSet<String>, catalog::error::Error> {
983 let mut labels = HashSet::new();
984 for target in targets {
985 if let Some(table) = manager
986 .table(
987 &target.catalog,
988 &target.schema,
989 &target.table,
990 Some(query_ctx),
991 )
992 .await?
993 && !is_internal_physical_metric_table(&table)
994 {
995 extend_tag_column_names(&mut labels, &table);
996 }
997 }
998 Ok(labels)
999}
1000
1001fn extend_tag_column_names(labels: &mut HashSet<String>, table: &TableRef) {
1002 labels.extend(table.primary_key_columns().filter_map(|column| {
1003 (column.name != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
1004 && column.name != DATA_SCHEMA_TSID_COLUMN_NAME)
1005 .then_some(column.name)
1006 }));
1007}
1008
1009async fn retrieve_series_from_query_result(
1010 result: Result<Output>,
1011 series: &mut Vec<HashMap<Column, String>>,
1012 query_ctx: &QueryContext,
1013 target: &PermissionTableTarget,
1014 manager: &CatalogManagerRef,
1015 metrics: &mut HashMap<String, u64>,
1016) -> Result<()> {
1017 let result = result?;
1018
1019 let table = manager
1021 .table(
1022 &target.catalog,
1023 &target.schema,
1024 &target.table,
1025 Some(query_ctx),
1026 )
1027 .await?
1028 .with_context(|| TableNotFoundSnafu {
1029 catalog: &target.catalog,
1030 schema: &target.schema,
1031 table: &target.table,
1032 })?;
1033 let tag_columns = table
1034 .primary_key_columns()
1035 .map(|c| c.name)
1036 .collect::<HashSet<_>>();
1037
1038 match result.data {
1039 OutputData::RecordBatches(batches) => {
1040 record_batches_to_series(batches, series, &target.table, &tag_columns)
1041 }
1042 OutputData::Stream(stream) => {
1043 let batches = RecordBatches::try_collect(stream)
1044 .await
1045 .context(CollectRecordbatchSnafu)?;
1046 record_batches_to_series(batches, series, &target.table, &tag_columns)
1047 }
1048 OutputData::AffectedRows(_) => Err(Error::UnexpectedResult {
1049 reason: "expected data result, but got affected rows".to_string(),
1050 location: Location::default(),
1051 }),
1052 }?;
1053
1054 if let Some(ref plan) = result.meta.plan {
1055 collect_plan_metrics(plan, &mut [metrics]);
1056 }
1057 Ok(())
1058}
1059
1060async fn retrieve_labels_name_from_query_result(
1062 result: Result<Output>,
1063 labels: &mut HashSet<String>,
1064 metrics: &mut HashMap<String, u64>,
1065) -> Result<()> {
1066 let result = result?;
1067 match result.data {
1068 OutputData::RecordBatches(batches) => record_batches_to_labels_name(batches, labels),
1069 OutputData::Stream(stream) => {
1070 let batches = RecordBatches::try_collect(stream)
1071 .await
1072 .context(CollectRecordbatchSnafu)?;
1073 record_batches_to_labels_name(batches, labels)
1074 }
1075 OutputData::AffectedRows(_) => UnexpectedResultSnafu {
1076 reason: "expected data result, but got affected rows".to_string(),
1077 }
1078 .fail(),
1079 }?;
1080 if let Some(ref plan) = result.meta.plan {
1081 collect_plan_metrics(plan, &mut [metrics]);
1082 }
1083 Ok(())
1084}
1085
1086fn record_batches_to_series(
1087 batches: RecordBatches,
1088 series: &mut Vec<HashMap<Column, String>>,
1089 table_name: &str,
1090 tag_columns: &HashSet<String>,
1091) -> Result<()> {
1092 for batch in batches.iter() {
1093 let projection = batch
1095 .schema
1096 .column_schemas()
1097 .iter()
1098 .enumerate()
1099 .filter_map(|(idx, col)| {
1100 if tag_columns.contains(&col.name) {
1101 Some(idx)
1102 } else {
1103 None
1104 }
1105 })
1106 .collect::<Vec<_>>();
1107 let batch = batch
1108 .try_project(&projection)
1109 .context(CollectRecordbatchSnafu)?;
1110
1111 let mut writer = RowWriter::new(&batch.schema, table_name);
1112 writer.write(batch, series)?;
1113 }
1114 Ok(())
1115}
1116
1117struct RowWriter {
1124 template: HashMap<Column, Option<String>>,
1127 current: Option<HashMap<Column, Option<String>>>,
1129}
1130
1131impl RowWriter {
1132 fn new(schema: &SchemaRef, table: &str) -> Self {
1133 let mut template = schema
1134 .column_schemas()
1135 .iter()
1136 .map(|x| (x.name.as_str().into(), None))
1137 .collect::<HashMap<Column, Option<String>>>();
1138 template.insert("__name__".into(), Some(table.to_string()));
1139 Self {
1140 template,
1141 current: None,
1142 }
1143 }
1144
1145 fn insert(&mut self, column: ColumnRef, value: impl ToString) {
1146 let current = self.current.get_or_insert_with(|| self.template.clone());
1147 match current.get_mut(&column as &dyn AsColumnRef) {
1148 Some(x) => {
1149 let _ = x.insert(value.to_string());
1150 }
1151 None => {
1152 let _ = current.insert(column.0.into(), Some(value.to_string()));
1153 }
1154 }
1155 }
1156
1157 fn insert_bytes(&mut self, column_schema: &ColumnSchema, bytes: &[u8]) -> Result<()> {
1158 let column_name = column_schema.name.as_str().into();
1159
1160 if column_schema.data_type.is_json() {
1161 let s = jsonb_to_string(bytes).context(ConvertScalarValueSnafu)?;
1162 self.insert(column_name, s);
1163 } else {
1164 let hex = bytes
1165 .iter()
1166 .map(|b| format!("{b:02x}"))
1167 .collect::<Vec<String>>()
1168 .join("");
1169 self.insert(column_name, hex);
1170 }
1171 Ok(())
1172 }
1173
1174 fn finish(&mut self) -> HashMap<Column, String> {
1175 let Some(current) = self.current.take() else {
1176 return HashMap::new();
1177 };
1178 current
1179 .into_iter()
1180 .filter_map(|(k, v)| v.map(|v| (k, v)))
1181 .collect()
1182 }
1183
1184 fn write(
1185 &mut self,
1186 record_batch: RecordBatch,
1187 series: &mut Vec<HashMap<Column, String>>,
1188 ) -> Result<()> {
1189 let schema = record_batch.schema.clone();
1190 let record_batch = record_batch.into_df_record_batch();
1191 for i in 0..record_batch.num_rows() {
1192 for (j, array) in record_batch.columns().iter().enumerate() {
1193 let column = schema.column_name_by_index(j).into();
1194
1195 if array.is_null(i) {
1196 self.insert(column, "Null");
1197 continue;
1198 }
1199
1200 match array.data_type() {
1201 DataType::Null => {
1202 self.insert(column, "Null");
1203 }
1204 DataType::Boolean => {
1205 let array = array.as_boolean();
1206 let v = array.value(i);
1207 self.insert(column, v);
1208 }
1209 DataType::UInt8 => {
1210 let array = array.as_primitive::<UInt8Type>();
1211 let v = array.value(i);
1212 self.insert(column, v);
1213 }
1214 DataType::UInt16 => {
1215 let array = array.as_primitive::<UInt16Type>();
1216 let v = array.value(i);
1217 self.insert(column, v);
1218 }
1219 DataType::UInt32 => {
1220 let array = array.as_primitive::<UInt32Type>();
1221 let v = array.value(i);
1222 self.insert(column, v);
1223 }
1224 DataType::UInt64 => {
1225 let array = array.as_primitive::<UInt64Type>();
1226 let v = array.value(i);
1227 self.insert(column, v);
1228 }
1229 DataType::Int8 => {
1230 let array = array.as_primitive::<Int8Type>();
1231 let v = array.value(i);
1232 self.insert(column, v);
1233 }
1234 DataType::Int16 => {
1235 let array = array.as_primitive::<Int16Type>();
1236 let v = array.value(i);
1237 self.insert(column, v);
1238 }
1239 DataType::Int32 => {
1240 let array = array.as_primitive::<Int32Type>();
1241 let v = array.value(i);
1242 self.insert(column, v);
1243 }
1244 DataType::Int64 => {
1245 let array = array.as_primitive::<Int64Type>();
1246 let v = array.value(i);
1247 self.insert(column, v);
1248 }
1249 DataType::Float32 => {
1250 let array = array.as_primitive::<Float32Type>();
1251 let v = array.value(i);
1252 self.insert(column, v);
1253 }
1254 DataType::Float64 => {
1255 let array = array.as_primitive::<Float64Type>();
1256 let v = array.value(i);
1257 self.insert(column, v);
1258 }
1259 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
1260 let v = datatypes::arrow_array::string_array_value(array, i);
1261 self.insert(column, v);
1262 }
1263 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
1264 let v = datatypes::arrow_array::binary_array_value(array, i);
1265 let column_schema = &schema.column_schemas()[j];
1266 self.insert_bytes(column_schema, v)?;
1267 }
1268 DataType::Date32 => {
1269 let array = array.as_primitive::<Date32Type>();
1270 let v = Date::new(array.value(i));
1271 self.insert(column, v);
1272 }
1273 DataType::Date64 => {
1274 let array = array.as_primitive::<Date64Type>();
1275 let v = Date::new((array.value(i) / 86_400_000) as i32);
1279 self.insert(column, v);
1280 }
1281 DataType::Timestamp(_, _) => {
1282 let v = datatypes::arrow_array::timestamp_array_value(array, i);
1283 self.insert(column, v.to_iso8601_string());
1284 }
1285 DataType::Time32(_) | DataType::Time64(_) => {
1286 let v = datatypes::arrow_array::time_array_value(array, i);
1287 self.insert(column, v.to_iso8601_string());
1288 }
1289 DataType::Interval(interval_unit) => match interval_unit {
1290 IntervalUnit::YearMonth => {
1291 let array = array.as_primitive::<IntervalYearMonthType>();
1292 let v: IntervalYearMonth = array.value(i).into();
1293 self.insert(column, v.to_iso8601_string());
1294 }
1295 IntervalUnit::DayTime => {
1296 let array = array.as_primitive::<IntervalDayTimeType>();
1297 let v: IntervalDayTime = array.value(i).into();
1298 self.insert(column, v.to_iso8601_string());
1299 }
1300 IntervalUnit::MonthDayNano => {
1301 let array = array.as_primitive::<IntervalMonthDayNanoType>();
1302 let v: IntervalMonthDayNano = array.value(i).into();
1303 self.insert(column, v.to_iso8601_string());
1304 }
1305 },
1306 DataType::Duration(_) => {
1307 let d = datatypes::arrow_array::duration_array_value(array, i);
1308 self.insert(column, d);
1309 }
1310 DataType::List(_) => {
1311 let v = ScalarValue::try_from_array(array, i).context(DataFusionSnafu)?;
1312 self.insert(column, v);
1313 }
1314 DataType::Struct(_) => {
1315 let v = ScalarValue::try_from_array(array, i).context(DataFusionSnafu)?;
1316 self.insert(column, v);
1317 }
1318 DataType::Decimal128(precision, scale) => {
1319 let array = array.as_primitive::<Decimal128Type>();
1320 let v = Decimal128::new(array.value(i), *precision, *scale);
1321 self.insert(column, v);
1322 }
1323 _ => {
1324 return NotSupportedSnafu {
1325 feat: format!("convert {} to http value", array.data_type()),
1326 }
1327 .fail();
1328 }
1329 }
1330 }
1331
1332 series.push(self.finish())
1333 }
1334 Ok(())
1335 }
1336}
1337
1338#[derive(Debug, Clone, Copy, PartialEq)]
1339struct ColumnRef<'a>(&'a str);
1340
1341impl<'a> From<&'a str> for ColumnRef<'a> {
1342 fn from(s: &'a str) -> Self {
1343 Self(s)
1344 }
1345}
1346
1347trait AsColumnRef {
1348 fn as_ref(&self) -> ColumnRef<'_>;
1349}
1350
1351impl AsColumnRef for Column {
1352 fn as_ref(&self) -> ColumnRef<'_> {
1353 self.0.as_str().into()
1354 }
1355}
1356
1357impl AsColumnRef for ColumnRef<'_> {
1358 fn as_ref(&self) -> ColumnRef<'_> {
1359 *self
1360 }
1361}
1362
1363impl<'a> PartialEq for dyn AsColumnRef + 'a {
1364 fn eq(&self, other: &Self) -> bool {
1365 self.as_ref() == other.as_ref()
1366 }
1367}
1368
1369impl<'a> Eq for dyn AsColumnRef + 'a {}
1370
1371impl<'a> Hash for dyn AsColumnRef + 'a {
1372 fn hash<H: Hasher>(&self, state: &mut H) {
1373 self.as_ref().0.hash(state);
1374 }
1375}
1376
1377impl<'a> Borrow<dyn AsColumnRef + 'a> for Column {
1378 fn borrow(&self) -> &(dyn AsColumnRef + 'a) {
1379 self
1380 }
1381}
1382
1383fn record_batches_to_labels_name(
1385 batches: RecordBatches,
1386 labels: &mut HashSet<String>,
1387) -> Result<()> {
1388 let mut column_indices = Vec::new();
1389 let mut value_column_indices = Vec::new();
1390 for (i, column) in batches.schema().column_schemas().iter().enumerate() {
1391 if is_prometheus_value_column(&column.data_type) {
1392 value_column_indices.push(i);
1393 }
1394 column_indices.push(i);
1395 }
1396
1397 if value_column_indices.is_empty() {
1398 return Err(Error::Internal {
1399 err_msg: "no value column found".to_string(),
1400 });
1401 }
1402
1403 for batch in batches.iter() {
1404 let names = column_indices
1405 .iter()
1406 .map(|c| batches.schema().column_name_by_index(*c).to_string())
1407 .collect::<Vec<_>>();
1408
1409 let value_columns = value_column_indices
1410 .iter()
1411 .map(|i| batch.column(*i))
1412 .collect::<Vec<_>>();
1413
1414 for row_index in 0..batch.num_rows() {
1415 if value_columns.iter().all(|c| c.is_null(row_index)) {
1417 continue;
1418 }
1419
1420 names.iter().for_each(|name| {
1422 let _ = labels.insert(name.clone());
1423 });
1424 return Ok(());
1425 }
1426 }
1427 Ok(())
1428}
1429
1430fn is_prometheus_value_column(data_type: &ConcreteDataType) -> bool {
1431 matches!(data_type, ConcreteDataType::Float64(_)) || is_native_histogram_value_type(data_type)
1432}
1433
1434pub(crate) fn retrieve_metric_name_and_result_type(
1435 promql_expr: &PromqlExpr,
1436) -> (Option<String>, ValueType) {
1437 (
1438 promql_expr_to_metric_name(promql_expr),
1439 promql_expr.value_type(),
1440 )
1441}
1442
1443fn current_schema_metric_targets(
1444 query_ctx: &QueryContext,
1445 metric_names: &[String],
1446) -> PermissionTableTargets {
1447 PermissionTableTargets::resolved(
1448 metric_names
1449 .iter()
1450 .map(|metric| {
1451 PermissionTableTarget::new(
1452 query_ctx.current_catalog(),
1453 query_ctx.current_schema(),
1454 metric,
1455 )
1456 })
1457 .collect(),
1458 )
1459}
1460
1461fn is_internal_physical_metric_table(table: &TableRef) -> bool {
1462 table.table_info().is_physical_table()
1463}
1464
1465fn promql_expr_to_metric_name(expr: &PromqlExpr) -> Option<String> {
1466 let mut metric_names = HashSet::new();
1467 collect_metric_names(expr, &mut metric_names);
1468
1469 if metric_names.len() == 1 {
1471 metric_names.into_iter().next()
1472 } else {
1473 None
1474 }
1475}
1476
1477fn binary_keeps_metric_name(op: &TokenType, modifier: Option<&BinModifier>) -> bool {
1487 if op.is_set_operator() {
1488 return true;
1489 }
1490
1491 if !op.is_comparison_operator() {
1492 return false;
1493 }
1494
1495 let Some(m) = modifier else {
1496 return true;
1497 };
1498
1499 if m.return_bool {
1500 return false;
1501 }
1502
1503 if m.card == VectorMatchCardinality::OneToOne {
1504 match &m.matching {
1505 Some(LabelModifier::Include(labels))
1506 if !labels.labels.contains(&METRIC_NAME.to_string()) =>
1507 {
1508 return false;
1509 }
1510 Some(LabelModifier::Exclude(labels))
1511 if labels.labels.contains(&METRIC_NAME.to_string()) =>
1512 {
1513 return false;
1514 }
1515 _ => {}
1516 }
1517 }
1518
1519 true
1520}
1521
1522fn binary_metric_name_side<'a>(
1530 lhs: &'a PromqlExpr,
1531 rhs: &'a PromqlExpr,
1532 modifier: Option<&BinModifier>,
1533) -> &'a PromqlExpr {
1534 let group_right =
1535 modifier.is_some_and(|m| matches!(m.card, VectorMatchCardinality::OneToMany(_)));
1536 if group_right || matches!(lhs.value_type(), ValueType::Scalar) {
1537 rhs
1538 } else {
1539 lhs
1540 }
1541}
1542
1543fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet<String>) {
1545 match expr {
1546 PromqlExpr::Aggregate(AggregateExpr { modifier, expr, .. }) => {
1547 match modifier {
1548 Some(LabelModifier::Include(labels))
1549 if !labels.labels.contains(&METRIC_NAME.to_string()) =>
1550 {
1551 metric_names.clear();
1552 return;
1553 }
1554 Some(LabelModifier::Exclude(labels))
1555 if labels.labels.contains(&METRIC_NAME.to_string()) =>
1556 {
1557 metric_names.clear();
1558 return;
1559 }
1560 _ => {}
1561 }
1562 collect_metric_names(expr, metric_names)
1563 }
1564 PromqlExpr::Unary(UnaryExpr { .. }) => metric_names.clear(),
1565 PromqlExpr::Binary(BinaryExpr {
1566 lhs,
1567 op,
1568 rhs,
1569 modifier,
1570 ..
1571 }) => {
1572 if binary_keeps_metric_name(op, modifier.as_ref()) {
1573 let side = binary_metric_name_side(lhs, rhs, modifier.as_ref());
1574 collect_metric_names(side, metric_names)
1575 } else {
1576 metric_names.clear()
1577 }
1578 }
1579 PromqlExpr::Paren(ParenExpr { expr }) => collect_metric_names(expr, metric_names),
1580 PromqlExpr::Subquery(SubqueryExpr { expr, .. }) => collect_metric_names(expr, metric_names),
1581 PromqlExpr::VectorSelector(VectorSelector { name, matchers, .. }) => {
1582 if let Some(name) = name {
1583 metric_names.insert(name.clone());
1584 } else if let Some(matcher) = matchers.find_matchers(METRIC_NAME).into_iter().next() {
1585 metric_names.insert(matcher.value);
1586 }
1587 }
1588 PromqlExpr::MatrixSelector(MatrixSelector { vs, .. }) => {
1589 let VectorSelector { name, matchers, .. } = vs;
1590 if let Some(name) = name {
1591 metric_names.insert(name.clone());
1592 } else if let Some(matcher) = matchers.find_matchers(METRIC_NAME).into_iter().next() {
1593 metric_names.insert(matcher.value);
1594 }
1595 }
1596 PromqlExpr::Call(Call { args, .. }) => {
1597 args.args
1598 .iter()
1599 .for_each(|e| collect_metric_names(e, metric_names));
1600 }
1601 PromqlExpr::NumberLiteral(_) | PromqlExpr::StringLiteral(_) | PromqlExpr::Extension(_) => {}
1602 }
1603}
1604
1605fn find_metric_name_and_matchers<E, F>(expr: &PromqlExpr, f: F) -> Option<E>
1606where
1607 F: Fn(&Option<String>, &Matchers) -> Option<E> + Clone,
1608{
1609 match expr {
1610 PromqlExpr::Aggregate(AggregateExpr { expr, .. }) => find_metric_name_and_matchers(expr, f),
1611 PromqlExpr::Unary(UnaryExpr { expr }) => find_metric_name_and_matchers(expr, f),
1612 PromqlExpr::Binary(BinaryExpr { lhs, rhs, .. }) => {
1613 find_metric_name_and_matchers(lhs, f.clone()).or(find_metric_name_and_matchers(rhs, f))
1614 }
1615 PromqlExpr::Paren(ParenExpr { expr }) => find_metric_name_and_matchers(expr, f),
1616 PromqlExpr::Subquery(SubqueryExpr { expr, .. }) => find_metric_name_and_matchers(expr, f),
1617 PromqlExpr::NumberLiteral(_) => None,
1618 PromqlExpr::StringLiteral(_) => None,
1619 PromqlExpr::Extension(_) => None,
1620 PromqlExpr::VectorSelector(VectorSelector { name, matchers, .. }) => f(name, matchers),
1621 PromqlExpr::MatrixSelector(MatrixSelector { vs, .. }) => {
1622 let VectorSelector { name, matchers, .. } = vs;
1623
1624 f(name, matchers)
1625 }
1626 PromqlExpr::Call(Call { args, .. }) => args
1627 .args
1628 .iter()
1629 .find_map(|e| find_metric_name_and_matchers(e, f.clone())),
1630 }
1631}
1632
1633struct MetricNameDiscovery {
1634 name_matchers: Vec<Matcher>,
1635 schema: Option<String>,
1636}
1637
1638fn find_metric_name_not_equal_matchers(expr: &PromqlExpr) -> Result<Option<MetricNameDiscovery>> {
1640 if find_metric_name_and_matchers(expr, |name, matchers| {
1641 (name.is_none() && !matchers.or_matchers.is_empty()).then_some(())
1642 })
1643 .is_some()
1644 {
1645 return Ok(None);
1646 }
1647
1648 find_metric_name_and_matchers(expr, |name, matchers| {
1649 if name.is_some() {
1651 return None;
1652 }
1653
1654 let name_matchers = matchers.find_matchers(METRIC_NAME);
1655 if name_matchers.len() != 1 || name_matchers[0].op == MatchOp::Equal {
1656 return None;
1657 }
1658
1659 Some(
1660 resolve_schema_from_matchers(&matchers.matchers).map(|schema| MetricNameDiscovery {
1661 name_matchers,
1662 schema,
1663 }),
1664 )
1665 })
1666 .transpose()
1667}
1668
1669fn expand_metric_name_queries(
1670 query: &ParsedPromQuery,
1671 metric_names: Vec<String>,
1672) -> Vec<ParsedPromQuery> {
1673 metric_names
1674 .into_iter()
1675 .map(|metric| {
1676 let mut query = query.clone();
1677 query.update_expr(|expr| update_metric_name_matcher(expr, &metric));
1678 query
1679 })
1680 .collect()
1681}
1682
1683fn static_promql_targets(
1684 expr: &PromqlExpr,
1685 query_ctx: &QueryContextRef,
1686) -> Result<(PermissionTableTargets, usize)> {
1687 let mut targets = Vec::new();
1688 let mut unresolved_selectors = 0;
1689 collect_static_promql_targets(expr, query_ctx, &mut targets, &mut unresolved_selectors)?;
1690 Ok((
1691 PermissionTableTargets::resolved(targets),
1692 unresolved_selectors,
1693 ))
1694}
1695
1696fn resolved_promql_targets(
1697 queries: &[ParsedPromQuery],
1698 query_ctx: &QueryContextRef,
1699) -> Option<Vec<PermissionTableTarget>> {
1700 let mut targets = Vec::new();
1701 for query in queries {
1702 let (PermissionTableTargets::Resolved(query_targets), unresolved_selectors) =
1703 static_promql_targets(query.expr(), query_ctx).ok()?
1704 else {
1705 return None;
1706 };
1707 if unresolved_selectors != 0 {
1708 return None;
1709 }
1710 for target in query_targets {
1711 if !targets.contains(&target) {
1712 targets.push(target);
1713 }
1714 }
1715 }
1716 Some(targets)
1717}
1718
1719fn collect_static_promql_targets(
1720 expr: &PromqlExpr,
1721 query_ctx: &QueryContextRef,
1722 targets: &mut Vec<PermissionTableTarget>,
1723 unresolved_selectors: &mut usize,
1724) -> Result<()> {
1725 match expr {
1726 PromqlExpr::Aggregate(AggregateExpr { expr, .. })
1727 | PromqlExpr::Unary(UnaryExpr { expr })
1728 | PromqlExpr::Paren(ParenExpr { expr })
1729 | PromqlExpr::Subquery(SubqueryExpr { expr, .. }) => {
1730 collect_static_promql_targets(expr, query_ctx, targets, unresolved_selectors)?
1731 }
1732 PromqlExpr::Binary(BinaryExpr { lhs, rhs, .. }) => {
1733 collect_static_promql_targets(lhs, query_ctx, targets, unresolved_selectors)?;
1734 collect_static_promql_targets(rhs, query_ctx, targets, unresolved_selectors)?;
1735 }
1736 PromqlExpr::VectorSelector(selector) => {
1737 collect_static_vector_target(selector, query_ctx, targets, unresolved_selectors)?
1738 }
1739 PromqlExpr::MatrixSelector(MatrixSelector { vs, .. }) => {
1740 collect_static_vector_target(vs, query_ctx, targets, unresolved_selectors)?
1741 }
1742 PromqlExpr::Call(Call { args, .. }) => {
1743 for expr in &args.args {
1744 collect_static_promql_targets(expr, query_ctx, targets, unresolved_selectors)?;
1745 }
1746 }
1747 PromqlExpr::NumberLiteral(_) | PromqlExpr::StringLiteral(_) | PromqlExpr::Extension(_) => {}
1748 }
1749 Ok(())
1750}
1751
1752fn collect_static_vector_target(
1753 selector: &VectorSelector,
1754 query_ctx: &QueryContextRef,
1755 targets: &mut Vec<PermissionTableTarget>,
1756 unresolved_selectors: &mut usize,
1757) -> Result<()> {
1758 if selector.name.is_none() && !selector.matchers.or_matchers.is_empty() {
1759 *unresolved_selectors += 1;
1760 return Ok(());
1761 }
1762
1763 let metric = selector.name.clone().or_else(|| {
1764 let mut matchers = selector.matchers.find_matchers(METRIC_NAME);
1765 if matchers.len() == 1 && matchers[0].op == MatchOp::Equal {
1766 matchers.pop().map(|matcher| matcher.value)
1767 } else {
1768 None
1769 }
1770 });
1771 let Some(metric) = metric else {
1772 *unresolved_selectors += 1;
1773 return Ok(());
1774 };
1775
1776 let schema = resolve_schema_from_matchers(&selector.matchers.matchers)?
1777 .unwrap_or_else(|| query_ctx.current_schema());
1778 targets.push(PermissionTableTarget::new(
1779 query_ctx.current_catalog(),
1780 schema,
1781 metric,
1782 ));
1783 Ok(())
1784}
1785
1786fn update_metric_name_matcher(expr: &mut PromqlExpr, metric_name: &str) {
1789 match expr {
1790 PromqlExpr::Aggregate(AggregateExpr { expr, .. }) => {
1791 update_metric_name_matcher(expr, metric_name)
1792 }
1793 PromqlExpr::Unary(UnaryExpr { expr }) => update_metric_name_matcher(expr, metric_name),
1794 PromqlExpr::Binary(BinaryExpr { lhs, rhs, .. }) => {
1795 update_metric_name_matcher(lhs, metric_name);
1796 update_metric_name_matcher(rhs, metric_name);
1797 }
1798 PromqlExpr::Paren(ParenExpr { expr }) => update_metric_name_matcher(expr, metric_name),
1799 PromqlExpr::Subquery(SubqueryExpr { expr, .. }) => {
1800 update_metric_name_matcher(expr, metric_name)
1801 }
1802 PromqlExpr::VectorSelector(VectorSelector { name, matchers, .. }) => {
1803 if name.is_some() {
1804 return;
1805 }
1806
1807 for m in &mut matchers.matchers {
1808 if m.name == METRIC_NAME && m.op != MatchOp::Equal {
1809 m.op = MatchOp::Equal;
1810 m.value = metric_name.to_string();
1811 }
1812 }
1813 }
1814 PromqlExpr::MatrixSelector(MatrixSelector { vs, .. }) => {
1815 let VectorSelector { name, matchers, .. } = vs;
1816 if name.is_some() {
1817 return;
1818 }
1819
1820 for m in &mut matchers.matchers {
1821 if m.name == METRIC_NAME && m.op != MatchOp::Equal {
1822 m.op = MatchOp::Equal;
1823 m.value = metric_name.to_string();
1824 }
1825 }
1826 }
1827 PromqlExpr::Call(Call { args, .. }) => {
1828 args.args.iter_mut().for_each(|e| {
1829 update_metric_name_matcher(e, metric_name);
1830 });
1831 }
1832 PromqlExpr::NumberLiteral(_) | PromqlExpr::StringLiteral(_) | PromqlExpr::Extension(_) => {}
1833 }
1834}
1835
1836#[derive(Debug, Default, Serialize, Deserialize)]
1837pub struct LabelValueQuery {
1838 start: Option<String>,
1839 end: Option<String>,
1840 lookback: Option<String>,
1841 #[serde(flatten)]
1842 matches: Matches,
1843 db: Option<String>,
1844 limit: Option<usize>,
1845}
1846
1847#[axum_macros::debug_handler]
1848#[tracing::instrument(
1849 skip_all,
1850 fields(protocol = "prometheus", request_type = "label_values_query")
1851)]
1852pub async fn label_values_query(
1853 user_provider: Option<Extension<UserProviderRef>>,
1854 State(handler): State<PrometheusHandlerRef>,
1855 Path(label_name): Path<String>,
1856 Extension(mut query_ctx): Extension<QueryContext>,
1857 Query(params): Query<LabelValueQuery>,
1858) -> PrometheusJsonResponse {
1859 if let Err(error) = select_database(
1860 &mut query_ctx,
1861 params.db.as_deref(),
1862 user_provider.as_ref().map(|Extension(provider)| provider),
1863 )
1864 .await
1865 {
1866 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
1867 }
1868 let query_ctx = Arc::new(query_ctx);
1869
1870 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
1871 .with_label_values(&[query_ctx.get_db_string().as_str(), "label_values_query"])
1872 .start_timer();
1873
1874 let matches = params.matches.0;
1875 if label_name == METRIC_NAME_LABEL {
1876 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
1877 let exact_metric_names = matches
1878 .iter()
1879 .filter_map(|selector| retrieve_exact_metric_name_from_promql(selector))
1880 .collect::<Vec<_>>();
1881 try_call_return_response!(
1882 handler
1883 .check_query_target_permission(
1884 current_schema_metric_targets(&query_ctx, &exact_metric_names),
1885 &query_ctx,
1886 )
1887 .await
1888 );
1889 let catalog_manager = handler.catalog_manager();
1890
1891 let enumerate_all = matches.is_empty();
1894 let (metadata_selectors, label_selectors) =
1895 try_call_return_response!(split_selectors_by_label_use(&matches));
1896
1897 let mut table_names = if enumerate_all || !metadata_selectors.is_empty() {
1898 try_call_return_response!(
1899 retrieve_table_names(&query_ctx, catalog_manager, metadata_selectors).await
1900 )
1901 } else {
1902 Vec::new()
1903 };
1904
1905 if !label_selectors.is_empty() {
1906 table_names.extend(try_call_return_response!(
1907 retrieve_table_names_by_labels(
1908 &handler,
1909 label_selectors,
1910 params.start.as_deref(),
1911 params.end.as_deref(),
1912 &query_ctx,
1913 )
1914 .await
1915 ));
1916 table_names.sort_unstable();
1917 table_names.dedup();
1918 }
1919
1920 table_names = try_call_return_response!(
1921 handler
1922 .filter_metadata_metric_names(
1923 table_names,
1924 query_ctx.current_schema().as_str(),
1925 &query_ctx,
1926 )
1927 .await
1928 );
1929
1930 truncate_results(&mut table_names, params.limit);
1931 return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(table_names));
1932 } else if label_name == FIELD_NAME_LABEL {
1933 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
1934 let field_columns =
1935 handle_schema_err!(retrieve_field_names(&query_ctx, &handler, matches).await)
1936 .unwrap_or_default();
1937 let mut field_columns = field_columns.into_iter().collect::<Vec<_>>();
1938 field_columns.sort_unstable();
1939 truncate_results(&mut field_columns, params.limit);
1940 return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(field_columns));
1941 } else if is_database_selection_label(&label_name) {
1942 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
1943 let mut schema_names =
1944 try_call_return_response!(retrieve_schema_names(&query_ctx, &handler, matches).await);
1945 truncate_results(&mut schema_names, params.limit);
1946 return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(schema_names));
1947 }
1948
1949 let queries = matches;
1950 if queries.is_empty() {
1951 return PrometheusJsonResponse::error(
1952 StatusCode::InvalidArguments,
1953 "match[] parameter is required",
1954 );
1955 }
1956
1957 let start = params.start.unwrap_or_else(yesterday_rfc3339);
1958 let end = params.end.unwrap_or_else(current_time_rfc3339);
1959 let lookback = params
1960 .lookback
1961 .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string());
1962 let prom_queries = queries
1963 .into_iter()
1964 .map(|query| PromQuery {
1965 query,
1966 start: start.clone(),
1967 end: end.clone(),
1968 step: DEFAULT_LOOKBACK_STRING.to_string(),
1969 lookback: lookback.clone(),
1970 alias: None,
1971 })
1972 .collect::<Vec<_>>();
1973 let prom_queries = try_call_return_response!(
1974 prom_queries
1975 .into_iter()
1976 .map(|query| ParsedPromQuery::parse(query, &query_ctx))
1977 .collect::<Result<Vec<_>>>()
1978 );
1979 try_call_return_response!(
1980 handler
1981 .check_query_permission_parsed(&prom_queries, &query_ctx)
1982 .await
1983 );
1984
1985 let mut label_values = HashSet::new();
1986
1987 let Some(first_query) = prom_queries.first() else {
1988 unreachable!("empty match[] is rejected above")
1989 };
1990 let QueryStatement::Promql(eval_stmt, _) = first_query.statement() else {
1991 unreachable!("query is parsed from PromQL")
1992 };
1993 let start = eval_stmt.start;
1994 let end = eval_stmt.end;
1995
1996 for prom_query in prom_queries {
1997 let (_, statement) = prom_query.into_parts();
1998 let QueryStatement::Promql(eval_stmt, _) = statement else {
1999 unreachable!("query is parsed from PromQL")
2000 };
2001 let PromqlExpr::VectorSelector(mut vector_selector) = eval_stmt.expr else {
2002 return PrometheusJsonResponse::error(
2003 StatusCode::InvalidArguments,
2004 "expected vector selector",
2005 );
2006 };
2007 let Some(name) = take_metric_name(&mut vector_selector) else {
2008 return PrometheusJsonResponse::error(
2009 StatusCode::InvalidArguments,
2010 "expected metric name",
2011 );
2012 };
2013 let VectorSelector { matchers, .. } = vector_selector;
2014 let matchers = matchers.matchers;
2016 let result = handler
2017 .query_label_values(name, label_name.clone(), matchers, start, end, &query_ctx)
2018 .await;
2019 if let Some(result) = handle_schema_err!(result) {
2020 label_values.extend(result.into_iter());
2021 }
2022 }
2023
2024 let mut label_values: Vec<_> = label_values.into_iter().collect();
2025 label_values.sort_unstable();
2026 truncate_results(&mut label_values, params.limit);
2027
2028 PrometheusJsonResponse::success(PrometheusResponse::LabelValues(label_values))
2029}
2030
2031fn truncate_results(label_values: &mut Vec<String>, limit: Option<usize>) {
2032 if let Some(limit) = limit
2033 && limit > 0
2034 && label_values.len() >= limit
2035 {
2036 label_values.truncate(limit);
2037 }
2038}
2039
2040fn take_metric_name(selector: &mut VectorSelector) -> Option<String> {
2043 if let Some(name) = selector.name.take() {
2044 return Some(name);
2045 }
2046
2047 let (pos, matcher) = selector
2048 .matchers
2049 .matchers
2050 .iter()
2051 .find_position(|matcher| matcher.name == "__name__" && matcher.op == MatchOp::Equal)?;
2052 let name = matcher.value.clone();
2053 selector.matchers.matchers.remove(pos);
2055
2056 Some(name)
2057}
2058
2059fn take_metric_name_matchers(selector: &mut VectorSelector) -> Vec<Matcher> {
2063 let mut taken = Vec::new();
2064 if let Some(name) = selector.name.take() {
2065 taken.push(Matcher::new(MatchOp::Equal, METRIC_NAME_LABEL, &name));
2066 }
2067
2068 let (name_matchers, rest) = std::mem::take(&mut selector.matchers.matchers)
2069 .into_iter()
2070 .partition(|matcher| matcher.name == METRIC_NAME_LABEL);
2071 selector.matchers.matchers = rest;
2072 taken.extend(name_matchers);
2073
2074 taken
2075}
2076
2077fn metric_name_matches(table_name: &str, matchers: &[Matcher]) -> bool {
2084 matchers.iter().all(|matcher| match &matcher.op {
2085 MatchOp::Equal => table_name == matcher.value,
2086 MatchOp::NotEqual => table_name != matcher.value,
2087 MatchOp::Re(re) => re.is_match(table_name),
2088 MatchOp::NotRe(re) => !re.is_match(table_name),
2089 })
2090}
2091
2092fn is_ordinary_label_matcher(matcher: &Matcher) -> bool {
2095 matcher.name != METRIC_NAME_LABEL
2096 && matcher.name != FIELD_NAME_LABEL
2097 && !is_database_selection_label(&matcher.name)
2098}
2099
2100fn split_selectors_by_label_use(matches: &[String]) -> Result<(Vec<String>, Vec<VectorSelector>)> {
2107 let mut metadata_only = Vec::new();
2108 let mut with_labels = Vec::new();
2109
2110 for selector in matches {
2111 let expr = promql_parser::parser::parse(selector)
2112 .map_err(|reason| InvalidQuerySnafu { reason }.build())?;
2113 let PromqlExpr::VectorSelector(vector_selector) = expr else {
2114 return InvalidQuerySnafu {
2115 reason: "expected vector selector".to_string(),
2116 }
2117 .fail();
2118 };
2119
2120 let constrains_labels = vector_selector.matchers.or_matchers.is_empty()
2121 && vector_selector
2122 .matchers
2123 .matchers
2124 .iter()
2125 .any(is_ordinary_label_matcher);
2126 if constrains_labels {
2127 with_labels.push(vector_selector);
2128 } else {
2129 metadata_only.push(selector.clone());
2130 }
2131 }
2132
2133 Ok((metadata_only, with_labels))
2134}
2135
2136async fn retrieve_table_names_by_labels(
2140 handler: &PrometheusHandlerRef,
2141 selectors: Vec<VectorSelector>,
2142 start: Option<&str>,
2143 end: Option<&str>,
2144 query_ctx: &QueryContextRef,
2145) -> Result<Vec<String>> {
2146 let start_arg = start.map(str::to_string).unwrap_or_else(yesterday_rfc3339);
2147 let end_arg = end.map(str::to_string).unwrap_or_else(current_time_rfc3339);
2148 let start = QueryLanguageParser::parse_promql_timestamp(&start_arg).with_context(|_| {
2149 ParseTimestampSnafu {
2150 timestamp: start_arg.clone(),
2151 }
2152 })?;
2153 let end = QueryLanguageParser::parse_promql_timestamp(&end_arg).with_context(|_| {
2154 ParseTimestampSnafu {
2155 timestamp: end_arg.clone(),
2156 }
2157 })?;
2158
2159 let schema = query_ctx.current_schema();
2160 let mut table_names = Vec::new();
2161 for mut selector in selectors {
2162 let name_matchers = take_metric_name_matchers(&mut selector);
2163 let label_matchers = selector
2166 .matchers
2167 .matchers
2168 .into_iter()
2169 .filter(is_ordinary_label_matcher)
2170 .collect();
2171 let matched = handler
2172 .query_metric_names_by_labels(label_matchers, &schema, start, end, query_ctx)
2173 .await?;
2174 table_names.extend(
2175 matched
2176 .into_iter()
2177 .filter(|name| metric_name_matches(name, &name_matchers)),
2178 );
2179 }
2180
2181 Ok(table_names)
2182}
2183
2184async fn retrieve_table_names(
2185 query_ctx: &QueryContext,
2186 catalog_manager: CatalogManagerRef,
2187 matches: Vec<String>,
2188) -> Result<Vec<String>> {
2189 let catalog = query_ctx.current_catalog();
2190 let schema = query_ctx.current_schema();
2191
2192 let mut tables_stream = catalog_manager.tables(catalog, &schema, Some(query_ctx));
2193 let mut table_names = Vec::new();
2194
2195 let name_matchers = matches
2196 .iter()
2197 .map(|selector| {
2198 let expr = promql_parser::parser::parse(selector)
2199 .map_err(|reason| InvalidQuerySnafu { reason }.build())?;
2200 let PromqlExpr::VectorSelector(selector) = expr else {
2201 return InvalidQuerySnafu {
2202 reason: "expected vector selector".to_string(),
2203 }
2204 .fail();
2205 };
2206 if !selector.matchers.or_matchers.is_empty() {
2207 return Ok(None);
2208 }
2209 if let Some(name) = selector.name {
2210 return Ok(Some(Matcher::new(MatchOp::Equal, METRIC_NAME_LABEL, &name)));
2211 }
2212 Ok(selector
2213 .matchers
2214 .matchers
2215 .into_iter()
2216 .find(|matcher| matcher.name == METRIC_NAME_LABEL))
2217 })
2218 .collect::<Result<Vec<_>>>()?;
2219
2220 while let Some(table) = tables_stream.next().await {
2221 let table = table?;
2222 if !table
2223 .table_info()
2224 .meta
2225 .options
2226 .extra_options
2227 .contains_key(LOGICAL_TABLE_METADATA_KEY)
2228 || is_internal_physical_metric_table(&table)
2229 {
2230 continue;
2232 }
2233
2234 let table_name = &table.table_info().name;
2235
2236 if name_matchers.is_empty()
2237 || name_matchers.iter().any(|matcher| match matcher {
2238 None => true,
2239 Some(matcher) => match &matcher.op {
2240 MatchOp::Equal => table_name == &matcher.value,
2241 MatchOp::Re(reg) => reg.is_match(table_name),
2242 _ => true,
2245 },
2246 })
2247 {
2248 table_names.push(table_name.clone());
2249 }
2250 }
2251
2252 table_names.sort_unstable();
2253 Ok(table_names)
2254}
2255
2256async fn retrieve_metric_metadata(
2257 query_ctx: &QueryContext,
2258 manager: CatalogManagerRef,
2259 params: &MetadataQuery,
2260) -> Result<BTreeMap<String, Vec<PromMetadata>>> {
2261 let mut metadata = BTreeMap::new();
2262 if params.limit == Some(0) {
2263 return Ok(metadata);
2264 }
2265
2266 let catalog = query_ctx.current_catalog();
2267 let schema = query_ctx.current_schema();
2268
2269 if let Some(metric) = ¶ms.metric {
2270 let Some(table) = manager
2271 .table(catalog, &schema, metric, Some(query_ctx))
2272 .await?
2273 else {
2274 return Ok(metadata);
2275 };
2276 let table_info = table.table_info();
2277 if is_prometheus_metric_table(table_info.as_ref()) {
2278 metadata.insert(
2279 table_info.name.clone(),
2280 vec![prometheus_metadata_from_table(table_info.as_ref())],
2281 );
2282 }
2283 return Ok(metadata);
2284 }
2285
2286 let mut tables_stream = manager.tables(catalog, &schema, Some(query_ctx));
2287
2288 while let Some(table) = tables_stream.next().await {
2289 let table = table?;
2290 let table_info = table.table_info();
2291 if !is_prometheus_metric_table(table_info.as_ref()) {
2292 continue;
2293 }
2294
2295 metadata.insert(
2296 table_info.name.clone(),
2297 vec![prometheus_metadata_from_table(table_info.as_ref())],
2298 );
2299 }
2300
2301 Ok(metadata)
2302}
2303
2304fn is_prometheus_metric_table(table_info: &TableInfo) -> bool {
2305 table_info
2306 .meta
2307 .options
2308 .extra_options
2309 .contains_key(LOGICAL_TABLE_METADATA_KEY)
2310}
2311
2312fn prometheus_metadata_from_table(table_info: &TableInfo) -> PromMetadata {
2313 let options = &table_info.meta.options.extra_options;
2314 let metric_type = match options.get(SEMANTIC_METRIC_TYPE) {
2315 Some(metric_type)
2316 if options
2317 .get(SEMANTIC_METRIC_TEMPORALITY)
2318 .is_some_and(|temporality| {
2319 matches!(
2320 temporality.as_str(),
2321 METRIC_TEMPORALITY_DELTA | SEMANTIC_VALUE_MIXED
2322 )
2323 })
2324 && matches!(
2325 metric_type.as_str(),
2326 "counter" | "histogram" | "updown_counter"
2327 ) =>
2328 {
2329 "unknown".to_string()
2330 }
2331 Some(metric_type) => match metric_type.as_str() {
2332 "updown_counter" => "gauge".to_string(),
2333 "gauge_histogram" => "gaugehistogram".to_string(),
2334 "mixed" => "unknown".to_string(),
2335 metric_type => metric_type.to_string(),
2336 },
2337 None if table_has_native_histogram_value(table_info) => "histogram".to_string(),
2338 None => String::new(),
2339 };
2340 let unit = options
2341 .get(SEMANTIC_METRIC_UNIT)
2342 .map(|unit| ucum_to_openmetrics_unit(unit))
2343 .unwrap_or_default();
2344
2345 PromMetadata {
2346 metric_type,
2347 unit,
2348 help: String::new(),
2350 }
2351}
2352
2353fn table_has_native_histogram_value(table_info: &TableInfo) -> bool {
2354 table_info
2355 .meta
2356 .schema
2357 .column_schemas()
2358 .iter()
2359 .any(|column| is_native_histogram_value_type(&column.data_type))
2360}
2361
2362fn metadata_name_selectors(matches: &[String]) -> Result<Vec<Vec<Matcher>>> {
2363 matches
2364 .iter()
2365 .map(|query| {
2366 let expr = promql_parser::parser::parse(query)
2367 .map_err(|reason| InvalidQuerySnafu { reason }.build())?;
2368 let PromqlExpr::VectorSelector(mut selector) = expr else {
2369 return InvalidQuerySnafu {
2370 reason: "expected vector selector".to_string(),
2371 }
2372 .fail();
2373 };
2374 ensure!(
2375 selector.matchers.or_matchers.is_empty(),
2376 InvalidQuerySnafu {
2377 reason: "OR matchers are not supported in schema and field metadata queries"
2378 .to_string(),
2379 }
2380 );
2381 Ok(take_metric_name_matchers(&mut selector))
2382 })
2383 .collect()
2384}
2385
2386async fn readable_metadata_tables(
2387 query_ctx: &QueryContextRef,
2388 handler: &PrometheusHandlerRef,
2389 schema: &str,
2390 selectors: &[Vec<Matcher>],
2391) -> Result<Vec<TableRef>> {
2392 let manager = handler.catalog_manager();
2393 let exact_names = selectors
2394 .iter()
2395 .map(|matchers| {
2396 matchers
2397 .iter()
2398 .find(|matcher| matcher.op == MatchOp::Equal)
2399 .map(|matcher| &matcher.value)
2400 })
2401 .collect::<Option<HashSet<_>>>();
2402 let mut candidates = Vec::new();
2403 if let Some(names) = exact_names.filter(|names| !names.is_empty()) {
2404 for name in names {
2405 if let Some(table) = manager
2406 .table(query_ctx.current_catalog(), schema, name, Some(query_ctx))
2407 .await?
2408 {
2409 candidates.push(table);
2410 }
2411 }
2412 } else {
2413 let mut tables = manager.tables(query_ctx.current_catalog(), schema, Some(query_ctx));
2414 while let Some(table) = tables.next().await {
2415 candidates.push(table?);
2416 }
2417 }
2418 candidates.retain(|table| {
2419 !is_internal_physical_metric_table(table)
2420 && (selectors.is_empty()
2421 || selectors
2422 .iter()
2423 .any(|matchers| metric_name_matches(&table.table_info().name, matchers)))
2424 });
2425 let names = candidates
2426 .iter()
2427 .map(|table| table.table_info().name.clone())
2428 .collect();
2429 let allowed: HashSet<_> = handler
2430 .filter_metadata_metric_names(names, schema, query_ctx)
2431 .await?
2432 .into_iter()
2433 .collect();
2434 candidates.retain(|table| allowed.contains(&table.table_info().name));
2435 Ok(candidates)
2436}
2437
2438async fn retrieve_field_names(
2439 query_ctx: &QueryContextRef,
2440 handler: &PrometheusHandlerRef,
2441 matches: Vec<String>,
2442) -> Result<HashSet<String>> {
2443 let mut selectors = metadata_name_selectors(&matches)?;
2444 selectors.retain(|matchers| !matchers.is_empty());
2446 if !matches.is_empty() && selectors.is_empty() {
2447 return Ok(HashSet::new());
2448 }
2449 let tables =
2450 readable_metadata_tables(query_ctx, handler, &query_ctx.current_schema(), &selectors)
2451 .await?;
2452 Ok(tables
2453 .iter()
2454 .flat_map(|table| table.field_columns().map(|column| column.name))
2455 .collect())
2456}
2457
2458async fn retrieve_schema_names(
2459 query_ctx: &QueryContextRef,
2460 handler: &PrometheusHandlerRef,
2461 matches: Vec<String>,
2462) -> Result<Vec<String>> {
2463 let selectors = metadata_name_selectors(&matches)?;
2464 let candidate_schemas = handler
2465 .catalog_manager()
2466 .schema_names(query_ctx.current_catalog(), Some(query_ctx))
2467 .await?;
2468 let candidate_schemas = handler
2469 .filter_metadata_schema_names(candidate_schemas, query_ctx)
2470 .await?;
2471 let mut schemas = Vec::new();
2472 for schema in candidate_schemas {
2473 if selectors.is_empty() {
2474 schemas.push(schema);
2475 continue;
2476 }
2477 let tables = readable_metadata_tables(query_ctx, handler, &schema, &selectors).await?;
2478 if selectors.iter().all(|matchers| {
2479 tables
2480 .iter()
2481 .any(|table| metric_name_matches(&table.table_info().name, matchers))
2482 }) {
2483 schemas.push(schema);
2484 }
2485 }
2486 schemas.sort_unstable();
2487 Ok(schemas)
2488}
2489
2490fn retrieve_exact_metric_name_from_promql(query: &str) -> Option<String> {
2491 let PromqlExpr::VectorSelector(selector) = promql_parser::parser::parse(query).ok()? else {
2492 return None;
2493 };
2494 if let Some(name) = selector.name {
2495 return (!name.is_empty()).then_some(name);
2496 }
2497
2498 let mut name_matchers = selector.matchers.find_matchers(METRIC_NAME);
2499 if name_matchers.len() != 1 {
2500 return None;
2501 }
2502 let matcher = name_matchers.pop()?;
2503 (matcher.op == MatchOp::Equal && !matcher.value.is_empty()).then_some(matcher.value)
2504}
2505
2506#[derive(Debug, Default, Serialize, Deserialize)]
2507pub struct SeriesQuery {
2508 start: Option<String>,
2509 end: Option<String>,
2510 lookback: Option<String>,
2511 #[serde(flatten)]
2512 matches: Matches,
2513 db: Option<String>,
2514}
2515
2516#[axum_macros::debug_handler]
2517#[tracing::instrument(
2518 skip_all,
2519 fields(protocol = "prometheus", request_type = "series_query")
2520)]
2521pub async fn series_query(
2522 user_provider: Option<Extension<UserProviderRef>>,
2523 State(handler): State<PrometheusHandlerRef>,
2524 Query(params): Query<SeriesQuery>,
2525 Extension(mut query_ctx): Extension<QueryContext>,
2526 Form(form_params): Form<SeriesQuery>,
2527) -> PrometheusJsonResponse {
2528 let mut queries: Vec<String> = params.matches.0;
2529 if queries.is_empty() {
2530 queries = form_params.matches.0;
2531 }
2532 if queries.is_empty() {
2533 return PrometheusJsonResponse::error(
2534 StatusCode::Unsupported,
2535 "match[] parameter is required",
2536 );
2537 }
2538 let start = params
2539 .start
2540 .or(form_params.start)
2541 .unwrap_or_else(yesterday_rfc3339);
2542 let end = params
2543 .end
2544 .or(form_params.end)
2545 .unwrap_or_else(current_time_rfc3339);
2546 let lookback = params
2547 .lookback
2548 .or(form_params.lookback)
2549 .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string());
2550
2551 if let Err(error) = select_database(
2553 &mut query_ctx,
2554 params.db.as_deref(),
2555 user_provider.as_ref().map(|Extension(provider)| provider),
2556 )
2557 .await
2558 {
2559 return PrometheusJsonResponse::error(error.status_code(), error.output_msg());
2560 }
2561 let query_ctx = Arc::new(query_ctx);
2562
2563 let _timer = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
2564 .with_label_values(&[query_ctx.get_db_string().as_str(), "series_query"])
2565 .start_timer();
2566
2567 let prom_queries = queries
2568 .into_iter()
2569 .map(|query| PromQuery {
2570 query,
2571 start: start.clone(),
2572 end: end.clone(),
2573 step: DEFAULT_LOOKBACK_STRING.to_string(),
2575 lookback: lookback.clone(),
2576 alias: None,
2577 })
2578 .collect::<Vec<_>>();
2579 let mut expanded_queries = Vec::new();
2580 for prom_query in prom_queries {
2581 let prom_query = try_call_return_response!(
2582 @output_msg ParsedPromQuery::parse(prom_query, &query_ctx),
2583 StatusCode::InvalidArguments
2584 );
2585 let promql_expr = prom_query.expr();
2586 let Some(discovery) =
2587 try_call_return_response!(find_metric_name_not_equal_matchers(promql_expr))
2588 else {
2589 expanded_queries.push(prom_query);
2590 continue;
2591 };
2592
2593 try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await);
2594 let (static_targets, unresolved_selectors) =
2595 try_call_return_response!(static_promql_targets(promql_expr, &query_ctx));
2596 try_call_return_response!(
2597 handler
2598 .check_query_target_permission(static_targets, &query_ctx)
2599 .await
2600 );
2601 if unresolved_selectors != 1 {
2602 try_call_return_response!(
2603 handler
2604 .check_query_target_permission(PermissionTableTargets::Unresolved, &query_ctx)
2605 .await
2606 );
2607 expanded_queries.push(prom_query);
2608 continue;
2609 }
2610
2611 let schema = discovery
2612 .schema
2613 .unwrap_or_else(|| query_ctx.current_schema());
2614 let metric_names = try_call_return_response!(
2615 handler
2616 .query_metric_names(discovery.name_matchers, &schema, &query_ctx)
2617 .await
2618 );
2619 expanded_queries.extend(expand_metric_name_queries(&prom_query, metric_names));
2620 }
2621 let prom_queries = expanded_queries;
2622 try_call_return_response!(
2623 handler
2624 .check_query_permission_parsed(&prom_queries, &query_ctx)
2625 .await
2626 );
2627
2628 let mut series = Vec::new();
2629 let mut merge_map = HashMap::new();
2630 for prom_query in prom_queries {
2631 let Some(target) = resolved_promql_targets(std::slice::from_ref(&prom_query), &query_ctx)
2632 .and_then(|targets| match targets.as_slice() {
2633 [target] => Some(target.clone()),
2634 _ => None,
2635 })
2636 else {
2637 return PrometheusJsonResponse::error(
2638 StatusCode::InvalidArguments,
2639 "series selector must resolve exactly one metric table",
2640 );
2641 };
2642 let result = handler.do_query_parsed(prom_query, query_ctx.clone()).await;
2643
2644 handle_schema_err!(
2645 retrieve_series_from_query_result(
2646 result,
2647 &mut series,
2648 &query_ctx,
2649 &target,
2650 &handler.catalog_manager(),
2651 &mut merge_map,
2652 )
2653 .await
2654 );
2655 }
2656 let merge_map = merge_map
2657 .into_iter()
2658 .map(|(k, v)| (k, Value::from(v)))
2659 .collect();
2660 let mut resp = PrometheusJsonResponse::success(PrometheusResponse::Series(series));
2661 resp.resp_metrics = merge_map;
2662 resp
2663}
2664
2665#[derive(Debug, Default, Serialize, Deserialize)]
2666pub struct ParseQuery {
2667 query: Option<String>,
2668 db: Option<String>,
2669}
2670
2671#[axum_macros::debug_handler]
2672#[tracing::instrument(
2673 skip_all,
2674 fields(protocol = "prometheus", request_type = "parse_query")
2675)]
2676pub async fn parse_query(
2677 State(_handler): State<PrometheusHandlerRef>,
2678 Query(params): Query<ParseQuery>,
2679 Extension(_query_ctx): Extension<QueryContext>,
2680 Form(form_params): Form<ParseQuery>,
2681) -> PrometheusJsonResponse {
2682 if let Some(query) = params.query.or(form_params.query) {
2683 let ast = try_call_return_response!(
2684 promql_parser::parser::parse(&query),
2685 StatusCode::InvalidArguments
2686 );
2687 PrometheusJsonResponse::success(PrometheusResponse::ParseResult(ast))
2688 } else {
2689 PrometheusJsonResponse::error(StatusCode::InvalidArguments, "query is required")
2690 }
2691}
2692
2693#[cfg(test)]
2694mod tests {
2695 use std::collections::HashSet;
2696 use std::sync::{Arc, Mutex};
2697
2698 use catalog::memory::MemoryCatalogManager;
2699 use catalog::{RegisterSchemaRequest, RegisterTableRequest};
2700 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2701 use common_query::native_histogram::{
2702 CounterResetHint, NativeHistogram, build_histogram_array, native_histogram_value_type,
2703 };
2704 use common_query::prelude::greptime_native_histogram;
2705 use datatypes::prelude::ConcreteDataType;
2706 use datatypes::schema::{ColumnSchema, Schema};
2707 use datatypes::vectors::{StringVector, StructVector};
2708 use promql_parser::parser::value::ValueType;
2709 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType, TableVersion};
2710 use table::requests::{
2711 SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, TableOptions,
2712 };
2713 use table::test_util::EmptyTable;
2714 use table::test_util::table_info::test_table_info;
2715
2716 use super::*;
2717 use crate::prometheus_handler::PrometheusHandler;
2718
2719 struct TestCase {
2720 name: &'static str,
2721 promql: &'static str,
2722 expected_metric: Option<&'static str>,
2723 expected_type: ValueType,
2724 should_error: bool,
2725 }
2726
2727 struct TestPrometheusHandler {
2728 catalog_manager: CatalogManagerRef,
2729 deny_operation: bool,
2730 denied_table: Option<&'static str>,
2731 metric_names: Vec<String>,
2732 label_metric_names: Vec<String>,
2735 label_lookups: Mutex<Vec<Vec<Matcher>>>,
2736 queries: Mutex<Vec<String>>,
2737 ordered_outputs: Mutex<Vec<bool>>,
2738 }
2739
2740 #[async_trait::async_trait]
2741 impl PrometheusHandler for TestPrometheusHandler {
2742 async fn query_exemplars(
2743 &self,
2744 _: crate::exemplar::query::Query,
2745 _: QueryContextRef,
2746 ) -> Result<Vec<crate::exemplar::query::Series>> {
2747 struct Scan<'a>(&'a Mutex<Vec<String>>);
2748 impl Drop for Scan<'_> {
2749 fn drop(&mut self) {
2750 self.0.lock().unwrap().push("cancelled".into());
2751 }
2752 }
2753 let _scan = Scan(&self.queries);
2754 self.queries.lock().unwrap().push("started".into());
2755 futures::future::pending().await
2756 }
2757
2758 async fn do_query(&self, _: &PromQuery, _: QueryContextRef) -> Result<Output> {
2759 unreachable!("HTTP handlers should execute the parsed query")
2760 }
2761
2762 async fn do_query_parsed(
2763 &self,
2764 query: ParsedPromQuery,
2765 _: QueryContextRef,
2766 ) -> Result<Output> {
2767 self.ordered_outputs
2768 .lock()
2769 .unwrap()
2770 .push(query.requires_output_ordering());
2771 self.queries
2772 .lock()
2773 .unwrap()
2774 .push(query.query().query.clone());
2775 Ok(Output::new_with_record_batches(RecordBatches::empty()))
2776 }
2777
2778 async fn check_query_permission(&self, _: &[PromQuery], _: &QueryContextRef) -> Result<()> {
2779 if self.deny_operation {
2780 return auth::error::PermissionDeniedSnafu
2781 .fail()
2782 .context(crate::error::AuthSnafu);
2783 }
2784 Ok(())
2785 }
2786
2787 async fn check_query_target_permission(
2788 &self,
2789 targets: PermissionTableTargets,
2790 _: &QueryContextRef,
2791 ) -> Result<()> {
2792 let denied = match targets {
2793 PermissionTableTargets::Unresolved => true,
2794 PermissionTableTargets::Resolved(targets) => self
2795 .denied_table
2796 .is_some_and(|denied| targets.iter().any(|target| target.table == denied)),
2797 };
2798 if denied {
2799 return auth::error::PermissionDeniedSnafu
2800 .fail()
2801 .context(crate::error::AuthSnafu);
2802 }
2803 Ok(())
2804 }
2805
2806 async fn filter_metadata_metric_names(
2807 &self,
2808 metric_names: Vec<String>,
2809 _: &str,
2810 _: &QueryContextRef,
2811 ) -> Result<Vec<String>> {
2812 Ok(metric_names
2813 .into_iter()
2814 .filter(|metric| metric != "denied")
2815 .collect())
2816 }
2817
2818 async fn filter_metadata_schema_names(
2819 &self,
2820 schema_names: Vec<String>,
2821 _: &QueryContextRef,
2822 ) -> Result<Vec<String>> {
2823 Ok(schema_names
2824 .into_iter()
2825 .filter(|schema| schema != "denied_schema")
2826 .collect())
2827 }
2828
2829 async fn query_metric_names(
2830 &self,
2831 _: Vec<Matcher>,
2832 _: &str,
2833 _: &QueryContextRef,
2834 ) -> Result<Vec<String>> {
2835 Ok(self.metric_names.clone())
2836 }
2837
2838 async fn query_metric_names_by_labels(
2839 &self,
2840 matchers: Vec<Matcher>,
2841 _: &str,
2842 _: std::time::SystemTime,
2843 _: std::time::SystemTime,
2844 _: &QueryContextRef,
2845 ) -> Result<Vec<String>> {
2846 self.label_lookups.lock().unwrap().push(matchers);
2847 Ok(self.label_metric_names.clone())
2848 }
2849
2850 async fn query_label_values(
2851 &self,
2852 _: String,
2853 _: String,
2854 _: Vec<Matcher>,
2855 _: std::time::SystemTime,
2856 _: std::time::SystemTime,
2857 _: &QueryContextRef,
2858 ) -> Result<Vec<String>> {
2859 unreachable!()
2860 }
2861
2862 fn catalog_manager(&self) -> CatalogManagerRef {
2863 self.catalog_manager.clone()
2864 }
2865 }
2866
2867 fn permission_denied_response() -> PrometheusJsonResponse {
2868 let result: auth::error::Result<()> = auth::error::PermissionDeniedSnafu.fail();
2869 try_call_return_response!(result);
2870 unreachable!()
2871 }
2872
2873 fn invalid_promql_response() -> PrometheusJsonResponse {
2874 let result = ParsedPromQuery::parse(
2875 PromQuery {
2876 query: "up{".to_string(),
2877 ..Default::default()
2878 },
2879 &QueryContext::arc(),
2880 );
2881 try_call_return_response!(@output_msg result, StatusCode::InvalidArguments);
2882 unreachable!()
2883 }
2884
2885 #[test]
2886 fn test_try_call_return_response_preserves_status() {
2887 use axum::response::IntoResponse;
2888
2889 let response = permission_denied_response();
2890 assert_eq!(Some(StatusCode::PermissionDenied), response.status_code);
2891 assert_eq!(
2892 axum::http::StatusCode::FORBIDDEN,
2893 response.into_response().status()
2894 );
2895 }
2896
2897 #[test]
2898 fn test_try_call_return_response_preserves_error_source() {
2899 let response = invalid_promql_response();
2900 assert_eq!(Some(StatusCode::InvalidArguments), response.status_code);
2901 let error = response.error.unwrap();
2902 assert!(
2903 error.contains("unexpected end of input inside braces"),
2904 "{error}"
2905 );
2906 }
2907
2908 #[tokio::test]
2909 async fn test_prometheus_endpoints_authorize_effective_database() {
2910 use axum::http::{Request, StatusCode as HttpStatusCode};
2911 use axum::response::IntoResponse;
2912 use tower::ServiceExt;
2913
2914 for authorization in ["Basic cmVhZGVyOnB3ZA==", "Bearer test-token"] {
2915 for (path, form_db, encoded_db) in [
2916 ("/v1/prometheus/api/v1/query", false, false),
2917 ("/v1/prometheus/api/v1/query_range", false, false),
2918 ("/v1/prometheus/api/v1/metadata", false, false),
2919 ("/v1/prometheus/api/v1/labels", false, false),
2920 ("/v1/prometheus/api/v1/label/job/values", false, false),
2921 ("/v1/prometheus/api/v1/series", false, false),
2922 ("/v1/prometheus/api/v1/query_exemplars", false, false),
2923 ("/v1/prometheus/api/v1/query_exemplars", true, false),
2924 ("/v1/prometheus/api/v1/query", false, true),
2925 ] {
2926 let handler = Arc::new(TestPrometheusHandler {
2927 catalog_manager: MemoryCatalogManager::new(),
2928 deny_operation: false,
2929 denied_table: None,
2930 metric_names: Vec::new(),
2931 label_metric_names: Vec::new(),
2932 label_lookups: Mutex::new(Vec::new()),
2933 queries: Mutex::new(Vec::new()),
2934 ordered_outputs: Mutex::new(Vec::new()),
2935 });
2936 let app = axum::Router::new()
2937 .route(
2938 "/v1/prometheus/api/v1/query",
2939 axum::routing::get(instant_query),
2940 )
2941 .route(
2942 "/v1/prometheus/api/v1/query_range",
2943 axum::routing::get(range_query),
2944 )
2945 .route(
2946 "/v1/prometheus/api/v1/metadata",
2947 axum::routing::get(metadata_query),
2948 )
2949 .route(
2950 "/v1/prometheus/api/v1/labels",
2951 axum::routing::get(labels_query),
2952 )
2953 .route(
2954 "/v1/prometheus/api/v1/label/{name}/values",
2955 axum::routing::get(label_values_query),
2956 )
2957 .route(
2958 "/v1/prometheus/api/v1/series",
2959 axum::routing::get(series_query),
2960 )
2961 .route(
2962 "/v1/prometheus/api/v1/query_exemplars",
2963 axum::routing::get(query_exemplars).post(query_exemplars),
2964 )
2965 .with_state(handler.clone() as PrometheusHandlerRef)
2966 .layer(axum::middleware::from_fn_with_state(
2967 crate::http::authorize::AuthState::new(Some(Arc::new(
2968 crate::http::authorize::tests::ScopedUserProvider,
2969 ))),
2970 crate::http::authorize::check_http_auth,
2971 ));
2972 let db = if encoded_db { "%64enied" } else { "denied" };
2973 let query = if form_db {
2974 String::new()
2975 } else {
2976 format!("&db={db}")
2977 };
2978 let mut request = Request::builder()
2979 .method(if form_db { "POST" } else { "GET" })
2980 .uri(format!(
2981 "{path}?query=requests&start=0&end=1&step=1&match[]=requests{query}"
2982 ))
2983 .header("authorization", authorization)
2984 .header("content-type", "application/x-www-form-urlencoded");
2985 if !encoded_db {
2986 request = request.header("x-greptime-db-name", "public");
2987 }
2988 let body = if form_db { "db=denied" } else { "" };
2989 let response = tokio::time::timeout(
2990 std::time::Duration::from_secs(5),
2991 app.oneshot(request.body(axum::body::Body::from(body)).unwrap()),
2992 )
2993 .await
2994 .unwrap()
2995 .unwrap();
2996 assert_eq!(
2997 HttpStatusCode::FORBIDDEN,
2998 response.into_response().status(),
2999 "{path}"
3000 );
3001 assert!(handler.queries.lock().unwrap().is_empty());
3002 assert!(handler.label_lookups.lock().unwrap().is_empty());
3003 }
3004 }
3005 }
3006
3007 #[tokio::test]
3008 async fn test_exemplar_timeout_cancels_scan() {
3009 for explicit in [false, true] {
3010 let handler = Arc::new(TestPrometheusHandler {
3011 catalog_manager: MemoryCatalogManager::new(),
3012 deny_operation: false,
3013 denied_table: None,
3014 metric_names: Vec::new(),
3015 label_metric_names: Vec::new(),
3016 label_lookups: Mutex::new(Vec::new()),
3017 queries: Mutex::new(Vec::new()),
3018 ordered_outputs: Mutex::new(Vec::new()),
3019 });
3020 let ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3021 if !explicit {
3022 ctx.set_query_timeout(std::time::Duration::from_millis(1));
3023 }
3024 let response = query_exemplars(
3025 None,
3026 State(handler.clone()),
3027 Query(ExemplarQuery {
3028 query: Some("requests".into()),
3029 start: Some("0".into()),
3030 end: Some("1".into()),
3031 timeout: explicit.then(|| "1ms".into()),
3032 ..Default::default()
3033 }),
3034 Extension(ctx),
3035 Form(ExemplarQuery::default()),
3036 )
3037 .await;
3038 assert_eq!(response.status_code, Some(StatusCode::Cancelled));
3039 assert_eq!(*handler.queries.lock().unwrap(), ["started", "cancelled"]);
3040 }
3041 }
3042
3043 #[tokio::test]
3044 async fn range_query_does_not_require_execution_output_ordering() {
3045 let handler = Arc::new(TestPrometheusHandler {
3046 catalog_manager: MemoryCatalogManager::new(),
3047 deny_operation: false,
3048 denied_table: None,
3049 metric_names: Vec::new(),
3050 label_metric_names: Vec::new(),
3051 label_lookups: Mutex::new(Vec::new()),
3052 queries: Mutex::new(Vec::new()),
3053 ordered_outputs: Mutex::new(Vec::new()),
3054 });
3055 let state: PrometheusHandlerRef = handler.clone();
3056 instant_query(
3057 None,
3058 State(state.clone()),
3059 Query(InstantQuery {
3060 query: Some("sort(vector(1))".to_string()),
3061 time: Some("0".to_string()),
3062 ..Default::default()
3063 }),
3064 Extension(QueryContext::with(
3065 DEFAULT_CATALOG_NAME,
3066 DEFAULT_SCHEMA_NAME,
3067 )),
3068 Form(InstantQuery::default()),
3069 )
3070 .await;
3071
3072 for end in ["0", "1"] {
3074 range_query(
3075 None,
3076 State(state.clone()),
3077 Query(RangeQuery {
3078 query: Some("sort(vector(1))".to_string()),
3079 start: Some("0".to_string()),
3080 end: Some(end.to_string()),
3081 step: Some("1s".to_string()),
3082 ..Default::default()
3083 }),
3084 Extension(QueryContext::with(
3085 DEFAULT_CATALOG_NAME,
3086 DEFAULT_SCHEMA_NAME,
3087 )),
3088 Form(RangeQuery::default()),
3089 )
3090 .await;
3091 }
3092
3093 assert_eq!(
3095 *handler.ordered_outputs.lock().unwrap(),
3096 vec![true, false, false]
3097 );
3098 }
3099
3100 #[tokio::test]
3101 async fn range_query_rejects_range_vector_with_equal_endpoints() {
3102 let handler = Arc::new(TestPrometheusHandler {
3103 catalog_manager: MemoryCatalogManager::new(),
3104 deny_operation: false,
3105 denied_table: None,
3106 metric_names: Vec::new(),
3107 label_metric_names: Vec::new(),
3108 label_lookups: Mutex::new(Vec::new()),
3109 queries: Mutex::new(Vec::new()),
3110 ordered_outputs: Mutex::new(Vec::new()),
3111 });
3112 let response = range_query(
3113 None,
3114 State(handler.clone() as PrometheusHandlerRef),
3115 Query(RangeQuery {
3116 query: Some("vector(1)[5m:]".to_string()),
3117 start: Some("300".to_string()),
3118 end: Some("300".to_string()),
3119 step: Some("1".to_string()),
3120 ..Default::default()
3121 }),
3122 Extension(QueryContext::with(
3123 DEFAULT_CATALOG_NAME,
3124 DEFAULT_SCHEMA_NAME,
3125 )),
3126 Form(RangeQuery::default()),
3127 )
3128 .await;
3129
3130 assert_eq!(Some(StatusCode::InvalidArguments), response.status_code);
3131 assert_eq!(
3132 response.error.as_deref(),
3133 Some(
3134 "invalid expression type \"range vector\" for range query, must be Scalar or instant Vector"
3135 )
3136 );
3137 assert!(handler.queries.lock().unwrap().is_empty());
3138 }
3139
3140 #[tokio::test]
3141 async fn test_promql_timer_records_parse_errors() {
3142 let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler {
3143 catalog_manager: MemoryCatalogManager::new(),
3144 deny_operation: false,
3145 denied_table: None,
3146 metric_names: Vec::new(),
3147 label_metric_names: Vec::new(),
3148 label_lookups: Mutex::new(Vec::new()),
3149 queries: Mutex::new(Vec::new()),
3150 ordered_outputs: Mutex::new(Vec::new()),
3151 });
3152 let query_ctx = QueryContext::with("promql_timer_test", "parse_error");
3153 let db = query_ctx.get_db_string();
3154
3155 let instant_histogram = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
3156 .with_label_values(&[db.as_str(), "instant_query"]);
3157 let instant_count = instant_histogram.get_sample_count();
3158 let response = instant_query(
3159 None,
3160 State(handler.clone()),
3161 Query(InstantQuery {
3162 query: Some("up{".to_string()),
3163 ..Default::default()
3164 }),
3165 Extension(query_ctx),
3166 Form(InstantQuery::default()),
3167 )
3168 .await;
3169 assert_eq!(Some(StatusCode::InvalidArguments), response.status_code);
3170 assert_eq!(instant_count + 1, instant_histogram.get_sample_count());
3171
3172 let range_histogram = crate::metrics::METRIC_HTTP_PROMETHEUS_PROMQL_ELAPSED
3173 .with_label_values(&[db.as_str(), "range_query"]);
3174 let range_count = range_histogram.get_sample_count();
3175 let response = range_query(
3176 None,
3177 State(handler),
3178 Query(RangeQuery {
3179 query: Some("up{".to_string()),
3180 start: Some("0".to_string()),
3181 end: Some("1".to_string()),
3182 step: Some("1s".to_string()),
3183 ..Default::default()
3184 }),
3185 Extension(QueryContext::with("promql_timer_test", "parse_error")),
3186 Form(RangeQuery::default()),
3187 )
3188 .await;
3189 assert_eq!(Some(StatusCode::InvalidArguments), response.status_code);
3190 assert_eq!(range_count + 1, range_histogram.get_sample_count());
3191 }
3192
3193 fn metadata_test_handler(manager: CatalogManagerRef) -> PrometheusHandlerRef {
3194 Arc::new(TestPrometheusHandler {
3195 catalog_manager: manager,
3196 deny_operation: false,
3197 denied_table: Some("denied"),
3198 metric_names: Vec::new(),
3199 label_metric_names: Vec::new(),
3200 label_lookups: Mutex::new(Vec::new()),
3201 queries: Mutex::new(Vec::new()),
3202 ordered_outputs: Mutex::new(Vec::new()),
3203 })
3204 }
3205
3206 #[tokio::test]
3207 async fn test_metadata_enumeration_filters_inaccessible_candidates() {
3208 let manager = MemoryCatalogManager::with_default_setup();
3209 for schema in ["empty", "denied_schema"] {
3210 manager
3211 .register_schema_sync(RegisterSchemaRequest {
3212 catalog: DEFAULT_CATALOG_NAME.to_string(),
3213 schema: schema.to_string(),
3214 })
3215 .unwrap();
3216 }
3217 for (id, schema, name, field) in [
3218 (1024, DEFAULT_SCHEMA_NAME, "allowed", "allowed_field"),
3219 (1025, DEFAULT_SCHEMA_NAME, "denied", "secret_field"),
3220 (1026, "denied_schema", "allowed", "hidden_field"),
3221 ] {
3222 let info = test_table_info(
3223 id,
3224 name,
3225 schema,
3226 DEFAULT_CATALOG_NAME,
3227 Arc::new(Schema::new(vec![ColumnSchema::new(
3228 field,
3229 ConcreteDataType::float64_datatype(),
3230 true,
3231 )])),
3232 );
3233 manager
3234 .register_table_sync(RegisterTableRequest {
3235 catalog: DEFAULT_CATALOG_NAME.to_string(),
3236 schema: schema.to_string(),
3237 table_name: name.to_string(),
3238 table_id: id,
3239 table: EmptyTable::from_table_info(&info),
3240 })
3241 .unwrap();
3242 }
3243 let handler = metadata_test_handler(manager);
3244 let query_ctx = Arc::new(QueryContext::with(
3245 DEFAULT_CATALOG_NAME,
3246 DEFAULT_SCHEMA_NAME,
3247 ));
3248 for selectors in [
3249 vec![],
3250 vec![r#"{__name__=~"allowed|denied"}"#.to_string()],
3251 vec!["allowed".to_string(), "denied".to_string()],
3252 vec!["allowed".to_string(), r#"{job="nonexistent"}"#.to_string()],
3253 vec![r#"allowed{__schema__="other"}"#.to_string()],
3254 ] {
3255 let fields = retrieve_field_names(&query_ctx, &handler, selectors)
3256 .await
3257 .unwrap();
3258 assert_eq!(fields, HashSet::from(["allowed_field".to_string()]));
3259 }
3260 for selectors in [
3261 vec![r#"{job="nonexistent"}"#.to_string()],
3262 vec![r#"{__database__="nonexistent"}"#.to_string()],
3263 vec![r#"{__field__="nonexistent"}"#.to_string()],
3264 vec!["missing".to_string(), r#"{job="nonexistent"}"#.to_string()],
3265 ] {
3266 assert!(
3267 retrieve_field_names(&query_ctx, &handler, selectors)
3268 .await
3269 .unwrap()
3270 .is_empty()
3271 );
3272 }
3273 assert!(
3274 retrieve_field_names(&query_ctx, &handler, vec!["denied".to_string()])
3275 .await
3276 .unwrap()
3277 .is_empty()
3278 );
3279 assert!(
3280 retrieve_field_names(&query_ctx, &handler, vec!["{".to_string()])
3281 .await
3282 .is_err()
3283 );
3284 for selector in ["allowed", r#"{__name__=~"allowed|denied"}"#] {
3285 assert_eq!(
3286 retrieve_schema_names(&query_ctx, &handler, vec![selector.to_string()])
3287 .await
3288 .unwrap(),
3289 vec![DEFAULT_SCHEMA_NAME]
3290 );
3291 }
3292 assert!(
3293 retrieve_schema_names(&query_ctx, &handler, vec!["denied".to_string()])
3294 .await
3295 .unwrap()
3296 .is_empty()
3297 );
3298 let schemas = retrieve_schema_names(&query_ctx, &handler, vec![])
3299 .await
3300 .unwrap();
3301 assert!(schemas.contains(&"empty".to_string()));
3302 assert!(!schemas.contains(&"denied_schema".to_string()));
3303 let response = label_values_query(
3304 None,
3305 State(handler),
3306 Path(FIELD_NAME_LABEL.to_string()),
3307 Extension(QueryContext::with(
3308 DEFAULT_CATALOG_NAME,
3309 DEFAULT_SCHEMA_NAME,
3310 )),
3311 Query(LabelValueQuery::default()),
3312 )
3313 .await;
3314 assert_eq!(None, response.status_code);
3315 }
3316
3317 fn label_values_handler(label_metric_names: Vec<&str>) -> Arc<TestPrometheusHandler> {
3321 let mut cpu_user = test_table_info(
3322 1024,
3323 "cpu_user",
3324 DEFAULT_SCHEMA_NAME,
3325 DEFAULT_CATALOG_NAME,
3326 Arc::new(Schema::new(vec![])),
3327 );
3328 cpu_user.meta.options.extra_options.insert(
3329 LOGICAL_TABLE_METADATA_KEY.to_string(),
3330 "physical_metrics".to_string(),
3331 );
3332 let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&cpu_user));
3333 let mut cpu_system = cpu_user.clone();
3334 cpu_system.ident.table_id = 1025;
3335 cpu_system.name = "cpu_system".to_string();
3336 manager
3337 .register_table_sync(RegisterTableRequest {
3338 catalog: DEFAULT_CATALOG_NAME.to_string(),
3339 schema: DEFAULT_SCHEMA_NAME.to_string(),
3340 table_name: cpu_system.name.clone(),
3341 table_id: cpu_system.table_id(),
3342 table: EmptyTable::from_table_info(&cpu_system),
3343 })
3344 .unwrap();
3345
3346 Arc::new(TestPrometheusHandler {
3347 catalog_manager: manager,
3348 deny_operation: false,
3349 denied_table: None,
3350 metric_names: Vec::new(),
3351 label_metric_names: label_metric_names.into_iter().map(String::from).collect(),
3352 label_lookups: Mutex::new(Vec::new()),
3353 queries: Mutex::new(Vec::new()),
3354 ordered_outputs: Mutex::new(Vec::new()),
3355 })
3356 }
3357
3358 async fn query_metric_name_values(
3359 handler: Arc<TestPrometheusHandler>,
3360 matches: Vec<&str>,
3361 ) -> Vec<String> {
3362 let state: PrometheusHandlerRef = handler;
3363 let response = label_values_query(
3364 None,
3365 State(state),
3366 Path(METRIC_NAME_LABEL.to_string()),
3367 Extension(QueryContext::with(
3368 DEFAULT_CATALOG_NAME,
3369 DEFAULT_SCHEMA_NAME,
3370 )),
3371 Query(LabelValueQuery {
3372 matches: Matches(matches.into_iter().map(String::from).collect()),
3373 ..Default::default()
3374 }),
3375 )
3376 .await;
3377
3378 assert!(
3379 response.status_code.is_none(),
3380 "status={:?}, error={:?}",
3381 response.status_code,
3382 response.error
3383 );
3384 match response.data {
3385 PrometheusResponse::LabelValues(values) => values,
3386 other => panic!("expected label values, got {other:?}"),
3387 }
3388 }
3389
3390 #[tokio::test]
3391 async fn label_matchers_resolve_metric_names_from_data() {
3392 let handler = label_values_handler(vec!["cpu_user"]);
3393 let values = query_metric_name_values(handler.clone(), vec![r#"{pod="abc"}"#]).await;
3394
3395 assert_eq!(vec!["cpu_user".to_string()], values);
3397
3398 let lookups = handler.label_lookups.lock().unwrap();
3399 assert_eq!(1, lookups.len());
3400 assert_eq!(
3401 vec!["pod".to_string()],
3402 lookups[0]
3403 .iter()
3404 .map(|matcher| matcher.name.clone())
3405 .collect::<Vec<_>>()
3406 );
3407 }
3408
3409 #[tokio::test]
3410 async fn metric_name_matchers_narrow_data_resolved_names() {
3411 let handler = label_values_handler(vec!["cpu_user", "cpu_system"]);
3412 let values =
3413 query_metric_name_values(handler, vec![r#"{__name__=~"cpu_u.*", pod="abc"}"#]).await;
3414
3415 assert_eq!(vec!["cpu_user".to_string()], values);
3416 }
3417
3418 #[tokio::test]
3419 async fn special_matchers_are_stripped_before_the_data_lookup() {
3420 let handler = label_values_handler(vec!["cpu_user"]);
3421 let values = query_metric_name_values(
3422 handler.clone(),
3423 vec![r#"{pod="abc", __field__="value", __database__="public"}"#],
3424 )
3425 .await;
3426
3427 assert_eq!(vec!["cpu_user".to_string()], values);
3428 let lookups = handler.label_lookups.lock().unwrap();
3429 assert_eq!(
3430 vec!["pod".to_string()],
3431 lookups[0]
3432 .iter()
3433 .map(|matcher| matcher.name.clone())
3434 .collect::<Vec<_>>()
3435 );
3436 }
3437
3438 #[tokio::test]
3439 async fn database_and_field_matchers_stay_on_the_metadata_path() {
3440 let handler = label_values_handler(vec!["never_returned"]);
3441 let values = query_metric_name_values(
3442 handler.clone(),
3443 vec![r#"{__name__=~"cpu_.*", __field__="value", __database__="other"}"#],
3444 )
3445 .await;
3446
3447 assert_eq!(
3448 vec!["cpu_system".to_string(), "cpu_user".to_string()],
3449 values
3450 );
3451 assert!(handler.label_lookups.lock().unwrap().is_empty());
3452 }
3453
3454 #[tokio::test]
3455 async fn empty_match_still_enumerates_every_metric() {
3456 let handler = label_values_handler(Vec::new());
3457 let values = query_metric_name_values(handler, Vec::new()).await;
3458
3459 assert_eq!(
3460 vec!["cpu_system".to_string(), "cpu_user".to_string()],
3461 values
3462 );
3463 }
3464
3465 #[tokio::test]
3466 async fn test_series_query_expands_metric_name_regex() {
3467 let cpu_user = test_table_info(
3468 1024,
3469 "cpu_user",
3470 DEFAULT_SCHEMA_NAME,
3471 DEFAULT_CATALOG_NAME,
3472 Arc::new(Schema::new(vec![])),
3473 );
3474 let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&cpu_user));
3475 let mut cpu_system = cpu_user.clone();
3476 cpu_system.ident.table_id = 1025;
3477 cpu_system.name = "cpu_system".to_string();
3478 manager
3479 .register_table_sync(RegisterTableRequest {
3480 catalog: DEFAULT_CATALOG_NAME.to_string(),
3481 schema: DEFAULT_SCHEMA_NAME.to_string(),
3482 table_name: cpu_system.name.clone(),
3483 table_id: cpu_system.table_id(),
3484 table: EmptyTable::from_table_info(&cpu_system),
3485 })
3486 .unwrap();
3487
3488 let handler = Arc::new(TestPrometheusHandler {
3489 catalog_manager: manager,
3490 deny_operation: false,
3491 denied_table: None,
3492 metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()],
3493 label_metric_names: Vec::new(),
3494 label_lookups: Mutex::new(Vec::new()),
3495 queries: Mutex::new(Vec::new()),
3496 ordered_outputs: Mutex::new(Vec::new()),
3497 });
3498 let state: PrometheusHandlerRef = handler.clone();
3499 let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3500 let target_ctx = Arc::new(query_ctx.clone());
3501 let response = series_query(
3502 None,
3503 State(state),
3504 Query(SeriesQuery {
3505 matches: Matches(vec![r#"{__name__=~"cpu_.*"}"#.to_string()]),
3506 ..Default::default()
3507 }),
3508 Extension(query_ctx),
3509 Form(SeriesQuery::default()),
3510 )
3511 .await;
3512
3513 assert!(
3514 response.status_code.is_none(),
3515 "status={:?}, error={:?}",
3516 response.status_code,
3517 response.error
3518 );
3519 let mut targets = handler
3520 .queries
3521 .lock()
3522 .unwrap()
3523 .iter()
3524 .map(|query| {
3525 let query = ParsedPromQuery::parse(
3526 PromQuery {
3527 query: query.clone(),
3528 ..Default::default()
3529 },
3530 &target_ctx,
3531 )
3532 .unwrap();
3533 resolved_promql_targets(std::slice::from_ref(&query), &target_ctx)
3534 .unwrap()
3535 .pop()
3536 .unwrap()
3537 .table
3538 })
3539 .collect_vec();
3540 targets.sort_unstable();
3541 assert_eq!(vec!["cpu_system", "cpu_user"], targets);
3542 }
3543
3544 #[test]
3545 fn test_retrieve_metric_name_and_result_type() {
3546 let test_cases = &[
3547 TestCase {
3549 name: "simple metric",
3550 promql: "cpu_usage",
3551 expected_metric: Some("cpu_usage"),
3552 expected_type: ValueType::Vector,
3553 should_error: false,
3554 },
3555 TestCase {
3556 name: "metric with selector",
3557 promql: r#"cpu_usage{instance="localhost"}"#,
3558 expected_metric: Some("cpu_usage"),
3559 expected_type: ValueType::Vector,
3560 should_error: false,
3561 },
3562 TestCase {
3563 name: "metric with range selector",
3564 promql: "cpu_usage[5m]",
3565 expected_metric: Some("cpu_usage"),
3566 expected_type: ValueType::Matrix,
3567 should_error: false,
3568 },
3569 TestCase {
3570 name: "metric with __name__ matcher",
3571 promql: r#"{__name__="cpu_usage"}"#,
3572 expected_metric: Some("cpu_usage"),
3573 expected_type: ValueType::Vector,
3574 should_error: false,
3575 },
3576 TestCase {
3577 name: "metric with unary operator",
3578 promql: "-cpu_usage",
3579 expected_metric: None,
3580 expected_type: ValueType::Vector,
3581 should_error: false,
3582 },
3583 TestCase {
3585 name: "metric with aggregation",
3586 promql: "sum(cpu_usage)",
3587 expected_metric: Some("cpu_usage"),
3588 expected_type: ValueType::Vector,
3589 should_error: false,
3590 },
3591 TestCase {
3592 name: "complex aggregation",
3593 promql: r#"sum by (instance) (cpu_usage{job="node"})"#,
3594 expected_metric: None,
3595 expected_type: ValueType::Vector,
3596 should_error: false,
3597 },
3598 TestCase {
3599 name: "complex aggregation",
3600 promql: r#"sum by (__name__) (cpu_usage{job="node"})"#,
3601 expected_metric: Some("cpu_usage"),
3602 expected_type: ValueType::Vector,
3603 should_error: false,
3604 },
3605 TestCase {
3606 name: "complex aggregation",
3607 promql: r#"sum without (instance) (cpu_usage{job="node"})"#,
3608 expected_metric: Some("cpu_usage"),
3609 expected_type: ValueType::Vector,
3610 should_error: false,
3611 },
3612 TestCase {
3614 name: "same metric addition",
3615 promql: "cpu_usage + cpu_usage",
3616 expected_metric: None,
3617 expected_type: ValueType::Vector,
3618 should_error: false,
3619 },
3620 TestCase {
3621 name: "metric with scalar addition",
3622 promql: r#"sum(rate(cpu_usage{job="node"}[5m])) + 100"#,
3623 expected_metric: None,
3624 expected_type: ValueType::Vector,
3625 should_error: false,
3626 },
3627 TestCase {
3629 name: "different metrics addition",
3630 promql: "cpu_usage + memory_usage",
3631 expected_metric: None,
3632 expected_type: ValueType::Vector,
3633 should_error: false,
3634 },
3635 TestCase {
3636 name: "different metrics subtraction",
3637 promql: "network_in - network_out",
3638 expected_metric: None,
3639 expected_type: ValueType::Vector,
3640 should_error: false,
3641 },
3642 TestCase {
3644 name: "unless with different metrics",
3645 promql: "cpu_usage unless memory_usage",
3646 expected_metric: Some("cpu_usage"),
3647 expected_type: ValueType::Vector,
3648 should_error: false,
3649 },
3650 TestCase {
3651 name: "unless with same metric",
3652 promql: "cpu_usage unless cpu_usage",
3653 expected_metric: Some("cpu_usage"),
3654 expected_type: ValueType::Vector,
3655 should_error: false,
3656 },
3657 TestCase {
3659 name: "arithmetic between metrics",
3660 promql: "a + b",
3661 expected_metric: None,
3662 expected_type: ValueType::Vector,
3663 should_error: false,
3664 },
3665 TestCase {
3666 name: "arithmetic with scalar",
3667 promql: "a * 2",
3668 expected_metric: None,
3669 expected_type: ValueType::Vector,
3670 should_error: false,
3671 },
3672 TestCase {
3681 name: "bool comparison between metrics",
3682 promql: "a > bool b",
3683 expected_metric: None,
3684 expected_type: ValueType::Vector,
3685 should_error: false,
3686 },
3687 TestCase {
3688 name: "bool comparison with scalar",
3689 promql: "a > bool 0.5",
3690 expected_metric: None,
3691 expected_type: ValueType::Vector,
3692 should_error: false,
3693 },
3694 TestCase {
3695 name: "comparison with scalar",
3696 promql: "a > 0.5",
3697 expected_metric: Some("a"),
3698 expected_type: ValueType::Vector,
3699 should_error: false,
3700 },
3701 TestCase {
3702 name: "comparison with scalar on the left",
3703 promql: "0.5 < a",
3704 expected_metric: Some("a"),
3705 expected_type: ValueType::Vector,
3706 should_error: false,
3707 },
3708 TestCase {
3709 name: "equality comparison with scalar",
3710 promql: "up == 0",
3711 expected_metric: Some("up"),
3712 expected_type: ValueType::Vector,
3713 should_error: false,
3714 },
3715 TestCase {
3716 name: "comparison between metrics",
3717 promql: "a > b",
3718 expected_metric: Some("a"),
3719 expected_type: ValueType::Vector,
3720 should_error: false,
3721 },
3722 TestCase {
3723 name: "comparison with ignoring",
3724 promql: "a > ignoring(x) b",
3725 expected_metric: Some("a"),
3726 expected_type: ValueType::Vector,
3727 should_error: false,
3728 },
3729 TestCase {
3730 name: "comparison with ignoring __name__",
3731 promql: "a > ignoring(__name__) b",
3732 expected_metric: None,
3733 expected_type: ValueType::Vector,
3734 should_error: false,
3735 },
3736 TestCase {
3737 name: "comparison with ignoring __name__ and group_left",
3738 promql: "a > ignoring(__name__) group_left b",
3739 expected_metric: Some("a"),
3740 expected_type: ValueType::Vector,
3741 should_error: false,
3742 },
3743 TestCase {
3744 name: "comparison with ignoring __name__ and group_right",
3745 promql: "a > ignoring(__name__) group_right b",
3746 expected_metric: Some("b"),
3747 expected_type: ValueType::Vector,
3748 should_error: false,
3749 },
3750 TestCase {
3751 name: "unless with ignoring __name__",
3752 promql: "a unless ignoring(__name__) b",
3753 expected_metric: Some("a"),
3754 expected_type: ValueType::Vector,
3755 should_error: false,
3756 },
3757 TestCase {
3758 name: "comparison with on",
3759 promql: "a > on(x) b",
3760 expected_metric: None,
3761 expected_type: ValueType::Vector,
3762 should_error: false,
3763 },
3764 TestCase {
3765 name: "comparison with on keeping __name__",
3766 promql: "a > on(__name__) b",
3767 expected_metric: Some("a"),
3768 expected_type: ValueType::Vector,
3769 should_error: false,
3770 },
3771 TestCase {
3772 name: "comparison with on and group_left",
3773 promql: "a > on(x) group_left b",
3774 expected_metric: Some("a"),
3775 expected_type: ValueType::Vector,
3776 should_error: false,
3777 },
3778 TestCase {
3779 name: "comparison with on and group_right",
3780 promql: "a > on(x) group_right b",
3781 expected_metric: Some("b"),
3782 expected_type: ValueType::Vector,
3783 should_error: false,
3784 },
3785 TestCase {
3787 name: "basic subquery",
3788 promql: "cpu_usage[5m:1m]",
3789 expected_metric: Some("cpu_usage"),
3790 expected_type: ValueType::Matrix,
3791 should_error: false,
3792 },
3793 TestCase {
3794 name: "subquery with multiple metrics",
3795 promql: "(cpu_usage + memory_usage)[5m:1m]",
3796 expected_metric: None,
3797 expected_type: ValueType::Matrix,
3798 should_error: false,
3799 },
3800 TestCase {
3802 name: "scalar value",
3803 promql: "42",
3804 expected_metric: None,
3805 expected_type: ValueType::Scalar,
3806 should_error: false,
3807 },
3808 TestCase {
3809 name: "string literal",
3810 promql: r#""hello world""#,
3811 expected_metric: None,
3812 expected_type: ValueType::String,
3813 should_error: false,
3814 },
3815 TestCase {
3817 name: "invalid syntax",
3818 promql: "cpu_usage{invalid=",
3819 expected_metric: None,
3820 expected_type: ValueType::Vector,
3821 should_error: true,
3822 },
3823 TestCase {
3824 name: "empty query",
3825 promql: "",
3826 expected_metric: None,
3827 expected_type: ValueType::Vector,
3828 should_error: true,
3829 },
3830 TestCase {
3831 name: "malformed brackets",
3832 promql: "cpu_usage[5m",
3833 expected_metric: None,
3834 expected_type: ValueType::Vector,
3835 should_error: true,
3836 },
3837 ];
3838
3839 for test_case in test_cases {
3840 let result = promql_parser::parser::parse(test_case.promql)
3841 .map(|expr| retrieve_metric_name_and_result_type(&expr));
3842
3843 if test_case.should_error {
3844 assert!(
3845 result.is_err(),
3846 "Test '{}' should have failed but succeeded with: {:?}",
3847 test_case.name,
3848 result
3849 );
3850 } else {
3851 let (metric_name, value_type) = result.unwrap_or_else(|e| {
3852 panic!(
3853 "Test '{}' should have succeeded but failed with error: {}",
3854 test_case.name, e
3855 )
3856 });
3857
3858 let expected_metric_name = test_case.expected_metric.map(|s| s.to_string());
3859 assert_eq!(
3860 metric_name, expected_metric_name,
3861 "Test '{}': metric name mismatch. Expected: {:?}, Got: {:?}",
3862 test_case.name, expected_metric_name, metric_name
3863 );
3864
3865 assert_eq!(
3866 value_type, test_case.expected_type,
3867 "Test '{}': value type mismatch. Expected: {:?}, Got: {:?}",
3868 test_case.name, test_case.expected_type, value_type
3869 );
3870 }
3871 }
3872 }
3873
3874 #[tokio::test]
3875 async fn test_get_all_column_names_uses_tag_columns() {
3876 let schema = Arc::new(Schema::new(vec![
3877 ColumnSchema::new(
3878 "greptime_timestamp",
3879 ConcreteDataType::timestamp_millisecond_datatype(),
3880 false,
3881 )
3882 .with_time_index(true),
3883 ColumnSchema::new("host", ConcreteDataType::string_datatype(), false),
3884 ColumnSchema::new("region", ConcreteDataType::string_datatype(), false),
3885 ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true),
3886 ColumnSchema::new(
3887 DATA_SCHEMA_TSID_COLUMN_NAME,
3888 ConcreteDataType::uint64_datatype(),
3889 true,
3890 ),
3891 ]));
3892 let mut options = TableOptions::default();
3893 options.extra_options.insert(
3894 LOGICAL_TABLE_METADATA_KEY.to_string(),
3895 "physical_metrics".to_string(),
3896 );
3897 let meta = TableMetaBuilder::empty()
3898 .schema(schema)
3899 .primary_key_indices(vec![1, 2, 4])
3900 .engine("metric".to_string())
3901 .next_column_id(5)
3902 .options(options)
3903 .build()
3904 .unwrap();
3905 let table_info = TableInfoBuilder::default()
3906 .table_id(1024)
3907 .table_version(0 as TableVersion)
3908 .name("cpu_usage")
3909 .catalog_name(DEFAULT_CATALOG_NAME)
3910 .schema_name(DEFAULT_SCHEMA_NAME)
3911 .table_type(TableType::Base)
3912 .meta(meta)
3913 .build()
3914 .unwrap();
3915 let manager =
3916 MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&table_info));
3917
3918 let physical_schema = Arc::new(Schema::new(vec![
3919 ColumnSchema::new(
3920 "greptime_timestamp",
3921 ConcreteDataType::timestamp_millisecond_datatype(),
3922 false,
3923 )
3924 .with_time_index(true),
3925 ColumnSchema::new(
3926 "physical_only_tag",
3927 ConcreteDataType::string_datatype(),
3928 false,
3929 ),
3930 ColumnSchema::new(
3931 "physical_only_value",
3932 ConcreteDataType::float64_datatype(),
3933 true,
3934 ),
3935 ]));
3936 let mut physical_options = TableOptions::default();
3937 physical_options.extra_options.insert(
3938 store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(),
3939 String::new(),
3940 );
3941 let physical_meta = TableMetaBuilder::empty()
3942 .schema(physical_schema)
3943 .primary_key_indices(vec![1])
3944 .engine("metric".to_string())
3945 .next_column_id(3)
3946 .options(physical_options)
3947 .build()
3948 .unwrap();
3949 let physical_table_info = TableInfoBuilder::default()
3950 .table_id(1025)
3951 .table_version(0 as TableVersion)
3952 .name("physical_metrics")
3953 .catalog_name(DEFAULT_CATALOG_NAME)
3954 .schema_name(DEFAULT_SCHEMA_NAME)
3955 .table_type(TableType::Base)
3956 .meta(physical_meta)
3957 .build()
3958 .unwrap();
3959 let physical_table = EmptyTable::from_table_info(&physical_table_info);
3960 manager
3961 .register_table_sync(RegisterTableRequest {
3962 catalog: DEFAULT_CATALOG_NAME.to_string(),
3963 schema: DEFAULT_SCHEMA_NAME.to_string(),
3964 table_name: "physical_metrics".to_string(),
3965 table_id: 1025,
3966 table: physical_table,
3967 })
3968 .unwrap();
3969 let manager: CatalogManagerRef = manager;
3970
3971 let (labels, table_names) =
3972 get_all_column_names(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, &manager)
3973 .await
3974 .unwrap();
3975
3976 assert_eq!(
3977 labels,
3978 HashSet::from(["host".to_string(), "region".to_string()])
3979 );
3980 assert_eq!(table_names, vec!["cpu_usage"]);
3981
3982 let query_ctx = Arc::new(QueryContext::with(
3983 DEFAULT_CATALOG_NAME,
3984 DEFAULT_SCHEMA_NAME,
3985 ));
3986 let handler = metadata_test_handler(manager.clone());
3987 for selector in ["cpu_usage", r#"{__name__=~"cpu.*"}"#] {
3988 assert_eq!(
3989 retrieve_schema_names(&query_ctx, &handler, vec![selector.to_string()])
3990 .await
3991 .unwrap(),
3992 vec![DEFAULT_SCHEMA_NAME]
3993 );
3994 }
3995 assert!(
3996 retrieve_schema_names(&query_ctx, &handler, vec![])
3997 .await
3998 .unwrap()
3999 .contains(&DEFAULT_SCHEMA_NAME.to_string())
4000 );
4001 for selectors in [
4002 vec!["physical_metrics".to_string()],
4003 vec!["missing".to_string(), "physical_metrics".to_string()],
4004 ] {
4005 assert!(
4006 retrieve_schema_names(&query_ctx, &handler, selectors)
4007 .await
4008 .unwrap()
4009 .is_empty()
4010 );
4011 }
4012
4013 assert_eq!(
4014 vec!["cpu_usage".to_string()],
4015 retrieve_table_names(
4016 &query_ctx,
4017 manager.clone(),
4018 vec!["missing".to_string(), r#"{__name__=~"cpu.*"}"#.to_string(),],
4019 )
4020 .await
4021 .unwrap()
4022 );
4023 assert!(
4024 retrieve_table_names(
4025 &query_ctx,
4026 manager.clone(),
4027 vec!["cpu_usage".to_string(), "{".to_string()],
4028 )
4029 .await
4030 .is_err()
4031 );
4032
4033 let fields =
4034 retrieve_field_names(&query_ctx, &handler, vec!["physical_metrics".to_string()])
4035 .await
4036 .unwrap();
4037 assert!(fields.is_empty());
4038 let fields = retrieve_field_names(&query_ctx, &handler, vec![])
4039 .await
4040 .unwrap();
4041 assert_eq!(fields, HashSet::from(["value".to_string()]));
4042 }
4043
4044 #[tokio::test]
4045 async fn test_get_target_column_names_uses_selected_schema() {
4046 let manager = MemoryCatalogManager::with_default_setup();
4047 manager
4048 .register_schema_sync(RegisterSchemaRequest {
4049 catalog: DEFAULT_CATALOG_NAME.to_string(),
4050 schema: "private".to_string(),
4051 })
4052 .unwrap();
4053
4054 for (schema_name, tag_name, table_id) in [
4055 (DEFAULT_SCHEMA_NAME, "public_tag", 2048),
4056 ("private", "private_tag", 2049),
4057 ] {
4058 let schema = Arc::new(Schema::new(vec![
4059 ColumnSchema::new(
4060 "greptime_timestamp",
4061 ConcreteDataType::timestamp_millisecond_datatype(),
4062 false,
4063 )
4064 .with_time_index(true),
4065 ColumnSchema::new(tag_name, ConcreteDataType::string_datatype(), false),
4066 ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true),
4067 ]));
4068 let meta = TableMetaBuilder::empty()
4069 .schema(schema)
4070 .primary_key_indices(vec![1])
4071 .engine("metric".to_string())
4072 .next_column_id(3)
4073 .build()
4074 .unwrap();
4075 let table_info = TableInfoBuilder::default()
4076 .table_id(table_id)
4077 .table_version(0 as TableVersion)
4078 .name("cpu")
4079 .catalog_name(DEFAULT_CATALOG_NAME)
4080 .schema_name(schema_name)
4081 .table_type(TableType::Base)
4082 .meta(meta)
4083 .build()
4084 .unwrap();
4085 manager
4086 .register_table_sync(RegisterTableRequest {
4087 catalog: DEFAULT_CATALOG_NAME.to_string(),
4088 schema: schema_name.to_string(),
4089 table_name: "cpu".to_string(),
4090 table_id,
4091 table: EmptyTable::from_table_info(&table_info),
4092 })
4093 .unwrap();
4094 }
4095
4096 let query_ctx = Arc::new(QueryContext::with(
4097 DEFAULT_CATALOG_NAME,
4098 DEFAULT_SCHEMA_NAME,
4099 ));
4100 let queries = [ParsedPromQuery::parse(
4101 PromQuery {
4102 query: r#"{__name__="cpu",__schema__="private"}"#.to_string(),
4103 ..Default::default()
4104 },
4105 &query_ctx,
4106 )
4107 .unwrap()];
4108 let targets = resolved_promql_targets(&queries, &query_ctx).unwrap();
4109 let manager: CatalogManagerRef = manager;
4110
4111 assert_eq!(
4112 get_target_column_names(&targets, &manager, &query_ctx)
4113 .await
4114 .unwrap(),
4115 HashSet::from(["private_tag".to_string()])
4116 );
4117 }
4118
4119 #[test]
4120 fn test_metric_name_discovery_uses_selected_schema() {
4121 let expr =
4122 promql_parser::parser::parse(r#"{__name__=~"cpu.*",__database__="private"}"#).unwrap();
4123 let discovery = find_metric_name_not_equal_matchers(&expr).unwrap().unwrap();
4124
4125 assert_eq!(discovery.schema.as_deref(), Some("private"));
4126 assert_eq!(discovery.name_matchers.len(), 1);
4127 assert_eq!(discovery.name_matchers[0].value, "cpu.*");
4128 assert!(matches!(discovery.name_matchers[0].op, MatchOp::Re(_)));
4129
4130 let expr = promql_parser::parser::parse(
4131 r#"{__name__=~"cpu.*",__database__="private",__schema__="public"}"#,
4132 )
4133 .unwrap();
4134 let discovery = find_metric_name_not_equal_matchers(&expr).unwrap().unwrap();
4135 assert_eq!(discovery.schema.as_deref(), Some("public"));
4136 }
4137
4138 #[test]
4139 fn test_metric_name_discovery_rejects_unsafe_selector_shapes() {
4140 for promql in [
4141 r#"{__name__="denied",__name__=~"no_such_.*"}"#,
4142 r#"{__name__=~"cpu.*" or job="api"}"#,
4143 ] {
4144 let expr = promql_parser::parser::parse(promql).unwrap();
4145 assert!(
4146 find_metric_name_not_equal_matchers(&expr)
4147 .unwrap()
4148 .is_none(),
4149 "{promql}"
4150 );
4151 }
4152
4153 let expr = promql_parser::parser::parse(r#"{__name__=~"cpu.*" or job="api"}"#).unwrap();
4154 assert_eq!(
4155 (PermissionTableTargets::Resolved(vec![]), 1),
4156 static_promql_targets(&expr, &QueryContext::arc()).unwrap()
4157 );
4158
4159 let expr =
4160 promql_parser::parser::parse(r#"allowed{job="api" or instance="host"}"#).unwrap();
4161 assert_eq!(
4162 (
4163 PermissionTableTargets::Resolved(vec![PermissionTableTarget::new(
4164 "greptime", "public", "allowed",
4165 )]),
4166 0,
4167 ),
4168 static_promql_targets(&expr, &QueryContext::arc()).unwrap()
4169 );
4170 }
4171
4172 #[test]
4173 fn test_metric_name_discovery_ignores_or_matchers_on_named_selector() {
4174 let expr = promql_parser::parser::parse(
4175 r#"{__name__=~"cpu.*"} + allowed{job="api" or instance="host"}"#,
4176 )
4177 .unwrap();
4178
4179 let discovery = find_metric_name_not_equal_matchers(&expr).unwrap().unwrap();
4180 assert_eq!(discovery.name_matchers[0].value, "cpu.*");
4181 }
4182
4183 #[test]
4184 fn test_retrieve_exact_metric_name_from_selector() {
4185 assert_eq!(
4186 retrieve_exact_metric_name_from_promql(r#"cpu{host="a"}"#).as_deref(),
4187 Some("cpu")
4188 );
4189 assert_eq!(
4190 retrieve_exact_metric_name_from_promql(r#"{__name__="cpu",host="a"}"#).as_deref(),
4191 Some("cpu")
4192 );
4193 assert!(retrieve_exact_metric_name_from_promql(r#"{__name__=~"cpu.*"}"#).is_none());
4194 }
4195
4196 #[test]
4197 fn test_expand_metric_name_queries_preserves_schema_matcher() {
4198 let query = PromQuery {
4199 query: r#"{__name__=~"cpu.*",__database__="private"}"#.to_string(),
4200 ..Default::default()
4201 };
4202 let ctx = QueryContext::arc();
4203 let query = ParsedPromQuery::parse(query, &ctx).unwrap();
4204 let queries = expand_metric_name_queries(
4205 &query,
4206 vec!["cpu_user".to_string(), "cpu_system".to_string()],
4207 );
4208
4209 assert_eq!(queries.len(), 2);
4210 for (query, metric_name) in queries.iter().zip(["cpu_user", "cpu_system"]) {
4211 let PromqlExpr::VectorSelector(selector) = query.expr() else {
4212 panic!("expected vector selector");
4213 };
4214 assert!(selector.matchers.matchers.iter().any(|matcher| {
4215 matcher.name == METRIC_NAME
4216 && matcher.op == MatchOp::Equal
4217 && matcher.value == metric_name
4218 }));
4219 assert!(selector.matchers.matchers.iter().any(|matcher| {
4220 matcher.name == "__database__"
4221 && matcher.op == MatchOp::Equal
4222 && matcher.value == "private"
4223 }));
4224 }
4225 }
4226
4227 #[test]
4228 fn test_expand_metric_name_queries_preserves_exact_selector() {
4229 let query = PromQuery {
4230 query: r#"denied + {__name__=~"cpu.*"}"#.to_string(),
4231 ..Default::default()
4232 };
4233 let ctx = QueryContext::arc();
4234 let query = ParsedPromQuery::parse(query, &ctx).unwrap();
4235 let queries = expand_metric_name_queries(&query, vec!["cpu_user".to_string()]);
4236
4237 assert_eq!(queries.len(), 1);
4238 let expanded = queries[0].expr();
4239 let ctx = Arc::new(QueryContext::with("greptime", "public"));
4240 let (PermissionTableTargets::Resolved(mut targets), unresolved_selectors) =
4241 static_promql_targets(expanded, &ctx).unwrap()
4242 else {
4243 panic!("expected resolved targets");
4244 };
4245 assert_eq!(unresolved_selectors, 0);
4246 targets.sort_unstable_by(|left, right| left.table.cmp(&right.table));
4247 assert_eq!(
4248 targets,
4249 vec![
4250 PermissionTableTarget::new("greptime", "public", "cpu_user"),
4251 PermissionTableTarget::new("greptime", "public", "denied"),
4252 ]
4253 );
4254 }
4255
4256 #[test]
4257 fn test_static_promql_targets_keep_exact_selector_during_discovery() {
4258 let expr = promql_parser::parser::parse(
4259 r#"denied + {__name__=~"no_match_.*",__schema__="private"}"#,
4260 )
4261 .unwrap();
4262 let ctx = Arc::new(QueryContext::with("greptime", "public"));
4263
4264 assert_eq!(
4265 static_promql_targets(&expr, &ctx).unwrap(),
4266 (
4267 PermissionTableTargets::resolved(vec![PermissionTableTarget::new(
4268 "greptime", "public", "denied",
4269 )]),
4270 1,
4271 )
4272 );
4273 }
4274
4275 #[test]
4276 fn test_static_promql_targets_count_unresolved_selectors() {
4277 let expr =
4278 promql_parser::parser::parse(r#"{__name__=~"no_such_metric"} or {job="api"}"#).unwrap();
4279 let ctx = Arc::new(QueryContext::with("greptime", "public"));
4280
4281 assert_eq!(
4282 static_promql_targets(&expr, &ctx).unwrap(),
4283 (PermissionTableTargets::resolved(Vec::new()), 2)
4284 );
4285 }
4286
4287 #[test]
4288 fn test_record_batches_to_labels_name_accepts_native_histogram_value() {
4289 let schema = Arc::new(Schema::new(vec![
4290 ColumnSchema::new("host", ConcreteDataType::string_datatype(), false),
4291 ColumnSchema::new(
4292 greptime_native_histogram(),
4293 native_histogram_value_type().clone(),
4294 true,
4295 ),
4296 ]));
4297 let batch = RecordBatch::new_empty(schema.clone());
4298 let batches = RecordBatches::try_new(schema.clone(), vec![batch]).unwrap();
4299 let mut labels = HashSet::new();
4300
4301 record_batches_to_labels_name(batches, &mut labels).unwrap();
4302
4303 assert!(labels.is_empty());
4304
4305 let histogram = NativeHistogram {
4306 schema: 0,
4307 zero_threshold: 0.001,
4308 sum: 1.0,
4309 reset_hint: CounterResetHint::Unknown,
4310 start_timestamp: None,
4311 custom_values: vec![],
4312 positive_spans: vec![],
4313 negative_spans: vec![],
4314 count: 1.0,
4315 zero_count: 1.0,
4316 positive_buckets: vec![],
4317 negative_buckets: vec![],
4318 };
4319 let histogram_array = build_histogram_array(&[Some(histogram)]);
4320 let histogram_array = histogram_array
4321 .as_any()
4322 .downcast_ref::<arrow::array::StructArray>()
4323 .unwrap()
4324 .clone();
4325 let ConcreteDataType::Struct(histogram_type) = native_histogram_value_type().clone() else {
4326 unreachable!("native histogram type must be a struct")
4327 };
4328 let batch = RecordBatch::new(
4329 schema.clone(),
4330 vec![
4331 Arc::new(StringVector::from(vec![Some("localhost")])) as _,
4332 Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) as _,
4333 ],
4334 )
4335 .unwrap();
4336 let batches = RecordBatches::try_new(schema, vec![batch]).unwrap();
4337
4338 record_batches_to_labels_name(batches, &mut labels).unwrap();
4339
4340 assert_eq!(
4341 labels,
4342 HashSet::from(["host".to_string(), greptime_native_histogram().to_string(),])
4343 );
4344 }
4345
4346 fn prometheus_metric_table_info(table_id: u32, name: &str) -> TableInfo {
4347 let schema = Arc::new(Schema::new(vec![
4348 ColumnSchema::new(
4349 "greptime_timestamp",
4350 ConcreteDataType::timestamp_millisecond_datatype(),
4351 false,
4352 )
4353 .with_time_index(true),
4354 ColumnSchema::new("host", ConcreteDataType::string_datatype(), false),
4355 ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true),
4356 ]));
4357 let mut options = TableOptions::default();
4358 options.extra_options.insert(
4359 LOGICAL_TABLE_METADATA_KEY.to_string(),
4360 "greptime_physical_table".to_string(),
4361 );
4362 options
4363 .extra_options
4364 .insert(SEMANTIC_METRIC_TYPE.to_string(), "counter".to_string());
4365 options
4366 .extra_options
4367 .insert(SEMANTIC_METRIC_UNIT.to_string(), "By".to_string());
4368 let meta = TableMetaBuilder::empty()
4369 .schema(schema)
4370 .primary_key_indices(vec![1])
4371 .engine("metric".to_string())
4372 .next_column_id(3)
4373 .options(options)
4374 .build()
4375 .unwrap();
4376
4377 TableInfoBuilder::default()
4378 .table_id(table_id)
4379 .table_version(0 as TableVersion)
4380 .name(name)
4381 .catalog_name(DEFAULT_CATALOG_NAME)
4382 .schema_name(DEFAULT_SCHEMA_NAME)
4383 .table_type(TableType::Base)
4384 .meta(meta)
4385 .build()
4386 .unwrap()
4387 }
4388
4389 #[tokio::test]
4390 async fn test_retrieve_metric_metadata_uses_semantic_options() {
4391 let table_info = prometheus_metric_table_info(1025, "http_requests_total");
4392 let manager: CatalogManagerRef =
4393 MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&table_info));
4394 let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4395
4396 let metadata = retrieve_metric_metadata(&query_ctx, manager, &MetadataQuery::default())
4397 .await
4398 .unwrap();
4399
4400 assert_eq!(
4401 metadata.get("http_requests_total"),
4402 Some(&vec![PromMetadata {
4403 metric_type: "counter".to_string(),
4404 unit: "bytes".to_string(),
4405 help: String::new(),
4406 }])
4407 );
4408
4409 for (ucum, openmetrics) in [
4410 ("1", "ratios"),
4411 ("ms", "milliseconds"),
4412 ("m/s", "meters_per_second"),
4413 ("USD", "USD"),
4414 ] {
4415 let mut table_info = table_info.clone();
4416 table_info
4417 .meta
4418 .options
4419 .extra_options
4420 .insert(SEMANTIC_METRIC_UNIT.to_string(), ucum.to_string());
4421 assert_eq!(
4422 prometheus_metadata_from_table(&table_info).unit,
4423 openmetrics
4424 );
4425 }
4426
4427 for (metric_type, temporality, expected) in [
4428 ("updown_counter", None, "gauge"),
4429 ("gauge_histogram", None, "gaugehistogram"),
4430 ("summary", None, "summary"),
4431 ("mixed", None, "unknown"),
4432 ("counter", Some("delta"), "unknown"),
4433 ("histogram", Some("delta"), "unknown"),
4434 ("counter", Some("mixed"), "unknown"),
4435 ("histogram", Some("mixed"), "unknown"),
4436 ("updown_counter", Some("mixed"), "unknown"),
4437 ] {
4438 let mut table_info = table_info.clone();
4439 table_info
4440 .meta
4441 .options
4442 .extra_options
4443 .insert(SEMANTIC_METRIC_TYPE.to_string(), metric_type.to_string());
4444 if let Some(temporality) = temporality {
4445 table_info.meta.options.extra_options.insert(
4446 SEMANTIC_METRIC_TEMPORALITY.to_string(),
4447 temporality.to_string(),
4448 );
4449 }
4450 assert_eq!(
4451 prometheus_metadata_from_table(&table_info).metric_type,
4452 expected
4453 );
4454 }
4455 }
4456
4457 #[tokio::test]
4458 async fn test_metadata_query_checks_permissions_and_filters_tables() {
4459 let allowed = prometheus_metric_table_info(1025, "allowed");
4460 let denied = prometheus_metric_table_info(1026, "denied");
4461 let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&allowed));
4462 manager
4463 .register_table_sync(RegisterTableRequest {
4464 catalog: DEFAULT_CATALOG_NAME.to_string(),
4465 schema: DEFAULT_SCHEMA_NAME.to_string(),
4466 table_name: denied.name.clone(),
4467 table_id: denied.table_id(),
4468 table: EmptyTable::from_table_info(&denied),
4469 })
4470 .unwrap();
4471 let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4472
4473 let metadata = retrieve_metric_metadata(
4474 &query_ctx,
4475 manager.clone(),
4476 &MetadataQuery {
4477 metric: Some("allowed".to_string()),
4478 ..Default::default()
4479 },
4480 )
4481 .await
4482 .unwrap();
4483 assert_eq!(metadata.len(), 1);
4484 assert!(metadata.contains_key("allowed"));
4485
4486 let response = metadata_query(
4487 None,
4488 State(Arc::new(TestPrometheusHandler {
4489 catalog_manager: manager.clone(),
4490 deny_operation: true,
4491 denied_table: None,
4492 metric_names: Vec::new(),
4493 label_metric_names: Vec::new(),
4494 label_lookups: Mutex::new(Vec::new()),
4495 queries: Mutex::new(Vec::new()),
4496 ordered_outputs: Mutex::new(Vec::new()),
4497 })),
4498 Query(MetadataQuery::default()),
4499 Extension(query_ctx.clone()),
4500 )
4501 .await;
4502 assert_eq!(Some(StatusCode::PermissionDenied), response.status_code);
4503
4504 let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler {
4505 catalog_manager: manager,
4506 deny_operation: false,
4507 denied_table: Some("denied"),
4508 metric_names: Vec::new(),
4509 label_metric_names: Vec::new(),
4510 label_lookups: Mutex::new(Vec::new()),
4511 queries: Mutex::new(Vec::new()),
4512 ordered_outputs: Mutex::new(Vec::new()),
4513 });
4514 let response = metadata_query(
4515 None,
4516 State(handler.clone()),
4517 Query(MetadataQuery::default()),
4518 Extension(query_ctx.clone()),
4519 )
4520 .await;
4521 assert!(response.status_code.is_none());
4522 let PrometheusResponse::Metadata(metadata) = response.data else {
4523 panic!("expected metadata response");
4524 };
4525 assert_eq!(
4526 metadata.keys().cloned().collect::<Vec<_>>(),
4527 vec!["allowed".to_string()]
4528 );
4529
4530 let response = metadata_query(
4531 None,
4532 State(handler),
4533 Query(MetadataQuery {
4534 metric: Some("denied".to_string()),
4535 ..Default::default()
4536 }),
4537 Extension(query_ctx),
4538 )
4539 .await;
4540 assert_eq!(Some(StatusCode::PermissionDenied), response.status_code);
4541 }
4542}