Skip to content

Commit d3cc6a9

Browse files
CodeBlanchCodeBlanch
authored
[query-engine] Add support for defining map schema while allowing undefined keys (open-telemetry#1281)
Requested by @drewrelmas ## Changes * Add support in KQL Parser for defining map schema while still allowing undefined keys ## Details Currently if you define schema for a map then any unknown key will yield "key-not-found" errors. @drewrelmas has a user scenario where he wants to define schema for known keys but still allow custom keys to be added/removed/modified. This PR essentially adds a flag `allow_undefined_keys` to enable this behavior. Note: There will be a follow-up PR exposing this in the bridge options.
1 parent 457e032 commit d3cc6a9

File tree

5 files changed

+658
-67
lines changed

5 files changed

+658
-67
lines changed

‎rust/experimental/query_engine/kql-parser/src/kql_parser.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,19 @@ impl Parser for KqlParser {
2222
}
2323
}
2424

25+
pub(crate) fn map_kql_errors(error: ParserError) -> ParserError {
26+
match error {
27+
ParserError::KeyNotFound { location, key } => ParserError::QueryLanguageDiagnostic {
28+
location,
29+
diagnostic_id: "KS142",
30+
message: format!(
31+
"The name '{key}' does not refer to any known column, table, variable or function"
32+
),
33+
},
34+
e => e,
35+
}
36+
}
37+
2538
#[cfg(test)]
2639
mod tests {
2740
use super::*;

‎rust/experimental/query_engine/kql-parser/src/scalar_primitive_expressions.rs‎

Lines changed: 216 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use data_engine_expressions::*;
88
use data_engine_parser_abstractions::*;
99
use pest::iterators::Pair;
1010

11-
use crate::{Rule, scalar_expression::parse_scalar_expression};
11+
use crate::{scalar_expression::parse_scalar_expression, *};
1212

1313
pub(crate) fn parse_type_unary_expressions(
1414
type_unary_expressions_rule: Pair<Rule>,
@@ -449,7 +449,9 @@ pub(crate) fn parse_accessor_expression(
449449
Rule::scalar_expression => {
450450
let mut scalar = parse_scalar_expression(pair, scope)?;
451451

452-
let value_type = scope.try_resolve_value_type(&mut scalar)?;
452+
let value_type = scope
453+
.try_resolve_value_type(&mut scalar)
454+
.map_err(map_kql_errors)?;
453455

454456
if negate_location.is_some() {
455457
if let Some(t) = value_type
@@ -510,7 +512,8 @@ pub(crate) fn parse_accessor_expression(
510512
Some(ValueType::Map)
511513
} else if let Some(schema) = scope.get_source_schema() {
512514
schema
513-
.try_resolve_value_type(selectors, &scope.get_pipeline().get_resolution_scope())?
515+
.try_resolve_value_type(selectors, &scope.get_pipeline().get_resolution_scope())
516+
.map_err(map_kql_errors)?
514517
} else {
515518
None
516519
};
@@ -566,10 +569,12 @@ pub(crate) fn parse_accessor_expression(
566569
match key {
567570
ParserMapKeySchema::Map(inner_schema) => {
568571
if let Some(schema) = inner_schema {
569-
resolved_value_type = schema.try_resolve_value_type(
570-
value_accessor.get_selectors_mut(),
571-
&scope.get_pipeline().get_resolution_scope(),
572-
)?;
572+
resolved_value_type = schema
573+
.try_resolve_value_type(
574+
value_accessor.get_selectors_mut(),
575+
&scope.get_pipeline().get_resolution_scope(),
576+
)
577+
.map_err(map_kql_errors)?;
573578
}
574579
}
575580
ParserMapKeySchema::Array | ParserMapKeySchema::Any => {
@@ -644,10 +649,12 @@ pub(crate) fn parse_accessor_expression(
644649
}
645650

646651
if let Some(schema) = default_map_schema {
647-
resolved_value_type = schema.try_resolve_value_type(
648-
value_accessor.get_selectors_mut(),
649-
&scope.get_pipeline().get_resolution_scope(),
650-
)?;
652+
resolved_value_type = schema
653+
.try_resolve_value_type(
654+
value_accessor.get_selectors_mut(),
655+
&scope.get_pipeline().get_resolution_scope(),
656+
)
657+
.map_err(map_kql_errors)?;
651658
}
652659

653660
value_accessor.insert_selector(
@@ -656,10 +663,17 @@ pub(crate) fn parse_accessor_expression(
656663
StringScalarExpression::new(root_location, default_map_key),
657664
)),
658665
);
666+
} else if schema.get_allow_undefined_keys() {
667+
value_accessor.insert_selector(
668+
0,
669+
ScalarExpression::Static(StaticScalarExpression::String(
670+
root_accessor_identity,
671+
)),
672+
);
659673
} else {
660674
return Err(ParserError::QueryLanguageDiagnostic {
661675
location: root_accessor_identity.get_query_location().clone(),
662-
diagnostic_id: "KS109",
676+
diagnostic_id: "KS142",
663677
message: format!(
664678
"The name '{}' does not refer to any known column, table, variable or function",
665679
root_accessor_identity.get_value()
@@ -1708,7 +1722,7 @@ mod tests {
17081722

17091723
run_test_failure(
17101724
"unknown_key",
1711-
"KS109",
1725+
"KS142",
17121726
"The name 'unknown_key' does not refer to any known column, table, variable or function",
17131727
);
17141728
}
@@ -2016,7 +2030,7 @@ mod tests {
20162030
}
20172031

20182032
#[test]
2019-
fn test_parse_accessor_expression_with_map_schema() {
2033+
fn test_parse_accessor_expression_with_default_map_schema() {
20202034
let run_test_success = |input: &str, expected: ScalarExpression| {
20212035
let mut result = KqlPestParser::parse(Rule::accessor_expression, input).unwrap();
20222036

@@ -2262,14 +2276,14 @@ mod tests {
22622276

22632277
run_test_failure(
22642278
"map.unknown",
2265-
None,
2266-
"The name 'unknown' does not refer to any known key on the target map",
2279+
Some("KS142"),
2280+
"The name 'unknown' does not refer to any known column, table, variable or function",
22672281
);
22682282

22692283
run_test_failure(
22702284
"map.map_value_with_schema.unknown",
2271-
None,
2272-
"The name 'unknown' does not refer to any known key on the target map",
2285+
Some("KS142"),
2286+
"The name 'unknown' does not refer to any known column, table, variable or function",
22732287
);
22742288

22752289
run_test_failure(
@@ -2280,8 +2294,190 @@ mod tests {
22802294

22812295
run_test_failure(
22822296
"source.field",
2283-
None,
2284-
"The name 'field' does not refer to any known key on the target map",
2297+
Some("KS142"),
2298+
"The name 'field' does not refer to any known column, table, variable or function",
2299+
);
2300+
}
2301+
2302+
#[test]
2303+
fn test_parse_accessor_expression_with_default_map_schema_and_allow_undefined_keys() {
2304+
let run_test_success = |input: &str, expected: ScalarExpression| {
2305+
let mut result = KqlPestParser::parse(Rule::accessor_expression, input).unwrap();
2306+
2307+
let state = ParserState::new_with_options(
2308+
input,
2309+
ParserOptions::new().with_source_map_schema(
2310+
ParserMapSchema::new()
2311+
.set_default_map_key("map")
2312+
.with_key_definition(
2313+
"map",
2314+
ParserMapKeySchema::Map(Some(
2315+
ParserMapSchema::new()
2316+
.with_key_definition("double_value", ParserMapKeySchema::Double)
2317+
.set_allow_undefined_keys(),
2318+
)),
2319+
),
2320+
),
2321+
);
2322+
2323+
let expression =
2324+
parse_accessor_expression(result.next().unwrap(), &state, false).unwrap();
2325+
2326+
assert_eq!(expected, expression);
2327+
};
2328+
2329+
run_test_success(
2330+
"double_value",
2331+
ScalarExpression::Source(SourceScalarExpression::new_with_value_type(
2332+
QueryLocation::new_fake(),
2333+
ValueAccessor::new_with_selectors(vec![
2334+
ScalarExpression::Static(StaticScalarExpression::String(
2335+
StringScalarExpression::new(QueryLocation::new_fake(), "map"),
2336+
)),
2337+
ScalarExpression::Static(StaticScalarExpression::String(
2338+
StringScalarExpression::new(QueryLocation::new_fake(), "double_value"),
2339+
)),
2340+
]),
2341+
Some(ValueType::Double),
2342+
)),
2343+
);
2344+
2345+
run_test_success(
2346+
"unknown",
2347+
ScalarExpression::Source(SourceScalarExpression::new_with_value_type(
2348+
QueryLocation::new_fake(),
2349+
ValueAccessor::new_with_selectors(vec![
2350+
ScalarExpression::Static(StaticScalarExpression::String(
2351+
StringScalarExpression::new(QueryLocation::new_fake(), "map"),
2352+
)),
2353+
ScalarExpression::Static(StaticScalarExpression::String(
2354+
StringScalarExpression::new(QueryLocation::new_fake(), "unknown"),
2355+
)),
2356+
]),
2357+
None,
2358+
)),
2359+
);
2360+
}
2361+
2362+
#[test]
2363+
fn test_parse_accessor_expression_with_schema() {
2364+
let run_test_success = |input: &str, expected: ScalarExpression| {
2365+
let mut result = KqlPestParser::parse(Rule::accessor_expression, input).unwrap();
2366+
2367+
let state = ParserState::new_with_options(
2368+
input,
2369+
ParserOptions::new().with_source_map_schema(
2370+
ParserMapSchema::new()
2371+
.with_key_definition("int_value", ParserMapKeySchema::Integer),
2372+
),
2373+
);
2374+
2375+
let expression =
2376+
parse_accessor_expression(result.next().unwrap(), &state, false).unwrap();
2377+
2378+
assert_eq!(expected, expression);
2379+
};
2380+
2381+
let run_test_failure = |input: &str, expected_id: Option<&str>, expected_msg: &str| {
2382+
let mut result = KqlPestParser::parse(Rule::accessor_expression, input).unwrap();
2383+
2384+
let state = ParserState::new_with_options(
2385+
input,
2386+
ParserOptions::new().with_source_map_schema(
2387+
ParserMapSchema::new()
2388+
.with_key_definition("int_value", ParserMapKeySchema::Integer),
2389+
),
2390+
);
2391+
2392+
let error =
2393+
parse_accessor_expression(result.next().unwrap(), &state, true).unwrap_err();
2394+
2395+
if let Some(expected_id) = expected_id {
2396+
if let ParserError::QueryLanguageDiagnostic {
2397+
location: _,
2398+
diagnostic_id: id,
2399+
message: msg,
2400+
} = error
2401+
{
2402+
assert_eq!(expected_id, id);
2403+
assert_eq!(expected_msg, msg);
2404+
} else {
2405+
panic!("Expected QueryLanguageDiagnostic");
2406+
}
2407+
} else if let ParserError::SyntaxError(_, msg) = error {
2408+
assert_eq!(expected_msg, msg);
2409+
} else {
2410+
panic!("Expected SyntaxError");
2411+
}
2412+
};
2413+
2414+
run_test_success(
2415+
"int_value",
2416+
ScalarExpression::Source(SourceScalarExpression::new_with_value_type(
2417+
QueryLocation::new_fake(),
2418+
ValueAccessor::new_with_selectors(vec![ScalarExpression::Static(
2419+
StaticScalarExpression::String(StringScalarExpression::new(
2420+
QueryLocation::new_fake(),
2421+
"int_value",
2422+
)),
2423+
)]),
2424+
Some(ValueType::Integer),
2425+
)),
2426+
);
2427+
2428+
run_test_failure(
2429+
"unknown",
2430+
Some("KS142"),
2431+
"The name 'unknown' does not refer to any known column, table, variable or function",
2432+
);
2433+
}
2434+
2435+
#[test]
2436+
fn test_parse_accessor_expression_with_schema_allow_undefined_keys() {
2437+
let run_test_success = |input: &str, expected: ScalarExpression| {
2438+
let mut result = KqlPestParser::parse(Rule::accessor_expression, input).unwrap();
2439+
2440+
let state = ParserState::new_with_options(
2441+
input,
2442+
ParserOptions::new().with_source_map_schema(
2443+
ParserMapSchema::new()
2444+
.with_key_definition("int_value", ParserMapKeySchema::Integer)
2445+
.set_allow_undefined_keys(),
2446+
),
2447+
);
2448+
2449+
let expression =
2450+
parse_accessor_expression(result.next().unwrap(), &state, false).unwrap();
2451+
2452+
assert_eq!(expected, expression);
2453+
};
2454+
2455+
run_test_success(
2456+
"int_value",
2457+
ScalarExpression::Source(SourceScalarExpression::new_with_value_type(
2458+
QueryLocation::new_fake(),
2459+
ValueAccessor::new_with_selectors(vec![ScalarExpression::Static(
2460+
StaticScalarExpression::String(StringScalarExpression::new(
2461+
QueryLocation::new_fake(),
2462+
"int_value",
2463+
)),
2464+
)]),
2465+
Some(ValueType::Integer),
2466+
)),
2467+
);
2468+
2469+
run_test_success(
2470+
"unknown",
2471+
ScalarExpression::Source(SourceScalarExpression::new_with_value_type(
2472+
QueryLocation::new_fake(),
2473+
ValueAccessor::new_with_selectors(vec![ScalarExpression::Static(
2474+
StaticScalarExpression::String(StringScalarExpression::new(
2475+
QueryLocation::new_fake(),
2476+
"unknown",
2477+
)),
2478+
)]),
2479+
None,
2480+
)),
22852481
);
22862482
}
22872483
}

0 commit comments

Comments
 (0)