Skip to content

Commit 8274ba9

Browse files
committed
Report error when SessionState::sql_to_expr_with_alias does not consume all input
1 parent 44e63e0 commit 8274ba9

File tree

2 files changed

+106
-5
lines changed

2 files changed

+106
-5
lines changed

datafusion/core/src/execution/session_state.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -434,7 +434,7 @@ impl SessionState {
434434
.with_dialect(dialect.as_ref())
435435
.with_recursion_limit(recursion_limit)
436436
.build()?
437-
.parse_expr()?;
437+
.parse_into_expr()?;
438438

439439
Ok(expr)
440440
}

datafusion/sql/src/parser.rs

Lines changed: 105 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -413,13 +413,18 @@ impl<'a> DFParser<'a> {
413413
parser.parse_statements()
414414
}
415415

416+
pub fn parse_sql_into_expr(sql: &str) -> Result<ExprWithAlias, DataFusionError> {
417+
DFParserBuilder::new(sql).build()?.parse_into_expr()
418+
}
419+
416420
pub fn parse_sql_into_expr_with_dialect(
417421
sql: &str,
418422
dialect: &dyn Dialect,
419423
) -> Result<ExprWithAlias, DataFusionError> {
420-
let mut parser = DFParserBuilder::new(sql).with_dialect(dialect).build()?;
421-
422-
parser.parse_expr()
424+
DFParserBuilder::new(sql)
425+
.with_dialect(dialect)
426+
.build()?
427+
.parse_into_expr()
423428
}
424429

425430
/// Parse a sql string into one or [`Statement`]s
@@ -465,6 +470,19 @@ impl<'a> DFParser<'a> {
465470
)
466471
}
467472

473+
fn expect_token(
474+
&mut self,
475+
expected: &str,
476+
token: Token,
477+
) -> Result<(), DataFusionError> {
478+
let next_token = self.parser.peek_token_ref();
479+
if next_token.token != token {
480+
self.expected(expected, next_token.clone())
481+
} else {
482+
Ok(())
483+
}
484+
}
485+
468486
/// Parse a new expression
469487
pub fn parse_statement(&mut self) -> Result<Statement, DataFusionError> {
470488
match self.parser.peek_token().token {
@@ -514,6 +532,16 @@ impl<'a> DFParser<'a> {
514532
Ok(self.parser.parse_expr_with_alias()?)
515533
}
516534

535+
/// Parses the entire SQL string into an expression.
536+
///
537+
/// In contrast to [`parse_expr`], this function will report an error if the input
538+
/// contains any trailing, unparsed tokens.
539+
pub fn parse_into_expr(&mut self) -> Result<ExprWithAlias, DataFusionError> {
540+
let expr = self.parse_expr()?;
541+
self.expect_token("end of expression", Token::EOF)?;
542+
Ok(expr)
543+
}
544+
517545
/// Helper method to parse a statement and handle errors consistently, especially for recursion limits
518546
fn parse_and_handle_statement(&mut self) -> Result<Statement, DataFusionError> {
519547
self.parser
@@ -1021,7 +1049,7 @@ mod tests {
10211049
use super::*;
10221050
use datafusion_common::assert_contains;
10231051
use sqlparser::ast::Expr::Identifier;
1024-
use sqlparser::ast::{BinaryOperator, DataType, Expr, Ident};
1052+
use sqlparser::ast::{BinaryOperator, DataType, Expr, Ident, ValueWithSpan};
10251053
use sqlparser::dialect::SnowflakeDialect;
10261054
use sqlparser::tokenizer::Span;
10271055

@@ -1783,4 +1811,77 @@ mod tests {
17831811
"SQL error: RecursionLimitExceeded (current limit: 1)"
17841812
);
17851813
}
1814+
1815+
fn expect_parse_expr_ok(sql: &str, expected: ExprWithAlias) {
1816+
let expr = DFParser::parse_sql_into_expr(sql).unwrap();
1817+
assert_eq!(expr, expected, "actual:\n{:#?}", expr);
1818+
}
1819+
1820+
/// Parses sql and asserts that the expected error message was found
1821+
fn expect_parse_expr_error(sql: &str, expected_error: &str) {
1822+
match DFParser::parse_sql_into_expr(sql) {
1823+
Ok(expr) => {
1824+
panic!("Expected parse error for '{sql}', but was successful: {expr:?}");
1825+
}
1826+
Err(e) => {
1827+
let error_message = e.to_string();
1828+
assert!(
1829+
error_message.contains(expected_error),
1830+
"Expected error '{expected_error}' not found in actual error '{error_message}'"
1831+
);
1832+
}
1833+
}
1834+
}
1835+
1836+
#[test]
1837+
fn literal() {
1838+
expect_parse_expr_ok(
1839+
"1234",
1840+
ExprWithAlias {
1841+
expr: Expr::Value(ValueWithSpan::from(Value::Number(
1842+
"1234".to_string(),
1843+
false,
1844+
))),
1845+
alias: None,
1846+
},
1847+
)
1848+
}
1849+
1850+
#[test]
1851+
fn literal_with_alias() {
1852+
expect_parse_expr_ok(
1853+
"1234 as foo",
1854+
ExprWithAlias {
1855+
expr: Expr::Value(ValueWithSpan::from(Value::Number(
1856+
"1234".to_string(),
1857+
false,
1858+
))),
1859+
alias: Some(Ident::from("foo")),
1860+
},
1861+
)
1862+
}
1863+
1864+
#[test]
1865+
fn literal_with_alias_and_trailing_tokens() {
1866+
expect_parse_expr_error("1234 as foo.bar", "Unexpected token .")
1867+
}
1868+
1869+
#[test]
1870+
fn literal_with_alias_and_trailing_whitespace() {
1871+
expect_parse_expr_ok(
1872+
"1234 as foo ",
1873+
ExprWithAlias {
1874+
expr: Expr::Value(ValueWithSpan::from(Value::Number(
1875+
"1234".to_string(),
1876+
false,
1877+
))),
1878+
alias: Some(Ident::from("foo")),
1879+
},
1880+
)
1881+
}
1882+
1883+
#[test]
1884+
fn literal_with_alias_and_trailing_whitespace_and_token() {
1885+
expect_parse_expr_error("1234 as foo bar", "Unexpected token bar")
1886+
}
17861887
}

0 commit comments

Comments
 (0)