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 names from non-physical 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    /// Removes schemas the user cannot access from metadata enumeration results.
180    async fn filter_metadata_schema_names(
181        &self,
182        schema_names: Vec<String>,
183        query_ctx: &QueryContextRef,
184    ) -> Result<Vec<String>>;
185
186    /// Query metric table names by the `__name__` matchers.
187    async fn query_metric_names(
188        &self,
189        matchers: Vec<Matcher>,
190        schema: &str,
191        ctx: &QueryContextRef,
192    ) -> Result<Vec<String>>;
193
194    /// Query metric table names that carry data matching `matchers` in the time
195    /// range. `matchers` must hold only ordinary label matchers: `__name__`
196    /// names a table and the database and field matchers name no column, so the
197    /// caller resolves all three before calling.
198    ///
199    /// Only metric engine tables are covered.
200    async fn query_metric_names_by_labels(
201        &self,
202        matchers: Vec<Matcher>,
203        schema: &str,
204        start: SystemTime,
205        end: SystemTime,
206        ctx: &QueryContextRef,
207    ) -> Result<Vec<String>>;
208
209    async fn query_label_values(
210        &self,
211        metric: String,
212        label_name: String,
213        matchers: Vec<Matcher>,
214        start: SystemTime,
215        end: SystemTime,
216        ctx: &QueryContextRef,
217    ) -> Result<Vec<String>>;
218
219    fn catalog_manager(&self) -> CatalogManagerRef;
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn test_resolve_schema_from_matchers() {
228        assert_eq!(resolve_schema_from_matchers(&[]).unwrap(), None);
229
230        for label in ["__database__", "__schema__"] {
231            let matchers = [Matcher::new(MatchOp::Equal, label, "private")];
232            assert_eq!(
233                resolve_schema_from_matchers(&matchers).unwrap(),
234                Some("private".to_string())
235            );
236        }
237
238        let matchers = [Matcher::new(
239            MatchOp::Equal,
240            "x_greptime_database",
241            "private",
242        )];
243        assert_eq!(resolve_schema_from_matchers(&matchers).unwrap(), None);
244
245        let matchers = [
246            Matcher::new(MatchOp::Equal, "__database__", "private"),
247            Matcher::new(MatchOp::Equal, "__schema__", "public"),
248        ];
249        assert_eq!(
250            resolve_schema_from_matchers(&matchers).unwrap(),
251            Some("public".to_string())
252        );
253    }
254
255    #[test]
256    fn test_resolve_schema_from_matchers_rejects_unsafe_matchers() {
257        assert!(
258            resolve_schema_from_matchers(&[Matcher::new(
259                MatchOp::NotEqual,
260                "__database__",
261                "private",
262            )])
263            .is_err()
264        );
265        assert!(
266            resolve_schema_from_matchers(&[Matcher::new(MatchOp::Equal, "__database__", "",)])
267                .is_err()
268        );
269    }
270}
servers/prometheus_handler.rs

servers/
prometheus_handler.rs
