frontend/instance/
promql.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
15use std::time::SystemTime;
16
17use auth::PermissionTableTarget;
18use catalog::information_schema::TABLES;
19use client::OutputData;
20use common_catalog::consts::INFORMATION_SCHEMA_NAME;
21use common_catalog::format_full_table_name;
22use common_recordbatch::util;
23use common_telemetry::tracing;
24use promql_parser::label::{Matcher, Matchers};
25use query::promql;
26use query::promql::planner::PromPlanner;
27use servers::prometheus;
28use session::context::QueryContextRef;
29use snafu::{OptionExt, ResultExt};
30
31use crate::error::{
32    CatalogSnafu, CollectRecordbatchSnafu, ExecLogicalPlanSnafu,
33    PrometheusLabelValuesQueryPlanSnafu, PrometheusMetricNamesQueryPlanSnafu, ReadTableSnafu,
34    Result, TableNotFoundSnafu, TableSnafu,
35};
36use crate::instance::Instance;
37
38impl Instance {
39    /// Handles metric names query request, returns the names.
40    #[tracing::instrument(skip_all)]
41    pub(crate) async fn handle_query_metric_names(
42        &self,
43        matchers: Vec<Matcher>,
44        schema: &str,
45        ctx: &QueryContextRef,
46    ) -> Result<Vec<String>> {
47        let _timer = crate::metrics::PROMQL_QUERY_METRICS_ELAPSED
48            .with_label_values(&[ctx.get_db_string().as_str()])
49            .start_timer();
50
51        let table = self
52            .catalog_manager
53            .table(
54                ctx.current_catalog(),
55                INFORMATION_SCHEMA_NAME,
56                TABLES,
57                Some(ctx),
58            )
59            .await
60            .context(CatalogSnafu)?
61            .with_context(|| TableNotFoundSnafu {
62                table_name: "greptime.information_schema.tables",
63            })?;
64
65        let dataframe = self
66            .query_engine
67            .read_table(table)
68            .with_context(|_| ReadTableSnafu {
69                table_name: "greptime.information_schema.tables",
70            })?;
71
72        let logical_plan =
73            prometheus::metric_name_matchers_to_plan(dataframe, matchers, schema, ctx)
74                .context(PrometheusMetricNamesQueryPlanSnafu)?;
75
76        let results = self
77            .query_engine
78            .execute(logical_plan, ctx.clone())
79            .await
80            .context(ExecLogicalPlanSnafu)?;
81
82        let batches = match results.data {
83            OutputData::Stream(stream) => util::collect(stream)
84                .await
85                .context(CollectRecordbatchSnafu)?,
86            OutputData::RecordBatches(rbs) => rbs.take(),
87            _ => unreachable!("should not happen"),
88        };
89
90        let mut results = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum());
91
92        for batch in batches {
93            // Only one column the results, ensured by `prometheus::metric_name_matchers_to_plan`.
94            batch
95                .iter_column_as_string(0)
96                .flatten()
97                .for_each(|x| results.push(x))
98        }
99
100        Ok(results)
101    }
102
103    /// Handles label values query request, returns the values.
104    #[tracing::instrument(skip_all)]
105    pub(crate) async fn handle_query_label_values(
106        &self,
107        target: PermissionTableTarget,
108        label_name: String,
109        matchers: Vec<Matcher>,
110        start: SystemTime,
111        end: SystemTime,
112        ctx: &QueryContextRef,
113    ) -> Result<Vec<String>> {
114        let full_table_name =
115            format_full_table_name(&target.catalog, &target.schema, &target.table);
116        let table = self
117            .catalog_manager
118            .table(&target.catalog, &target.schema, &target.table, Some(ctx))
119            .await
120            .context(CatalogSnafu)?
121            .with_context(|| TableNotFoundSnafu {
122                table_name: full_table_name.clone(),
123            })?;
124
125        // Check label column existence before building the query plan so a missing label can be
126        // reported as `TableColumnNotFound` and handled like Prometheus expects.
127        if table.schema().column_schema_by_name(&label_name).is_none() {
128            return table::error::ColumnNotExistsSnafu {
129                column_name: label_name,
130                table_name: full_table_name,
131            }
132            .fail()
133            .context(TableSnafu);
134        }
135
136        let dataframe = self
137            .query_engine
138            .read_table(table.clone())
139            .with_context(|_| ReadTableSnafu {
140                table_name: full_table_name,
141            })?;
142
143        let scan_plan = dataframe.into_unoptimized_plan();
144        let filter_conditions =
145            PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema())
146                .context(PrometheusLabelValuesQueryPlanSnafu)?;
147        let logical_plan = promql::label_values::rewrite_label_values_query(
148            table,
149            scan_plan,
150            filter_conditions,
151            label_name,
152            start,
153            end,
154        )
155        .context(PrometheusLabelValuesQueryPlanSnafu)?;
156
157        let results = self
158            .query_engine
159            .execute(logical_plan, ctx.clone())
160            .await
161            .context(ExecLogicalPlanSnafu)?;
162
163        let batches = match results.data {
164            OutputData::Stream(stream) => util::collect(stream)
165                .await
166                .context(CollectRecordbatchSnafu)?,
167            OutputData::RecordBatches(rbs) => rbs.take(),
168            _ => unreachable!("should not happen"),
169        };
170
171        let mut results = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum());
172        for batch in batches {
173            // Only one column in results, ensured by `prometheus::label_values_matchers_to_plan`.
174            batch
175                .iter_column_as_string(0)
176                .flatten()
177                .for_each(|x| results.push(x))
178        }
179
180        Ok(results)
181    }
182}
frontend/instance/promql.rs

frontend/instance/
promql.rs
