MSSQL: Add support for OUTPUT clause on INSERT/UPDATE/DELETE (apache#2228)

Signed-off-by: Guan-Ming Chiu <guanmingchiu@gmail.com>
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Co-authored-by: Ifeanyi Ubah <ifeanyi@validio.io>

Author: guan404ming

src/ast/dml.rs

@@ -79,6 +79,9 @@ pub struct Insert {
     pub on: Option<OnInsert>,
     /// RETURNING
     pub returning: Option<Vec<SelectItem>>,
+    /// OUTPUT (MSSQL)
+    /// See <https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql>
+    pub output: Option<OutputClause>,
     /// Only for mysql
     pub replace_into: bool,
     /// Only for mysql
@@ -203,6 +206,11 @@ impl Display for Insert {
             SpaceOrNewline.fmt(f)?;
         }
 
+        if let Some(output) = &self.output {
+            write!(f, "{output}")?;
+            SpaceOrNewline.fmt(f)?;
+        }
+
         if let Some(settings) = &self.settings {
             write!(f, "SETTINGS {}", display_comma_separated(settings))?;
             SpaceOrNewline.fmt(f)?;
@@ -289,6 +297,9 @@ pub struct Delete {
     pub selection: Option<Expr>,
     /// RETURNING
     pub returning: Option<Vec<SelectItem>>,
+    /// OUTPUT (MSSQL)
+    /// See <https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql>
+    pub output: Option<OutputClause>,
     /// ORDER BY (MySQL)
     pub order_by: Vec<OrderByExpr>,
     /// LIMIT (MySQL)
@@ -314,6 +325,10 @@ impl Display for Delete {
                 indented_list(f, from)?;
             }
         }
+        if let Some(output) = &self.output {
+            SpaceOrNewline.fmt(f)?;
+            write!(f, "{output}")?;
+        }
         if let Some(using) = &self.using {
             SpaceOrNewline.fmt(f)?;
             f.write_str("USING")?;
@@ -367,6 +382,9 @@ pub struct Update {
     pub selection: Option<Expr>,
     /// RETURNING
     pub returning: Option<Vec<SelectItem>>,
+    /// OUTPUT (MSSQL)
+    /// See <https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql>
+    pub output: Option<OutputClause>,
     /// SQLite-specific conflict resolution clause
     pub or: Option<SqliteOnConflict>,
     /// LIMIT
@@ -396,6 +414,10 @@ impl Display for Update {
             f.write_str("SET")?;
             indented_list(f, &self.assignments)?;
         }
+        if let Some(output) = &self.output {
+            SpaceOrNewline.fmt(f)?;
+            write!(f, "{output}")?;
+        }
         if let Some(UpdateTableFromKind::AfterSet(from)) = &self.from {
             SpaceOrNewline.fmt(f)?;
             f.write_str("FROM")?;
@@ -717,11 +739,11 @@ impl Display for MergeUpdateExpr {
     }
 }
 
-/// A `OUTPUT` Clause in the end of a `MERGE` Statement
+/// An `OUTPUT` clause on `MERGE`, `INSERT`, `UPDATE`, or `DELETE` (MSSQL).
 ///
 /// Example:
 /// OUTPUT $action, deleted.* INTO dbo.temp_products;
-/// [mssql](https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql)
+/// <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))]

src/ast/spans.rs

@@ -906,6 +906,7 @@ impl Spanned for Delete {
             using,
             selection,
             returning,
+            output,
             order_by,
             limit,
         } = self;
@@ -923,6 +924,7 @@ impl Spanned for Delete {
                     )
                     .chain(selection.iter().map(|i| i.span()))
                     .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
+                    .chain(output.iter().map(|i| i.span()))
                     .chain(order_by.iter().map(|i| i.span()))
                     .chain(limit.iter().map(|i| i.span())),
             ),
@@ -940,6 +942,7 @@ impl Spanned for Update {
             from,
             selection,
             returning,
+            output,
             or: _,
             limit,
         } = self;
@@ -951,6 +954,7 @@ impl Spanned for Update {
                 .chain(from.iter().map(|i| i.span()))
                 .chain(selection.iter().map(|i| i.span()))
                 .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
+                .chain(output.iter().map(|i| i.span()))
                 .chain(limit.iter().map(|i| i.span())),
         )
     }
@@ -1312,6 +1316,7 @@ impl Spanned for Insert {
             has_table_keyword: _, // bool
             on,
             returning,
+            output,
             replace_into: _, // bool
             priority: _,     // todo, mysql specific
             insert_alias: _, // todo, mysql specific
@@ -1334,7 +1339,8 @@ impl Spanned for Insert {
                 .chain(partitioned.iter().flat_map(|i| i.iter().map(|k| k.span())))
                 .chain(after_columns.iter().map(|i| i.span))
                 .chain(on.as_ref().map(|i| i.span()))
-                .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span()))),
+                .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
+                .chain(output.iter().map(|i| i.span())),
         )
     }
 }

src/dialect/snowflake.rs

@@ -1784,6 +1784,7 @@ fn parse_multi_table_insert(
         has_table_keyword: false,
         on: None,
         returning: None,
+        output: None,
         replace_into: false,
         priority: None,
         insert_alias: None,

src/keywords.rs

@@ -1210,6 +1210,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::ANTI,
     Keyword::SEMI,
     Keyword::RETURNING,
+    Keyword::OUTPUT,
     Keyword::ASOF,
     Keyword::MATCH_CONDITION,
     // for MSSQL-specific OUTER APPLY (seems reserved in most dialects)
@@ -1264,6 +1265,7 @@ pub const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[
     Keyword::CLUSTER,
     Keyword::DISTRIBUTE,
     Keyword::RETURNING,
+    Keyword::VALUES,
     // Reserved only as a column alias in the `SELECT` clause
     Keyword::FROM,
     Keyword::INTO,

src/parser/merge.rs

@@ -218,7 +218,21 @@ impl Parser<'_> {
         self.parse_parenthesized_qualified_column_list(IsOptional::Optional, allow_empty)
     }
 
-    fn parse_output(
+    /// Parses an `OUTPUT` clause if present (MSSQL).
+    pub(super) fn maybe_parse_output_clause(
+        &mut self,
+    ) -> Result<Option<OutputClause>, ParserError> {
+        if self.parse_keyword(Keyword::OUTPUT) {
+            Ok(Some(self.parse_output(
+                Keyword::OUTPUT,
+                self.get_current_token().clone(),
+            )?))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub(super) fn parse_output(
         &mut self,
         start_keyword: Keyword,
         start_token: TokenWithSpan,

src/parser/mod.rs

@@ -13309,6 +13309,9 @@ impl<'a> Parser<'a> {
         };
 
         let from = self.parse_comma_separated(Parser::parse_table_and_joins)?;
+
+        let output = self.maybe_parse_output_clause()?;
+
         let using = if self.parse_keyword(Keyword::USING) {
             Some(self.parse_comma_separated(Parser::parse_table_and_joins)?)
         } else {
@@ -13347,6 +13350,7 @@ impl<'a> Parser<'a> {
             using,
             selection,
             returning,
+            output,
             order_by,
             limit,
         }))
@@ -17275,10 +17279,10 @@ impl<'a> Parser<'a> {
 
             let is_mysql = dialect_of!(self is MySqlDialect);
 
-            let (columns, partitioned, after_columns, source, assignments) = if self
+            let (columns, partitioned, after_columns, output, source, assignments) = if self
                 .parse_keywords(&[Keyword::DEFAULT, Keyword::VALUES])
             {
-                (vec![], None, vec![], None, vec![])
+                (vec![], None, vec![], None, None, vec![])
             } else {
                 let (columns, partitioned, after_columns) = if !self.peek_subquery_start() {
                     let columns = self.parse_parenthesized_column_list(Optional, is_mysql)?;
@@ -17295,6 +17299,8 @@ impl<'a> Parser<'a> {
                     Default::default()
                 };
 
+                let output = self.maybe_parse_output_clause()?;
+
                 let (source, assignments) = if self.peek_keyword(Keyword::FORMAT)
                     || self.peek_keyword(Keyword::SETTINGS)
                 {
@@ -17305,7 +17311,14 @@ impl<'a> Parser<'a> {
                     (Some(self.parse_query()?), vec![])
                 };
 
-                (columns, partitioned, after_columns, source, assignments)
+                (
+                    columns,
+                    partitioned,
+                    after_columns,
+                    output,
+                    source,
+                    assignments,
+                )
             };
 
             let (format_clause, settings) = if self.dialect.supports_insert_format() {
@@ -17407,6 +17420,7 @@ impl<'a> Parser<'a> {
                 has_table_keyword: table,
                 on,
                 returning,
+                output,
                 replace_into,
                 priority,
                 insert_alias,
@@ -17512,6 +17526,9 @@ impl<'a> Parser<'a> {
         };
         self.expect_keyword(Keyword::SET)?;
         let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
+
+        let output = self.maybe_parse_output_clause()?;
+
         let from = if from_before_set.is_none() && self.parse_keyword(Keyword::FROM) {
             Some(UpdateTableFromKind::AfterSet(
                 self.parse_table_with_joins()?,
@@ -17542,6 +17559,7 @@ impl<'a> Parser<'a> {
             from,
             selection,
             returning,
+            output,
             or,
             limit,
         }

tests/sqlparser_common.rs

@@ -530,6 +530,7 @@ fn parse_update_set_from() {
                 ])),
             }),
             returning: None,
+            output: None,
             or: None,
             limit: None
         })
@@ -553,6 +554,7 @@ fn parse_update_with_table_alias() {
             limit: None,
             optimizer_hints,
             update_token: _,
+            output: _,
         }) if optimizer_hints.is_empty() => {
             assert_eq!(
                 TableWithJoins {

tests/sqlparser_mssql.rs

@@ -2806,3 +2806,45 @@ fn test_exec_dynamic_sql() {
         .expect("EXEC (@sql) followed by DROP TABLE should parse");
     assert_eq!(stmts.len(), 2);
 }
+
+// MSSQL OUTPUT clause on INSERT/UPDATE/DELETE
+// https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql
+#[test]
+fn parse_mssql_insert_with_output() {
+    ms_and_generic().verified_stmt(
+        "INSERT INTO customers (name, email) OUTPUT INSERTED.id, INSERTED.name VALUES ('John', 'john@example.com')",
+    );
+}
+
+#[test]
+fn parse_mssql_insert_with_output_into() {
+    ms_and_generic().verified_stmt(
+        "INSERT INTO customers (name, email) OUTPUT INSERTED.id, INSERTED.name INTO @new_ids VALUES ('John', 'john@example.com')",
+    );
+}
+
+#[test]
+fn parse_mssql_delete_with_output() {
+    ms_and_generic().verified_stmt("DELETE FROM customers OUTPUT DELETED.* WHERE id = 1");
+}
+
+#[test]
+fn parse_mssql_delete_with_output_into() {
+    ms_and_generic().verified_stmt(
+        "DELETE FROM customers OUTPUT DELETED.id, DELETED.name INTO @deleted_rows WHERE active = 0",
+    );
+}
+
+#[test]
+fn parse_mssql_update_with_output() {
+    ms_and_generic().verified_stmt(
+        "UPDATE employees SET salary = salary * 1.1 OUTPUT INSERTED.id, DELETED.salary, INSERTED.salary WHERE department = 'Engineering'",
+    );
+}
+
+#[test]
+fn parse_mssql_update_with_output_into() {
+    ms_and_generic().verified_stmt(
+        "UPDATE employees SET salary = salary * 1.1 OUTPUT INSERTED.id, DELETED.salary, INSERTED.salary INTO @changes WHERE department = 'Engineering'",
+    );
+}

tests/sqlparser_mysql.rs

@@ -2671,6 +2671,7 @@ fn parse_update_with_joins() {
             limit: None,
             optimizer_hints,
             update_token: _,
+            output: _,
         }) if optimizer_hints.is_empty() => {
             assert_eq!(
                 TableWithJoins {

tests/sqlparser_postgres.rs

@@ -5530,6 +5530,7 @@ fn test_simple_postgres_insert_with_alias() {
             has_table_keyword: false,
             on: None,
             returning: None,
+            output: None,
             replace_into: false,
             priority: None,
             insert_alias: None,
@@ -5612,6 +5613,7 @@ fn test_simple_postgres_insert_with_alias() {
             has_table_keyword: false,
             on: None,
             returning: None,
+            output: None,
             replace_into: false,
             priority: None,
             insert_alias: None,
@@ -5692,6 +5694,7 @@ fn test_simple_insert_with_quoted_alias() {
             has_table_keyword: false,
             on: None,
             returning: None,
+            output: None,
             replace_into: false,
             priority: None,
             insert_alias: None,