Moved the Update struct out of the Statement enum

LucaCappelletti94 · LucaCappelletti94 · commit 9a4d23563b60 · 2025-10-06T13:32:55.000+02:00
diff --git a/src/ast/dml.rs b/src/ast/dml.rs
@@ -29,7 +29,7 @@ use crate::display_utils::{indented_list, Indent, SpaceOrNewline};
 use super::{
     display_comma_separated, query::InputFormatClause, Assignment, Expr, FromTable, Ident,
     InsertAliases, MysqlInsertPriority, ObjectName, OnInsert, OrderByExpr, Query, SelectItem,
-    Setting, SqliteOnConflict, TableObject, TableWithJoins,
+    Setting, SqliteOnConflict, TableObject, TableWithJoins, UpdateTableFromKind,
 };
 
 /// INSERT statement.
@@ -240,3 +240,66 @@ impl Display for Delete {
         Ok(())
     }
 }
+
+/// UPDATE statement.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct Update {
+    /// TABLE
+    pub table: TableWithJoins,
+    /// Column assignments
+    pub assignments: Vec<Assignment>,
+    /// Table which provide value to be set
+    pub from: Option<UpdateTableFromKind>,
+    /// WHERE
+    pub selection: Option<Expr>,
+    /// RETURNING
+    pub returning: Option<Vec<SelectItem>>,
+    /// SQLite-specific conflict resolution clause
+    pub or: Option<SqliteOnConflict>,
+    /// LIMIT
+    pub limit: Option<Expr>,
+}
+
+impl Display for Update {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str("UPDATE ")?;
+        if let Some(or) = &self.or {
+            or.fmt(f)?;
+            f.write_str(" ")?;
+        }
+        self.table.fmt(f)?;
+        if let Some(UpdateTableFromKind::BeforeSet(from)) = &self.from {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("FROM")?;
+            indented_list(f, from)?;
+        }
+        if !self.assignments.is_empty() {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("SET")?;
+            indented_list(f, &self.assignments)?;
+        }
+        if let Some(UpdateTableFromKind::AfterSet(from)) = &self.from {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("FROM")?;
+            indented_list(f, from)?;
+        }
+        if let Some(selection) = &self.selection {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("WHERE")?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(selection).fmt(f)?;
+        }
+        if let Some(returning) = &self.returning {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("RETURNING")?;
+            indented_list(f, returning)?;
+        }
+        if let Some(limit) = &self.limit {
+            SpaceOrNewline.fmt(f)?;
+            write!(f, "LIMIT {limit}")?;
+        }
+        Ok(())
+    }
+}
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -43,7 +43,7 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::{
-    display_utils::{indented_list, SpaceOrNewline},
+    display_utils::SpaceOrNewline,
     tokenizer::{Span, Token},
 };
 use crate::{
@@ -72,7 +72,7 @@ pub use self::ddl::{
     TableConstraint, TagsColumnOption, Truncate, UserDefinedTypeCompositeAttributeDef,
     UserDefinedTypeRepresentation, ViewColumnDef,
 };
-pub use self::dml::{Delete, Insert};
+pub use self::dml::{Delete, Insert, Update};
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
@@ -3241,22 +3241,7 @@ pub enum Statement {
     /// ```sql
     /// UPDATE
     /// ```
-    Update {
-        /// TABLE
-        table: TableWithJoins,
-        /// Column assignments
-        assignments: Vec<Assignment>,
-        /// Table which provide value to be set
-        from: Option<UpdateTableFromKind>,
-        /// WHERE
-        selection: Option<Expr>,
-        /// RETURNING
-        returning: Option<Vec<SelectItem>>,
-        /// SQLite-specific conflict resolution clause
-        or: Option<SqliteOnConflict>,
-        /// LIMIT
-        limit: Option<Expr>,
-    },
+    Update(Update),
     /// ```sql
     /// DELETE
     /// ```
@@ -4676,53 +4661,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Update {
-                table,
-                assignments,
-                from,
-                selection,
-                returning,
-                or,
-                limit,
-            } => {
-                f.write_str("UPDATE ")?;
-                if let Some(or) = or {
-                    or.fmt(f)?;
-                    f.write_str(" ")?;
-                }
-                table.fmt(f)?;
-                if let Some(UpdateTableFromKind::BeforeSet(from)) = from {
-                    SpaceOrNewline.fmt(f)?;
-                    f.write_str("FROM")?;
-                    indented_list(f, from)?;
-                }
-                if !assignments.is_empty() {
-                    SpaceOrNewline.fmt(f)?;
-                    f.write_str("SET")?;
-                    indented_list(f, assignments)?;
-                }
-                if let Some(UpdateTableFromKind::AfterSet(from)) = from {
-                    SpaceOrNewline.fmt(f)?;
-                    f.write_str("FROM")?;
-                    indented_list(f, from)?;
-                }
-                if let Some(selection) = selection {
-                    SpaceOrNewline.fmt(f)?;
-                    f.write_str("WHERE")?;
-                    SpaceOrNewline.fmt(f)?;
-                    Indent(selection).fmt(f)?;
-                }
-                if let Some(returning) = returning {
-                    SpaceOrNewline.fmt(f)?;
-                    f.write_str("RETURNING")?;
-                    indented_list(f, returning)?;
-                }
-                if let Some(limit) = limit {
-                    SpaceOrNewline.fmt(f)?;
-                    write!(f, "LIMIT {limit}")?;
-                }
-                Ok(())
-            }
+            Statement::Update(update) => update.fmt(f),
             Statement::Delete(delete) => delete.fmt(f),
             Statement::Open(open) => open.fmt(f),
             Statement::Close { cursor } => {
@@ -10891,6 +10830,12 @@ impl From<Insert> for Statement {
     }
 }
 
+impl From<Update> for Statement {
+    fn from(u: Update) -> Self {
+        Self::Update(u)
+    }
+}
+
 impl From<CaseStatement> for Statement {
     fn from(c: CaseStatement) -> Self {
         Self::Case(c)
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -40,8 +40,8 @@ use super::{
     RaiseStatementValue, ReferentialAction, RenameSelectItem, ReplaceSelectElement,
     ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript,
     SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
-    TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef,
-    WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    TableOptionsClustered, TableWithJoins, Update, UpdateTableFromKind, Use, Value, Values,
+    ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -346,21 +346,7 @@ impl Spanned for Statement {
                 CloseCursor::All => Span::empty(),
                 CloseCursor::Specific { name } => name.span,
             },
-            Statement::Update {
-                table,
-                assignments,
-                from,
-                selection,
-                returning,
-                or: _,
-                limit: _,
-            } => union_spans(
-                core::iter::once(table.span())
-                    .chain(assignments.iter().map(|i| i.span()))
-                    .chain(from.iter().map(|i| i.span()))
-                    .chain(selection.iter().map(|i| i.span()))
-                    .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span()))),
-            ),
+            Statement::Update(update) => update.span(),
             Statement::Delete(delete) => delete.span(),
             Statement::CreateView {
                 or_alter: _,
@@ -997,6 +983,29 @@ impl Spanned for Delete {
     }
 }
 
+impl Spanned for Update {
+    fn span(&self) -> Span {
+        let Update {
+            table,
+            assignments,
+            from,
+            selection,
+            returning,
+            or: _,
+            limit,
+        } = self;
+
+        union_spans(
+            core::iter::once(table.span())
+                .chain(assignments.iter().map(|i| i.span()))
+                .chain(from.iter().map(|i| i.span()))
+                .chain(selection.iter().map(|i| i.span()))
+                .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
+                .chain(limit.iter().map(|i| i.span())),
+        )
+    }
+}
+
 impl Spanned for FromTable {
     fn span(&self) -> Span {
         match self {
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -15640,15 +15640,16 @@ impl<'a> Parser<'a> {
         } else {
             None
         };
-        Ok(Statement::Update {
+        Ok(Update {
             table,
             assignments,
             from,
             selection,
             returning,
             or,
             limit,
-        })
+        }
+        .into())
     }
 
     /// Parse a `var = expr` assignment, used in an UPDATE statement
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -377,12 +377,12 @@ fn parse_insert_sqlite() {
 fn parse_update() {
     let sql = "UPDATE t SET a = 1, b = 2, c = 3 WHERE d";
     match verified_stmt(sql) {
-        Statement::Update {
+        Statement::Update(Update {
             table,
             assignments,
             selection,
             ..
-        } => {
+        }) => {
             assert_eq!(table.to_string(), "t".to_string());
             assert_eq!(
                 assignments,
@@ -439,7 +439,7 @@ fn parse_update_set_from() {
     let stmt = dialects.verified_stmt(sql);
     assert_eq!(
         stmt,
-        Statement::Update {
+        Statement::Update(Update {
             table: TableWithJoins {
                 relation: table_from_name(ObjectName::from(vec![Ident::new("t1")])),
                 joins: vec![],
@@ -516,7 +516,7 @@ fn parse_update_set_from() {
             returning: None,
             or: None,
             limit: None
-        }
+        })
     );
 
     let sql = "UPDATE T SET a = b FROM U, (SELECT foo FROM V) AS W WHERE 1 = 1";
@@ -527,15 +527,15 @@ fn parse_update_set_from() {
 fn parse_update_with_table_alias() {
     let sql = "UPDATE users AS u SET u.username = 'new_user' WHERE u.username = 'old_user'";
     match verified_stmt(sql) {
-        Statement::Update {
+        Statement::Update(Update {
             table,
             assignments,
             from: _from,
             selection,
             returning,
             or: None,
             limit: None,
-        } => {
+        }) => {
             assert_eq!(
                 TableWithJoins {
                     relation: TableFactor::Table {
@@ -591,7 +591,7 @@ fn parse_update_with_table_alias() {
 #[test]
 fn parse_update_or() {
     let expect_or_clause = |sql: &str, expected_action: SqliteOnConflict| match verified_stmt(sql) {
-        Statement::Update { or, .. } => assert_eq!(or, Some(expected_action)),
+        Statement::Update(Update { or, .. }) => assert_eq!(or, Some(expected_action)),
         other => unreachable!("Expected update with or, got {:?}", other),
     };
     expect_or_clause(
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -2601,15 +2601,15 @@ fn parse_insert_with_numeric_prefix_column_name() {
 fn parse_update_with_joins() {
     let sql = "UPDATE orders AS o JOIN customers AS c ON o.customer_id = c.id SET o.completed = true WHERE c.firstname = 'Peter'";
     match mysql().verified_stmt(sql) {
-        Statement::Update {
+        Statement::Update(Update {
             table,
             assignments,
             from: _from,
             selection,
             returning,
             or: None,
             limit: None,
-        } => {
+        }) => {
             assert_eq!(
                 TableWithJoins {
                     relation: TableFactor::Table {
@@ -4177,7 +4177,7 @@ fn test_variable_assignment_using_colon_equal() {
     let stmt = mysql().verified_stmt(sql_update);
 
     match stmt {
-        Statement::Update { assignments, .. } => {
+        Statement::Update(Update { assignments, .. }) => {
             assert_eq!(
                 assignments,
                 vec![Assignment {
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -2008,7 +2008,7 @@ fn parse_pg_returning() {
              RETURNING temp_lo AS lo, temp_hi AS hi, prcp",
     );
     match stmt {
-        Statement::Update { returning, .. } => {
+        Statement::Update(Update { returning, .. }) => {
             assert_eq!(
                 Some(vec![
                     SelectItem::ExprWithAlias {
diff --git a/tests/sqlparser_sqlite.rs b/tests/sqlparser_sqlite.rs
@@ -467,7 +467,7 @@ fn parse_update_tuple_row_values() {
     // See https://github.com/sqlparser-rs/sqlparser-rs/issues/1311
     assert_eq!(
         sqlite().verified_stmt("UPDATE x SET (a, b) = (1, 2)"),
-        Statement::Update {
+        Statement::Update(Update {
             or: None,
             assignments: vec![Assignment {
                 target: AssignmentTarget::Tuple(vec![
@@ -487,7 +487,7 @@ fn parse_update_tuple_row_values() {
             from: None,
             returning: None,
             limit: None
-        }
+        })
     );
 }
 
@@ -596,7 +596,7 @@ fn test_regexp_operator() {
 #[test]
 fn test_update_delete_limit() {
     match sqlite().verified_stmt("UPDATE foo SET bar = 1 LIMIT 99") {
-        Statement::Update { limit, .. } => {
+        Statement::Update(Update { limit, .. }) => {
             assert_eq!(limit, Some(Expr::value(number("99"))));
         }
         _ => unreachable!(),
