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 query_exemplars(
132        &self,
133        _query: crate::exemplar::query::Query,
134        _ctx: QueryContextRef,
135    ) -> Result<Vec<crate::exemplar::query::Series>> {
136        crate::error::NotSupportedSnafu {
137            feat: "query_exemplars",
138        }
139        .fail()
140    }
141
142    async fn do_query(&self, query: &PromQuery, query_ctx: QueryContextRef) -> Result<Output>;
143
144    /// Executes a query whose statement has already been parsed.
145    async fn do_query_parsed(
146        &self,
147        query: ParsedPromQuery,
148        query_ctx: QueryContextRef,
149    ) -> Result<Output> {
150        self.do_query(query.query(), query_ctx).await
151    }
152
153    /// Checks access to every table referenced by a batch of PromQL queries.
154    ///
155    /// An empty batch still checks the operation privilege.
156    async fn check_query_permission(
157        &self,
158        queries: &[PromQuery],
159        query_ctx: &QueryContextRef,
160    ) -> Result<()>;
161
162    /// Checks access without parsing the queries again.
163    async fn check_query_permission_parsed(
164        &self,
165        queries: &[ParsedPromQuery],
166        query_ctx: &QueryContextRef,
167    ) -> Result<()> {
168        let queries = queries
169            .iter()
170            .map(|query| query.query().clone())
171            .collect::<Vec<_>>();
172        self.check_query_permission(&queries, query_ctx).await
173    }
174
175    /// Checks access to targets resolved by Prometheus metadata APIs.
176    async fn check_query_target_permission(
177        &self,
178        targets: PermissionTableTargets,
179        query_ctx: &QueryContextRef,
180    ) -> Result<()>;
181
182    /// Removes inaccessible names from non-physical table metadata enumeration results.
183    async fn filter_metadata_metric_names(
184        &self,
185        metric_names: Vec<String>,
186        schema: &str,
187        query_ctx: &QueryContextRef,
188    ) -> Result<Vec<String>>;
189
190    /// Removes schemas the user cannot access from metadata enumeration results.
191    async fn filter_metadata_schema_names(
192        &self,
193        schema_names: Vec<String>,
194        query_ctx: &QueryContextRef,
195    ) -> Result<Vec<String>>;
196
197    /// Query metric table names by the `__name__` matchers.
198    async fn query_metric_names(
199        &self,
200        matchers: Vec<Matcher>,
201        schema: &str,
202        ctx: &QueryContextRef,
203    ) -> Result<Vec<String>>;
204
205    /// Query metric table names that carry data matching `matchers` in the time
206    /// range. `matchers` must hold only ordinary label matchers: `__name__`
207    /// names a table and the database and field matchers name no column, so the
208    /// caller resolves all three before calling.
209    ///
210    /// Only metric engine tables are covered.
211    async fn query_metric_names_by_labels(
212        &self,
213        matchers: Vec<Matcher>,
214        schema: &str,
215        start: SystemTime,
216        end: SystemTime,
217        ctx: &QueryContextRef,
218    ) -> Result<Vec<String>>;
219
220    async fn query_label_values(
221        &self,
222        metric: String,
223        label_name: String,
224        matchers: Vec<Matcher>,
225        start: SystemTime,
226        end: SystemTime,
227        ctx: &QueryContextRef,
228    ) -> Result<Vec<String>>;
229
230    fn catalog_manager(&self) -> CatalogManagerRef;
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn test_resolve_schema_from_matchers() {
239        assert_eq!(resolve_schema_from_matchers(&[]).unwrap(), None);
240
241        for label in ["__database__", "__schema__"] {
242            let matchers = [Matcher::new(MatchOp::Equal, label, "private")];
243            assert_eq!(
244                resolve_schema_from_matchers(&matchers).unwrap(),
245                Some("private".to_string())
246            );
247        }
248
249        let matchers = [Matcher::new(
250            MatchOp::Equal,
251            "x_greptime_database",
252            "private",
253        )];
254        assert_eq!(resolve_schema_from_matchers(&matchers).unwrap(), None);
255
256        let matchers = [
257            Matcher::new(MatchOp::Equal, "__database__", "private"),
258            Matcher::new(MatchOp::Equal, "__schema__", "public"),
259        ];
260        assert_eq!(
261            resolve_schema_from_matchers(&matchers).unwrap(),
262            Some("public".to_string())
263        );
264    }
265
266    #[test]
267    fn test_resolve_schema_from_matchers_rejects_unsafe_matchers() {
268        assert!(
269            resolve_schema_from_matchers(&[Matcher::new(
270                MatchOp::NotEqual,
271                "__database__",
272                "private",
273            )])
274            .is_err()
275        );
276        assert!(
277            resolve_schema_from_matchers(&[Matcher::new(MatchOp::Equal, "__database__", "",)])
278                .is_err()
279        );
280    }
281}
servers/prometheus_handler.rs

servers/
prometheus_handler.rs
