diff --git a/rust/experimental/query_engine/kql-parser/src/if_else_expression.rs b/rust/experimental/query_engine/kql-parser/src/if_else_expression.rs new file mode 100644 index 000000000..5532658be --- /dev/null +++ b/rust/experimental/query_engine/kql-parser/src/if_else_expression.rs @@ -0,0 +1,322 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +use data_engine_expressions::{ + ConditionalDataExpression, ConditionalDataExpressionBranch, DataExpression, Expression, + LogicalExpression, QueryLocation, +}; +use data_engine_parser_abstractions::{ParserError, ParserScope, to_query_location}; +use pest::iterators::Pair; + +use crate::{ + Rule, logical_expressions::parse_logical_expression, + tabular_expressions::parse_tabular_expression_rule, +}; + +pub fn parse_if_else_expression( + if_else_expression: Pair, + scope: &dyn ParserScope, +) -> Result { + let query_location = to_query_location(&if_else_expression); + let mut conditional_expr = ConditionalDataExpression::new(query_location); + + // keep track of the location of the current branch + let mut branch_location_start = 0; + let mut branch_location_line = 0; + let mut branch_location_col = 0; + + let mut next_condition: Option = None; + let mut next_branch: Vec = Vec::new(); + + let inner = if_else_expression.into_inner(); + for rule in inner { + match rule.as_rule() { + // parse the condition for the branch + Rule::if_condition_expression => { + branch_location_start = rule.as_span().start(); + let (line_number, column_number) = rule.line_col(); + branch_location_line = line_number; + branch_location_col = column_number; + + let condition_rule = rule.into_inner().next().ok_or_else(|| { + // under normal invocation of this function this shouldn't happen as this + // missing expression should be caught by the parser + ParserError::SyntaxError( + conditional_expr.get_query_location().clone(), + "expected if_condition_expression to contain one inner logical_expression" + .to_string(), + ) + })?; + next_condition = Some(parse_logical_expression(condition_rule, scope)?); + } + + // parse the pipeline of data expressions for this branch + Rule::if_else_branch_expression => { + let branch_loc_end = rule.as_span().end(); + + // parse all the rules + for tabular_rule in rule.into_inner() { + let mut expr = parse_tabular_expression_rule(tabular_rule, scope)?; + next_branch.append(&mut expr); + } + + // take the data expressions for the branch and reset next_branch + let curr_branch = next_branch; + next_branch = Vec::new(); + + let query_location = QueryLocation::new( + branch_location_start, + branch_loc_end, + branch_location_line, + branch_location_col, + ) + .map_err(|e| { + ParserError::SyntaxError( + conditional_expr.get_query_location().clone(), + format!("invalid query location {e}"), + ) + })?; + + let condition = next_condition.take().ok_or_else(|| { + // next_condition should always be Some here as we should see the + // if_condition_expression first.Under normal invocation of this function + // the order should be enforced by the parser + ParserError::SyntaxError( + conditional_expr.get_query_location().clone(), + "expected if_condition_expression before if_else_branch_expression" + .to_string(), + ) + })?; + + conditional_expr = conditional_expr.with_branch( + ConditionalDataExpressionBranch::new(query_location, condition, curr_branch), + ); + } + + // parse the data expressions for the else branch + Rule::else_expression => { + let mut else_branch_exprs = Vec::new(); + let branch_rules = rule.into_inner().next().ok_or_else(|| { + // under normal invocation of this function this shouldn't happen as this + // missing expression should be caught by the parser + ParserError::SyntaxError( + conditional_expr.get_query_location().clone(), + "expected else_expression to contain one inner if_else_branch_expression" + .to_string(), + ) + })?; + + for tabular_rule in branch_rules.into_inner() { + let mut expr = parse_tabular_expression_rule(tabular_rule, scope)?; + else_branch_exprs.append(&mut expr); + } + conditional_expr = conditional_expr.with_default_branch(else_branch_exprs); + } + _ => { + return Err(ParserError::SyntaxError( + conditional_expr.get_query_location().clone(), + format!("invalid rule found in if_else_expression {rule}"), + )); + } + } + } + + Ok(DataExpression::Conditional(conditional_expr)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kql_parser::KqlParser; + use data_engine_expressions::{ + EqualToLogicalExpression, MutableValueExpression, ScalarExpression, SetTransformExpression, + SourceScalarExpression, StaticScalarExpression, StringScalarExpression, + TransformExpression, ValueAccessor, + }; + use data_engine_parser_abstractions::Parser; + + fn equals_logical_expr(field_name: &'static str, value: &'static str) -> LogicalExpression { + LogicalExpression::EqualTo(EqualToLogicalExpression::new( + QueryLocation::new_fake(), + ScalarExpression::Source(SourceScalarExpression::new( + QueryLocation::new_fake(), + ValueAccessor::new_with_selectors(vec![ScalarExpression::Static( + StaticScalarExpression::String(StringScalarExpression::new( + QueryLocation::new_fake(), + field_name, + )), + )]), + )), + ScalarExpression::Static(StaticScalarExpression::String(StringScalarExpression::new( + QueryLocation::new_fake(), + value, + ))), + false, + )) + } + + fn assign_attribute_expression( + attr_key: &'static str, + attr_val: &'static str, + ) -> DataExpression { + DataExpression::Transform(TransformExpression::Set(SetTransformExpression::new( + QueryLocation::new_fake(), + ScalarExpression::Static(StaticScalarExpression::String(StringScalarExpression::new( + QueryLocation::new_fake(), + attr_val, + ))), + MutableValueExpression::Source(SourceScalarExpression::new( + QueryLocation::new_fake(), + ValueAccessor::new_with_selectors(vec![ + ScalarExpression::Static(StaticScalarExpression::String( + StringScalarExpression::new(QueryLocation::new_fake(), "attributes"), + )), + ScalarExpression::Static(StaticScalarExpression::String( + StringScalarExpression::new(QueryLocation::new_fake(), attr_key), + )), + ]), + )), + ))) + } + + #[test] + pub fn test_parse_if_else_expression() { + let query = r#" + logs | if (severity_text == "ERROR") { + extend attributes["important"] = "very" | extend attributes["triggers_alarm"] = "true" + } else if (severity_text == "WARN") { + extend attributes["important"] = "somewhat" + } else if (severity_text == "INFO") { + extend attributes["important"] = "rarely" + } else { + extend attributes["important"] = "no" + } + "#; + + let result = KqlParser::parse(query); + assert!(result.is_ok()); + + let pipeline = result.unwrap().pipeline; + let expressions = pipeline.get_expressions(); + assert_eq!(expressions.len(), 1); + + let expected = DataExpression::Conditional( + ConditionalDataExpression::new(QueryLocation::new_fake()) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "ERROR"), + vec![ + assign_attribute_expression("important", "very"), + assign_attribute_expression("triggers_alarm", "true"), + ], + )) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "WARN"), + vec![assign_attribute_expression("important", "somewhat")], + )) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "INFO"), + vec![assign_attribute_expression("important", "rarely")], + )) + .with_default_branch(vec![assign_attribute_expression("important", "no")]), + ); + assert_eq!(expressions[0], expected); + } + + #[test] + pub fn test_parse_if_else_expression_no_else() { + let query = r#" + logs | if (severity_text == "ERROR") { + extend attributes["important"] = "very" | extend attributes["triggers_alarm"] = "true" + } else if (severity_text == "WARN") { + extend attributes["important"] = "somewhat" + } + "#; + + let result = KqlParser::parse(query); + assert!(result.is_ok()); + + let pipeline = result.unwrap().pipeline; + let expressions = pipeline.get_expressions(); + assert_eq!(expressions.len(), 1); + + let expected = DataExpression::Conditional( + ConditionalDataExpression::new(QueryLocation::new_fake()) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "ERROR"), + vec![ + assign_attribute_expression("important", "very"), + assign_attribute_expression("triggers_alarm", "true"), + ], + )) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "WARN"), + vec![assign_attribute_expression("important", "somewhat")], + )), + ); + assert_eq!(expressions[0], expected); + } + + #[test] + pub fn test_parse_if_else_expression_no_elseif() { + let query = r#" + logs | if (severity_text == "ERROR") { + extend attributes["important"] = "very" | extend attributes["triggers_alarm"] = "true" + } else { + extend attributes["important"] = "no" + } + "#; + + let result = KqlParser::parse(query); + assert!(result.is_ok()); + + let pipeline = result.unwrap().pipeline; + let expressions = pipeline.get_expressions(); + assert_eq!(expressions.len(), 1); + + let expected = DataExpression::Conditional( + ConditionalDataExpression::new(QueryLocation::new_fake()) + .with_branch(ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "ERROR"), + vec![ + assign_attribute_expression("important", "very"), + assign_attribute_expression("triggers_alarm", "true"), + ], + )) + .with_default_branch(vec![assign_attribute_expression("important", "no")]), + ); + assert_eq!(expressions[0], expected); + } + + #[test] + pub fn test_parse_if_else_expression_if_only() { + let query = r#" + logs | if (severity_text == "ERROR") { + extend attributes["triggers_alarm"] = "true" + } + "#; + + let result = KqlParser::parse(query); + assert!(result.is_ok()); + + let pipeline = result.unwrap().pipeline; + let expressions = pipeline.get_expressions(); + assert_eq!(expressions.len(), 1); + + let expected = DataExpression::Conditional( + ConditionalDataExpression::new(QueryLocation::new_fake()).with_branch( + ConditionalDataExpressionBranch::new( + QueryLocation::new_fake(), + equals_logical_expr("severity_text", "ERROR"), + vec![assign_attribute_expression("triggers_alarm", "true")], + ), + ), + ); + assert_eq!(expressions[0], expected); + } +} diff --git a/rust/experimental/query_engine/kql-parser/src/kql.pest b/rust/experimental/query_engine/kql-parser/src/kql.pest index c2bc4cedb..4b0f992ac 100644 --- a/rust/experimental/query_engine/kql-parser/src/kql.pest +++ b/rust/experimental/query_engine/kql-parser/src/kql.pest @@ -320,6 +320,24 @@ group_by_expression = { | scalar_expression } +if_condition_expression = { + "if" ~ "(" ~ logical_expression ~ ")" +} + +if_else_branch_expression = { + tabular_expressions ~ ("|" ~ tabular_expressions)* +} + +else_expression = { + "else" ~ "{" ~ if_else_branch_expression ~ "}" +} + +if_else_expression = { + if_condition_expression ~ "{" ~ if_else_branch_expression ~ "}" + ~ ("else" ~ if_condition_expression ~ "{" ~ if_else_branch_expression ~ "}")* + ~ (else_expression)? +} + summarize_expression = { "summarize" ~ (aggregate_expression ~ ("," ~ aggregate_expression)*)? @@ -335,6 +353,7 @@ tabular_expressions = _{ | project_rename_expression | where_expression | summarize_expression + | if_else_expression } tabular_expression = { diff --git a/rust/experimental/query_engine/kql-parser/src/lib.rs b/rust/experimental/query_engine/kql-parser/src/lib.rs index 79297c8c2..53718c978 100644 --- a/rust/experimental/query_engine/kql-parser/src/lib.rs +++ b/rust/experimental/query_engine/kql-parser/src/lib.rs @@ -3,6 +3,7 @@ pub(crate) mod aggregate_expressions; pub(crate) mod date_utils; +pub(crate) mod if_else_expression; pub(crate) mod kql_parser; pub(crate) mod logical_expressions; pub(crate) mod query_expression; diff --git a/rust/experimental/query_engine/kql-parser/src/tabular_expressions.rs b/rust/experimental/query_engine/kql-parser/src/tabular_expressions.rs index ea0ef2664..73dd1b90b 100644 --- a/rust/experimental/query_engine/kql-parser/src/tabular_expressions.rs +++ b/rust/experimental/query_engine/kql-parser/src/tabular_expressions.rs @@ -11,6 +11,7 @@ use regex::Regex; use crate::{ Rule, aggregate_expressions::parse_aggregate_expression, + if_else_expression::parse_if_else_expression, logical_expressions::parse_logical_expression, scalar_expression::{parse_scalar_expression, try_resolve_identifier}, scalar_primitive_expressions::{parse_accessor_expression, parse_string_literal}, @@ -555,6 +556,9 @@ pub(crate) fn parse_tabular_expression_rule( Rule::summarize_expression => { expressions.push(parse_summarize_expression(tabular_expression_rule, scope)?) } + Rule::if_else_expression => { + expressions.push(parse_if_else_expression(tabular_expression_rule, scope)?); + } _ => panic!("Unexpected rule in tabular_expression: {tabular_expression_rule}"), }