diff --git a/src/ast/mod.rs b/src/ast/mod.rs index cd937857..71cb6c5e 100644 --- a/src/ast/mod.rs +++ b/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, - pub into_table: SelectInto, +pub enum OutputClause { + Output { + select_items: Vec, + into_table: Option, + }, + Returning { + select_items: Vec, + }, } 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) + } + } } } diff --git a/src/ast/query.rs b/src/ast/query.rs index 2ef456b1..967af85c 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -161,6 +161,7 @@ pub enum SetExpr { Insert(Statement), Update(Statement), Delete(Statement), + Merge(Statement), Table(Box), } @@ -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, diff --git a/src/ast/spans.rs b/src/ast/spans.rs index add6c390..39761751 100644 --- a/src/ast/spans.rs +++ b/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(), } } } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index c4c72e9c..6c559eed 100644 --- a/src/parser/mod.rs +++ b/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, ParserError> { + Ok(Box::new(SetExpr::Merge(self.parse_merge()?))) + } + 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. @@ -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 { - self.expect_keyword_is(Keyword::OUTPUT)?; + fn parse_output(&mut self, start_keyword: Keyword) -> Result { 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 { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 54ad1732..8b99bb1d 100644 --- a/tests/sqlparser_common.rs +++ b/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 \