Support Update Set From statement (#450)

src/ast/mod.rs

@@ -758,6 +758,8 @@ pub enum Statement {
         table: TableWithJoins,
         /// Column assignments
         assignments: Vec<Assignment>,
+        /// Table which provide value to be set
+        from: Option<TableWithJoins>,
         /// WHERE
         selection: Option<Expr>,
     },
@@ -1191,12 +1193,16 @@ impl fmt::Display for Statement {
             Statement::Update {
                 table,
                 assignments,
+                from,
                 selection,
             } => {
                 write!(f, "UPDATE {}", table)?;
                 if !assignments.is_empty() {
                     write!(f, " SET {}", display_comma_separated(assignments))?;
                 }
+                if let Some(from) = from {
+                    write!(f, " FROM {}", from)?;
+                }
                 if let Some(selection) = selection {
                     write!(f, " WHERE {}", selection)?;
                 }

src/parser.rs

@@ -3611,6 +3611,11 @@ impl<'a> Parser<'a> {
         let table = self.parse_table_and_joins()?;
         self.expect_keyword(Keyword::SET)?;
         let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
+        let from = if self.parse_keyword(Keyword::FROM) && dialect_of!(self is PostgreSqlDialect) {
+            Some(self.parse_table_and_joins()?)
+        } else {
+            None
+        };
         let selection = if self.parse_keyword(Keyword::WHERE) {
             Some(self.parse_expr()?)
         } else {
@@ -3619,6 +3624,7 @@ impl<'a> Parser<'a> {
         Ok(Statement::Update {
             table,
             assignments,
+            from,
             selection,
         })
     }

tests/sqlparser_common.rs

@@ -195,6 +195,7 @@ fn parse_update_with_table_alias() {
         Statement::Update {
             table,
             assignments,
+            from: _from,
             selection,
         } => {
             assert_eq!(

tests/sqlparser_mysql.rs

@@ -614,6 +614,7 @@ fn parse_update_with_joins() {
         Statement::Update {
             table,
             assignments,
+            from: _from,
             selection,
         } => {
             assert_eq!(

tests/sqlparser_postgres.rs

@@ -401,6 +401,84 @@ PHP	₱ USD $
     //assert_eq!(sql, ast.to_string());
 }
 
+#[test]
+fn parse_update_set_from() {
+    let sql = "UPDATE t1 SET name = t2.name FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 WHERE t1.id = t2.id";
+    let stmt = pg().verified_stmt(sql);
+    assert_eq!(
+        stmt,
+        Statement::Update {
+            table: TableWithJoins {
+                relation: TableFactor::Table {
+                    name: ObjectName(vec![Ident::new("t1")]),
+                    alias: None,
+                    args: vec![],
+                    with_hints: vec![],
+                },
+                joins: vec![],
+            },
+            assignments: vec![Assignment {
+                id: vec![Ident::new("name")],
+                value: Expr::CompoundIdentifier(vec![Ident::new("t2"), Ident::new("name")])
+            }],
+            from: Some(TableWithJoins {
+                relation: TableFactor::Derived {
+                    lateral: false,
+                    subquery: Box::new(Query {
+                        with: None,
+                        body: SetExpr::Select(Box::new(Select {
+                            distinct: false,
+                            top: None,
+                            projection: vec![
+                                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("name"))),
+                                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("id"))),
+                            ],
+                            into: None,
+                            from: vec![TableWithJoins {
+                                relation: TableFactor::Table {
+                                    name: ObjectName(vec![Ident::new("t1")]),
+                                    alias: None,
+                                    args: vec![],
+                                    with_hints: vec![],
+                                },
+                                joins: vec![],
+                            }],
+                            lateral_views: vec![],
+                            selection: None,
+                            group_by: vec![Expr::Identifier(Ident::new("id"))],
+                            cluster_by: vec![],
+                            distribute_by: vec![],
+                            sort_by: vec![],
+                            having: None,
+                        })),
+                        order_by: vec![],
+                        limit: None,
+                        offset: None,
+                        fetch: None,
+                        lock: None,
+                    }),
+                    alias: Some(TableAlias {
+                        name: Ident::new("t2"),
+                        columns: vec![],
+                    })
+                },
+                joins: vec![],
+            }),
+            selection: Some(Expr::BinaryOp {
+                left: Box::new(Expr::CompoundIdentifier(vec![
+                    Ident::new("t1"),
+                    Ident::new("id")
+                ])),
+                op: BinaryOperator::Eq,
+                right: Box::new(Expr::CompoundIdentifier(vec![
+                    Ident::new("t2"),
+                    Ident::new("id")
+                ])),
+            }),
+        }
+    );
+}
+
 #[test]
 fn test_copy_from() {
     let stmt = pg().verified_stmt("COPY users FROM 'data.csv'");