Databricks: Add support for UPDATE SET * and INSERT * in MERGE statements

Add parsing support for the Databricks star shorthand syntax in MERGE
statements, allowing `UPDATE SET *` and `INSERT *` to reference all columns
from the source.

- Add `MergeUpdateKind` enum (`Set(Vec<Assignment>)` | `Star`) to replace
  the previous `assignments` field on `MergeUpdateExpr`, enforcing valid
  states at the type level
- Add `MergeInsertKind::Star` variant for `INSERT *` shorthand
- Add `supports_merge_star_syntax()` dialect method for Databricks and Generic

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 `*` shorthand to insert all columns.
+    ///
+    /// Example:
+    /// ```sql
+    /// INSERT *
+    /// ```
+    /// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html)
+    Star,
 }
 
 impl Display for MergeInsertKind {
@@ -660,6 +668,9 @@ impl Display for MergeInsertKind {
             MergeInsertKind::Row => {
                 write!(f, "ROW")
             }
+            MergeInsertKind::Star => {
+                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 `*` shorthand to update all columns from the source.
+    ///
+    /// Example:
+    /// ```sql
+    /// UPDATE SET *
+    /// ```
+    /// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html)
+    Star,
+}
+
+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::Star => {
+                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,7 +85,8 @@ pub use self::ddl::{
 };
 pub use self::dml::{
     Delete, Insert, Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr,
-    MergeInsertKind, MergeUpdateExpr, MultiTableInsertIntoClause, MultiTableInsertType,
+    MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, MultiTableInsertIntoClause,
+    MultiTableInsertType,
     MultiTableInsertValue, MultiTableInsertValues, MultiTableInsertWhenClause, OutputClause,
     Update,
 };

src/ast/spans.rs

@@ -39,7 +39,8 @@ 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,
+    MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, NamedParenthesizedList,
+    NamedWindowDefinition, ObjectName,
     ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy,
     OrderByExpr, OrderByKind, OutputClause, Parens, Partition, PartitionBoundValue,
     PivotValueSource, ProjectionSelect, Query, RaiseStatement, RaiseStatementValue,
@@ -2531,7 +2532,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::Star => Span::empty(),
                 },
             ]
             .into_iter()
@@ -2543,9 +2544,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::Star => 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)),
         )
@@ -2927,7 +2932,7 @@ WHERE id = 1
         );
         if let MergeAction::Update(MergeUpdateExpr {
             update_token,
-            assignments: _,
+            kind: _,
             update_predicate: _,
             delete_predicate: _,
         }) = &clauses[1].action

src/dialect/databricks.rs

@@ -108,4 +108,9 @@ impl Dialect for DatabricksDialect {
     fn supports_select_item_multi_column_alias(&self) -> bool {
         true
     }
+
+    /// See <https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html>
+    fn supports_merge_star_syntax(&self) -> bool {
+        true
+    }
 }

src/dialect/generic.rs

@@ -304,4 +304,8 @@ impl Dialect for GenericDialect {
     fn supports_xml_expressions(&self) -> bool {
         true
     }
+
+    fn supports_merge_star_syntax(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -1164,6 +1164,13 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports `UPDATE SET *` and `INSERT *`
+    /// shorthand syntax in `MERGE` statements.
+    /// <https://docs.databricks.com/en/sql/language-manual/delta-merge-into.html>
+    fn supports_merge_star_syntax(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports the `LISTEN`, `UNLISTEN` and `NOTIFY` statements
     fn supports_listen_notify(&self) -> bool {
         false

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,13 @@ 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.dialect.supports_merge_star_syntax()
+                        && self.consume_token(&Token::Mul)
+                    {
+                        MergeUpdateKind::Star
+                    } 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 +141,7 @@ impl Parser<'_> {
                     };
                     MergeAction::Update(MergeUpdateExpr {
                         update_token: update_token.into(),
-                        assignments,
+                        kind,
                         update_predicate,
                         delete_predicate,
                     })
@@ -167,32 +174,44 @@ 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)
+                    if self.dialect.supports_merge_star_syntax() && self.consume_token(&Token::Mul)
                     {
-                        (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()?)
+                        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::Star,
+                            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,
                         }),

tests/sqlparser_databricks.rs

@@ -738,3 +738,61 @@ fn parse_cte_without_as() {
         .parse_sql_statements("WITH cte (SELECT 1) SELECT * FROM cte")
         .is_err());
 }
+
+#[test]
+fn test_merge_update_set_star_and_insert_star() {
+    let sql = "MERGE INTO target USING source ON target.id = source.id WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT *";
+    databricks_and_generic().verified_stmt(sql);
+
+    match databricks().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::Star));
+                }
+                _ => panic!("Expected UPDATE action"),
+            }
+
+            match &merge.clauses[1].action {
+                MergeAction::Insert(insert_expr) => {
+                    assert!(matches!(insert_expr.kind, MergeInsertKind::Star));
+                    assert!(insert_expr.columns.is_empty());
+                }
+                _ => panic!("Expected INSERT action"),
+            }
+        }
+        _ => panic!("Expected MERGE statement"),
+    }
+}
+
+#[test]
+fn test_merge_update_set_star_with_predicate() {
+    let sql = "MERGE INTO target USING source ON target.id = source.id WHEN MATCHED AND source.active = true THEN UPDATE SET *";
+    databricks_and_generic().verified_stmt(sql);
+}
+
+#[test]
+fn test_merge_insert_star_with_not_matched_by_target() {
+    let sql = "MERGE INTO target USING source ON target.id = source.id WHEN NOT MATCHED BY TARGET THEN INSERT *";
+    databricks_and_generic().verified_stmt(sql);
+}
+
+#[test]
+fn test_merge_mixed_star_and_explicit() {
+    let sql = "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)";
+    databricks_and_generic().verified_stmt(sql);
+}
+
+#[test]
+fn test_merge_star_with_subquery_source() {
+    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 *"
+    );
+    databricks_and_generic().verified_stmt(sql);
+}