feat: MERGE statements: add RETURNING and OUTPUT without INTO (apache#2011)

Author: lovasoa

src/ast/mod.rs

@@ -9107,24 +9107,36 @@ impl Display for MergeClause {
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub struct OutputClause {
-    pub select_items: Vec<SelectItem>,
-    pub into_table: SelectInto,
+pub enum OutputClause {
+    Output {
+        select_items: Vec<SelectItem>,
+        into_table: Option<SelectInto>,
+    },
+    Returning {
+        select_items: Vec<SelectItem>,
+    },
 }
 
 impl fmt::Display for OutputClause {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let OutputClause {
-            select_items,
-            into_table,
-        } = self;
-
-        write!(
-            f,
-            "OUTPUT {} {}",
-            display_comma_separated(select_items),
-            into_table
-        )
+        match self {
+            OutputClause::Output {
+                select_items,
+                into_table,
+            } => {
+                f.write_str("OUTPUT ")?;
+                display_comma_separated(select_items).fmt(f)?;
+                if let Some(into_table) = into_table {
+                    f.write_str(" ")?;
+                    into_table.fmt(f)?;
+                }
+                Ok(())
+            }
+            OutputClause::Returning { select_items } => {
+                f.write_str("RETURNING ")?;
+                display_comma_separated(select_items).fmt(f)
+            }
+        }
     }
 }
 

src/ast/query.rs

@@ -161,6 +161,7 @@ pub enum SetExpr {
     Insert(Statement),
     Update(Statement),
     Delete(Statement),
+    Merge(Statement),
     Table(Box<Table>),
 }
 
@@ -188,6 +189,7 @@ impl fmt::Display for SetExpr {
             SetExpr::Insert(v) => v.fmt(f),
             SetExpr::Update(v) => v.fmt(f),
             SetExpr::Delete(v) => v.fmt(f),
+            SetExpr::Merge(v) => v.fmt(f),
             SetExpr::Table(t) => t.fmt(f),
             SetExpr::SetOperation {
                 left,

src/ast/spans.rs

@@ -214,6 +214,7 @@ impl Spanned for SetExpr {
             SetExpr::Table(_) => Span::empty(),
             SetExpr::Update(statement) => statement.span(),
             SetExpr::Delete(statement) => statement.span(),
+            SetExpr::Merge(statement) => statement.span(),
         }
     }
 }

src/parser/mod.rs

@@ -11508,6 +11508,13 @@ impl<'a> Parser<'a> {
         Ok(Box::new(SetExpr::Delete(self.parse_delete()?)))
     }
 
+    /// Parse a MERGE statement, returning a `Box`ed SetExpr
+    ///
+    /// This is used to reduce the size of the stack frames in debug builds
+    fn parse_merge_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
+        Ok(Box::new(SetExpr::Merge(self.parse_merge()?)))
+    }
+
     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.
@@ -11719,6 +11726,20 @@ impl<'a> Parser<'a> {
                 pipe_operators: vec![],
             }
             .into())
+        } else if self.parse_keyword(Keyword::MERGE) {
+            Ok(Query {
+                with,
+                body: self.parse_merge_setexpr_boxed()?,
+                limit_clause: None,
+                order_by: None,
+                fetch: None,
+                locks: vec![],
+                for_clause: None,
+                settings: None,
+                format_clause: None,
+                pipe_operators: vec![],
+            }
+            .into())
         } else {
             let body = self.parse_query_body(self.dialect.prec_unknown())?;
 
@@ -16571,15 +16592,22 @@ impl<'a> Parser<'a> {
         Ok(clauses)
     }
 
-    fn parse_output(&mut self) -> Result<OutputClause, ParserError> {
-        self.expect_keyword_is(Keyword::OUTPUT)?;
+    fn parse_output(&mut self, start_keyword: Keyword) -> Result<OutputClause, ParserError> {
         let select_items = self.parse_projection()?;
-        self.expect_keyword_is(Keyword::INTO)?;
-        let into_table = self.parse_select_into()?;
+        let into_table = if start_keyword == Keyword::OUTPUT && self.peek_keyword(Keyword::INTO) {
+            self.expect_keyword_is(Keyword::INTO)?;
+            Some(self.parse_select_into()?)
+        } else {
+            None
+        };
 
-        Ok(OutputClause {
-            select_items,
-            into_table,
+        Ok(if start_keyword == Keyword::OUTPUT {
+            OutputClause::Output {
+                select_items,
+                into_table,
+            }
+        } else {
+            OutputClause::Returning { select_items }
         })
     }
 
@@ -16609,10 +16637,9 @@ impl<'a> Parser<'a> {
         self.expect_keyword_is(Keyword::ON)?;
         let on = self.parse_expr()?;
         let clauses = self.parse_merge_clauses()?;
-        let output = if self.peek_keyword(Keyword::OUTPUT) {
-            Some(self.parse_output()?)
-        } else {
-            None
+        let output = match self.parse_one_of_keywords(&[Keyword::OUTPUT, Keyword::RETURNING]) {
+            Some(start_keyword) => Some(self.parse_output(start_keyword)?),
+            None => None,
         };
 
         Ok(Statement::Merge {

tests/sqlparser_common.rs

@@ -9902,6 +9902,29 @@ fn parse_merge() {
     verified_stmt(sql);
 }
 
+#[test]
+fn test_merge_in_cte() {
+    verified_only_select(
+        "WITH x AS (\
+            MERGE INTO t USING (VALUES (1)) ON 1 = 1 \
+            WHEN MATCHED THEN DELETE \
+            RETURNING *\
+        ) SELECT * FROM x",
+    );
+}
+
+#[test]
+fn test_merge_with_returning() {
+    let sql = "MERGE INTO wines AS w \
+    USING wine_stock_changes AS s \
+        ON s.winename = w.winename \
+    WHEN NOT MATCHED AND s.stock_delta > 0 THEN INSERT VALUES (s.winename, s.stock_delta) \
+    WHEN MATCHED AND w.stock + s.stock_delta > 0 THEN UPDATE SET stock = w.stock + s.stock_delta \
+    WHEN MATCHED THEN DELETE \
+    RETURNING merge_action(), w.*";
+    verified_stmt(sql);
+}
+
 #[test]
 fn test_merge_with_output() {
     let sql = "MERGE INTO target_table USING source_table \
@@ -9915,6 +9938,14 @@ fn test_merge_with_output() {
     verified_stmt(sql);
 }
 
+#[test]
+fn test_merge_with_output_without_into() {
+    let sql = "MERGE INTO a USING b ON a.id = b.id \
+        WHEN MATCHED THEN DELETE \
+        OUTPUT inserted.*";
+    verified_stmt(sql);
+}
+
 #[test]
 fn test_merge_into_using_table() {
     let sql = "MERGE INTO target_table USING source_table \