Add support for mysql table hints (apache#1675)

Author: AvivDavid-Satori

src/ast/mod.rs

@@ -69,11 +69,11 @@ pub use self::query::{
     OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
     SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
-    TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableSample,
-    TableSampleBucket, TableSampleKind, TableSampleMethod, TableSampleModifier,
-    TableSampleQuantity, TableSampleSeed, TableSampleSeedModifier, TableSampleUnit, TableVersion,
-    TableWithJoins, Top, TopQuantity, UpdateTableFromKind, ValueTableMode, Values,
-    WildcardAdditionalOptions, With, WithFill,
+    TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableIndexHintForClause,
+    TableIndexHintType, TableIndexHints, TableIndexType, TableSample, TableSampleBucket,
+    TableSampleKind, TableSampleMethod, TableSampleModifier, TableSampleQuantity, TableSampleSeed,
+    TableSampleSeedModifier, TableSampleUnit, TableVersion, TableWithJoins, Top, TopQuantity,
+    UpdateTableFromKind, ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill,
 };
 
 pub use self::trigger::{

src/ast/query.rs

@@ -975,6 +975,81 @@ pub struct TableFunctionArgs {
     pub settings: Option<Vec<Setting>>,
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableIndexHintType {
+    Use,
+    Ignore,
+    Force,
+}
+
+impl fmt::Display for TableIndexHintType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            TableIndexHintType::Use => "USE",
+            TableIndexHintType::Ignore => "IGNORE",
+            TableIndexHintType::Force => "FORCE",
+        })
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableIndexType {
+    Index,
+    Key,
+}
+
+impl fmt::Display for TableIndexType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            TableIndexType::Index => "INDEX",
+            TableIndexType::Key => "KEY",
+        })
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableIndexHintForClause {
+    Join,
+    OrderBy,
+    GroupBy,
+}
+
+impl fmt::Display for TableIndexHintForClause {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            TableIndexHintForClause::Join => "JOIN",
+            TableIndexHintForClause::OrderBy => "ORDER BY",
+            TableIndexHintForClause::GroupBy => "GROUP BY",
+        })
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableIndexHints {
+    pub hint_type: TableIndexHintType,
+    pub index_type: TableIndexType,
+    pub for_clause: Option<TableIndexHintForClause>,
+    pub index_names: Vec<Ident>,
+}
+
+impl fmt::Display for TableIndexHints {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{} {} ", self.hint_type, self.index_type)?;
+        if let Some(for_clause) = &self.for_clause {
+            write!(f, "FOR {} ", for_clause)?;
+        }
+        write!(f, "({})", display_comma_separated(&self.index_names))
+    }
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1009,6 +1084,9 @@ pub enum TableFactor {
         /// Optional table sample modifier
         /// See: <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#sample-clause>
         sample: Option<TableSampleKind>,
+        /// Optional index hints(mysql)
+        /// See: <https://dev.mysql.com/doc/refman/8.4/en/index-hints.html>
+        index_hints: Vec<TableIndexHints>,
     },
     Derived {
         lateral: bool,
@@ -1590,6 +1668,7 @@ impl fmt::Display for TableFactor {
                 with_ordinality,
                 json_path,
                 sample,
+                index_hints,
             } => {
                 write!(f, "{name}")?;
                 if let Some(json_path) = json_path {
@@ -1618,6 +1697,9 @@ impl fmt::Display for TableFactor {
                 if let Some(alias) = alias {
                     write!(f, " AS {alias}")?;
                 }
+                if !index_hints.is_empty() {
+                    write!(f, " {}", display_separated(index_hints, " "))?;
+                }
                 if !with_hints.is_empty() {
                     write!(f, " WITH ({})", display_comma_separated(with_hints))?;
                 }

src/ast/spans.rs

@@ -1766,6 +1766,7 @@ impl Spanned for TableFactor {
                 partitions: _,
                 json_path: _,
                 sample: _,
+                index_hints: _,
             } => union_spans(
                 name.0
                     .iter()

src/dialect/mod.rs

@@ -854,6 +854,10 @@ pub trait Dialect: Debug + Any {
     fn supports_string_escape_constant(&self) -> bool {
         false
     }
+    /// Returns true if the dialect supports the table hints in the `FROM` clause.
+    fn supports_table_hints(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/dialect/mysql.rs

@@ -25,6 +25,10 @@ use crate::{
     parser::{Parser, ParserError},
 };
 
+use super::keywords;
+
+const RESERVED_FOR_TABLE_ALIAS_MYSQL: &[Keyword] = &[Keyword::USE, Keyword::IGNORE, Keyword::FORCE];
+
 /// A [`Dialect`] for [MySQL](https://www.mysql.com/)
 #[derive(Debug)]
 pub struct MySqlDialect {}
@@ -111,6 +115,16 @@ impl Dialect for MySqlDialect {
     fn supports_user_host_grantee(&self) -> bool {
         true
     }
+
+    fn is_table_factor_alias(&self, explicit: bool, kw: &Keyword, _parser: &mut Parser) -> bool {
+        explicit
+            || (!keywords::RESERVED_FOR_TABLE_ALIAS.contains(kw)
+                && !RESERVED_FOR_TABLE_ALIAS_MYSQL.contains(kw))
+    }
+
+    fn supports_table_hints(&self) -> bool {
+        true
+    }
 }
 
 /// `LOCK TABLES`

src/parser/mod.rs

@@ -8910,6 +8910,64 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn parse_table_index_hints(&mut self) -> Result<Vec<TableIndexHints>, ParserError> {
+        let mut hints = vec![];
+        while let Some(hint_type) =
+            self.parse_one_of_keywords(&[Keyword::USE, Keyword::IGNORE, Keyword::FORCE])
+        {
+            let hint_type = match hint_type {
+                Keyword::USE => TableIndexHintType::Use,
+                Keyword::IGNORE => TableIndexHintType::Ignore,
+                Keyword::FORCE => TableIndexHintType::Force,
+                _ => {
+                    return self.expected(
+                        "expected to match USE/IGNORE/FORCE keyword",
+                        self.peek_token(),
+                    )
+                }
+            };
+            let index_type = match self.parse_one_of_keywords(&[Keyword::INDEX, Keyword::KEY]) {
+                Some(Keyword::INDEX) => TableIndexType::Index,
+                Some(Keyword::KEY) => TableIndexType::Key,
+                _ => {
+                    return self.expected("expected to match INDEX/KEY keyword", self.peek_token())
+                }
+            };
+            let for_clause = if self.parse_keyword(Keyword::FOR) {
+                let clause = if self.parse_keyword(Keyword::JOIN) {
+                    TableIndexHintForClause::Join
+                } else if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
+                    TableIndexHintForClause::OrderBy
+                } else if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
+                    TableIndexHintForClause::GroupBy
+                } else {
+                    return self.expected(
+                        "expected to match FOR/ORDER BY/GROUP BY table hint in for clause",
+                        self.peek_token(),
+                    );
+                };
+                Some(clause)
+            } else {
+                None
+            };
+
+            self.expect_token(&Token::LParen)?;
+            let index_names = if self.peek_token().token != Token::RParen {
+                self.parse_comma_separated(Parser::parse_identifier)?
+            } else {
+                vec![]
+            };
+            self.expect_token(&Token::RParen)?;
+            hints.push(TableIndexHints {
+                hint_type,
+                index_type,
+                for_clause,
+                index_names,
+            });
+        }
+        Ok(hints)
+    }
+
     /// Wrapper for parse_optional_alias_inner, left for backwards-compatibility
     /// but new flows should use the context-specific methods such as `maybe_parse_select_item_alias`
     /// and `maybe_parse_table_alias`.
@@ -11257,6 +11315,14 @@ impl<'a> Parser<'a> {
 
             let alias = self.maybe_parse_table_alias()?;
 
+            // MYSQL-specific table hints:
+            let index_hints = if self.dialect.supports_table_hints() {
+                self.maybe_parse(|p| p.parse_table_index_hints())?
+                    .unwrap_or(vec![])
+            } else {
+                vec![]
+            };
+
             // MSSQL-specific table hints:
             let mut with_hints = vec![];
             if self.parse_keyword(Keyword::WITH) {
@@ -11285,6 +11351,7 @@ impl<'a> Parser<'a> {
                 with_ordinality,
                 json_path,
                 sample,
+                index_hints,
             };
 
             while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {

src/test_utils.rs

@@ -362,6 +362,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         with_ordinality: false,
         json_path: None,
         sample: None,
+        index_hints: vec![],
     }
 }
 
@@ -376,6 +377,7 @@ pub fn table_from_name(name: ObjectName) -> TableFactor {
         with_ordinality: false,
         json_path: None,
         sample: None,
+        index_hints: vec![],
     }
 }
 
@@ -393,6 +395,7 @@ pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> Ta
         with_ordinality: false,
         json_path: None,
         sample: None,
+        index_hints: vec![],
     }
 }
 

tests/sqlparser_bigquery.rs

@@ -1565,6 +1565,7 @@ fn parse_table_time_travel() {
                 with_ordinality: false,
                 json_path: None,
                 sample: None,
+                index_hints: vec![],
             },
             joins: vec![]
         },]
@@ -1665,6 +1666,7 @@ fn parse_merge() {
                     with_ordinality: false,
                     json_path: None,
                     sample: None,
+                    index_hints: vec![],
                 },
                 table
             );
@@ -1682,6 +1684,7 @@ fn parse_merge() {
                     with_ordinality: false,
                     json_path: None,
                     sample: None,
+                    index_hints: vec![],
                 },
                 source
             );