Repository: apache/datafusion-sqlparser-rs (Rust)

Issue:
Add support for OUTPUT Clause in TSQL (MS SQL) Merge statements
Add the end of a merge statement in T-SQL we can use an OUTPUT clause to write the changed rows to a table variable. This is often used in Data Warehouse workloads were we would like to track slowly changing dimensions. 

A merge statement with the output clause could look like this:
```sql

MERGE dso.products AS t
USING dsi.products AS s
	ON s.ProductID = t.ProductID
WHEN MATCHED 
	AND NOT (t.[ProductName] = s.[ProductName] OR (ISNULL(t.[ProductName], s.[ProductName]) IS NULL))
THEN
	UPDATE SET t.[ProductName] = s.[ProductName]

WHEN NOT MATCHED BY TARGET THEN
INSERT ([ProductID]
           ,[ProductName])
     VALUES
           (s.[ProductID]
           ,s.[ProductName])

WHEN NOT MATCHED BY SOURCE  THEN
	DELETE

OUTPUT $action, deleted.* INTO dsi.temp_products;
```

Currently this fails with the error `ParserError("Expected: WHEN, found: OUTPUT at Line: 112, Column: 1")`.

I will try to create a PR solving this issue, but any pointers on a good approach would be highly appreciated. 

Thanks for the awesome crate btw. 



Fix (patch):
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 07a22798e..8537bd858 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -3817,6 +3817,7 @@ pub enum Statement {
     /// ```
     /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/merge)
     /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement)
+    /// [MSSQL](https://learn.microsoft.com/en-us/sql/t-sql/statements/merge-transact-sql?view=sql-server-ver16)
     Merge {
         /// optional INTO keyword
         into: bool,
@@ -3828,6 +3829,8 @@ pub enum Statement {
         on: Box<Expr>,
         /// Specifies the actions to perform when values match or do not match.
         clauses: Vec<MergeClause>,
+        // Specifies the output to save changes in MSSQL
+        output: Option<OutputClause>,
     },
     /// ```sql
     /// CACHE [ FLAG ] TABLE <table_name> [ OPTIONS('K1' = 'V1', 'K2' = V2) ] [ AS ] [ <query> ]
@@ -5407,6 +5410,7 @@ impl fmt::Display for Statement {
                 source,
                 on,
                 clauses,
+                output,
             } => {
                 write!(
                     f,
@@ -5414,7 +5418,11 @@ impl fmt::Display for Statement {
                     int = if *into { " INTO" } else { "" }
                 )?;
                 write!(f, "ON {on} ")?;
-                write!(f, "{}", display_separated(clauses, " "))
+                write!(f, "{}", display_separated(clauses, " "))?;
+                if let Some(output) = output {
+                    write!(f, " {output}")?;
+                }
+                Ok(())
             }
             Statement::Cache {
                 table_name,
@@ -7945,6 +7953,35 @@ impl Display for MergeClause {
     }
 }
 
+/// A Output Clause in the end of a 'MERGE' Statement
+///
+/// Example:
+/// OUTPUT $action, deleted.* INTO dbo.temp_products;
+/// [mssql](https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql)
+#[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,
+}
+
+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
+        )
+    }
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
diff --git a/src/keywords.rs b/src/keywords.rs
index bf1206f6a..a1b4c0c3a 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -632,6 +632,7 @@ define_keywords!(
     ORGANIZATION,
     OUT,
     OUTER,
+    OUTPUT,
     OUTPUTFORMAT,
     OVER,
     OVERFLOW,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 40d6b0acd..50a9aae84 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -10910,18 +10910,7 @@ impl<'a> Parser<'a> {
             };
 
         let into = if self.parse_keyword(Keyword::INTO) {
-            let temporary = self
-                .parse_one_of_keywords(&[Keyword::TEMP, Keyword::TEMPORARY])
-                .is_some();
-            let unlogged = self.parse_keyword(Keyword::UNLOGGED);
-            let table = self.parse_keyword(Keyword::TABLE);
-            let name = self.parse_object_name(false)?;
-            Some(SelectInto {
-                temporary,
-                unlogged,
-                table,
-                name,
-            })
+            Some(self.parse_select_into()?)
         } else {
             None
         };
@@ -14513,10 +14502,9 @@ impl<'a> Parser<'a> {
     pub fn parse_merge_clauses(&mut self) -> Result<Vec<MergeClause>, ParserError> {
         let mut clauses = vec![];
         loop {
-            if self.peek_token() == Token::EOF || self.peek_token() == Token::SemiColon {
+            if !(self.parse_keyword(Keyword::WHEN)) {
                 break;
             }
-            self.expect_keyword_is(Keyword::WHEN)?;
 
             let mut clause_kind = MergeClauseKind::Matched;
             if self.parse_keyword(Keyword::NOT) {
@@ -14610,6 +14598,34 @@ impl<'a> Parser<'a> {
         Ok(clauses)
     }
 
+    fn parse_output(&mut self) -> Result<OutputClause, ParserError> {
+        self.expect_keyword_is(Keyword::OUTPUT)?;
+        let select_items = self.parse_projection()?;
+        self.expect_keyword_is(Keyword::INTO)?;
+        let into_table = self.parse_select_into()?;
+
+        Ok(OutputClause {
+            select_items,
+            into_table,
+        })
+    }
+
+    fn parse_select_into(&mut self) -> Result<SelectInto, ParserError> {
+        let temporary = self
+            .parse_one_of_keywords(&[Keyword::TEMP, Keyword::TEMPORARY])
+            .is_some();
+        let unlogged = self.parse_keyword(Keyword::UNLOGGED);
+        let table = self.parse_keyword(Keyword::TABLE);
+        let name = self.parse_object_name(false)?;
+
+        Ok(SelectInto {
+            temporary,
+            unlogged,
+            table,
+            name,
+        })
+    }
+
     pub fn parse_merge(&mut self) -> Result<Statement, ParserError> {
         let into = self.parse_keyword(Keyword::INTO);
 
@@ -14620,6 +14636,11 @@ 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
+        };
 
         Ok(Statement::Merge {
             into,
@@ -14627,6 +14648,7 @@ impl<'a> Parser<'a> {
             source,
             on: Box::new(on),
             clauses,
+            output,
         })
     }