fix(select): classify function argument planner errors

This commit is contained in:
overtrue
2026-08-11 23:14:23 +08:00
parent 7c1d9dec8f
commit 3f994c59eb
2 changed files with 62 additions and 1 deletions
+15 -1
View File
@@ -68,7 +68,7 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
match stmt {
Statement::Query(_) => {
validate_s3_select_statement(&stmt)?;
let df_plan = self.df_planner.sql_statement_to_plan(stmt)?;
let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?;
let plan = Plan::Query(QueryPlan {
df_plan,
is_tag_scan: false,
@@ -81,6 +81,20 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
}
}
fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryError {
if matches!(
&error,
datafusion::common::DataFusionError::Plan(message)
if message.starts_with("Failed to coerce arguments to satisfy a call to")
|| (message.starts_with("Internal error: Function '")
&& message.contains("' failed to match any signature, errors:"))
) {
return SelectError::IncorrectSqlFunctionArgumentType.into();
}
error.into()
}
fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> {
let Statement::Query(query) = statement else {
return Err(unsupported_structure("only SELECT queries are supported"));
@@ -127,6 +127,53 @@ mod error_handling_tests {
assert!(matches!(error.select_error(), SelectError::UnsupportedSqlStructure { .. }));
}
#[tokio::test]
async fn test_function_argument_coercion_failure_is_typed() {
for sql in ["SELECT ROUND(3.14, 1.1) FROM S3Object", "SELECT SQRT(1, 2) FROM S3Object"] {
let input = create_test_input_with_sql(sql);
let db = get_global_db(input.clone(), true)
.await
.expect("test database should initialize");
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let error = match db.execute(&query).await {
Err(error) => error,
Ok(_) => panic!("invalid function arguments must fail during planning: {sql}"),
};
assert_eq!(
error.select_error(),
SelectError::IncorrectSqlFunctionArgumentType,
"unexpected planner error for {sql}: {error:?}"
);
}
}
#[tokio::test]
async fn test_other_planner_failures_remain_invalid_query() {
for sql in [
"SELECT DEFINITELY_UNKNOWN_FUNCTION(1) FROM S3Object",
"SELECT 1 + 'text' FROM S3Object",
] {
let input = create_test_input_with_sql(sql);
let db = get_global_db(input.clone(), true)
.await
.expect("test database should initialize");
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let error = match db.execute(&query).await {
Err(error) => error,
Ok(_) => panic!("invalid query must fail during planning: {sql}"),
};
assert_eq!(
error.select_error(),
SelectError::InvalidQuery,
"unexpected planner error for {sql}: {error:?}"
);
}
}
#[tokio::test]
async fn test_invalid_column_references() {
let invalid_column_sqls = vec![