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}
44
45impl ParsedPromQuery {
46    /// Parses a Prometheus query once for permission checking and execution.
47    pub fn parse(query: PromQuery, query_ctx: &QueryContextRef) -> Result<Self> {
48        let statement =
49            QueryLanguageParser::parse_promql(&query, query_ctx).with_context(|_| {
50                ParsePromQLSnafu {
51                    query: query.clone(),
52                }
53            })?;
54        Ok(Self { query, statement })
55    }
56
57    /// Returns the original query parameters.
58    pub fn query(&self) -> &PromQuery {
59        &self.query
60    }
61
62    /// Returns the parsed query statement.
63    pub fn statement(&self) -> &QueryStatement {
64        &self.statement
65    }
66
67    /// Returns the parsed PromQL expression.
68    pub fn expr(&self) -> &PromqlExpr {
69        let QueryStatement::Promql(eval_stmt, _) = &self.statement else {
70            unreachable!("query is parsed from PromQL")
71        };
72        &eval_stmt.expr
73    }
74
75    pub(crate) fn update_expr(&mut self, update: impl FnOnce(&mut PromqlExpr)) {
76        let QueryStatement::Promql(eval_stmt, _) = &mut self.statement else {
77            unreachable!("query is parsed from PromQL")
78        };
79        update(&mut eval_stmt.expr);
80        self.query.query = eval_stmt.expr.to_string();
81    }
82
83    /// Splits the parsed query into its original parameters and statement.
84    pub fn into_parts(self) -> (PromQuery, QueryStatement) {
85        (self.query, self.statement)
86    }
87}
88
89/// Resolves the optional schema selector shared by Prometheus query paths.
90///
91/// Like the PromQL planner, the last equality matcher selects the schema.
92pub fn resolve_schema_from_matchers(matchers: &[Matcher]) -> Result<Option<String>> {
93    let mut schema = None;
94    for matcher in matchers
95        .iter()
96        .filter(|matcher| PROMQL_SCHEMA_MATCHERS.contains(&matcher.name.as_str()))
97    {
98        if matcher.op != MatchOp::Equal || matcher.value.is_empty() {
99            return Err(InvalidQuerySnafu {
100                reason: format!(
101                    "expected a non-empty equality matcher for '{}'",
102                    matcher.name
103                ),
104            }
105            .build());
106        }
107        schema = Some(matcher.value.clone());
108    }
109    Ok(schema)
110}
111
112#[async_trait]
113pub trait PrometheusHandler {
114    async fn do_query(&self, query: &PromQuery, query_ctx: QueryContextRef) -> Result<Output>;
115
116    /// Executes a query whose statement has already been parsed.
117    async fn do_query_parsed(
118        &self,
119        query: ParsedPromQuery,
120        query_ctx: QueryContextRef,
121    ) -> Result<Output> {
122        self.do_query(query.query(), query_ctx).await
123    }
124
125    /// Checks access to every table referenced by a batch of PromQL queries.
126    ///
127    /// An empty batch still checks the operation privilege.
128    async fn check_query_permission(
129        &self,
130        queries: &[PromQuery],
131        query_ctx: &QueryContextRef,
132    ) -> Result<()>;
133
134    /// Checks access without parsing the queries again.
135    async fn check_query_permission_parsed(
136        &self,
137        queries: &[ParsedPromQuery],
138        query_ctx: &QueryContextRef,
139    ) -> Result<()> {
140        let queries = queries
141            .iter()
142            .map(|query| query.query().clone())
143            .collect::<Vec<_>>();
144        self.check_query_permission(&queries, query_ctx).await
145    }
146
147    /// Checks access to targets resolved by Prometheus metadata APIs.
148    async fn check_query_target_permission(
149        &self,
150        targets: PermissionTableTargets,
151        query_ctx: &QueryContextRef,
152    ) -> Result<()>;
153
154    /// Removes inaccessible metric names from logical-table metadata enumeration results.
155    async fn filter_metadata_metric_names(
156        &self,
157        metric_names: Vec<String>,
158        schema: &str,
159        query_ctx: &QueryContextRef,
160    ) -> Result<Vec<String>>;
161
162    /// Query metric table names by the `__name__` matchers.
163    async fn query_metric_names(
164        &self,
165        matchers: Vec<Matcher>,
166        schema: &str,
167        ctx: &QueryContextRef,
168    ) -> Result<Vec<String>>;
169
170    async fn query_label_values(
171        &self,
172        metric: String,
173        label_name: String,
174        matchers: Vec<Matcher>,
175        start: SystemTime,
176        end: SystemTime,
177        ctx: &QueryContextRef,
178    ) -> Result<Vec<String>>;
179
180    fn catalog_manager(&self) -> CatalogManagerRef;
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn test_resolve_schema_from_matchers() {
189        assert_eq!(resolve_schema_from_matchers(&[]).unwrap(), None);
190
191        for label in ["__database__", "__schema__"] {
192            let matchers = [Matcher::new(MatchOp::Equal, label, "private")];
193            assert_eq!(
194                resolve_schema_from_matchers(&matchers).unwrap(),
195                Some("private".to_string())
196            );
197        }
198
199        let matchers = [Matcher::new(
200            MatchOp::Equal,
201            "x_greptime_database",
202            "private",
203        )];
204        assert_eq!(resolve_schema_from_matchers(&matchers).unwrap(), None);
205
206        let matchers = [
207            Matcher::new(MatchOp::Equal, "__database__", "private"),
208            Matcher::new(MatchOp::Equal, "__schema__", "public"),
209        ];
210        assert_eq!(
211            resolve_schema_from_matchers(&matchers).unwrap(),
212            Some("public".to_string())
213        );
214    }
215
216    #[test]
217    fn test_resolve_schema_from_matchers_rejects_unsafe_matchers() {
218        assert!(
219            resolve_schema_from_matchers(&[Matcher::new(
220                MatchOp::NotEqual,
221                "__database__",
222                "private",
223            )])
224            .is_err()
225        );
226        assert!(
227            resolve_schema_from_matchers(&[Matcher::new(MatchOp::Equal, "__database__", "",)])
228                .is_err()
229        );
230    }
231}
servers/prometheus_handler.rs

servers/
prometheus_handler.rs
