MySQL: Add support for ORDER BY on UPDATE

Author: tpyo

src/ast/dml.rs

@@ -387,6 +387,9 @@ pub struct Update {
     pub output: Option<OutputClause>,
     /// SQLite-specific conflict resolution clause
     pub or: Option<SqliteOnConflict>,
+    /// ORDER BY (MySQL extension for single-table UPDATE)
+    /// See <https://dev.mysql.com/doc/refman/8.4/en/update.html>
+    pub order_by: Vec<OrderByExpr>,
     /// LIMIT
     pub limit: Option<Expr>,
 }
@@ -434,6 +437,11 @@ impl Display for Update {
             f.write_str("RETURNING")?;
             indented_list(f, returning)?;
         }
+        if !self.order_by.is_empty() {
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("ORDER BY")?;
+            indented_list(f, &self.order_by)?;
+        }
         if let Some(limit) = &self.limit {
             SpaceOrNewline.fmt(f)?;
             write!(f, "LIMIT {limit}")?;

src/ast/spans.rs

@@ -952,6 +952,7 @@ impl Spanned for Update {
             returning,
             output,
             or: _,
+            order_by,
             limit,
         } = self;
 
@@ -963,6 +964,7 @@ impl Spanned for Update {
                 .chain(selection.iter().map(|i| i.span()))
                 .chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
                 .chain(output.iter().map(|i| i.span()))
+                .chain(order_by.iter().map(|i| i.span()))
                 .chain(limit.iter().map(|i| i.span())),
         )
     }

src/dialect/mod.rs

@@ -520,6 +520,16 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports `ORDER BY` in `UPDATE` statements.
+    ///
+    /// ```sql
+    /// UPDATE foo SET bar = false WHERE foo = true ORDER BY foo ASC;
+    /// ```
+    /// See <https://dev.mysql.com/doc/refman/8.4/en/update.html>
+    fn supports_update_order_by(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports an `EXCEPT` clause following a
     /// wildcard in a select list.
     ///

src/dialect/mysql.rs

@@ -179,6 +179,11 @@ impl Dialect for MySqlDialect {
         true
     }
 
+    /// See: <https://dev.mysql.com/doc/refman/8.4/en/update.html>
+    fn supports_update_order_by(&self) -> bool {
+        true
+    }
+
     fn supports_data_type_signed_suffix(&self) -> bool {
         true
     }

src/parser/mod.rs

@@ -17745,6 +17745,13 @@ impl<'a> Parser<'a> {
         } else {
             None
         };
+        let order_by = if self.dialect.supports_update_order_by()
+            && self.parse_keywords(&[Keyword::ORDER, Keyword::BY])
+        {
+            self.parse_comma_separated(Parser::parse_order_by_expr)?
+        } else {
+            vec![]
+        };
         let limit = if self.parse_keyword(Keyword::LIMIT) {
             Some(self.parse_expr()?)
         } else {
@@ -17760,6 +17767,7 @@ impl<'a> Parser<'a> {
             returning,
             output,
             or,
+            order_by,
             limit,
         }
         .into())

tests/sqlparser_common.rs

@@ -535,6 +535,7 @@ fn parse_update_set_from() {
             returning: None,
             output: None,
             or: None,
+            order_by: vec![],
             limit: None
         })
     );
@@ -554,6 +555,7 @@ fn parse_update_with_table_alias() {
             selection,
             returning,
             or: None,
+            order_by: _,
             limit: None,
             optimizer_hints,
             update_token: _,

tests/sqlparser_mysql.rs

@@ -2707,6 +2707,7 @@ fn parse_update_with_joins() {
             selection,
             returning,
             or: None,
+            order_by: _,
             limit: None,
             optimizer_hints,
             update_token: _,
@@ -2784,6 +2785,57 @@ fn parse_update_with_joins() {
     }
 }
 
+#[test]
+fn parse_update_with_order_by() {
+    let sql = "UPDATE foo SET bar = false WHERE foo = true ORDER BY foo ASC";
+    match mysql().verified_stmt(sql) {
+        Statement::Update(Update { order_by, .. }) => {
+            assert_eq!(
+                vec![OrderByExpr {
+                    expr: Expr::Identifier(Ident {
+                        value: "foo".to_owned(),
+                        quote_style: None,
+                        span: Span::empty(),
+                    }),
+                    options: OrderByOptions {
+                        asc: Some(true),
+                        nulls_first: None,
+                    },
+                    with_fill: None,
+                }],
+                order_by
+            );
+        }
+        _ => unreachable!(),
+    }
+}
+
+#[test]
+fn parse_update_with_order_by_and_limit() {
+    let sql = "UPDATE foo SET bar = false WHERE foo = true ORDER BY foo ASC LIMIT 10";
+    match mysql().verified_stmt(sql) {
+        Statement::Update(Update { order_by, limit, .. }) => {
+            assert_eq!(
+                vec![OrderByExpr {
+                    expr: Expr::Identifier(Ident {
+                        value: "foo".to_owned(),
+                        quote_style: None,
+                        span: Span::empty(),
+                    }),
+                    options: OrderByOptions {
+                        asc: Some(true),
+                        nulls_first: None,
+                    },
+                    with_fill: None,
+                }],
+                order_by
+            );
+            assert_eq!(Some(Expr::value(number("10"))), limit);
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_delete_with_order_by() {
     let sql = "DELETE FROM customers ORDER BY id DESC";

tests/sqlparser_sqlite.rs

@@ -497,6 +497,7 @@ fn parse_update_tuple_row_values() {
             from: None,
             returning: None,
             output: None,
+            order_by: vec![],
             limit: None,
             update_token: AttachedToken::empty()
         })