feat(binder): support update from (#17464)

* feat(binder): support update from

* chore(parser): fix token

* chore(test): add sqllogictest

* chore(test): update stmt unit-test

Author: Dousir9

src/query/ast/src/ast/statements/merge_into.rs

@@ -30,13 +30,13 @@ use crate::ast::TableReference;
 use crate::ast::WithOptions;
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
-pub struct MergeUpdateExpr {
+pub struct MutationUpdateExpr {
     pub table: Option<Identifier>,
     pub name: Identifier,
     pub expr: Expr,
 }
 
-impl Display for MergeUpdateExpr {
+impl Display for MutationUpdateExpr {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         if self.table.is_some() {
             write!(f, "{}.", self.table.clone().unwrap())?;
@@ -49,7 +49,7 @@ impl Display for MergeUpdateExpr {
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum MatchOperation {
     Update {
-        update_list: Vec<MergeUpdateExpr>,
+        update_list: Vec<MutationUpdateExpr>,
         is_star: bool,
     },
     Delete,
@@ -86,7 +86,7 @@ pub struct MergeIntoStmt {
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table_ident: Identifier,
-    pub source: MergeSource,
+    pub source: MutationSource,
     // target_alias is belong to target
     pub target_alias: Option<TableAlias>,
     pub join_expr: Expr,
@@ -167,7 +167,7 @@ impl Display for MergeIntoStmt {
 }
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
-pub enum MergeSource {
+pub enum MutationSource {
     StreamingV2 {
         settings: FileFormatOptions,
         on_error_mode: Option<String>,
@@ -187,7 +187,7 @@ pub enum MergeSource {
     },
 }
 
-impl MergeSource {
+impl MutationSource {
     pub fn transform_table_reference(&self) -> TableReference {
         match self {
             Self::StreamingV2 {
@@ -229,10 +229,10 @@ impl MergeSource {
     }
 }
 
-impl Display for MergeSource {
+impl Display for MutationSource {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         match self {
-            MergeSource::StreamingV2 {
+            MutationSource::StreamingV2 {
                 settings,
                 on_error_mode,
                 start: _,
@@ -245,12 +245,12 @@ impl Display for MergeSource {
                 )
             }
 
-            MergeSource::Select {
+            MutationSource::Select {
                 query,
                 source_alias,
             } => write!(f, "({query}) AS {source_alias}"),
 
-            MergeSource::Table {
+            MutationSource::Table {
                 catalog,
                 database,
                 table,

src/query/ast/src/ast/statements/update.rs

@@ -19,17 +19,24 @@ use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
 use crate::ast::write_comma_separated_list;
+use crate::ast::write_dot_separated_list;
 use crate::ast::Expr;
 use crate::ast::Hint;
 use crate::ast::Identifier;
-use crate::ast::TableReference;
+use crate::ast::MutationSource;
+use crate::ast::MutationUpdateExpr;
+use crate::ast::TableAlias;
 use crate::ast::With;
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct UpdateStmt {
     pub hints: Option<Hint>,
-    pub table: TableReference,
-    pub update_list: Vec<UpdateExpr>,
+    pub catalog: Option<Identifier>,
+    pub database: Option<Identifier>,
+    pub table: Identifier,
+    pub table_alias: Option<TableAlias>,
+    pub update_list: Vec<MutationUpdateExpr>,
+    pub from: Option<MutationSource>,
     pub selection: Option<Expr>,
     // With clause, common table expression
     pub with: Option<With>,
@@ -44,8 +51,21 @@ impl Display for UpdateStmt {
         if let Some(hints) = &self.hints {
             write!(f, "{} ", hints)?;
         }
-        write!(f, "{} SET ", self.table)?;
+        write_dot_separated_list(
+            f,
+            self.catalog
+                .iter()
+                .chain(&self.database)
+                .chain(Some(&self.table)),
+        )?;
+        if let Some(alias) = &self.table_alias {
+            write!(f, " AS {}", alias.name)?;
+        }
+        write!(f, " SET ")?;
         write_comma_separated_list(f, &self.update_list)?;
+        if let Some(from) = &self.from {
+            write!(f, " FROM {} ", from)?;
+        }
         if let Some(conditions) = &self.selection {
             write!(f, " WHERE {conditions}")?;
         }

src/query/ast/src/parser/statement.rs

@@ -227,7 +227,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         rule! {
             MERGE ~ #hint?
             ~ INTO ~ #dot_separated_idents_1_to_3 ~ #table_alias?
-            ~ USING ~ #merge_source
+            ~ USING ~ #mutation_source
             ~ ON ~ #expr ~ (#match_clause | #unmatch_clause)*
         },
         |(
@@ -271,15 +271,30 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let update = map(
         rule! {
-            #with? ~ UPDATE ~ #hint? ~ #table_reference_only
-            ~ SET ~ ^#comma_separated_list1(update_expr)
+            #with? ~ UPDATE ~ #hint? ~ #dot_separated_idents_1_to_3 ~ #table_alias?
+            ~ SET ~ ^#comma_separated_list1(mutation_update_expr)
+            ~ ( FROM ~ #mutation_source )?
             ~ ( WHERE ~ ^#expr )?
         },
-        |(with, _, hints, table, _, update_list, opt_selection)| {
+        |(
+            with,
+            _,
+            hints,
+            (catalog, database, table),
+            table_alias,
+            _,
+            update_list,
+            from,
+            opt_selection,
+        )| {
             Statement::Update(UpdateStmt {
                 hints,
+                catalog,
+                database,
                 table,
+                table_alias,
                 update_list,
+                from: from.map(|(_, table)| table),
                 selection: opt_selection.map(|(_, selection)| selection),
                 with,
             })
@@ -2899,28 +2914,28 @@ pub fn raw_insert_source(i: Input) -> IResult<InsertSource> {
     )(i)
 }
 
-pub fn merge_source(i: Input) -> IResult<MergeSource> {
+pub fn mutation_source(i: Input) -> IResult<MutationSource> {
     let streaming_v2 = map(
         rule! {
            #file_format_clause  ~ (ON_ERROR ~ ^"=" ~ ^#ident)? ~  #rest_str
         },
-        |(options, on_error_opt, (_, start))| MergeSource::StreamingV2 {
+        |(options, on_error_opt, (_, start))| MutationSource::StreamingV2 {
             settings: options,
             on_error_mode: on_error_opt.map(|v| v.2.to_string()),
             start,
         },
     );
 
     let query = map(rule! {#query ~ #table_alias}, |(query, source_alias)| {
-        MergeSource::Select {
+        MutationSource::Select {
             query: Box::new(query),
             source_alias,
         }
     });
 
     let source_table = map(
         rule!(#dot_separated_idents_1_to_3 ~ #with_options? ~ #table_alias?),
-        |((catalog, database, table), with_options, alias)| MergeSource::Table {
+        |((catalog, database, table), with_options, alias)| MutationSource::Table {
             catalog,
             database,
             table,
@@ -3813,7 +3828,7 @@ fn match_operation(i: Input) -> IResult<MatchOperation> {
         value(MatchOperation::Delete, rule! { DELETE }),
         map(
             rule! {
-                UPDATE ~ SET ~ ^#comma_separated_list1(merge_update_expr)
+                UPDATE ~ SET ~ ^#comma_separated_list1(mutation_update_expr)
             },
             |(_, _, update_list)| MatchOperation::Update {
                 update_list,
@@ -4640,10 +4655,10 @@ pub fn udf_definition(i: Input) -> IResult<UDFDefinition> {
     )(i)
 }
 
-pub fn merge_update_expr(i: Input) -> IResult<MergeUpdateExpr> {
+pub fn mutation_update_expr(i: Input) -> IResult<MutationUpdateExpr> {
     map(
         rule! { #dot_separated_idents_1_to_2 ~ "=" ~ ^#expr },
-        |((table, name), _, expr)| MergeUpdateExpr { table, name, expr },
+        |((table, name), _, expr)| MutationUpdateExpr { table, name, expr },
     )(i)
 }
 

src/query/ast/src/parser/token.rs

@@ -1711,6 +1711,7 @@ impl TokenKind {
             // | TokenKind::SIMILAR
             | TokenKind::SOME
             | TokenKind::SEMI
+            | TokenKind::SET
             | TokenKind::SAMPLE
             // | TokenKind::SYMMETRIC
             // | TokenKind::TABLESAMPLE

src/query/ast/tests/it/testdata/stmt.txt

@@ -16489,38 +16489,29 @@ UPDATE db1.tb1 SET a = a + 1, b = 2 WHERE c > 3
 Update(
     UpdateStmt {
         hints: None,
-        table: Table {
-            span: Some(
-                7..14,
-            ),
-            catalog: None,
-            database: Some(
-                Identifier {
-                    span: Some(
-                        7..10,
-                    ),
-                    name: "db1",
-                    quote: None,
-                    ident_type: None,
-                },
-            ),
-            table: Identifier {
+        catalog: None,
+        database: Some(
+            Identifier {
                 span: Some(
-                    11..14,
+                    7..10,
                 ),
-                name: "tb1",
+                name: "db1",
                 quote: None,
                 ident_type: None,
             },
-            alias: None,
-            temporal: None,
-            with_options: None,
-            pivot: None,
-            unpivot: None,
-            sample: None,
+        ),
+        table: Identifier {
+            span: Some(
+                11..14,
+            ),
+            name: "tb1",
+            quote: None,
+            ident_type: None,
         },
+        table_alias: None,
         update_list: [
-            UpdateExpr {
+            MutationUpdateExpr {
+                table: None,
                 name: Identifier {
                     span: Some(
                         19..20,
@@ -16563,7 +16554,8 @@ Update(
                     },
                 },
             },
-            UpdateExpr {
+            MutationUpdateExpr {
+                table: None,
                 name: Identifier {
                     span: Some(
                         30..31,
@@ -16582,6 +16574,7 @@ Update(
                 },
             },
         ],
+        from: None,
         selection: Some(
             BinaryOp {
                 span: Some(

src/query/sql/src/planner/binder/bind_mutation/delete.rs

@@ -68,6 +68,7 @@ impl Binder {
             target_table_identifier,
             expression: MutationExpression::Delete {
                 target: table.clone(),
+                from: None,
                 filter: selection.clone(),
             },
             strategy: MutationStrategy::MatchedOnly,