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