servers/
prometheus.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 catalog::system_schema::information_schema::tables::{
16    CREATE_OPTIONS, ENGINE as TABLE_ENGINE, TABLE_CATALOG, TABLE_NAME, TABLE_SCHEMA,
17};
18use common_telemetry::tracing;
19use datafusion::dataframe::DataFrame;
20use datafusion::prelude::{Expr, col, lit, regexp_match};
21use datafusion_expr::LogicalPlan;
22use promql_parser::label::{MatchOp, Matcher};
23use session::context::QueryContextRef;
24use snafu::ResultExt;
25use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, PHYSICAL_TABLE_METADATA_KEY};
26
27use crate::error::{self, Result};
28
29/// The maximum number of metrics at one time.
30const MAX_METRICS_NUM: usize = 1024;
31
32/// Create a DataFrame from promql `__name__` matchers.
33/// # Panics
34///  Panic when the machers contains `MatchOp::Equal`.
35#[tracing::instrument(skip_all)]
36pub fn metric_name_matchers_to_plan(
37    dataframe: DataFrame,
38    matchers: Vec<Matcher>,
39    schema: &str,
40    ctx: &QueryContextRef,
41) -> Result<LogicalPlan> {
42    assert!(!matchers.is_empty());
43
44    let mut conditions = Vec::with_capacity(matchers.len() + 5);
45
46    conditions.push(col(TABLE_CATALOG).eq(lit(ctx.current_catalog())));
47    conditions.push(col(TABLE_SCHEMA).eq(lit(schema)));
48    // Must be metric engine
49    conditions.push(col(TABLE_ENGINE).eq(lit("metric")));
50    // Physical metric tables are internal. PromQL queries user-visible logical tables.
51    conditions.push(
52        regexp_match(
53            col(CREATE_OPTIONS),
54            lit(format!("(^| ){LOGICAL_TABLE_METADATA_KEY}=")),
55            None,
56        )
57        .is_not_null(),
58    );
59    conditions.push(
60        regexp_match(
61            col(CREATE_OPTIONS),
62            lit(format!("(^| ){PHYSICAL_TABLE_METADATA_KEY}=")),
63            None,
64        )
65        .is_null(),
66    );
67
68    for m in matchers {
69        let value = &m.value;
70
71        match &m.op {
72            MatchOp::NotEqual => {
73                conditions.push(col(TABLE_NAME).not_eq(lit(value)));
74            }
75            // Case sensitive regexp match
76            MatchOp::Re(regex) => {
77                conditions.push(
78                    regexp_match(col(TABLE_NAME), lit(regex.to_string()), None).is_not_null(),
79                );
80            }
81            // Case sensitive regexp not match
82            MatchOp::NotRe(regex) => {
83                conditions
84                    .push(regexp_match(col(TABLE_NAME), lit(regex.to_string()), None).is_null());
85            }
86            _ => unreachable!("checked outside"),
87        }
88    }
89
90    // Safety: conditions MUST not be empty, reduce always return Some(expr).
91    let conditions = conditions.into_iter().reduce(Expr::and).unwrap();
92
93    let dataframe = dataframe
94        .filter(conditions)
95        .context(error::DataFrameSnafu)?
96        .select(vec![col(TABLE_NAME)])
97        .context(error::DataFrameSnafu)?
98        .limit(0, Some(MAX_METRICS_NUM))
99        .context(error::DataFrameSnafu)?;
100
101    Ok(dataframe.into_parts().1)
102}
103
104#[cfg(test)]
105mod tests {
106    use std::sync::Arc;
107
108    use arrow::array::StringArray;
109    use arrow::record_batch::RecordBatch;
110    use arrow_schema::{DataType, Field, Schema};
111    use datafusion::prelude::SessionContext;
112    use promql_parser::label::{MatchOp, Matcher};
113    use session::context::QueryContext;
114
115    use super::*;
116
117    #[tokio::test]
118    async fn test_metric_name_plan_only_returns_logical_metric_tables() {
119        let schema = Arc::new(Schema::new(vec![
120            Field::new(TABLE_CATALOG, DataType::Utf8, false),
121            Field::new(TABLE_SCHEMA, DataType::Utf8, false),
122            Field::new(TABLE_NAME, DataType::Utf8, false),
123            Field::new(TABLE_ENGINE, DataType::Utf8, false),
124            Field::new(CREATE_OPTIONS, DataType::Utf8, false),
125        ]));
126        let batch = RecordBatch::try_new(
127            schema,
128            vec![
129                Arc::new(StringArray::from(vec!["greptime"; 5])),
130                Arc::new(StringArray::from(vec![
131                    "public", "private", "private", "private", "private",
132                ])),
133                Arc::new(StringArray::from(vec![
134                    "logical_public",
135                    "logical_private",
136                    "physical",
137                    "misleading",
138                    "mito",
139                ])),
140                Arc::new(StringArray::from(vec![
141                    "metric", "metric", "metric", "metric", "mito",
142                ])),
143                Arc::new(StringArray::from(vec![
144                    "on_physical_table=physical",
145                    "ttl=1d on_physical_table=physical",
146                    "physical_metric_table=",
147                    "foo=x on_physical_table=physical physical_metric_table=",
148                    "on_physical_table=physical",
149                ])),
150            ],
151        )
152        .unwrap();
153
154        let session = SessionContext::new();
155        let dataframe = session.read_batch(batch).unwrap();
156        let query_ctx = Arc::new(QueryContext::with("greptime", "public"));
157        let plan = metric_name_matchers_to_plan(
158            dataframe,
159            vec![Matcher::new(MatchOp::NotEqual, "__name__", "")],
160            "private",
161            &query_ctx,
162        )
163        .unwrap();
164        let batches = session
165            .execute_logical_plan(plan)
166            .await
167            .unwrap()
168            .collect()
169            .await
170            .unwrap();
171
172        let mut names = Vec::new();
173        for batch in batches {
174            let column = batch
175                .column(0)
176                .as_any()
177                .downcast_ref::<StringArray>()
178                .unwrap();
179            names.extend(column.iter().flatten().map(str::to_owned));
180        }
181        assert_eq!(vec!["logical_private"], names);
182    }
183}
servers/prometheus.rs

servers/
prometheus.rs
