mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-09-14 01:46:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Support subquery expression in SET expressions (#1343)

Browse source
This commit is contained in:
Ifeanyi Ubah 2024-07-20 12:55:24 +02:00 • committed by GitHub
parent 845a1aaddd
commit 028ada8350
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 57 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

42
src/parser/mod.rs
View file

@ -1208,20 +1208,18 @@ impl<'a> Parser<'a> {
Ok(Expr::Value(self.parse_value()?)) Ok(Expr::Value(self.parse_value()?))
} }
Token::LParen => { Token::LParen => {
let expr = let expr = if let Some(expr) = self.try_parse_expr_sub_query()? {
if self.parse_keyword(Keyword::SELECT) || self.parse_keyword(Keyword::WITH) { expr
self.prev_token(); } else if let Some(lambda) = self.try_parse_lambda() {
Expr::Subquery(self.parse_boxed_query()?) return Ok(lambda);
} else if let Some(lambda) = self.try_parse_lambda() { } else {
return Ok(lambda); let exprs = self.parse_comma_separated(Parser::parse_expr)?;
} else { match exprs.len() {
let exprs = self.parse_comma_separated(Parser::parse_expr)?; 0 => unreachable!(), // parse_comma_separated ensures 1 or more
match exprs.len() { 1 => Expr::Nested(Box::new(exprs.into_iter().next().unwrap())),
0 => unreachable!(), // parse_comma_separated ensures 1 or more _ => Expr::Tuple(exprs),
1 => Expr::Nested(Box::new(exprs.into_iter().next().unwrap())), }
_ => Expr::Tuple(exprs), };
}
};
self.expect_token(&Token::RParen)?; self.expect_token(&Token::RParen)?;
if !self.consume_token(&Token::Period) { if !self.consume_token(&Token::Period) {
Ok(expr) Ok(expr)
@ -1263,6 +1261,18 @@ impl<'a> Parser<'a> {
} }
} }
fn try_parse_expr_sub_query(&mut self) -> Result<Option<Expr>, ParserError> {
if self
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_none()
{
return Ok(None);
}
self.prev_token();
Ok(Some(Expr::Subquery(self.parse_boxed_query()?)))
}
fn try_parse_lambda(&mut self) -> Option<Expr> { fn try_parse_lambda(&mut self) -> Option<Expr> {
if !self.dialect.supports_lambda_functions() { if !self.dialect.supports_lambda_functions() {
return None; return None;
@ -8709,7 +8719,9 @@ impl<'a> Parser<'a> {
let mut values = vec![]; let mut values = vec![];
loop { loop {
let value = if let Ok(expr) = self.parse_expr() { let value = if let Some(expr) = self.try_parse_expr_sub_query()? {
expr
} else if let Ok(expr) = self.parse_expr() {
expr expr
} else { } else {
self.expected("variable value", self.peek_token())? self.expected("variable value", self.peek_token())?

30
tests/sqlparser_common.rs
View file

@ -7135,9 +7135,39 @@ fn parse_set_variable() {
_ => unreachable!(), _ => unreachable!(),
} }
// Subquery expression
for (sql, canonical) in [
(
"SET (a) = (SELECT 22 FROM tbl1)",
"SET (a) = ((SELECT 22 FROM tbl1))",
),
(
"SET (a) = (SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2))",
"SET (a) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)))",
),
(
"SET (a) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)))",
"SET (a) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)))",
),
(
"SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), SELECT 33 FROM tbl3)",
"SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), (SELECT 33 FROM tbl3))",
),
] {
multi_variable_dialects.one_statement_parses_to(sql, canonical);
}
let error_sqls = [ let error_sqls = [
("SET (a, b, c) = (1, 2, 3", "Expected: ), found: EOF"), ("SET (a, b, c) = (1, 2, 3", "Expected: ), found: EOF"),
("SET (a, b, c) = 1, 2, 3", "Expected: (, found: 1"), ("SET (a, b, c) = 1, 2, 3", "Expected: (, found: 1"),
(
"SET (a) = ((SELECT 22 FROM tbl1)",
"Expected: ), found: EOF",
),
(
"SET (a) = ((SELECT 22 FROM tbl1) (SELECT 22 FROM tbl1))",
"Expected: ), found: (",
),
]; ];
for (sql, error) in error_sqls { for (sql, error) in error_sqls {
assert_eq!( assert_eq!(