Repository: rcoh/angle-grinder
Pull request #152 (2021-08-27 14:42:35)

[lang] add the "if" operator


Patch:
diff --git a/README.md b/README.md
index c167d9e..47287da 100644
--- a/README.md
+++ b/README.md
@@ -284,6 +284,19 @@ Multiply `value` by 100 to get the percentage
 * | json | value * 100 as percentage
 ```
 
+##### `if` Operator
+
+`if(<condition>, <value-if-true>, <value-if-false>)`: Choose between two values
+based on the provided condition.
+
+*Examples*
+
+To get byte counts for successful requests:
+
+```agrind
+* | json | if(status == 200, sc_bytes, 0) as ok_bytes
+```
+
 #### Aggregate Operators
 Aggregate operators group and combine your data by 0 or more key fields. The same query can include multiple aggregates.
 The general syntax is:
diff --git a/src/data.rs b/src/data.rs
index ccaeda4..fc370c4 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -511,6 +511,12 @@ impl Record {
                     found: "function call".to_string(),
                 })
             }
+            Expr::IfOp { .. } => {
+                return Err(EvalError::ExpectedXYZ {
+                    expected: "valid expr".to_string(),
+                    found: "if operator".to_string(),
+                })
+            }
             Expr::Value(_) => {
                 return Err(EvalError::ExpectedXYZ {
                     expected: "valid expr".to_string(),
diff --git a/src/lang.rs b/src/lang.rs
index 2360a68..dece754 100644
--- a/src/lang.rs
+++ b/src/lang.rs
@@ -194,6 +194,11 @@ pub enum Expr {
         name: String,
         args: Vec<Expr>,
     },
+    IfOp {
+        cond: Box<Expr>,
+        value_if_true: Box<Expr>,
+        value_if_false: Box<Expr>,
+    },
     Value(data::Value),
     Error,
 }
@@ -649,6 +654,30 @@ fn fcall(input: Span) -> IResult<Span, Expr> {
         .parse(input)
 }
 
+fn if_op(input: Span) -> IResult<Span, Expr> {
+    tag("if")
+        .precedes(with_pos(arg_list))
+        .map(|Positioned { range, value: args }| match args.as_slice() {
+            [cond, value_if_true, value_if_false] => Expr::IfOp {
+                cond: Box::new(cond.clone()),
+                value_if_true: Box::new(value_if_true.clone()),
+                value_if_false: Box::new(value_if_false.clone()),
+            },
+            _ => {
+                input
+                    .extra
+                    .report_error_for(
+                        "the 'if' operator expects exactly 3 arguments, the condition, \
+                        value-if-true, and value-if-false",
+                    )
+                    .with_code_range(range, format!("supplied {} arguments", args.len()))
+                    .send_report();
+                Expr::Error
+            }
+        })
+        .parse(input)
+}
+
 fn filter_atom(input: Span) -> IResult<Span, Option<Search>> {
     let keyword = take_while1(is_keyword).map(|i: Span| Search::from_keyword_input(i.fragment()));
 
@@ -869,7 +898,7 @@ fn atomic(input: Span) -> IResult<Span, Expr> {
             .send_report()
     });
 
-    alt((fcall, value, column_ref, parens)).parse(input)
+    alt((if_op, fcall, value, column_ref, parens)).parse(input)
 }
 
 /// Parses an atomic expression with an optional unary prefix
@@ -1838,6 +1867,36 @@ mod tests {
 
     #[test]
     fn parse_expr() {
+        check_query(
+            "* | if() as i",
+            expect![[r#"
+                Query {
+                    search: And(
+                        [],
+                    ),
+                    operators: [],
+                }
+                error: the 'if' operator expects exactly 3 arguments, the condition, value-if-true, and value-if-false
+                  |
+                1 | * | if() as i
+                  |       ^^ supplied 0 arguments
+                  |"#]],
+        );
+        check_query(
+            "* | if(a, b, c, d) as i",
+            expect![[r#"
+                Query {
+                    search: And(
+                        [],
+                    ),
+                    operators: [],
+                }
+                error: the 'if' operator expects exactly 3 arguments, the condition, value-if-true, and value-if-false
+                  |
+                1 | * | if(a, b, c, d) as i
+                  |       ^^^^^^^^^^^^ supplied 4 arguments
+                  |"#]],
+        );
         check_query(
             "* | now() - 1w2d as yesterday",
             expect![[r#"
diff --git a/src/operator.rs b/src/operator.rs
index 3edb1da..cbfd271 100644
--- a/src/operator.rs
+++ b/src/operator.rs
@@ -236,6 +236,11 @@ pub enum Expr {
         func: &'static funcs::FunctionContainer,
         args: Vec<Expr>,
     },
+    IfOp {
+        cond: Box<Expr>,
+        value_if_true: Box<Expr>,
+        value_if_false: Box<Expr>,
+    },
     Value(&'static data::Value),
 }
 
@@ -408,6 +413,19 @@ impl Evaluatable<data::Value> for Expr {
 
                 func.eval_func(&evaluated_args?)
             }
+            Expr::IfOp {
+                ref cond,
+                ref value_if_true,
+                ref value_if_false,
+            } => {
+                let evaluated_cond: bool = (*cond).eval(record)?;
+
+                if evaluated_cond {
+                    (*value_if_true).eval(record)
+                } else {
+                    (*value_if_false).eval(record)
+                }
+            }
             Expr::Value(v) => Ok(v.clone()),
         }
     }
diff --git a/src/typecheck.rs b/src/typecheck.rs
index 34b5445..f015d10 100644
--- a/src/typecheck.rs
+++ b/src/typecheck.rs
@@ -124,6 +124,15 @@ impl TypeCheck<operator::Expr> for lang::Expr {
                     Err(TypeError::UnknownFunction { name })
                 }
             }
+            lang::Expr::IfOp {
+                cond,
+                value_if_true,
+                value_if_false,
+            } => Ok(operator::Expr::IfOp {
+                cond: Box::new(cond.type_check(error_builder)?),
+                value_if_true: Box::new(value_if_true.type_check(error_builder)?),
+                value_if_false: Box::new(value_if_false.type_check(error_builder)?),
+            }),
             lang::Expr::Value(value) => {
                 let boxed = Box::new(value);
                 let static_value: &'static mut Value = Box::leak(boxed);


Test patch:
diff --git a/tests/structured_tests/if-1.toml b/tests/structured_tests/if-1.toml
new file mode 100644
index 0000000..7feaace
--- /dev/null
+++ b/tests/structured_tests/if-1.toml
@@ -0,0 +1,13 @@
+query = """* | json | if(kind == 1, v1, v2) as vany"""
+input = """
+{"kind": 1, "v1": 5}
+{"kind": 1, "v1": 6}
+{"kind": 2, "v2": 7}
+{"kind": 1, "v1": 8}
+"""
+output = """
+[kind=1]           [v1=5]         [vany=5]
+[kind=1]           [v1=6]         [vany=6]
+[kind=2]                          [vany=7]           [v2=7]
+[kind=1]           [v1=8]         [vany=8]
+"""
diff --git a/tests/structured_tests/if-2.toml b/tests/structured_tests/if-2.toml
new file mode 100644
index 0000000..76d13c7
--- /dev/null
+++ b/tests/structured_tests/if-2.toml
@@ -0,0 +1,15 @@
+query = """* | json | if(kind, v1, v2) as vany"""
+input = """
+{"kind": 1, "v1": 5}
+{"kind": 1, "v1": 6}
+{"kind": 2, "v2": 7}
+{"kind": 1, "v1": 8}
+"""
+output = """
+"""
+error = """
+error: Expected boolean, found 1
+error: Expected boolean, found 1
+error: Expected boolean, found 2
+error: Expected boolean, found 1
+"""