[MySQL, Oracle] Parse optimizer hints for UPDATEs

Author: xitep

src/ast/dml.rs

@@ -265,6 +265,11 @@ impl Display for Delete {
 pub struct Update {
     /// Token for the `UPDATE` keyword
     pub update_token: AttachedToken,
+    /// A query optimizer hint
+    ///
+    /// [MySQL](https://dev.mysql.com/doc/refman/8.4/en/optimizer-hints.html)
+    /// [Oracle](https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/Comments.html#GUID-D316D545-89E2-4D54-977F-FC97815CD62E)
+    pub optimizer_hint: Option<OptimizerHint>,
     /// TABLE
     pub table: TableWithJoins,
     /// Column assignments
@@ -284,6 +289,10 @@ pub struct Update {
 impl Display for Update {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.write_str("UPDATE ")?;
+        if let Some(hint) = self.optimizer_hint.as_ref() {
+            hint.fmt(f)?;
+            f.write_str(" ")?;
+        }
         if let Some(or) = &self.or {
             or.fmt(f)?;
             f.write_str(" ")?;

src/ast/spans.rs

@@ -927,6 +927,7 @@ impl Spanned for Update {
     fn span(&self) -> Span {
         let Update {
             update_token,
+            optimizer_hint: _,
             table,
             assignments,
             from,

src/parser/mod.rs

@@ -17011,6 +17011,7 @@ impl<'a> Parser<'a> {
 
     /// Parse an `UPDATE` statement and return `Statement::Update`.
     pub fn parse_update(&mut self, update_token: TokenWithSpan) -> Result<Statement, ParserError> {
+        let optimizer_hint = self.parse_optional_optimizer_hint()?;
         let or = self.parse_conflict_clause();
         let table = self.parse_table_and_joins()?;
         let from_before_set = if self.parse_keyword(Keyword::FROM) {
@@ -17046,6 +17047,7 @@ impl<'a> Parser<'a> {
         };
         Ok(Update {
             update_token: update_token.into(),
+            optimizer_hint,
             table,
             assignments,
             from,

tests/sqlparser_common.rs

@@ -457,6 +457,7 @@ fn parse_update_set_from() {
         stmt,
         Statement::Update(Update {
             update_token: AttachedToken::empty(),
+            optimizer_hint: None,
             table: TableWithJoins {
                 relation: table_from_name(ObjectName::from(vec![Ident::new("t1")])),
                 joins: vec![],
@@ -550,6 +551,7 @@ fn parse_update_with_table_alias() {
             returning,
             or: None,
             limit: None,
+            optimizer_hint: None,
             update_token: _,
         }) => {
             assert_eq!(

tests/sqlparser_mysql.rs

@@ -2635,6 +2635,7 @@ fn parse_update_with_joins() {
             returning,
             or: None,
             limit: None,
+            optimizer_hint: None,
             update_token: _,
         }) => {
             assert_eq!(
@@ -4385,4 +4386,10 @@ fn test_optimizer_hints() {
     mysql_dialect.verified_stmt("\
        REPLACE /*+ foobar */ INTO test \
        VALUES (1, 'Old', '2014-08-20 18:47:00')");
+
+    // ~ updates
+    mysql_dialect.verified_stmt("\
+       UPDATE /*+ quux */ table_name \
+       SET column1 = 1 \
+       WHERE 1 = 1");
 }

tests/sqlparser_oracle.rs

@@ -373,4 +373,7 @@ fn test_optimizer_hints() {
     oracle_dialect.verified_stmt(
         "INSERT /*+ append */ INTO t1 SELECT * FROM all_objects");
 
+    // ~ updates
+    oracle_dialect.verified_stmt(
+        "UPDATE /*+ DISABLE_PARALLEL_DML */ table_name SET column1 = 1");
 }

tests/sqlparser_sqlite.rs

@@ -477,6 +477,7 @@ fn parse_update_tuple_row_values() {
     assert_eq!(
         sqlite().verified_stmt("UPDATE x SET (a, b) = (1, 2)"),
         Statement::Update(Update {
+            optimizer_hint: None,
             or: None,
             assignments: vec![Assignment {
                 target: AssignmentTarget::Tuple(vec![