Merge pull request #7925 from zhyass/feature_cluster

feat(planner): support update ast and planner

src/query/ast/src/ast/format/ast_format.rs

@@ -994,6 +994,27 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         self.children.push(node);
     }
 
+    fn visit_update(&mut self, update: &'ast UpdateStmt<'ast>) {
+        let mut children = Vec::new();
+        self.visit_table_reference(&update.table);
+        children.push(self.children.pop().unwrap());
+
+        for update_expr in update.update_list.iter() {
+            self.visit_identifier(&update_expr.name);
+            children.push(self.children.pop().unwrap());
+            self.visit_expr(&update_expr.expr);
+            children.push(self.children.pop().unwrap());
+        }
+        if let Some(selection) = &update.selection {
+            self.visit_expr(selection);
+            children.push(self.children.pop().unwrap());
+        }
+        let name = "Update".to_string();
+        let format_ctx = AstFormatContext::with_children(name, children.len());
+        let node = FormatTreeNode::with_children(format_ctx, children);
+        self.children.push(node);
+    }
+
     fn visit_show_databases(&mut self, stmt: &'ast ShowDatabasesStmt<'ast>) {
         let mut children = Vec::new();
         if let Some(limit) = &stmt.limit {

src/query/ast/src/ast/format/syntax/dml.rs

@@ -27,6 +27,8 @@ use crate::ast::Expr;
 use crate::ast::InsertSource;
 use crate::ast::InsertStmt;
 use crate::ast::TableReference;
+use crate::ast::UpdateExpr;
+use crate::ast::UpdateStmt;
 
 pub(crate) fn pretty_insert(insert_stmt: InsertStmt) -> RcDoc {
     RcDoc::text("INSERT")
@@ -108,6 +110,46 @@ pub(crate) fn pretty_delete<'a>(
         })
 }
 
+pub(crate) fn pretty_update(update_stmt: UpdateStmt) -> RcDoc {
+    RcDoc::text("UPDATE")
+        .append(
+            RcDoc::line()
+                .nest(NEST_FACTOR)
+                .append(pretty_table(update_stmt.table)),
+        )
+        .append(RcDoc::line().append(RcDoc::text("SET")))
+        .append(pretty_update_list(update_stmt.update_list))
+        .append(if let Some(selection) = update_stmt.selection {
+            RcDoc::line().append(RcDoc::text("WHERE")).append(
+                RcDoc::line()
+                    .nest(NEST_FACTOR)
+                    .append(pretty_expr(selection).nest(NEST_FACTOR).group()),
+            )
+        } else {
+            RcDoc::nil()
+        })
+}
+
+fn pretty_update_list(update_list: Vec<UpdateExpr>) -> RcDoc {
+    if update_list.len() > 1 {
+        RcDoc::line()
+    } else {
+        RcDoc::space()
+    }
+    .nest(NEST_FACTOR)
+    .append(
+        interweave_comma(update_list.into_iter().map(|update_expr| {
+            RcDoc::text(update_expr.name.to_string())
+                .append(RcDoc::space())
+                .append(RcDoc::text("="))
+                .append(RcDoc::space())
+                .append(pretty_expr(update_expr.expr))
+        }))
+        .nest(NEST_FACTOR)
+        .group(),
+    )
+}
+
 pub(crate) fn pretty_copy(copy_stmt: CopyStmt) -> RcDoc {
     RcDoc::text("COPY")
         .append(RcDoc::line().append(RcDoc::text("INTO ")))

src/query/ast/src/ast/format/syntax/mod.rs

@@ -35,6 +35,7 @@ pub fn pretty_statement(stmt: Statement, max_width: usize) -> Result<String> {
             selection,
         } => pretty_delete(table_reference, selection),
         Statement::Copy(copy_stmt) => pretty_copy(copy_stmt),
+        Statement::Update(update_stmt) => pretty_update(update_stmt),
         Statement::CreateTable(create_table_stmt) => pretty_create_table(create_table_stmt),
         Statement::AlterTable(alter_table_stmt) => pretty_alter_table(alter_table_stmt),
         Statement::CreateView(create_view_stmt) => pretty_create_view(create_view_stmt),

src/query/ast/src/ast/statements/mod.rs

@@ -24,6 +24,7 @@ mod show;
 mod stage;
 mod statement;
 mod table;
+mod update;
 mod user;
 mod view;
 
@@ -39,5 +40,6 @@ pub use show::*;
 pub use stage::*;
 pub use statement::*;
 pub use table::*;
+pub use update::*;
 pub use user::*;
 pub use view::*;

src/query/ast/src/ast/statements/statement.rs

@@ -66,6 +66,9 @@ pub enum Statement<'a> {
         table_reference: TableReference<'a>,
         selection: Option<Expr<'a>>,
     },
+
+    Update(UpdateStmt<'a>),
+
     // Databases
     ShowDatabases(ShowDatabasesStmt<'a>),
     ShowCreateDatabase(ShowCreateDatabaseStmt<'a>),
@@ -205,6 +208,7 @@ impl<'a> Display for Statement<'a> {
                     write!(f, "WHERE {conditions} ")?;
                 }
             }
+            Statement::Update(update) => write!(f, "{update}")?,
             Statement::Copy(stmt) => write!(f, "{stmt}")?,
             Statement::ShowSettings { like } => {
                 write!(f, "SHOW SETTINGS")?;

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

@@ -0,0 +1,51 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use crate::ast::write_comma_separated_list;
+use crate::ast::Expr;
+use crate::ast::Identifier;
+use crate::ast::TableReference;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct UpdateStmt<'a> {
+    pub table: TableReference<'a>,
+    pub update_list: Vec<UpdateExpr<'a>>,
+    pub selection: Option<Expr<'a>>,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct UpdateExpr<'a> {
+    pub name: Identifier<'a>,
+    pub expr: Expr<'a>,
+}
+
+impl Display for UpdateStmt<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "UPDATE {} SET ", self.table)?;
+        write_comma_separated_list(f, &self.update_list)?;
+        if let Some(conditions) = &self.selection {
+            write!(f, " WHERE {conditions}")?;
+        }
+        Ok(())
+    }
+}
+
+impl Display for UpdateExpr<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{} = {}", self.name, self.expr)
+    }
+}

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

@@ -111,6 +111,21 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
         },
     );
 
+    let update = map(
+        rule! {
+            UPDATE ~ #table_reference_only
+            ~ SET ~ ^#comma_separated_list1(update_expr)
+            ~ ( WHERE ~ ^#expr )?
+        },
+        |(_, table, _, update_list, opt_selection)| {
+            Statement::Update(UpdateStmt {
+                table,
+                update_list,
+                selection: opt_selection.map(|(_, selection)| selection),
+            })
+        },
+    );
+
     let show_settings = map(
         rule! {
             SHOW ~ SETTINGS ~ (LIKE ~ #literal_string)?
@@ -893,6 +908,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
             | #explain : "`EXPLAIN [PIPELINE | GRAPH] <statement>`"
             | #insert : "`INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
             | #delete : "`DELETE FROM <table> [WHERE ...]`"
+            | #update : "`UPDATE <table> SET <column> = <expr> [, <column> = <expr> , ... ] [WHERE ...]`"
             | #show_settings : "`SHOW SETTINGS [<show_limit>]`"
             | #show_stages : "`SHOW STAGES`"
             | #show_engines : "`SHOW ENGINES`"
@@ -1656,3 +1672,9 @@ pub fn table_reference_only(i: Input) -> IResult<TableReference> {
         },
     )(i)
 }
+
+pub fn update_expr(i: Input) -> IResult<UpdateExpr> {
+    map(rule! { ( #ident ~ "=" ~ ^#expr ) }, |(name, _, expr)| {
+        UpdateExpr { name, expr }
+    })(i)
+}

src/query/ast/src/visitors/visitor.rs

@@ -383,6 +383,8 @@ pub trait Visitor<'ast>: Sized {
     ) {
     }
 
+    fn visit_update(&mut self, _update: &'ast UpdateStmt<'ast>) {}
+
     fn visit_show_databases(&mut self, _stmt: &'ast ShowDatabasesStmt<'ast>) {}
 
     fn visit_show_create_databases(&mut self, _stmt: &'ast ShowCreateDatabaseStmt<'ast>) {}

src/query/ast/src/visitors/visitor_mut.rs

@@ -386,6 +386,8 @@ pub trait VisitorMut: Sized {
     ) {
     }
 
+    fn visit_update(&mut self, _update: &mut UpdateStmt<'_>) {}
+
     fn visit_show_databases(&mut self, _stmt: &mut ShowDatabasesStmt<'_>) {}
 
     fn visit_show_create_databases(&mut self, _stmt: &mut ShowCreateDatabaseStmt<'_>) {}

src/query/ast/src/visitors/walk.rs

@@ -294,6 +294,7 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
             selection,
             ..
         } => visitor.visit_delete(table_reference, selection),
+        Statement::Update(update) => visitor.visit_update(update),
         Statement::Copy(stmt) => visitor.visit_copy(stmt),
         Statement::ShowSettings { like } => visitor.visit_show_settings(like),
         Statement::ShowProcessList => visitor.visit_show_process_list(),

src/query/ast/src/visitors/walk_mut.rs

@@ -300,6 +300,7 @@ pub fn walk_statement_mut<'a, V: VisitorMut>(visitor: &mut V, statement: &mut St
             selection,
             ..
         } => visitor.visit_delete(table_reference, selection),
+        Statement::Update(update) => visitor.visit_update(update),
         Statement::Copy(stmt) => visitor.visit_copy(stmt),
         Statement::ShowSettings { like } => visitor.visit_show_settings(like),
         Statement::ShowProcessList => visitor.visit_show_process_list(),

src/query/ast/tests/it/parser.rs

@@ -295,6 +295,7 @@ fn test_statement() {
         r#"SHOW GRANTS ON TABLE db1.tb1;"#,
         r#"SHOW GRANTS ON DATABASE db;"#,
         r#"SHOW GRANTS OF SHARE t;"#,
+        r#"UPDATE db1.tb1 set a = a + 1, b = 2 WHERE c > 3;"#,
     ];
 
     for case in cases {

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

@@ -7394,3 +7394,114 @@ ShowGrantsOfShare(
 )
 
 
+---------- Input ----------
+UPDATE db1.tb1 set a = a + 1, b = 2 WHERE c > 3;
+---------- Output ---------
+UPDATE db1.tb1 SET a = a + 1, b = 2 WHERE c > 3
+---------- AST ------------
+Update(
+    UpdateStmt {
+        table: Table {
+            span: [
+                Ident(7..10),
+                Period(10..11),
+                Ident(11..14),
+            ],
+            catalog: None,
+            database: Some(
+                Identifier {
+                    name: "db1",
+                    quote: None,
+                    span: Ident(7..10),
+                },
+            ),
+            table: Identifier {
+                name: "tb1",
+                quote: None,
+                span: Ident(11..14),
+            },
+            alias: None,
+            travel_point: None,
+        },
+        update_list: [
+            UpdateExpr {
+                name: Identifier {
+                    name: "a",
+                    quote: None,
+                    span: Ident(19..20),
+                },
+                expr: BinaryOp {
+                    span: [
+                        Plus(25..26),
+                    ],
+                    op: Plus,
+                    left: ColumnRef {
+                        span: [
+                            Ident(23..24),
+                        ],
+                        database: None,
+                        table: None,
+                        column: Identifier {
+                            name: "a",
+                            quote: None,
+                            span: Ident(23..24),
+                        },
+                    },
+                    right: Literal {
+                        span: [
+                            LiteralInteger(27..28),
+                        ],
+                        lit: Integer(
+                            1,
+                        ),
+                    },
+                },
+            },
+            UpdateExpr {
+                name: Identifier {
+                    name: "b",
+                    quote: None,
+                    span: Ident(30..31),
+                },
+                expr: Literal {
+                    span: [
+                        LiteralInteger(34..35),
+                    ],
+                    lit: Integer(
+                        2,
+                    ),
+                },
+            },
+        ],
+        selection: Some(
+            BinaryOp {
+                span: [
+                    Gt(44..45),
+                ],
+                op: Gt,
+                left: ColumnRef {
+                    span: [
+                        Ident(42..43),
+                    ],
+                    database: None,
+                    table: None,
+                    column: Identifier {
+                        name: "c",
+                        quote: None,
+                        span: Ident(42..43),
+                    },
+                },
+                right: Literal {
+                    span: [
+                        LiteralInteger(46..47),
+                    ],
+                    lit: Integer(
+                        3,
+                    ),
+                },
+            },
+        ),
+    },
+)
+
+

src/query/service/src/interpreters/access/privilege_access.rs

@@ -149,6 +149,18 @@ impl AccessChecker for PrivilegeAccess {
             // Others.
             Plan::Insert(_) => {}
             Plan::Delete(_) => {}
+            Plan::Update(plan) => {
+                session
+                    .validate_privilege(
+                        &GrantObject::Table(
+                            plan.catalog.clone(),
+                            plan.database.clone(),
+                            plan.table.clone(),
+                        ),
+                        UserPrivilegeType::Update,
+                    )
+                    .await?;
+            }
             Plan::CreateView(plan) => {
                 session
                     .validate_privilege(