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