fix: Do not panic when bad arguments are passed to selectors (#6441)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
pull/24376/head
Andrew Lamb 2022-12-20 05:37:56 -05:00 committed by GitHub
parent 96fe89bd3e
commit 3e2362ea9b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 38 additions and 7 deletions

View File

@ -464,15 +464,24 @@ fn make_uda(name: &str, factory_builder: FactoryBuilder) -> AggregateUDF {
//
// The inputs are (value, time) and the output is a struct with a
// 'value' and 'time' field of the same time.
let captured_name = name.to_string();
let return_type_func: ReturnTypeFunction = Arc::new(move |arg_types| {
assert_eq!(
arg_types.len(),
2,
"selector expected exactly 2 arguments, got {}",
arg_types.len()
);
if arg_types.len() != 2 {
return Err(DataFusionError::Plan(format!(
"{} requires exactly 2 arguments, got {}",
captured_name,
arg_types.len()
)));
}
let input_type = &arg_types[0];
assert_eq!(&arg_types[1], &TIME_DATA_TYPE());
let time_type = &arg_types[1];
if time_type != &TIME_DATA_TYPE() {
return Err(DataFusionError::Plan(format!(
"{} second argument must be a timestamp, but got {}",
captured_name, time_type
)));
}
let return_type = output_type.return_type(input_type);
Ok(Arc::new(return_type))

View File

@ -362,4 +362,26 @@ async fn sql_create_schema() {
.await;
}
#[tokio::test]
async fn bad_selector_num_args() {
let expected_error = "selector_last requires exactly 2 arguments, got 1";
run_sql_error_test_case(
scenarios::delete::NoDeleteOneChunk {},
"select selector_last(time)['bar'] from cpu;",
expected_error,
)
.await;
}
#[tokio::test]
async fn bad_selector_arg_types() {
let expected_error = "selector_last second argument must be a timestamp, but got Float64";
run_sql_error_test_case(
scenarios::delete::NoDeleteOneChunk {},
"select selector_last(time, bar)['value'] from cpu;",
expected_error,
)
.await;
}
// --------------------------------------------------------