MSSQL: Add support for functionality MERGE output clause (apache#1790)

Author: dilovancelik

src/ast/mod.rs

@@ -3835,6 +3835,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,
@@ -3846,6 +3847,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> ]
@@ -5425,14 +5428,19 @@ impl fmt::Display for Statement {
                 source,
                 on,
                 clauses,
+                output,
             } => {
                 write!(
                     f,
                     "MERGE{int} {table} USING {source} ",
                     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,
@@ -7963,6 +7971,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))]

src/keywords.rs

@@ -632,6 +632,7 @@ define_keywords!(
     ORGANIZATION,
     OUT,
     OUTER,
+    OUTPUT,
     OUTPUTFORMAT,
     OVER,
     OVERFLOW,

src/parser/mod.rs

@@ -10920,18 +10920,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
         };
@@ -14523,10 +14512,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) {
@@ -14620,6 +14608,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);
 
@@ -14630,13 +14646,19 @@ 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,
             table,
             source,
             on: Box::new(on),
             clauses,
+            output,
         })
     }
 

tests/sqlparser_bigquery.rs

@@ -1735,13 +1735,15 @@ fn parse_merge() {
             },
         ],
     };
+
     match bigquery_and_generic().verified_stmt(sql) {
         Statement::Merge {
             into,
             table,
             source,
             on,
             clauses,
+            ..
         } => {
             assert!(!into);
             assert_eq!(

tests/sqlparser_common.rs

@@ -9360,13 +9360,15 @@ fn parse_merge() {
                 source,
                 on,
                 clauses,
+                ..
             },
             Statement::Merge {
                 into: no_into,
                 table: table_no_into,
                 source: source_no_into,
                 on: on_no_into,
                 clauses: clauses_no_into,
+                ..
             },
         ) => {
             assert!(into);
@@ -9559,6 +9561,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 \

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);
+}

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);
 }