impl Spanned for MERGE statements (#2100)

src/ast/mod.rs

@@ -4064,6 +4064,8 @@ pub enum Statement {
     /// [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 {
+        /// The `MERGE` token that starts the statement.
+        merge_token: AttachedToken,
         /// optional INTO keyword
         into: bool,
         /// Specifies the table to merge
@@ -4088,7 +4090,6 @@ pub enum Statement {
         /// Table flag
         table_flag: Option<ObjectName>,
         /// Table name
-
         #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
         table_name: ObjectName,
         has_as: bool,
@@ -5488,6 +5489,7 @@ impl fmt::Display for Statement {
                 write!(f, "RELEASE SAVEPOINT {name}")
             }
             Statement::Merge {
+                merge_token: _,
                 into,
                 table,
                 source,
@@ -8620,6 +8622,8 @@ impl Display for MergeInsertKind {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct MergeInsertExpr {
+    /// The `INSERT` token that starts the sub-expression.
+    pub insert_token: AttachedToken,
     /// Columns (if any) specified by the insert.
     ///
     /// Example:
@@ -8628,6 +8632,8 @@ pub struct MergeInsertExpr {
     /// INSERT (product, quantity) ROW
     /// ```
     pub columns: Vec<Ident>,
+    /// The token, `[VALUES | ROW]` starting `kind`.
+    pub kind_token: AttachedToken,
     /// The insert type used by the statement.
     pub kind: MergeInsertKind,
 }
@@ -8667,9 +8673,16 @@ pub enum MergeAction {
     /// ```sql
     /// UPDATE SET quantity = T.quantity + S.quantity
     /// ```
-    Update { assignments: Vec<Assignment> },
+    Update {
+        /// The `UPDATE` token that starts the sub-expression.
+        update_token: AttachedToken,
+        assignments: Vec<Assignment>,
+    },
     /// A plain `DELETE` clause
-    Delete,
+    Delete {
+        /// The `DELETE` token that starts the sub-expression.
+        delete_token: AttachedToken,
+    },
 }
 
 impl Display for MergeAction {
@@ -8678,10 +8691,10 @@ impl Display for MergeAction {
             MergeAction::Insert(insert) => {
                 write!(f, "INSERT {insert}")
             }
-            MergeAction::Update { assignments } => {
+            MergeAction::Update { assignments, .. } => {
                 write!(f, "UPDATE SET {}", display_comma_separated(assignments))
             }
-            MergeAction::Delete => {
+            MergeAction::Delete { .. } => {
                 write!(f, "DELETE")
             }
         }
@@ -8700,6 +8713,8 @@ impl Display for MergeAction {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct MergeClause {
+    /// The `WHEN` token that starts the sub-expression.
+    pub when_token: AttachedToken,
     pub clause_kind: MergeClauseKind,
     pub predicate: Option<Expr>,
     pub action: MergeAction,
@@ -8708,6 +8723,7 @@ pub struct MergeClause {
 impl Display for MergeClause {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let MergeClause {
+            when_token: _,
             clause_kind,
             predicate,
             action,
@@ -8731,10 +8747,12 @@ impl Display for MergeClause {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum OutputClause {
     Output {
+        output_token: AttachedToken,
         select_items: Vec<SelectItem>,
         into_table: Option<SelectInto>,
     },
     Returning {
+        returning_token: AttachedToken,
         select_items: Vec<SelectItem>,
     },
 }
@@ -8743,6 +8761,7 @@ impl fmt::Display for OutputClause {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             OutputClause::Output {
+                output_token: _,
                 select_items,
                 into_table,
             } => {
@@ -8754,7 +8773,10 @@ impl fmt::Display for OutputClause {
                 }
                 Ok(())
             }
-            OutputClause::Returning { select_items } => {
+            OutputClause::Returning {
+                returning_token: _,
+                select_items,
+            } => {
                 f.write_str("RETURNING ")?;
                 display_comma_separated(select_items).fmt(f)
             }

src/ast/spans.rs

@@ -35,14 +35,15 @@ use super::{
     FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
     IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr, Join,
     JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
-    MatchRecognizePattern, Measure, NamedParenthesizedList, NamedWindowDefinition, ObjectName,
-    ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy,
-    OrderByExpr, OrderByKind, Partition, PivotValueSource, ProjectionSelect, Query, RaiseStatement,
-    RaiseStatementValue, ReferentialAction, RenameSelectItem, ReplaceSelectElement,
-    ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript,
-    SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
-    TableOptionsClustered, TableWithJoins, Update, UpdateTableFromKind, Use, Value, Values,
-    ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    MatchRecognizePattern, Measure, MergeAction, MergeClause, MergeInsertExpr, MergeInsertKind,
+    NamedParenthesizedList, NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict,
+    OnConflictAction, OnInsert, OpenStatement, OrderBy, OrderByExpr, OrderByKind, OutputClause,
+    Partition, PivotValueSource, ProjectionSelect, Query, RaiseStatement, RaiseStatementValue,
+    ReferentialAction, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select,
+    SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias,
+    TableAliasColumnDef, TableConstraint, TableFactor, TableObject, TableOptionsClustered,
+    TableWithJoins, Update, UpdateTableFromKind, Use, Value, Values, ViewColumnDef, WhileStatement,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -287,7 +288,6 @@ impl Spanned for Values {
 /// - [Statement::Explain]
 /// - [Statement::Savepoint]
 /// - [Statement::ReleaseSavepoint]
-/// - [Statement::Merge]
 /// - [Statement::Cache]
 /// - [Statement::UNCache]
 /// - [Statement::CreateSequence]
@@ -439,7 +439,20 @@ impl Spanned for Statement {
             Statement::Explain { .. } => Span::empty(),
             Statement::Savepoint { .. } => Span::empty(),
             Statement::ReleaseSavepoint { .. } => Span::empty(),
-            Statement::Merge { .. } => Span::empty(),
+            Statement::Merge {
+                merge_token,
+                into: _,
+                table: _,
+                source: _,
+                on,
+                clauses,
+                output,
+            } => union_spans(
+                [merge_token.0.span, on.span()]
+                    .into_iter()
+                    .chain(clauses.iter().map(Spanned::span))
+                    .chain(output.iter().map(Spanned::span)),
+            ),
             Statement::Cache { .. } => Span::empty(),
             Statement::UNCache { .. } => Span::empty(),
             Statement::CreateSequence { .. } => Span::empty(),
@@ -2381,11 +2394,72 @@ impl Spanned for CreateOperatorClass {
     }
 }
 
+impl Spanned for MergeClause {
+    fn span(&self) -> Span {
+        union_spans([self.when_token.0.span, self.action.span()].into_iter())
+    }
+}
+
+impl Spanned for MergeAction {
+    fn span(&self) -> Span {
+        match self {
+            MergeAction::Insert(expr) => expr.span(),
+            MergeAction::Update {
+                update_token,
+                assignments,
+            } => union_spans(
+                core::iter::once(update_token.0.span).chain(assignments.iter().map(Spanned::span)),
+            ),
+            MergeAction::Delete { delete_token } => delete_token.0.span,
+        }
+    }
+}
+
+impl Spanned for MergeInsertExpr {
+    fn span(&self) -> Span {
+        union_spans(
+            [
+                self.insert_token.0.span,
+                self.kind_token.0.span,
+                match self.kind {
+                    MergeInsertKind::Values(ref values) => values.span(),
+                    MergeInsertKind::Row => Span::empty(), // ~ covered by `kind_token`
+                },
+            ]
+            .into_iter()
+            .chain(self.columns.iter().map(|i| i.span)),
+        )
+    }
+}
+
+impl Spanned for OutputClause {
+    fn span(&self) -> Span {
+        match self {
+            OutputClause::Output {
+                output_token,
+                select_items,
+                into_table,
+            } => union_spans(
+                core::iter::once(output_token.0.span)
+                    .chain(into_table.iter().map(Spanned::span))
+                    .chain(select_items.iter().map(Spanned::span)),
+            ),
+            OutputClause::Returning {
+                returning_token,
+                select_items,
+            } => union_spans(
+                core::iter::once(returning_token.0.span)
+                    .chain(select_items.iter().map(Spanned::span)),
+            ),
+        }
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use crate::dialect::{Dialect, GenericDialect, SnowflakeDialect};
     use crate::parser::Parser;
-    use crate::tokenizer::Span;
+    use crate::tokenizer::{Location, Span};
 
     use super::*;
 
@@ -2647,4 +2721,203 @@ WHERE id = 1
         assert_eq!(stmt_span.start, (2, 7).into());
         assert_eq!(stmt_span.end, (4, 24).into());
     }
+
+    #[test]
+    fn test_merge_statement_spans() {
+        let sql = r#"
+        -- plain merge statement; no RETURNING, no OUTPUT
+
+        MERGE INTO target_table USING source_table
+                ON target_table.id = source_table.oooid
+
+        /* an inline comment */ WHEN NOT MATCHED THEN
+            INSERT (ID, description)
+               VALUES (source_table.id, source_table.description)
+
+            -- another one
+                WHEN MATCHED AND target_table.x = 'X' THEN
+            UPDATE SET target_table.description = source_table.description
+
+              WHEN MATCHED AND target_table.x != 'X' THEN   DELETE
+        WHEN NOT MATCHED AND 1 THEN INSERT (product, quantity) ROW 
+        "#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        // ~ assert the span of the whole statement
+        let stmt_span = r[0].span();
+        assert_eq!(stmt_span.start, (4, 9).into());
+        assert_eq!(stmt_span.end, (16, 67).into());
+
+        // ~ individual tokens within the statement
+        let Statement::Merge {
+            merge_token,
+            into: _,
+            table: _,
+            source: _,
+            on: _,
+            clauses,
+            output,
+        } = &r[0]
+        else {
+            panic!("not a MERGE statement");
+        };
+        assert_eq!(
+            merge_token.0.span,
+            Span::new(Location::new(4, 9), Location::new(4, 14))
+        );
+        assert_eq!(clauses.len(), 4);
+
+        // ~ the INSERT clause's TOKENs
+        assert_eq!(
+            clauses[0].when_token.0.span,
+            Span::new(Location::new(7, 33), Location::new(7, 37))
+        );
+        if let MergeAction::Insert(MergeInsertExpr {
+            insert_token,
+            kind_token,
+            ..
+        }) = &clauses[0].action
+        {
+            assert_eq!(
+                insert_token.0.span,
+                Span::new(Location::new(8, 13), Location::new(8, 19))
+            );
+            assert_eq!(
+                kind_token.0.span,
+                Span::new(Location::new(9, 16), Location::new(9, 22))
+            );
+        } else {
+            panic!("not a MERGE INSERT clause");
+        }
+
+        // ~ the UPDATE token(s)
+        assert_eq!(
+            clauses[1].when_token.0.span,
+            Span::new(Location::new(12, 17), Location::new(12, 21))
+        );
+        if let MergeAction::Update {
+            update_token,
+            assignments: _,
+        } = &clauses[1].action
+        {
+            assert_eq!(
+                update_token.0.span,
+                Span::new(Location::new(13, 13), Location::new(13, 19))
+            );
+        } else {
+            panic!("not a MERGE UPDATE clause");
+        }
+
+        // the DELETE token(s)
+        assert_eq!(
+            clauses[2].when_token.0.span,
+            Span::new(Location::new(15, 15), Location::new(15, 19))
+        );
+        if let MergeAction::Delete { delete_token } = &clauses[2].action {
+            assert_eq!(
+                delete_token.0.span,
+                Span::new(Location::new(15, 61), Location::new(15, 67))
+            );
+        } else {
+            panic!("not a MERGE DELETE clause");
+        }
+
+        // ~ an INSERT clause's ROW token
+        assert_eq!(
+            clauses[3].when_token.0.span,
+            Span::new(Location::new(16, 9), Location::new(16, 13))
+        );
+        if let MergeAction::Insert(MergeInsertExpr {
+            insert_token,
+            kind_token,
+            ..
+        }) = &clauses[3].action
+        {
+            assert_eq!(
+                insert_token.0.span,
+                Span::new(Location::new(16, 37), Location::new(16, 43))
+            );
+            assert_eq!(
+                kind_token.0.span,
+                Span::new(Location::new(16, 64), Location::new(16, 67))
+            );
+        } else {
+            panic!("not a MERGE INSERT clause");
+        }
+
+        assert!(output.is_none());
+    }
+
+    #[test]
+    fn test_merge_statement_spans_with_returning() {
+        let sql = r#"
+    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.*
+        "#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        // ~ assert the span of the whole statement
+        let stmt_span = r[0].span();
+        assert_eq!(
+            stmt_span,
+            Span::new(Location::new(2, 5), Location::new(8, 34))
+        );
+
+        // ~ individual tokens within the statement
+        if let Statement::Merge { output, .. } = &r[0] {
+            if let Some(OutputClause::Returning {
+                returning_token, ..
+            }) = output
+            {
+                assert_eq!(
+                    returning_token.0.span,
+                    Span::new(Location::new(8, 5), Location::new(8, 14))
+                );
+            } else {
+                panic!("unexpected MERGE output clause");
+            }
+        } else {
+            panic!("not a MERGE statement");
+        };
+    }
+
+    #[test]
+    fn test_merge_statement_spans_with_output() {
+        let sql = r#"MERGE INTO a USING b ON a.id = b.id
+        WHEN MATCHED THEN DELETE
+              OUTPUT inserted.*"#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        // ~ assert the span of the whole statement
+        let stmt_span = r[0].span();
+        assert_eq!(
+            stmt_span,
+            Span::new(Location::new(1, 1), Location::new(3, 32))
+        );
+
+        // ~ individual tokens within the statement
+        if let Statement::Merge { output, .. } = &r[0] {
+            if let Some(OutputClause::Output { output_token, .. }) = output {
+                assert_eq!(
+                    output_token.0.span,
+                    Span::new(Location::new(3, 15), Location::new(3, 21))
+                );
+            } else {
+                panic!("unexpected MERGE output clause");
+            }
+        } else {
+            panic!("not a MERGE statement");
+        };
+    }
 }

src/parser/mod.rs

@@ -622,7 +622,7 @@ impl<'a> Parser<'a> {
                 Keyword::DEALLOCATE => self.parse_deallocate(),
                 Keyword::EXECUTE | Keyword::EXEC => self.parse_execute(),
                 Keyword::PREPARE => self.parse_prepare(),
-                Keyword::MERGE => self.parse_merge(),
+                Keyword::MERGE => self.parse_merge(next_token),
                 // `LISTEN`, `UNLISTEN` and `NOTIFY` are Postgres-specific
                 // syntaxes. They are used for Postgres statement.
                 Keyword::LISTEN if self.dialect.supports_listen_notify() => self.parse_listen(),
@@ -12125,8 +12125,11 @@ impl<'a> Parser<'a> {
     /// 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<Box<SetExpr>, ParserError> {
-        Ok(Box::new(SetExpr::Merge(self.parse_merge()?)))
+    fn parse_merge_setexpr_boxed(
+        &mut self,
+        merge_token: TokenWithSpan,
+    ) -> Result<Box<SetExpr>, ParserError> {
+        Ok(Box::new(SetExpr::Merge(self.parse_merge(merge_token)?)))
     }
 
     pub fn parse_delete(&mut self, delete_token: TokenWithSpan) -> Result<Statement, ParserError> {
@@ -12344,7 +12347,7 @@ impl<'a> Parser<'a> {
         } else if self.parse_keyword(Keyword::MERGE) {
             Ok(Query {
                 with,
-                body: self.parse_merge_setexpr_boxed()?,
+                body: self.parse_merge_setexpr_boxed(self.get_current_token().clone())?,
                 limit_clause: None,
                 order_by: None,
                 fetch: None,
@@ -17200,6 +17203,7 @@ impl<'a> Parser<'a> {
             if !(self.parse_keyword(Keyword::WHEN)) {
                 break;
             }
+            let when_token = self.get_current_token().clone();
 
             let mut clause_kind = MergeClauseKind::Matched;
             if self.parse_keyword(Keyword::NOT) {
@@ -17235,12 +17239,16 @@ impl<'a> Parser<'a> {
                         clause_kind,
                         MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
                     ) {
-                        return Err(ParserError::ParserError(format!(
-                            "UPDATE is not allowed in a {clause_kind} merge clause"
-                        )));
+                        return parser_err!(
+                            format_args!("UPDATE is not allowed in a {clause_kind} merge clause"),
+                            self.get_current_token().span.start
+                        );
                     }
+
+                    let update_token = self.get_current_token().clone();
                     self.expect_keyword_is(Keyword::SET)?;
                     MergeAction::Update {
+                        update_token: update_token.into(),
                         assignments: self.parse_comma_separated(Parser::parse_assignment)?,
                     }
                 }
@@ -17249,42 +17257,58 @@ impl<'a> Parser<'a> {
                         clause_kind,
                         MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
                     ) {
-                        return Err(ParserError::ParserError(format!(
-                            "DELETE is not allowed in a {clause_kind} merge clause"
-                        )));
+                        return parser_err!(
+                            format_args!("DELETE is not allowed in a {clause_kind} merge clause"),
+                            self.get_current_token().span.start
+                        );
+                    };
+
+                    let delete_token = self.get_current_token().clone();
+                    MergeAction::Delete {
+                        delete_token: delete_token.into(),
                     }
-                    MergeAction::Delete
                 }
                 Some(Keyword::INSERT) => {
                     if !matches!(
                         clause_kind,
                         MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget
                     ) {
-                        return Err(ParserError::ParserError(format!(
-                            "INSERT is not allowed in a {clause_kind} merge clause"
-                        )));
-                    }
+                        return parser_err!(
+                            format_args!("INSERT is not allowed in a {clause_kind} merge clause"),
+                            self.get_current_token().span.start
+                        );
+                    };
+
+                    let insert_token = self.get_current_token().clone();
                     let is_mysql = dialect_of!(self is MySqlDialect);
 
                     let columns = self.parse_parenthesized_column_list(Optional, is_mysql)?;
-                    let kind = if dialect_of!(self is BigQueryDialect | GenericDialect)
+                    let (kind, kind_token) = if dialect_of!(self is BigQueryDialect | GenericDialect)
                         && self.parse_keyword(Keyword::ROW)
                     {
-                        MergeInsertKind::Row
+                        (MergeInsertKind::Row, self.get_current_token().clone())
                     } else {
                         self.expect_keyword_is(Keyword::VALUES)?;
+                        let values_token = self.get_current_token().clone();
                         let values = self.parse_values(is_mysql, false)?;
-                        MergeInsertKind::Values(values)
+                        (MergeInsertKind::Values(values), values_token)
                     };
-                    MergeAction::Insert(MergeInsertExpr { columns, kind })
+                    MergeAction::Insert(MergeInsertExpr {
+                        insert_token: insert_token.into(),
+                        columns,
+                        kind_token: kind_token.into(),
+                        kind,
+                    })
                 }
                 _ => {
-                    return Err(ParserError::ParserError(
-                        "expected UPDATE, DELETE or INSERT in merge clause".to_string(),
-                    ));
+                    return parser_err!(
+                        "expected UPDATE, DELETE or INSERT in merge clause",
+                        self.peek_token_ref().span.start
+                    );
                 }
             };
             clauses.push(MergeClause {
+                when_token: when_token.into(),
                 clause_kind,
                 predicate,
                 action: merge_clause,
@@ -17293,7 +17317,11 @@ impl<'a> Parser<'a> {
         Ok(clauses)
     }
 
-    fn parse_output(&mut self, start_keyword: Keyword) -> Result<OutputClause, ParserError> {
+    fn parse_output(
+        &mut self,
+        start_keyword: Keyword,
+        start_token: TokenWithSpan,
+    ) -> Result<OutputClause, ParserError> {
         let select_items = self.parse_projection()?;
         let into_table = if start_keyword == Keyword::OUTPUT && self.peek_keyword(Keyword::INTO) {
             self.expect_keyword_is(Keyword::INTO)?;
@@ -17304,11 +17332,15 @@ impl<'a> Parser<'a> {
 
         Ok(if start_keyword == Keyword::OUTPUT {
             OutputClause::Output {
+                output_token: start_token.into(),
                 select_items,
                 into_table,
             }
         } else {
-            OutputClause::Returning { select_items }
+            OutputClause::Returning {
+                returning_token: start_token.into(),
+                select_items,
+            }
         })
     }
 
@@ -17328,7 +17360,7 @@ impl<'a> Parser<'a> {
         })
     }
 
-    pub fn parse_merge(&mut self) -> Result<Statement, ParserError> {
+    pub fn parse_merge(&mut self, merge_token: TokenWithSpan) -> Result<Statement, ParserError> {
         let into = self.parse_keyword(Keyword::INTO);
 
         let table = self.parse_table_factor()?;
@@ -17339,11 +17371,12 @@ impl<'a> Parser<'a> {
         let on = self.parse_expr()?;
         let clauses = self.parse_merge_clauses()?;
         let output = match self.parse_one_of_keywords(&[Keyword::OUTPUT, Keyword::RETURNING]) {
-            Some(start_keyword) => Some(self.parse_output(start_keyword)?),
+            Some(keyword) => Some(self.parse_output(keyword, self.get_current_token().clone())?),
             None => None,
         };
 
         Ok(Statement::Merge {
+            merge_token: merge_token.into(),
             into,
             table,
             source,