Databricks: Add support for UPDATE SET * and INSERT * in MERGE statements (#2325)

Co-authored-by: Ifeanyi Ubah <ify1992@yahoo.com>

Author: finchxxia

src/ast/dml.rs

@@ -649,6 +649,14 @@ pub enum MergeInsertKind {
     /// ```
     /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement)
     Row,
+    /// The insert expression uses the `*` wildcard to insert all columns.
+    ///
+    /// Example:
+    /// ```sql
+    /// INSERT *
+    /// ```
+    /// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html)
+    Wildcard,
 }
 
 impl Display for MergeInsertKind {
@@ -660,6 +668,9 @@ impl Display for MergeInsertKind {
             MergeInsertKind::Row => {
                 write!(f, "ROW")
             }
+            MergeInsertKind::Wildcard => {
+                write!(f, "*")
+            }
         }
     }
 }
@@ -710,33 +721,70 @@ impl Display for MergeInsertExpr {
     }
 }
 
+/// The kind of update used within a `MERGE` statement.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum MergeUpdateKind {
+    /// Standard update with explicit assignments.
+    ///
+    /// Example:
+    /// ```sql
+    /// UPDATE SET quantity = source.quantity, name = source.name
+    /// ```
+    Set(Vec<Assignment>),
+    /// The `*` wildcard to update all columns from the source.
+    ///
+    /// Example:
+    /// ```sql
+    /// UPDATE SET *
+    /// ```
+    /// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html)
+    Wildcard,
+}
+
+impl Display for MergeUpdateKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            MergeUpdateKind::Set(assignments) => {
+                write!(f, "SET {}", display_comma_separated(assignments))
+            }
+            MergeUpdateKind::Wildcard => {
+                write!(f, "SET *")
+            }
+        }
+    }
+}
+
 /// The expression used to update rows within a `MERGE` statement.
 ///
 /// Examples
 /// ```sql
 /// UPDATE SET quantity = T.quantity + S.quantity
+/// UPDATE SET *
 /// ```
 ///
 /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/merge)
 /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement)
 /// [Oracle](https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/MERGE.html)
+/// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html)
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct MergeUpdateExpr {
     /// The `UPDATE` token that starts the sub-expression.
     pub update_token: AttachedToken,
-    /// The update assiment expressions
-    pub assignments: Vec<Assignment>,
-    /// `where_clause` for the update (Oralce specific)
+    /// The kind of update: explicit assignments or `*` shorthand.
+    pub kind: MergeUpdateKind,
+    /// `where_clause` for the update (Oracle specific)
     pub update_predicate: Option<Expr>,
     /// `delete_clause` for the update "delete where" (Oracle specific)
     pub delete_predicate: Option<Expr>,
 }
 
 impl Display for MergeUpdateExpr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "SET {}", display_comma_separated(&self.assignments))?;
+        write!(f, "{}", self.kind)?;
         if let Some(predicate) = self.update_predicate.as_ref() {
             write!(f, " WHERE {predicate}")?;
         }

src/ast/mod.rs

@@ -85,9 +85,9 @@ pub use self::ddl::{
 };
 pub use self::dml::{
     Delete, Insert, Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr,
-    MergeInsertKind, MergeUpdateExpr, MultiTableInsertIntoClause, MultiTableInsertType,
-    MultiTableInsertValue, MultiTableInsertValues, MultiTableInsertWhenClause, OutputClause,
-    Update,
+    MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, MultiTableInsertIntoClause,
+    MultiTableInsertType, MultiTableInsertValue, MultiTableInsertValues,
+    MultiTableInsertWhenClause, OutputClause, Update,
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{

src/ast/spans.rs

@@ -39,15 +39,15 @@ use super::{
     IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr, Join,
     JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
     MatchRecognizePattern, Measure, Merge, MergeAction, MergeClause, MergeInsertExpr,
-    MergeInsertKind, MergeUpdateExpr, NamedParenthesizedList, NamedWindowDefinition, ObjectName,
-    ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy,
-    OrderByExpr, OrderByKind, OutputClause, Parens, Partition, PartitionBoundValue,
-    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, Values, ViewColumnDef, WhileStatement,
-    WildcardAdditionalOptions, With, WithFill,
+    MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, NamedParenthesizedList,
+    NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction,
+    OnInsert, OpenStatement, OrderBy, OrderByExpr, OrderByKind, OutputClause, Parens, Partition,
+    PartitionBoundValue, 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, Values, ViewColumnDef,
+    WhileStatement, WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -2534,7 +2534,7 @@ impl Spanned for MergeInsertExpr {
                 self.kind_token.0.span,
                 match self.kind {
                     MergeInsertKind::Values(ref values) => values.span(),
-                    MergeInsertKind::Row => Span::empty(), // ~ covered by `kind_token`
+                    MergeInsertKind::Row | MergeInsertKind::Wildcard => Span::empty(),
                 },
             ]
             .into_iter()
@@ -2546,9 +2546,13 @@ impl Spanned for MergeInsertExpr {
 
 impl Spanned for MergeUpdateExpr {
     fn span(&self) -> Span {
+        let kind_span = match &self.kind {
+            MergeUpdateKind::Set(assignments) => union_spans(assignments.iter().map(Spanned::span)),
+            MergeUpdateKind::Wildcard => Span::empty(),
+        };
         union_spans(
             core::iter::once(self.update_token.0.span)
-                .chain(self.assignments.iter().map(Spanned::span))
+                .chain(core::iter::once(kind_span))
                 .chain(self.update_predicate.iter().map(Spanned::span))
                 .chain(self.delete_predicate.iter().map(Spanned::span)),
         )
@@ -2930,7 +2934,7 @@ WHERE id = 1
         );
         if let MergeAction::Update(MergeUpdateExpr {
             update_token,
-            assignments: _,
+            kind: _,
             update_predicate: _,
             delete_predicate: _,
         }) = &clauses[1].action

src/parser/merge.rs

@@ -18,11 +18,12 @@ use alloc::{boxed::Box, format, vec, vec::Vec};
 use crate::{
     ast::{
         Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr, MergeInsertKind,
-        MergeUpdateExpr, ObjectName, OutputClause, SetExpr,
+        MergeUpdateExpr, MergeUpdateKind, ObjectName, OutputClause, SetExpr,
     },
     dialect::{BigQueryDialect, GenericDialect, MySqlDialect},
     keywords::Keyword,
     parser::IsOptional,
+    tokenizer::Token,
     tokenizer::TokenWithSpan,
 };
 
@@ -120,7 +121,11 @@ impl Parser<'_> {
 
                     let update_token = self.get_current_token().clone();
                     self.expect_keyword_is(Keyword::SET)?;
-                    let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
+                    let kind = if self.consume_token(&Token::Mul) {
+                        MergeUpdateKind::Wildcard
+                    } else {
+                        MergeUpdateKind::Set(self.parse_comma_separated(Parser::parse_assignment)?)
+                    };
                     let update_predicate = if self.parse_keyword(Keyword::WHERE) {
                         Some(self.parse_expr()?)
                     } else {
@@ -134,7 +139,7 @@ impl Parser<'_> {
                     };
                     MergeAction::Update(MergeUpdateExpr {
                         update_token: update_token.into(),
-                        assignments,
+                        kind,
                         update_predicate,
                         delete_predicate,
                     })
@@ -167,32 +172,43 @@ impl Parser<'_> {
                     };
 
                     let insert_token = self.get_current_token().clone();
-                    let is_mysql = dialect_of!(self is MySqlDialect);
 
-                    let columns = self.parse_merge_clause_insert_columns(is_mysql)?;
-                    let (kind, kind_token) = if dialect_of!(self is BigQueryDialect | GenericDialect)
-                        && self.parse_keyword(Keyword::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), values_token)
-                    };
-                    let insert_predicate = if self.parse_keyword(Keyword::WHERE) {
-                        Some(self.parse_expr()?)
+                    if self.consume_token(&Token::Mul) {
+                        let star_token = self.get_current_token().clone();
+                        MergeAction::Insert(MergeInsertExpr {
+                            insert_token: insert_token.into(),
+                            columns: vec![],
+                            kind_token: star_token.into(),
+                            kind: MergeInsertKind::Wildcard,
+                            insert_predicate: None,
+                        })
                     } else {
-                        None
-                    };
+                        let is_mysql = dialect_of!(self is MySqlDialect);
+                        let columns = self.parse_merge_clause_insert_columns(is_mysql)?;
+                        let (kind, kind_token) = if dialect_of!(self is BigQueryDialect | GenericDialect)
+                            && self.parse_keyword(Keyword::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), values_token)
+                        };
+                        let insert_predicate = if self.parse_keyword(Keyword::WHERE) {
+                            Some(self.parse_expr()?)
+                        } else {
+                            None
+                        };
 
-                    MergeAction::Insert(MergeInsertExpr {
-                        insert_token: insert_token.into(),
-                        columns,
-                        kind_token: kind_token.into(),
-                        kind,
-                        insert_predicate,
-                    })
+                        MergeAction::Insert(MergeInsertExpr {
+                            insert_token: insert_token.into(),
+                            columns,
+                            kind_token: kind_token.into(),
+                            kind,
+                            insert_predicate,
+                        })
+                    }
                 }
                 _ => {
                     return parser_err!(

tests/sqlparser_bigquery.rs

@@ -1830,7 +1830,7 @@ fn parse_merge() {
     });
     let update_action = MergeAction::Update(MergeUpdateExpr {
         update_token: AttachedToken::empty(),
-        assignments: vec![
+        kind: MergeUpdateKind::Set(vec![
             Assignment {
                 target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("a")])),
                 value: Expr::value(number("1")),
@@ -1839,7 +1839,7 @@ fn parse_merge() {
                 target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("b")])),
                 value: Expr::value(number("2")),
             },
-        ],
+        ]),
         update_predicate: None,
         delete_predicate: None,
     });

tests/sqlparser_common.rs

@@ -10157,7 +10157,7 @@ fn parse_merge() {
                         }),
                         action: MergeAction::Update(MergeUpdateExpr {
                             update_token: AttachedToken::empty(),
-                            assignments: vec![
+                            kind: MergeUpdateKind::Set(vec![
                                 Assignment {
                                     target: AssignmentTarget::ColumnName(ObjectName::from(vec![
                                         Ident::new("dest"),
@@ -10178,7 +10178,7 @@ fn parse_merge() {
                                         Ident::new("G"),
                                     ]),
                                 },
-                            ],
+                            ]),
                             update_predicate: None,
                             delete_predicate: None,
                         }),
@@ -10239,6 +10239,45 @@ WHEN NOT MATCHED THEN \
 INSERT (PLAYGROUND.FOO.ID, PLAYGROUND.FOO.NAME) \
 VALUES (1, 'abc')";
     all_dialects().verified_stmt(sql);
+
+    // MERGE with wildcard (UPDATE SET * and INSERT *)
+    let sql = "MERGE INTO target USING source ON target.id = source.id WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT *";
+    match verified_stmt(sql) {
+        Statement::Merge(merge) => {
+            assert_eq!(merge.clauses.len(), 2);
+
+            match &merge.clauses[0].action {
+                MergeAction::Update(update_expr) => {
+                    assert!(matches!(update_expr.kind, MergeUpdateKind::Wildcard));
+                }
+                _ => panic!("Expected UPDATE action"),
+            }
+
+            match &merge.clauses[1].action {
+                MergeAction::Insert(insert_expr) => {
+                    assert!(matches!(insert_expr.kind, MergeInsertKind::Wildcard));
+                    assert!(insert_expr.columns.is_empty());
+                }
+                _ => panic!("Expected INSERT action"),
+            }
+        }
+        _ => panic!("Expected MERGE statement"),
+    }
+
+    verified_stmt("MERGE INTO target USING source ON target.id = source.id WHEN MATCHED AND source.active = 1 THEN UPDATE SET *");
+
+    verified_stmt("MERGE INTO target USING source ON target.id = source.id WHEN NOT MATCHED BY TARGET THEN INSERT *");
+
+    verified_stmt("MERGE INTO target USING source ON target.id = source.id WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT (a, b) VALUES (source.a, source.b)");
+
+    let sql = concat!(
+        "MERGE INTO t1 AS target ",
+        "USING (SELECT * FROM t2) AS source ",
+        "ON target.id = source.id ",
+        "WHEN MATCHED THEN UPDATE SET * ",
+        "WHEN NOT MATCHED THEN INSERT *"
+    );
+    verified_stmt(sql);
 }
 
 #[test]