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. 

Pull request description:
MSSQL: Add support for functionality `MERGE` output clause
Hey I have added a feature to handle OUTPUT Statements in the end of merge statements, which is used in MS SQL fixes #1789  that I created. 

- I've added the OUTPUT Keyword
- I've added a struct in for Output in ast/mod.rs
- I've added an optional Output parameter to the Merge struct in ast/mod.rs
- added function for parsing the output parameter in parser/mod.rs
- updated parse merge clause function, to break on output keyword
- updated parse merge to handling of the output clause
- added and updated tests

This is my first Pull Request to an open source project so please let me know, if there is something missing, or something I did wrong. 

Thank you for a great crate.

Patch:
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 07a2279..8537bd8 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 bf1206f..a1b4c0c 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 40d6b0a..50a9aae 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,
         })
     }
 


Test patch:
diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs
index 3037d4a..5eb30d1 100644
--- a/tests/sqlparser_bigquery.rs
+++ b/tests/sqlparser_bigquery.rs
@@ -1735,6 +1735,7 @@ fn parse_merge() {
             },
         ],
     };
+
     match bigquery_and_generic().verified_stmt(sql) {
         Statement::Merge {
             into,
@@ -1742,6 +1743,7 @@ fn parse_merge() {
             source,
             on,
             clauses,
+            ..
         } => {
             assert!(!into);
             assert_eq!(
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 9fe6eae..3653321 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -9359,6 +9359,7 @@ fn parse_merge() {
                 source,
                 on,
                 clauses,
+                ..
             },
             Statement::Merge {
                 into: no_into,
@@ -9366,6 +9367,7 @@ fn parse_merge() {
                 source: source_no_into,
                 on: on_no_into,
                 clauses: clauses_no_into,
+                ..
             },
         ) => {
             assert!(into);
@@ -9558,6 +9560,19 @@ fn parse_merge() {
     verified_stmt(sql);
 }
 
+#[test]
+fn test_merge_with_output() {
+    let sql = "MERGE INTO target_table USING source_table \
+        ON target_table.id = source_table.oooid \
+        WHEN MATCHED THEN \
+            UPDATE SET target_table.description = source_table.description \
+        WHEN NOT MATCHED THEN \
+            INSERT (ID, description) VALUES (source_table.id, source_table.description) \
+        OUTPUT inserted.* INTO log_target";
+
+    verified_stmt(sql);
+}
+
 #[test]
 fn test_merge_into_using_table() {
     let sql = "MERGE INTO target_table USING source_table \
diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs
index 2bfc38a..5d76fd0 100644
--- a/tests/sqlparser_mssql.rs
+++ b/tests/sqlparser_mssql.rs
@@ -1921,3 +1921,19 @@ fn ms() -> TestedDialects {
 fn ms_and_generic() -> TestedDialects {
     TestedDialects::new(vec![Box::new(MsSqlDialect {}), Box::new(GenericDialect {})])
 }
+
+#[test]
+fn parse_mssql_merge_with_output() {
+    let stmt = "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.ProductID INTO dsi.temp_products";
+    ms_and_generic().verified_stmt(stmt);
+}
diff --git a/tests/sqlparser_redshift.rs b/tests/sqlparser_redshift.rs
index c75abe1..060e385 100644
--- a/tests/sqlparser_redshift.rs
+++ b/tests/sqlparser_redshift.rs
@@ -395,5 +395,5 @@ fn test_parse_nested_quoted_identifier() {
 #[test]
 fn parse_extract_single_quotes() {
     let sql = "SELECT EXTRACT('month' FROM my_timestamp) FROM my_table";
-    redshift().verified_stmt(&sql);
+    redshift().verified_stmt(sql);
 }