servers/
prometheus_handler.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::sync::Arc;
18use std::time::SystemTime;
19
20use async_trait::async_trait;
21use auth::PermissionTableTargets;
22use catalog::CatalogManagerRef;
23use common_query::Output;
24use promql_parser::label::{MatchOp, Matcher};
25use promql_parser::parser::Expr as PromqlExpr;
26use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
27use session::context::QueryContextRef;
28use snafu::ResultExt;
29
30use crate::error::{InvalidQuerySnafu, ParsePromQLSnafu, Result};
31
32pub const PROMETHEUS_API_VERSION: &str = "v1";
33
34const PROMQL_SCHEMA_MATCHERS: [&str; 2] = ["__database__", "__schema__"];
35
36pub type PrometheusHandlerRef = Arc<dyn PrometheusHandler + Send + Sync>;
37
38/// A Prometheus query paired with its parsed statement.
39#[derive(Debug, Clone)]
40pub struct ParsedPromQuery {
41    query: PromQuery,
42    statement: QueryStatement,
43    requires_output_ordering: bool,
44}
45
46impl ParsedPromQuery {
47    /// Parses a Prometheus query once for permission checking and execution.
48    pub fn parse(query: PromQuery, query_ctx: &QueryContextRef) -> Result<Self> {
49        let statement =
50            QueryLanguageParser::parse_promql(&query, query_ctx).with_context(|_| {
51                ParsePromQLSnafu {
52                    query: query.clone(),
53                }
54            })?;
55        Ok(Self {
56            query,
57            statement,
58            requires_output_ordering: true,
59        })
60    }
61
62    /// Returns the original query parameters.
63    pub fn query(&self) -> &PromQuery {
64        &self.query
65    }
66
67    /// Returns the parsed query statement.
68    pub fn statement(&self) -> &QueryStatement {
69        &self.statement
70    }
71
72    /// Returns whether the caller observes the query output in execution order.
73    /// When it does not, the executor may drop the plan's output sort.
74    pub fn requires_output_ordering(&self) -> bool {
75        self.requires_output_ordering
76    }
77
78    /// Marks the output order as irrelevant to the caller.
79    pub(crate) fn with_unordered_output(mut self) -> Self {
80        self.requires_output_ordering = false;
81        self
82    }
83
84    /// Returns the parsed PromQL expression.
85    pub fn expr(&self) -> &PromqlExpr {
86        let QueryStatement::Promql(eval_stmt, _) = &self.statement else {
87            unreachable!("query is parsed from PromQL")
88        };
89        &eval_stmt.expr
90    }
91
92    pub(crate) fn update_expr(&mut self, update: impl FnOnce(&mut PromqlExpr)) {
93        let QueryStatement::Promql(eval_stmt, _) = &mut self.statement else {
94            unreachable!("query is parsed from PromQL")
95        };
96        update(&mut eval_stmt.expr);
97        self.query.query = eval_stmt.expr.to_string();
98    }
99
100    /// Splits the parsed query into its original parameters and statement.
101    pub fn into_parts(self) -> (PromQuery, QueryStatement) {
102        (self.query, self.statement)
103    }
104}
105
106/// Resolves the optional schema selector shared by Prometheus query paths.
107///
108/// Like the PromQL planner, the last equality matcher selects the schema.
109pub fn resolve_schema_from_matchers(matchers: &[Matcher]) -> Result<Option<String>> {
110    let mut schema = None;
111    for matcher in matchers
112        .iter()
113        .filter(|matcher| PROMQL_SCHEMA_MATCHERS.contains(&matcher.name.as_str()))
114    {
115        if matcher.op != MatchOp::Equal || matcher.value.is_empty() {
116            return Err(InvalidQuerySnafu {
117                reason: format!(
118                    "expected a non-empty equality matcher for '{}'",
119                    matcher.name
120                ),
121            }
122            .build());
123        }
124        schema = Some(matcher.value.clone());
125    }
126    Ok(schema)
127}
128
129#[async_trait]
130pub trait PrometheusHandler {
131    async fn do_query(&self, query: &PromQuery, query_ctx: QueryContextRef) -> Result<Output>;
132
133    /// Executes a query whose statement has already been parsed.
134    async fn do_query_parsed(
135        &self,
136        query: ParsedPromQuery,
137        query_ctx: QueryContextRef,
138    ) -> Result<Output> {
139        self.do_query(query.query(), query_ctx).await
140    }
141
142    /// Checks access to every table referenced by a batch of PromQL queries.
143    ///
144    /// An empty batch still checks the operation privilege.
145    async fn check_query_permission(
146        &self,
147        queries: &[PromQuery],
148        query_ctx: &QueryContextRef,
149    ) -> Result<()>;
150
151    /// Checks access without parsing the queries again.
152    async fn check_query_permission_parsed(
153        &self,
154        queries: &[ParsedPromQuery],
155        query_ctx: &QueryContextRef,
156    ) -> Result<()> {
157        let queries = queries
158            .iter()
159            .map(|query| query.query().clone())
160            .collect::<Vec<_>>();
161        self.check_query_permission(&queries, query_ctx).await
162    }
163
164    /// Checks access to targets resolved by Prometheus metadata APIs.
165    async fn check_query_target_permission(
166        &self,
167        targets: PermissionTableTargets,
168        query_ctx: &QueryContextRef,
169    ) -> Result<()>;
170
171    /// Removes inaccessible metric names from logical-table metadata enumeration results.
172    async fn filter_metadata_metric_names(
173        &self,
174        metric_names: Vec<String>,
175        schema: &str,
176        query_ctx: &QueryContextRef,
177    ) -> Result<Vec<String>>;
178
179    /// Query metric table names by the `__name__` matchers.
180    async fn query_metric_names(
181        &self,
182        matchers: Vec<Matcher>,
183        schema: &str,
184        ctx: &QueryContextRef,
185    ) -> Result<Vec<String>>;
186
187    /// Query metric table names that carry data matching `matchers` in the time
188    /// range. `matchers` must hold only ordinary label matchers: `__name__`
189    /// names a table and the database and field matchers name no column, so the
190    /// caller resolves all three before calling.
191    ///
192    /// Only metric engine tables are covered.
193    async fn query_metric_names_by_labels(
194        &self,
195        matchers: Vec<Matcher>,
196        schema: &str,
197        start: SystemTime,
198        end: SystemTime,
199        ctx: &QueryContextRef,
200    ) -> Result<Vec<String>>;
201
202    async fn query_label_values(
203        &self,
204        metric: String,
205        label_name: String,
206        matchers: Vec<Matcher>,
207        start: SystemTime,
208        end: SystemTime,
209        ctx: &QueryContextRef,
210    ) -> Result<Vec<String>>;
211
212    fn catalog_manager(&self) -> CatalogManagerRef;
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[test]
220    fn test_resolve_schema_from_matchers() {
221        assert_eq!(resolve_schema_from_matchers(&[]).unwrap(), None);
222
223        for label in ["__database__", "__schema__"] {
224            let matchers = [Matcher::new(MatchOp::Equal, label, "private")];
225            assert_eq!(
226                resolve_schema_from_matchers(&matchers).unwrap(),
227                Some("private".to_string())
228            );
229        }
230
231        let matchers = [Matcher::new(
232            MatchOp::Equal,
233            "x_greptime_database",
234            "private",
235        )];
236        assert_eq!(resolve_schema_from_matchers(&matchers).unwrap(), None);
237
238        let matchers = [
239            Matcher::new(MatchOp::Equal, "__database__", "private"),
240            Matcher::new(MatchOp::Equal, "__schema__", "public"),
241        ];
242        assert_eq!(
243            resolve_schema_from_matchers(&matchers).unwrap(),
244            Some("public".to_string())
245        );
246    }
247
248    #[test]
249    fn test_resolve_schema_from_matchers_rejects_unsafe_matchers() {
250        assert!(
251            resolve_schema_from_matchers(&[Matcher::new(
252                MatchOp::NotEqual,
253                "__database__",
254                "private",
255            )])
256            .is_err()
257        );
258        assert!(
259            resolve_schema_from_matchers(&[Matcher::new(MatchOp::Equal, "__database__", "",)])
260                .is_err()
261        );
262    }
263}
servers/prometheus_handler.rs

servers/
prometheus_handler.rs
