From 8a03f85aeaf4e3f4b9e55043339bdb1374fe2571 Mon Sep 17 00:00:00 2001
From: ifeanyi <ifeanyi@validio.io>
Date: Thu, 23 Jan 2025 18:11:27 +0100
Subject: [PATCH] Support multiple tables in `UPDATE FROM` clause

Adds support for multiple tables in the `FROM` clause
of an update statement.

```sql
UPDATE A
SET ..
FROM T, U, V
```
---
 src/ast/mod.rs            |  4 ++--
 src/ast/query.rs          |  4 ++--
 src/ast/spans.rs          |  9 +++++----
 src/parser/mod.rs         |  6 ++++--
 tests/sqlparser_common.rs | 11 +++++++----
 5 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index b473dc11f..e64b7d3db 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -3872,13 +3872,13 @@ impl fmt::Display for Statement {
                 }
                 write!(f, "{table}")?;
                 if let Some(UpdateTableFromKind::BeforeSet(from)) = from {
-                    write!(f, " FROM {from}")?;
+                    write!(f, " FROM {}", display_comma_separated(from))?;
                 }
                 if !assignments.is_empty() {
                     write!(f, " SET {}", display_comma_separated(assignments))?;
                 }
                 if let Some(UpdateTableFromKind::AfterSet(from)) = from {
-                    write!(f, " FROM {from}")?;
+                    write!(f, " FROM {}", display_comma_separated(from))?;
                 }
                 if let Some(selection) = selection {
                     write!(f, " WHERE {selection}")?;
diff --git a/src/ast/query.rs b/src/ast/query.rs
index 4053dd239..e982c7f0d 100644
--- a/src/ast/query.rs
+++ b/src/ast/query.rs
@@ -2829,8 +2829,8 @@ impl fmt::Display for ValueTableMode {
 pub enum UpdateTableFromKind {
     /// Update Statement where the 'FROM' clause is before the 'SET' keyword (Supported by Snowflake)
     /// For Example: `UPDATE FROM t1 SET t1.name='aaa'`
-    BeforeSet(TableWithJoins),
+    BeforeSet(Vec<TableWithJoins>),
     /// Update Statement where the 'FROM' clause is after the 'SET' keyword (Which is the standard way)
     /// For Example: `UPDATE SET t1.name='aaa' FROM t1`
-    AfterSet(TableWithJoins),
+    AfterSet(Vec<TableWithJoins>),
 }
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
index 5316bfbda..aed1c6c2b 100644
--- a/src/ast/spans.rs
+++ b/src/ast/spans.rs
@@ -2138,10 +2138,11 @@ impl Spanned for SelectInto {
 
 impl Spanned for UpdateTableFromKind {
     fn span(&self) -> Span {
-        match self {
-            UpdateTableFromKind::BeforeSet(from) => from.span(),
-            UpdateTableFromKind::AfterSet(from) => from.span(),
-        }
+        let from = match self {
+            UpdateTableFromKind::BeforeSet(from) => from,
+            UpdateTableFromKind::AfterSet(from) => from,
+        };
+        union_spans(from.iter().map(|t| t.span()))
     }
 }
 
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 0d2973c7f..c6e1eb196 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -12617,7 +12617,7 @@ impl<'a> Parser<'a> {
         let table = self.parse_table_and_joins()?;
         let from_before_set = if self.parse_keyword(Keyword::FROM) {
             Some(UpdateTableFromKind::BeforeSet(
-                self.parse_table_and_joins()?,
+                self.parse_table_with_joins()?,
             ))
         } else {
             None
@@ -12625,7 +12625,9 @@ impl<'a> Parser<'a> {
         self.expect_keyword(Keyword::SET)?;
         let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
         let from = if from_before_set.is_none() && self.parse_keyword(Keyword::FROM) {
-            Some(UpdateTableFromKind::AfterSet(self.parse_table_and_joins()?))
+            Some(UpdateTableFromKind::AfterSet(
+                self.parse_table_with_joins()?,
+            ))
         } else {
             from_before_set
         };
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 5a1e812d0..5c11b2901 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -403,7 +403,7 @@ fn parse_update_set_from() {
                 target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("name")])),
                 value: Expr::CompoundIdentifier(vec![Ident::new("t2"), Ident::new("name")])
             }],
-            from: Some(UpdateTableFromKind::AfterSet(TableWithJoins {
+            from: Some(UpdateTableFromKind::AfterSet(vec![TableWithJoins {
                 relation: TableFactor::Derived {
                     lateral: false,
                     subquery: Box::new(Query {
@@ -455,7 +455,7 @@ fn parse_update_set_from() {
                     })
                 },
                 joins: vec![]
-            })),
+            }])),
             selection: Some(Expr::BinaryOp {
                 left: Box::new(Expr::CompoundIdentifier(vec![
                     Ident::new("t1"),
@@ -471,6 +471,9 @@ fn parse_update_set_from() {
             or: None,
         }
     );
+
+    let sql = "UPDATE T SET a = b FROM U, (SELECT foo FROM V) AS W WHERE 1 = 1";
+    dialects.verified_stmt(sql);
 }
 
 #[test]
@@ -13051,8 +13054,8 @@ fn parse_select_without_projection() {
 
 #[test]
 fn parse_update_from_before_select() {
-    all_dialects()
-    .verified_stmt("UPDATE t1 FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 SET name = t2.name WHERE t1.id = t2.id");
+    verified_stmt("UPDATE t1 FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 SET name = t2.name WHERE t1.id = t2.id");
+    verified_stmt("UPDATE t1 FROM U, (SELECT id FROM V) AS W SET a = b WHERE 1 = 1");
 
     let query =
     "UPDATE t1 FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 SET name = t2.name FROM (SELECT name from t2) AS t2";