cr

Author: yuval-illumex

src/ast/mod.rs

@@ -72,8 +72,8 @@ pub use self::query::{
     TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableSample,
     TableSampleBucket, TableSampleKind, TableSampleMethod, TableSampleModifier,
     TableSampleQuantity, TableSampleSeed, TableSampleSeedModifier, TableSampleUnit, TableVersion,
-    TableWithJoins, Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With,
-    WithFill,
+    TableWithJoins, Top, TopQuantity, UpdateTableFromKind, ValueTableMode, Values,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 pub use self::trigger::{
@@ -2479,7 +2479,7 @@ pub enum Statement {
         /// Column assignments
         assignments: Vec<Assignment>,
         /// Table which provide value to be set
-        from: Option<TableWithJoins>,
+        from: Option<UpdateTableFromKind>,
         /// WHERE
         selection: Option<Expr>,
         /// RETURNING
@@ -3751,10 +3751,13 @@ impl fmt::Display for Statement {
                     write!(f, "{or} ")?;
                 }
                 write!(f, "{table}")?;
+                if let Some(UpdateTableFromKind::BeforeSet(from)) = from {
+                    write!(f, " FROM {from}")?;
+                }
                 if !assignments.is_empty() {
                     write!(f, " SET {}", display_comma_separated(assignments))?;
                 }
-                if let Some(from) = from {
+                if let Some(UpdateTableFromKind::AfterSet(from)) = from {
                     write!(f, " FROM {from}")?;
                 }
                 if let Some(selection) = selection {

src/ast/query.rs

@@ -2788,3 +2788,14 @@ impl fmt::Display for ValueTableMode {
         }
     }
 }
+
+/// The update table from options
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum UpdateTableFromKind {
+    /// Table sample located before the table alias option
+    BeforeSet(TableWithJoins),
+    /// Table sample located after the table alias option
+    AfterSet(TableWithJoins),
+}

src/ast/spans.rs

@@ -32,8 +32,8 @@ use super::{
     PivotValueSource, ProjectionSelect, Query, ReferentialAction, RenameSelectItem,
     ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption,
     Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint,
-    TableFactor, TableOptionsClustered, TableWithJoins, Use, Value, Values, ViewColumnDef,
-    WildcardAdditionalOptions, With, WithFill,
+    TableFactor, TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use, Value, Values,
+    ViewColumnDef, WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -2098,6 +2098,15 @@ impl Spanned for SelectInto {
     }
 }
 
+impl Spanned for UpdateTableFromKind {
+    fn span(&self) -> Span {
+        match self {
+            UpdateTableFromKind::BeforeSet(vec) => vec.span(),
+            UpdateTableFromKind::AfterSet(vec) => vec.span(),
+        }
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use crate::dialect::{Dialect, GenericDialect, SnowflakeDialect};

src/parser/mod.rs

@@ -11699,14 +11699,16 @@ impl<'a> Parser<'a> {
         let or = self.parse_conflict_clause();
         let table = self.parse_table_and_joins()?;
         let from_before_set = if self.parse_keyword(Keyword::FROM) {
-            Some(self.parse_table_and_joins()?)
+            Some(UpdateTableFromKind::BeforeSet(
+                self.parse_table_and_joins()?,
+            ))
         } else {
             None
         };
         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(self.parse_table_and_joins()?)
+            Some(UpdateTableFromKind::AfterSet(self.parse_table_and_joins()?))
         } else {
             from_before_set
         };

tests/sqlparser_common.rs

@@ -366,7 +366,7 @@ fn parse_update_set_from() {
                 target: AssignmentTarget::ColumnName(ObjectName(vec![Ident::new("name")])),
                 value: Expr::CompoundIdentifier(vec![Ident::new("t2"), Ident::new("name")])
             }],
-            from: Some(TableWithJoins {
+            from: Some(UpdateTableFromKind::AfterSet(TableWithJoins {
                 relation: TableFactor::Derived {
                     lateral: false,
                     subquery: Box::new(Query {
@@ -417,8 +417,8 @@ fn parse_update_set_from() {
                         columns: vec![],
                     })
                 },
-                joins: vec![],
-            }),
+                joins: vec![]
+            })),
             selection: Some(Expr::BinaryOp {
                 left: Box::new(Expr::CompoundIdentifier(vec![
                     Ident::new("t1"),
@@ -12449,10 +12449,8 @@ fn overflow() {
 
 #[test]
 fn parse_update_from_before_select() {
-    let res = all_dialects()
-    .parse_sql_statements("UPDATE t1 FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 SET name = t2.name WHERE t1.id = t2.id");
-
-    assert!(res.is_ok(), "{res:?}");
+    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");
 
     let res = all_dialects().parse_sql_statements(
         "UPDATE t1 FROM (SELECT name, id FROM t1 GROUP BY id) AS t2 SET name = t2.name FROM (SELECT name from t2) AS t2",