[MySQL, Oracle] Parse optimizer hints for INSERTs

Author: xitep

src/ast/dml.rs

@@ -25,15 +25,12 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::{
-    ast::display_separated,
-    display_utils::{indented_list, Indent, SpaceOrNewline},
+    ast::{display_separated},
+    display_utils::{Indent, SpaceOrNewline, indented_list},
 };
 
 use super::{
-    display_comma_separated, helpers::attached_token::AttachedToken, query::InputFormatClause,
-    Assignment, Expr, FromTable, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnInsert,
-    OrderByExpr, Query, SelectInto, SelectItem, Setting, SqliteOnConflict, TableFactor,
-    TableObject, TableWithJoins, UpdateTableFromKind, Values,
+    Assignment, Expr, FromTable, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnInsert, OptimizerHint, OrderByExpr, Query, SelectInto, SelectItem, Setting, SqliteOnConflict, TableFactor, TableObject, TableWithJoins, UpdateTableFromKind, Values, display_comma_separated, helpers::attached_token::AttachedToken, query::InputFormatClause
 };
 
 /// INSERT statement.
@@ -43,6 +40,11 @@ use super::{
 pub struct Insert {
     /// Token for the `INSERT` keyword (or its substitutes)
     pub insert_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>,
     /// Only for Sqlite
     pub or: Option<SqliteOnConflict>,
     /// Only for mysql
@@ -102,7 +104,11 @@ impl Display for Insert {
         };
 
         if let Some(on_conflict) = self.or {
-            write!(f, "INSERT {on_conflict} INTO {table_name} ")?;
+            f.write_str("INSERT")?;
+            if let Some(hint) = self.optimizer_hint.as_ref() {
+                write!(f, " {hint}")?;
+            }
+            write!(f, " {on_conflict} INTO {table_name} ")?;
         } else {
             write!(
                 f,
@@ -111,8 +117,10 @@ impl Display for Insert {
                     "REPLACE"
                 } else {
                     "INSERT"
-                },
-            )?;
+                })?;
+            if let Some(hint) = self.optimizer_hint.as_ref() {
+                write!(f, " {hint}")?;
+            }
             if let Some(priority) = self.priority {
                 write!(f, " {priority}",)?;
             }

src/ast/spans.rs

@@ -1288,6 +1288,7 @@ impl Spanned for Insert {
     fn span(&self) -> Span {
         let Insert {
             insert_token,
+            optimizer_hint: _,
             or: _,     // enum, sqlite specific
             ignore: _, // bool
             into: _,   // bool

src/parser/mod.rs

@@ -16738,6 +16738,7 @@ impl<'a> Parser<'a> {
 
     /// Parse an INSERT statement
     pub fn parse_insert(&mut self, insert_token: TokenWithSpan) -> Result<Statement, ParserError> {
+        let optimizer_hint = self.parse_optional_optimizer_hint()?;
         let or = self.parse_conflict_clause();
         let priority = if !dialect_of!(self is MySqlDialect | GenericDialect) {
             None
@@ -16907,6 +16908,7 @@ impl<'a> Parser<'a> {
 
             Ok(Statement::Insert(Insert {
                 insert_token: insert_token.into(),
+                optimizer_hint,
                 or,
                 table: table_object,
                 table_alias,

tests/sqlparser_mysql.rs

@@ -4361,16 +4361,28 @@ fn test_create_index_options() {
 }
 
 #[test]
-fn test_select_optimizer_hints() {
-    mysql_and_generic().verified_stmt(
+fn test_optimizer_hints() {
+    let mysql_dialect = mysql_and_generic();
+
+    // ~ selects
+    mysql_dialect.verified_stmt(
         "\
        SELECT /*+ SET_VAR(optimizer_switch = 'mrr_cost_based=off') \
                   SET_VAR(max_heap_table_size = 1G) */ 1",
     );
 
-    mysql_and_generic().verified_stmt(
+    mysql_dialect.verified_stmt(
         "\
        SELECT /*+ SET_VAR(target_partitions=1) */ * FROM \
            (SELECT /*+ SET_VAR(target_partitions=8) */ * FROM t1 LIMIT 1) AS dt",
     );
+
+    // ~ inserts / replace
+    mysql_dialect.verified_stmt("\
+       INSERT /*+ RESOURCE_GROUP(Batch) */ \
+       INTO t2 VALUES (2)");
+
+    mysql_dialect.verified_stmt("\
+       REPLACE /*+ foobar */ INTO test \
+       VALUES (1, 'Old', '2014-08-20 18:47:00')");
 }

tests/sqlparser_oracle.rs

@@ -335,9 +335,10 @@ fn parse_national_quote_delimited_string_but_is_a_word() {
 }
 
 #[test]
-fn parse_optimizer_hints() {
+fn test_optimizer_hints() {
     let oracle_dialect = oracle();
 
+    // ~ selects
     let select = oracle_dialect.verified_only_select_with_canonical(
         "SELECT /*+one two three*/ /*+not a hint!*/ 1 FROM dual",
         "SELECT /*+one two three*/ 1 FROM dual",
@@ -367,4 +368,9 @@ fn parse_optimizer_hints() {
             .map(|hint| hint.text.as_str()),
         Some(" one two three /* asdf */\n")
     );
+
+    // ~ inserts
+    oracle_dialect.verified_stmt(
+        "INSERT /*+ append */ INTO t1 SELECT * FROM all_objects");
+
 }

tests/sqlparser_postgres.rs

@@ -5317,6 +5317,7 @@ fn test_simple_postgres_insert_with_alias() {
         statement,
         Statement::Insert(Insert {
             insert_token: AttachedToken::empty(),
+            optimizer_hint: None,
             or: None,
             ignore: false,
             into: true,
@@ -5388,6 +5389,7 @@ fn test_simple_postgres_insert_with_alias() {
         statement,
         Statement::Insert(Insert {
             insert_token: AttachedToken::empty(),
+            optimizer_hint: None,
             or: None,
             ignore: false,
             into: true,
@@ -5461,6 +5463,7 @@ fn test_simple_insert_with_quoted_alias() {
         statement,
         Statement::Insert(Insert {
             insert_token: AttachedToken::empty(),
+            optimizer_hint: None,
             or: None,
             ignore: false,
             into: true,