feat: add |-operator for metric expression #83

GreptimeTeam · May 14, 2024 · 76bc3d1 · 76bc3d1
1 parent ac502fe
commit 76bc3d1
Show file tree

Hide file tree

Showing 4 changed files with 160 additions and 5 deletions.
diff --git a/src/label/matcher.rs b/src/label/matcher.rs
@@ -70,6 +70,7 @@ pub struct Matcher {
     pub op: MatchOp,
     pub name: String,
     pub value: String,
+    pub is_or: bool,
 }
 
 impl Matcher {
@@ -78,6 +79,16 @@ impl Matcher {
             op,
             name: name.into(),
             value: value.into(),
+            is_or: false,
+        }
+    }
+
+    pub fn new_or(op: MatchOp, name: &str, value: &str) -> Self {
+        Self {
+            op,
+            name: name.into(),
+            value: value.into(),
+            is_or: true,
         }
     }
 
@@ -102,15 +113,24 @@ impl Matcher {
     }
 
     pub fn new_matcher(id: TokenId, name: String, value: String) -> Result<Matcher, String> {
+        let op = Self::find_matcher_op(id, &value)?;
+        op.map(|op| Matcher::new(op, name.as_str(), value.as_str()))
+    }
+
+    pub fn new_matcher_or(id: TokenId, name: String, value: String) -> Result<Matcher, String> {
+        let op = Self::find_matcher_op(id, &value)?;
+        op.map(|op| Matcher::new_or(op, name.as_str(), value.as_str()))
+    }
+
+    fn find_matcher_op(id: TokenId, value: &str) -> Result<Result<MatchOp, String>, String> {
         let op = match id {
             T_EQL => Ok(MatchOp::Equal),
             T_NEQ => Ok(MatchOp::NotEqual),
-            T_EQL_REGEX => Ok(MatchOp::Re(Matcher::try_parse_re(&value)?)),
-            T_NEQ_REGEX => Ok(MatchOp::NotRe(Matcher::try_parse_re(&value)?)),
+            T_EQL_REGEX => Ok(MatchOp::Re(Matcher::try_parse_re(value)?)),
+            T_NEQ_REGEX => Ok(MatchOp::NotRe(Matcher::try_parse_re(value)?)),
             _ => Err(format!("invalid match op {}", token_display(id))),
         };
-
-        op.map(|op| Matcher { op, name, value })
+        Ok(op)
     }
 }
 
@@ -234,7 +254,23 @@ impl Matchers {
 
 impl fmt::Display for Matchers {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", join_vector(&self.matchers, ",", true))
+        let not_contains_or = &self.matchers.iter().all(|matcher| !matcher.is_or);
+        if *not_contains_or {
+            write!(f, "{}", join_vector(&self.matchers, ",", true))
+        } else {
+            let matchers_str = self
+                .matchers
+                .iter()
+                .map(|matcher| {
+                    if matcher.is_or {
+                        format!(" or {}", matcher)
+                    } else {
+                        format!(",{}", matcher)
+                    }
+                })
+                .collect::<String>();
+            write!(f, "{}", matchers_str.trim_start_matches(','))
+        }
     }
 }
 

diff --git a/src/parser/lex.rs b/src/parser/lex.rs
@@ -564,6 +564,13 @@ impl Lexer {
         match self.pop() {
             Some('#') => State::LineComment,
             Some(',') => State::Lexeme(T_COMMA),
+            Some('o') => match self.peek() {
+                Some('r') => {
+                    self.pop();
+                    State::Lexeme(T_LOR)
+                }
+                _ => State::Identifier,
+            },
             Some(ch) if ch.is_ascii_whitespace() => State::Space,
             Some(ch) if is_alpha(ch) => State::Identifier,
             Some(ch) if STRING_SYMBOLS.contains(ch) => State::String(ch),

diff --git a/src/parser/parse.rs b/src/parser/parse.rs
@@ -37,6 +37,7 @@ mod tests {
     use regex::Regex;
 
     use crate::label::{Labels, MatchOp, Matcher, Matchers, METRIC_NAME};
+    use crate::parser;
     use crate::parser::function::get_function;
     use crate::parser::{
         token, AtModifier as At, BinModifier, Expr, FunctionArgs, LabelModifier, Offset,
@@ -2115,4 +2116,77 @@ mod tests {
         ];
         assert_cases(fail_cases);
     }
+
+    #[test]
+    fn test_label_or() {
+        let cases = vec![
+            (r#"foo{label1="1" or label2="2"}"#, {
+                let matchers = Matchers::new(vec![
+                    Matcher::new(MatchOp::Equal, "label1", "1"),
+                    Matcher::new_or(MatchOp::Equal, "label2", "2"),
+                ]);
+                Expr::new_vector_selector(Some(String::from("foo")), matchers)
+            }),
+            (r#"foo{label1="1", label2="2"}"#, {
+                let matchers = Matchers::new(vec![
+                    Matcher::new(MatchOp::Equal, "label1", "1"),
+                    Matcher::new(MatchOp::Equal, "label2", "2"),
+                ]);
+                Expr::new_vector_selector(Some(String::from("foo")), matchers)
+            }),
+            (r#"foo{label1="1" or label2="2", label3="3"}"#, {
+                let matchers = Matchers::new(vec![
+                    Matcher::new(MatchOp::Equal, "label1", "1"),
+                    Matcher::new_or(MatchOp::Equal, "label2", "2"),
+                    Matcher::new(MatchOp::Equal, "label3", "3"),
+                ]);
+                Expr::new_vector_selector(Some(String::from("foo")), matchers)
+            }),
+            (r#"foo{label1="1", label2="2" or label3="3"}"#, {
+                let matchers = Matchers::new(vec![
+                    Matcher::new(MatchOp::Equal, "label1", "1"),
+                    Matcher::new(MatchOp::Equal, "label2", "2"),
+                    Matcher::new_or(MatchOp::Equal, "label3", "3"),
+                ]);
+                Expr::new_vector_selector(Some(String::from("foo")), matchers)
+            }),
+            (
+                r#"foo{label1="1", label2="2" or label3="3", label4="4"}"#,
+                {
+                    let matchers = Matchers::new(vec![
+                        Matcher::new(MatchOp::Equal, "label1", "1"),
+                        Matcher::new(MatchOp::Equal, "label2", "2"),
+                        Matcher::new_or(MatchOp::Equal, "label3", "3"),
+                        Matcher::new(MatchOp::Equal, "label4", "4"),
+                    ]);
+                    Expr::new_vector_selector(Some(String::from("foo")), matchers)
+                },
+            ),
+        ];
+        assert_cases(Case::new_result_cases(cases));
+
+        let promql = r#"a{on="1" or label2="2"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+
+        let promql = r#"a{label1="1",label2="2"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+
+        let promql = r#"a{label1="1" or label2="2",label3="3"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+
+        let promql = r#"a{label1="1",label2="2" or label3="3"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+
+        let promql = r#"a{label1="1",label2="2" or label3="3",label4="4"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+
+        let promql = r#"a{o="1",o1="2" or o2="3",o3="4"}"#;
+        let expr = parser::parse(promql).unwrap();
+        assert_eq!(expr.to_string(), promql);
+    }
 }
diff --git a/src/parser/promql.y b/src/parser/promql.y
@@ -407,6 +407,7 @@ label_matchers -> Result<Matchers, String>:
 
 label_match_list -> Result<Matchers, String>:
                 label_match_list COMMA label_matcher { Ok($1?.append($3?)) }
+        |       label_match_list LOR label_matcher_or { Ok($1?.append($3?)) }
         |       label_matcher { Ok(Matchers::empty().append($1?)) }
 ;
 
@@ -447,6 +448,43 @@ label_matcher -> Result<Matcher, String>:
                 }
 ;
 
+label_matcher_or -> Result<Matcher, String>:
+                IDENTIFIER match_op STRING
+                {
+                        let name = lexeme_to_string($lexer, &$1)?;
+                        let value = lexeme_to_string($lexer, &$3)?;
+                        Matcher::new_matcher_or($2?.id(), name, value)
+                }
+        |       IDENTIFIER match_op match_op
+                {
+                        let op = $3?.val;
+                        Err(format!("unexpected '{op}' in label matching, expected string"))
+
+                }
+        |       IDENTIFIER match_op match_op STRING
+                {
+                        let op = $3?.val;
+                        Err(format!("unexpected '{op}' in label matching, expected string"))
+
+                }
+        |       IDENTIFIER match_op match_op IDENTIFIER
+                {
+                        let op = $3?.val;
+                        Err(format!("unexpected '{op}' in label matching, expected string"))
+
+                }
+        |       IDENTIFIER match_op IDENTIFIER
+                {
+                        let id = lexeme_to_string($lexer, &$3)?;
+                        Err(format!("unexpected identifier '{id}' in label matching, expected string"))
+                }
+        |       IDENTIFIER
+                {
+                        let id = lexeme_to_string($lexer, &$1)?;
+                        Err(format!("invalid label matcher, expected label matching operator after '{id}'"))
+                }
+;
+
 /*
  * Metric descriptions.
  */