add support for with clauses in delete statements

fixes apache#1763

Author: lovasoa

src/ast/query.rs

@@ -156,6 +156,7 @@ pub enum SetExpr {
     Values(Values),
     Insert(Statement),
     Update(Statement),
+    Delete(Statement),
     Table(Box<Table>),
 }
 
@@ -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,

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(),
         }
     }
 }

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<Box<SetExpr>, ParserError> {
+        Ok(Box::new(SetExpr::Delete(self.parse_delete()?)))
+    }
+
     pub fn parse_delete(&mut self) -> Result<Statement, ParserError> {
         let (tables, with_from_keyword) = if !self.parse_keyword(Keyword::FROM) {
             // `FROM` keyword is optional in BigQuery SQL.
@@ -10202,19 +10209,25 @@ 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<With, ParserError> {
+        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
     /// expect the initial keyword to be already consumed
     pub fn parse_query(&mut self) -> Result<Box<Query>, 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())?;
 

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";