Skip to main content

servers/http/
prometheus.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! prom supply the prometheus HTTP API Server compliance
16
17use 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/// For [ValueType::Vector] result type
87#[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/// For [ValueType::Matrix] result type
98#[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/// A native histogram sample in the Prometheus HTTP API JSON format.
109#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
110pub struct PromNativeHistogram {
111    /// Total number of observations, encoded as a string.
112    pub count: String,
113    /// Sum of all observations, encoded as a string.
114    pub sum: String,
115    /// Populated buckets as `(boundary rule, lower bound, upper bound, count)`.
116    ///
117    /// Boundary rules are `0` for lower-open/upper-closed, `1` for
118    /// lower-closed/upper-open, `2` for both bounds open, and `3` for both
119    /// bounds closed. Bounds and counts are encoded as strings.
120    pub buckets: Vec<(u8, String, String, String)>,
121}
122
123/// Variants corresponding to [ValueType]
124#[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/// Metadata for a Prometheus metric family.
147#[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/// A "holder" for the reference([Arc]) to a column name,
156/// to help avoiding cloning [String]s when used as a [HashMap] key.
157#[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    /// Append the other [`PrometheusResponse]`.
184    /// # NOTE
185    ///   Only append matrix and vector results, otherwise just ignore the other response.
186    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                // TODO(dennis): process other cases?
215            }
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/// Query parameters for the Prometheus metric metadata endpoint.
230#[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(&params.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) = &params.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(), &params).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
344/// Helper macro which try to evaluate the expression and return its results.
345/// If the evaluation fails, return a `PrometheusJsonResponse` early.
346macro_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    // Extract time from query string, or use current server time if not specified.
389    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    // update catalog and schema in query context if necessary
406    if let Some(db) = &params.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
489/// Executes a single instant query and returns response
490async 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    // update catalog and schema in query context if necessary
537    if let Some(db) = &params.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    // Prometheus rejects a range-vector expression in a range query even when start == end,
551    // which the planner cannot tell apart from an instant query.
552    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        // Safety: at least one responses, checked above
613        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
625/// Executes a single range query and returns response
626async 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    // Matrix serialization sorts samples and series, so execution order never
634    // reaches the response.
635    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
660// Custom Deserialize method to support parsing repeated match[]
661impl<'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
692/// Handles schema errors, transforming a Result into an Option.
693/// - If the input is `Ok(v)`, returns `Some(v)`
694/// - If the input is `Err(err)` and the error status code is `TableNotFound` or
695///   `TableColumnNotFound`, returns `None` (ignoring these specific errors)
696/// - If the input is `Err(err)` with any other error code, directly returns a
697///   `PrometheusJsonResponse::error`.
698macro_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                    // Prometheus won't report error if querying nonexist label and metric
707                    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(&params.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    // Fetch all columns if no query matcher is provided
782    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    // Fetch tag columns only from the tables checked above.
803    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    // intersect `fetched_labels` with `labels` to filter out non-tag columns
830    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
842/// Get all tag column name of the given schema
843async 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    // fetch tag list
904    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
944/// Retrieve labels name from query result
945async 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        // project record batch to only contains tag columns
978        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
1001/// Writer from a row in the record batch to a Prometheus time series:
1002///
1003/// `{__name__="<metric name>", <label name>="<label value>", ...}`
1004///
1005/// The metrics name is the table name; label names are the column names and
1006/// the label values are the corresponding row values (all are converted to strings).
1007struct RowWriter {
1008    /// The template that is to produce a Prometheus time series. It is pre-filled with metrics name
1009    /// and label names, waiting to be filled by row values afterward.
1010    template: HashMap<Column, Option<String>>,
1011    /// The current filling row.
1012    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                        // `Date64` values are milliseconds representation of `Date32` values,
1160                        // according to its specification. So we convert the `Date64` value here to
1161                        // the `Date32` value to process them unified.
1162                        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
1267/// Retrieve labels name from record batches
1268fn 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 all value columns are null, skip this row
1300            if value_columns.iter().all(|c| c.is_null(row_index)) {
1301                continue;
1302            }
1303
1304            // if a value is not null, record the tag name and return
1305            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
1327/// Tries to get catalog and schema from an optional db param. And retrieves
1328/// them from [QueryContext] if they don't present.
1329pub(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
1340/// Update catalog and schema in [QueryContext] if necessary.
1341pub(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    // Return the metric name only if there's exactly one unique metric name
1375    if metric_names.len() == 1 {
1376        metric_names.into_iter().next()
1377    } else {
1378        None
1379    }
1380}
1381
1382/// Whether the metric name survives a binary operation; see
1383/// [`binary_metric_name_side`] for the operand that supplies it.
1384///
1385/// Follows Prometheus' `shouldDropMetricName` and `resultMetric`: set operators
1386/// and comparisons keep the name, while arithmetic operators drop it. A
1387/// comparison also drops it when it returns a bool or is a one-to-one
1388/// `on(...)` match that does not keep `__name__` or an `ignoring(...)` match
1389/// that excludes it. Prometheus keeps only the `on(...)` labels or deletes the
1390/// `ignoring(...)` labels for one-to-one matches.
1391fn 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
1427/// The operand whose labels supply the result metric name of a binary
1428/// operation.
1429///
1430/// Prometheus builds the result from the side whose labels it carries, i.e. the
1431/// "many" side of the match: `group_right` makes the right-hand operand that
1432/// side, and a scalar left-hand operand (for example `0.5 < a`) leaves the
1433/// vector on the right. Every other case takes the left-hand operand.
1434fn 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
1448/// Recursively collect all metric names from a PromQL expression
1449fn 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
1543/// Finds non-equality `__name__` matchers and their unambiguous schema.
1544fn 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        // Has name, ignore the matchers
1555        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
1691/// Update the `__name__` matchers in expression into special value
1692/// Returns the updated expression.
1693fn 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(&params.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        // An empty `match[]` enumerates every metric; otherwise only the
1789        // selectors answerable from metadata go down that path.
1790        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        // Only use and filter matchers.
1956        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
1981/// Take metric name from the [VectorSelector].
1982/// It takes the name in the selector or removes the name matcher.
1983fn 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    // We need to remove the name matcher to avoid using it as a filter in query.
1995    selector.matchers.matchers.remove(pos);
1996
1997    Some(name)
1998}
1999
2000/// Removes every `__name__` matcher from the selector and returns them, so the
2001/// rest can be planned as column predicates. A name given as `VectorSelector::name`
2002/// comes back as an equality matcher, making both spellings filter alike.
2003fn 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
2018/// Whether a metric name satisfies every `__name__` matcher of one selector.
2019///
2020/// Negated matchers are honoured here, unlike in [`retrieve_table_names`] where
2021/// they keep every table so the caller authorizes the full candidate set: the
2022/// names reaching this point are already narrowed by the data, so filtering them
2023/// can only remove names, never widen what a caller gets to see.
2024fn 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
2033/// Whether a matcher constrains an ordinary label. The others name the metric,
2034/// the database or the field, none of which is a column to scan.
2035fn 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
2041/// Splits `match[]` selectors by whether they constrain an ordinary label. The
2042/// first group is answerable from table metadata; the second needs the data read
2043/// and is returned as parsed selectors.
2044///
2045/// `or` matchers stay in the metadata group, which ignores them, rather than
2046/// being silently dropped from a data scan that cannot express them.
2047fn 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
2077/// Resolves selectors constraining ordinary labels into metric names: the data
2078/// answers the label matchers, then each selector's `__name__` matchers narrow
2079/// the names it found.
2080async 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        // The database and field matchers name no column, and the metadata path
2105        // ignores them too.
2106        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            // skip non-prometheus (non-metricengine) tables for __name__ query
2172            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                    // != and !~ are not supported, so include every table and
2184                    // let the caller authorize the complete candidate set.
2185                    _ => 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) = &params.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        // TODO: Persist and return Prometheus help text and OTLP metric descriptions.
2290        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        // query all tables if no matcher is provided
2315        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    // update catalog and schema in query context if necessary
2460    if let Some(db) = &params.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            // TODO: find a better value for step
2477            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        /// Names the label-matcher path resolves, kept apart from `metric_names`
2636        /// so a test can tell which path answered.
2637        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        // Both a single-point and a multi-step range query take the same path.
2813        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        // `sort()` stays observable for instant queries, but not for range queries.
2833        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    /// A handler over the logical metric tables `cpu_user` and `cpu_system`,
2990    /// which is what the metadata path enumerates, with the label-matcher path
2991    /// answering `label_metric_names`.
2992    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        // The metadata path would have enumerated both metrics.
3067        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            // Single metric cases
3218            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            // Aggregation and function cases
3254            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            // Same metric binary operations
3283            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            // Multiple metrics cases
3298            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            // Unless operator cases
3313            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            // Arithmetic operators always drop the metric name.
3328            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            // Comparisons keep the metric name from the side whose labels the
3343            // result carries, i.e. the "many" side of the match: normally the
3344            // left operand, but `group_right` takes it from the right operand and
3345            // a scalar left operand leaves the vector on the right. The name is
3346            // dropped when the comparison returns a bool or is a one-to-one
3347            // `on(...)` match that does not keep `__name__` or an `ignoring(...)`
3348            // match that excludes it. A `group_left`/`group_right` modifier does
3349            // not reduce the labels, so the name is kept.
3350            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            // Subquery cases
3456            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            // Literal values
3471            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            // Error cases
3486            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}