From ab8b8724f180066146af85f976df9d3424f5cf30 Mon Sep 17 00:00:00 2001 From: lovasoa Date: Mon, 10 Mar 2025 22:55:36 +0100 Subject: [PATCH 1/2] add support for with clauses in delete statements fixes https://github.com/apache/datafusion-sqlparser-rs/issues/1763 --- src/ast/query.rs | 2 ++ src/ast/spans.rs | 1 + src/parser/mod.rs | 40 +++++++++++++++++++++++++++++++++------ tests/sqlparser_common.rs | 27 ++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 6 deletions(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index bed991114..12f72932f 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -156,6 +156,7 @@ pub enum SetExpr { Values(Values), Insert(Statement), Update(Statement), + Delete(Statement), Table(Box), } @@ -178,6 +179,7 @@ impl fmt::Display for SetExpr { SetExpr::Values(v) => write!(f, "{v}"), SetExpr::Insert(v) => write!(f, "{v}"), SetExpr::Update(v) => write!(f, "{v}"), + SetExpr::Delete(v) => write!(f, "{v}"), SetExpr::Table(t) => write!(f, "{t}"), SetExpr::SetOperation { left, diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 0a64fb8ea..f69dc014f 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -191,6 +191,7 @@ impl Spanned for SetExpr { SetExpr::Insert(statement) => statement.span(), SetExpr::Table(_) => Span::empty(), SetExpr::Update(statement) => statement.span(), + SetExpr::Delete(statement) => statement.span(), } } } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index b34415388..6a7303015 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -10049,6 +10049,13 @@ impl<'a> Parser<'a> { Ok(parent_type(inside_type.into())) } + /// Parse a DELETE statement, returning a `Box`ed SetExpr + /// + /// This is used to reduce the size of the stack frames in debug builds + fn parse_delete_setexpr_boxed(&mut self) -> Result, ParserError> { + Ok(Box::new(SetExpr::Delete(self.parse_delete()?))) + } + pub fn parse_delete(&mut self) -> Result { let (tables, with_from_keyword) = if !self.parse_keyword(Keyword::FROM) { // `FROM` keyword is optional in BigQuery SQL. @@ -10202,6 +10209,17 @@ impl<'a> Parser<'a> { } } + /// Parse a `WITH` clause, i.e. a `WITH` keyword followed by a `RECURSIVE` keyword + /// and a comma-separated list of CTE declarations. + fn parse_with_clause(&mut self) -> Result { + let with_token = self.get_current_token(); + Ok(With { + with_token: with_token.clone().into(), + recursive: self.parse_keyword(Keyword::RECURSIVE), + cte_tables: self.parse_comma_separated(Parser::parse_cte)?, + }) + } + /// Parse a query expression, i.e. a `SELECT` statement optionally /// preceded with some `WITH` CTE declarations and optionally followed /// by `ORDER BY`. Unlike some other parse_... methods, this one doesn't @@ -10209,12 +10227,7 @@ impl<'a> Parser<'a> { pub fn parse_query(&mut self) -> Result, ParserError> { let _guard = self.recursion_counter.try_decrease()?; let with = if self.parse_keyword(Keyword::WITH) { - let with_token = self.get_current_token(); - Some(With { - with_token: with_token.clone().into(), - recursive: self.parse_keyword(Keyword::RECURSIVE), - cte_tables: self.parse_comma_separated(Parser::parse_cte)?, - }) + Some(self.parse_with_clause()?) } else { None }; @@ -10248,6 +10261,21 @@ impl<'a> Parser<'a> { format_clause: None, } .into()) + } else if self.parse_keyword(Keyword::DELETE) { + Ok(Query { + with, + body: self.parse_delete_setexpr_boxed()?, + limit: None, + limit_by: vec![], + order_by: None, + offset: None, + fetch: None, + locks: vec![], + for_clause: None, + settings: None, + format_clause: None, + } + .into()) } else { let body = self.parse_query_body(self.dialect.prec_unknown())?; diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index a8ccd70a7..8225d3672 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -7383,6 +7383,33 @@ fn parse_recursive_cte() { assert_eq!(with.cte_tables.first().unwrap(), &expected); } +#[test] +fn parse_cte_in_data_modification_statements() { + match verified_stmt("WITH x AS (SELECT 1) UPDATE t SET bar = (SELECT * FROM x)") { + Statement::Query(query) => { + assert_eq!(query.with.unwrap().to_string(), "WITH x AS (SELECT 1)"); + assert!(matches!(*query.body, SetExpr::Update(_))); + } + other => panic!("Expected: UPDATE, got: {:?}", other), + } + + match verified_stmt("WITH t (x) AS (SELECT 9) DELETE FROM q WHERE id IN (SELECT x FROM t)") { + Statement::Query(query) => { + assert_eq!(query.with.unwrap().to_string(), "WITH t (x) AS (SELECT 9)"); + assert!(matches!(*query.body, SetExpr::Delete(_))); + } + other => panic!("Expected: DELETE, got: {:?}", other), + } + + match verified_stmt("WITH x AS (SELECT 42) INSERT INTO t SELECT foo FROM x") { + Statement::Query(query) => { + assert_eq!(query.with.unwrap().to_string(), "WITH x AS (SELECT 42)"); + assert!(matches!(*query.body, SetExpr::Insert(_))); + } + other => panic!("Expected: INSERT, got: {:?}", other), + } +} + #[test] fn parse_derived_tables() { let sql = "SELECT a.x, b.y FROM (SELECT x FROM foo) AS a CROSS JOIN (SELECT y FROM bar) AS b"; From fbc89a9bea51383a9bdfd4e024eadffdb5ad3470 Mon Sep 17 00:00:00 2001 From: lovasoa Date: Tue, 11 Mar 2025 20:31:17 +0100 Subject: [PATCH 2/2] inline extracted function https://github.com/apache/datafusion-sqlparser-rs/pull/1764#discussion_r1988453837 --- src/parser/mod.rs | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 6a7303015..88ed3b080 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -10209,17 +10209,6 @@ impl<'a> Parser<'a> { } } - /// Parse a `WITH` clause, i.e. a `WITH` keyword followed by a `RECURSIVE` keyword - /// and a comma-separated list of CTE declarations. - fn parse_with_clause(&mut self) -> Result { - let with_token = self.get_current_token(); - Ok(With { - with_token: with_token.clone().into(), - recursive: self.parse_keyword(Keyword::RECURSIVE), - cte_tables: self.parse_comma_separated(Parser::parse_cte)?, - }) - } - /// Parse a query expression, i.e. a `SELECT` statement optionally /// preceded with some `WITH` CTE declarations and optionally followed /// by `ORDER BY`. Unlike some other parse_... methods, this one doesn't @@ -10227,7 +10216,12 @@ impl<'a> Parser<'a> { pub fn parse_query(&mut self) -> Result, ParserError> { let _guard = self.recursion_counter.try_decrease()?; let with = if self.parse_keyword(Keyword::WITH) { - Some(self.parse_with_clause()?) + let with_token = self.get_current_token(); + Some(With { + with_token: with_token.clone().into(), + recursive: self.parse_keyword(Keyword::RECURSIVE), + cte_tables: self.parse_comma_separated(Parser::parse_cte)?, + }) } else { None };