MySQL: Add support for SELECT modifiers (apache#2172)

Author: mvzink

src/ast/mod.rs

@@ -97,14 +97,15 @@ pub use self::query::{
     OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions,
     PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
-    SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SetExpr, SetOperator,
-    SetQuantifier, Setting, SymbolDefinition, Table, 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, XmlNamespaceDefinition,
-    XmlPassingArgument, XmlPassingClause, XmlTableColumn, XmlTableColumnOption,
+    SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SelectModifiers,
+    SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, 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,
+    XmlNamespaceDefinition, XmlPassingArgument, XmlPassingClause, XmlTableColumn,
+    XmlTableColumnOption,
 };
 
 pub use self::trigger::{

src/ast/query.rs

@@ -334,6 +334,108 @@ pub enum SelectFlavor {
     FromFirstNoSelect,
 }
 
+/// MySQL-specific SELECT modifiers that appear after the SELECT keyword.
+///
+/// These modifiers affect query execution and optimization. They can appear in any order after
+/// SELECT and before the column list, can be repeated, and can be interleaved with
+/// DISTINCT/DISTINCTROW/ALL:
+///
+/// ```sql
+/// SELECT
+///     [ALL | DISTINCT | DISTINCTROW]
+///     [HIGH_PRIORITY]
+///     [STRAIGHT_JOIN]
+///     [SQL_SMALL_RESULT] [SQL_BIG_RESULT] [SQL_BUFFER_RESULT]
+///     [SQL_NO_CACHE] [SQL_CALC_FOUND_ROWS]
+///     select_expr [, select_expr] ...
+/// ```
+///
+/// See [MySQL SELECT](https://dev.mysql.com/doc/refman/8.4/en/select.html).
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Default)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SelectModifiers {
+    /// `HIGH_PRIORITY` gives the SELECT higher priority than statements that update a table.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub high_priority: bool,
+    /// `STRAIGHT_JOIN` forces the optimizer to join tables in the order listed in the FROM clause.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub straight_join: bool,
+    /// `SQL_SMALL_RESULT` hints that the result set is small, using in-memory temp tables.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub sql_small_result: bool,
+    /// `SQL_BIG_RESULT` hints that the result set is large, using disk-based temp tables.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub sql_big_result: bool,
+    /// `SQL_BUFFER_RESULT` forces the result to be put into a temporary table to release locks early.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub sql_buffer_result: bool,
+    /// `SQL_NO_CACHE` tells MySQL not to cache the query result. (Deprecated in 8.4+.)
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    pub sql_no_cache: bool,
+    /// `SQL_CALC_FOUND_ROWS` tells MySQL to calculate the total number of rows. (Deprecated in 8.0.17+.)
+    ///
+    /// - [MySQL SELECT modifiers](https://dev.mysql.com/doc/refman/8.4/en/select.html)
+    /// - [`FOUND_ROWS()`](https://dev.mysql.com/doc/refman/8.4/en/information-functions.html#function_found-rows)
+    pub sql_calc_found_rows: bool,
+}
+
+impl fmt::Display for SelectModifiers {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.high_priority {
+            f.write_str(" HIGH_PRIORITY")?;
+        }
+        if self.straight_join {
+            f.write_str(" STRAIGHT_JOIN")?;
+        }
+        if self.sql_small_result {
+            f.write_str(" SQL_SMALL_RESULT")?;
+        }
+        if self.sql_big_result {
+            f.write_str(" SQL_BIG_RESULT")?;
+        }
+        if self.sql_buffer_result {
+            f.write_str(" SQL_BUFFER_RESULT")?;
+        }
+        if self.sql_no_cache {
+            f.write_str(" SQL_NO_CACHE")?;
+        }
+        if self.sql_calc_found_rows {
+            f.write_str(" SQL_CALC_FOUND_ROWS")?;
+        }
+        Ok(())
+    }
+}
+
+impl SelectModifiers {
+    /// Returns true if any of the modifiers are set.
+    pub fn is_any_set(&self) -> bool {
+        // Using irrefutable destructuring to catch fields added in the future
+        let Self {
+            high_priority,
+            straight_join,
+            sql_small_result,
+            sql_big_result,
+            sql_buffer_result,
+            sql_no_cache,
+            sql_calc_found_rows,
+        } = self;
+        *high_priority
+            || *straight_join
+            || *sql_small_result
+            || *sql_big_result
+            || *sql_buffer_result
+            || *sql_no_cache
+            || *sql_calc_found_rows
+    }
+}
+
 /// A restricted variant of `SELECT` (without CTEs/`ORDER BY`), which may
 /// appear either as the only body item of a `Query`, or as an operand
 /// to a set operation like `UNION`.
@@ -350,6 +452,10 @@ pub struct Select {
     pub optimizer_hint: Option<OptimizerHint>,
     /// `SELECT [DISTINCT] ...`
     pub distinct: Option<Distinct>,
+    /// MySQL-specific SELECT modifiers.
+    ///
+    /// See [MySQL SELECT](https://dev.mysql.com/doc/refman/8.4/en/select.html).
+    pub select_modifiers: Option<SelectModifiers>,
     /// MSSQL syntax: `TOP (<N>) [ PERCENT ] [ WITH TIES ]`
     pub top: Option<Top>,
     /// Whether the top was located before `ALL`/`DISTINCT`
@@ -442,6 +548,10 @@ impl fmt::Display for Select {
             }
         }
 
+        if let Some(ref select_modifiers) = self.select_modifiers {
+            select_modifiers.fmt(f)?;
+        }
+
         if !self.projection.is_empty() {
             indented_list(f, &self.projection)?;
         }
@@ -3351,8 +3461,14 @@ impl fmt::Display for NonBlock {
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-/// `DISTINCT` or `DISTINCT ON (...)` modifiers for `SELECT` lists.
+/// `ALL`, `DISTINCT`, or `DISTINCT ON (...)` modifiers for `SELECT` lists.
 pub enum Distinct {
+    /// `ALL` (keep duplicate rows)
+    ///
+    /// Generally this is the default if omitted, but omission should be represented as
+    /// `None::<Option<Distinct>>`
+    All,
+
     /// `DISTINCT` (remove duplicate rows)
     Distinct,
 
@@ -3363,6 +3479,7 @@ pub enum Distinct {
 impl fmt::Display for Distinct {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
+            Distinct::All => write!(f, "ALL"),
             Distinct::Distinct => write!(f, "DISTINCT"),
             Distinct::On(col_names) => {
                 let col_names = display_comma_separated(col_names);

src/ast/spans.rs

@@ -2238,7 +2238,8 @@ impl Spanned for Select {
             select_token,
             optimizer_hint: _,
             distinct: _, // todo
-            top: _,      // todo, mysql specific
+            select_modifiers: _,
+            top: _, // todo, mysql specific
             projection,
             exclude: _,
             into,

src/dialect/mod.rs

@@ -698,6 +698,19 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports MySQL-specific SELECT modifiers
+    /// like `HIGH_PRIORITY`, `STRAIGHT_JOIN`, `SQL_SMALL_RESULT`, etc.
+    ///
+    /// For example:
+    /// ```sql
+    /// SELECT HIGH_PRIORITY STRAIGHT_JOIN SQL_SMALL_RESULT * FROM t1 JOIN t2 ON ...
+    /// ```
+    ///
+    /// [MySQL](https://dev.mysql.com/doc/refman/8.4/en/select.html)
+    fn supports_select_modifiers(&self) -> bool {
+        false
+    }
+
     /// Dialect-specific infix parser override
     ///
     /// This method is called to parse the next infix expression.

src/dialect/mysql.rs

@@ -156,6 +156,10 @@ impl Dialect for MySqlDialect {
         true
     }
 
+    fn supports_select_modifiers(&self) -> bool {
+        true
+    }
+
     fn supports_set_names(&self) -> bool {
         true
     }

src/keywords.rs

@@ -333,6 +333,7 @@ define_keywords!(
     DISCARD,
     DISCONNECT,
     DISTINCT,
+    DISTINCTROW,
     DISTRIBUTE,
     DIV,
     DO,
@@ -956,6 +957,11 @@ define_keywords!(
     SQLEXCEPTION,
     SQLSTATE,
     SQLWARNING,
+    SQL_BIG_RESULT,
+    SQL_BUFFER_RESULT,
+    SQL_CALC_FOUND_ROWS,
+    SQL_NO_CACHE,
+    SQL_SMALL_RESULT,
     SQRT,
     SRID,
     STABLE,

src/parser/mod.rs

@@ -4927,16 +4927,27 @@ impl<'a> Parser<'a> {
     /// and results in a [`ParserError`] if both `ALL` and `DISTINCT` are found.
     pub fn parse_all_or_distinct(&mut self) -> Result<Option<Distinct>, ParserError> {
         let loc = self.peek_token().span.start;
-        let all = self.parse_keyword(Keyword::ALL);
-        let distinct = self.parse_keyword(Keyword::DISTINCT);
-        if !distinct {
-            return Ok(None);
-        }
-        if all {
-            return parser_err!("Cannot specify both ALL and DISTINCT".to_string(), loc);
-        }
-        let on = self.parse_keyword(Keyword::ON);
-        if !on {
+        let distinct = match self.parse_one_of_keywords(&[Keyword::ALL, Keyword::DISTINCT]) {
+            Some(Keyword::ALL) => {
+                if self.peek_keyword(Keyword::DISTINCT) {
+                    return parser_err!("Cannot specify ALL then DISTINCT".to_string(), loc);
+                }
+                Some(Distinct::All)
+            }
+            Some(Keyword::DISTINCT) => {
+                if self.peek_keyword(Keyword::ALL) {
+                    return parser_err!("Cannot specify DISTINCT then ALL".to_string(), loc);
+                }
+                Some(Distinct::Distinct)
+            }
+            None => return Ok(None),
+            _ => return parser_err!("ALL or DISTINCT", loc),
+        };
+
+        let Some(Distinct::Distinct) = distinct else {
+            return Ok(distinct);
+        };
+        if !self.parse_keyword(Keyword::ON) {
             return Ok(Some(Distinct::Distinct));
         }
 
@@ -13861,6 +13872,7 @@ impl<'a> Parser<'a> {
                     select_token: AttachedToken(from_token),
                     optimizer_hint: None,
                     distinct: None,
+                    select_modifiers: None,
                     top: None,
                     top_before_distinct: false,
                     projection: vec![],
@@ -13890,13 +13902,26 @@ impl<'a> Parser<'a> {
         let optimizer_hint = self.maybe_parse_optimizer_hint()?;
         let value_table_mode = self.parse_value_table_mode()?;
 
+        let (select_modifiers, distinct_select_modifier) =
+            if self.dialect.supports_select_modifiers() {
+                self.parse_select_modifiers()?
+            } else {
+                (None, None)
+            };
+
         let mut top_before_distinct = false;
         let mut top = None;
         if self.dialect.supports_top_before_distinct() && self.parse_keyword(Keyword::TOP) {
             top = Some(self.parse_top()?);
             top_before_distinct = true;
         }
-        let distinct = self.parse_all_or_distinct()?;
+
+        let distinct = if distinct_select_modifier.is_some() {
+            distinct_select_modifier
+        } else {
+            self.parse_all_or_distinct()?
+        };
+
         if !self.dialect.supports_top_before_distinct() && self.parse_keyword(Keyword::TOP) {
             top = Some(self.parse_top()?);
         }
@@ -14044,6 +14069,7 @@ impl<'a> Parser<'a> {
             select_token: AttachedToken(select_token),
             optimizer_hint,
             distinct,
+            select_modifiers,
             top,
             top_before_distinct,
             projection,
@@ -14120,6 +14146,68 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parses MySQL SELECT modifiers and DISTINCT/ALL in any order.
+    ///
+    /// Manual testing shows odifiers can appear in any order, and modifiers other than DISTINCT/ALL
+    /// can be repeated.
+    ///
+    /// <https://dev.mysql.com/doc/refman/8.4/en/select.html>
+    fn parse_select_modifiers(
+        &mut self,
+    ) -> Result<(Option<SelectModifiers>, Option<Distinct>), ParserError> {
+        let mut modifiers = SelectModifiers::default();
+        let mut distinct = None;
+
+        let keywords = &[
+            Keyword::ALL,
+            Keyword::DISTINCT,
+            Keyword::DISTINCTROW,
+            Keyword::HIGH_PRIORITY,
+            Keyword::STRAIGHT_JOIN,
+            Keyword::SQL_SMALL_RESULT,
+            Keyword::SQL_BIG_RESULT,
+            Keyword::SQL_BUFFER_RESULT,
+            Keyword::SQL_NO_CACHE,
+            Keyword::SQL_CALC_FOUND_ROWS,
+        ];
+
+        while let Some(keyword) = self.parse_one_of_keywords(keywords) {
+            match keyword {
+                Keyword::ALL | Keyword::DISTINCT if distinct.is_none() => {
+                    self.prev_token();
+                    distinct = self.parse_all_or_distinct()?;
+                }
+                // DISTINCTROW is a MySQL-specific legacy (but not deprecated) alias for DISTINCT
+                Keyword::DISTINCTROW if distinct.is_none() => {
+                    distinct = Some(Distinct::Distinct);
+                }
+                Keyword::HIGH_PRIORITY => modifiers.high_priority = true,
+                Keyword::STRAIGHT_JOIN => modifiers.straight_join = true,
+                Keyword::SQL_SMALL_RESULT => modifiers.sql_small_result = true,
+                Keyword::SQL_BIG_RESULT => modifiers.sql_big_result = true,
+                Keyword::SQL_BUFFER_RESULT => modifiers.sql_buffer_result = true,
+                Keyword::SQL_NO_CACHE => modifiers.sql_no_cache = true,
+                Keyword::SQL_CALC_FOUND_ROWS => modifiers.sql_calc_found_rows = true,
+                _ => {
+                    self.prev_token();
+                    return self.expected(
+                        "HIGH_PRIORITY, STRAIGHT_JOIN, or other MySQL select modifier",
+                        self.peek_token(),
+                    );
+                }
+            }
+        }
+
+        // Avoid polluting the AST with `Some(SelectModifiers::default())` empty value unless there
+        // actually were some modifiers set.
+        let select_modifiers = if modifiers.is_any_set() {
+            Some(modifiers)
+        } else {
+            None
+        };
+        Ok((select_modifiers, distinct))
+    }
+
     fn parse_value_table_mode(&mut self) -> Result<Option<ValueTableMode>, ParserError> {
         if !dialect_of!(self is BigQueryDialect) {
             return Ok(None);