feat: support mysql partition selection

Author: chunshao90

src/ast/query.rs

@@ -665,6 +665,8 @@ pub enum TableFactor {
         /// Optional version qualifier to facilitate table time-travel, as
         /// supported by BigQuery and MSSQL.
         version: Option<TableVersion>,
+        /// Partition selection, supported by Mysql.
+        partitions: Vec<Expr>,
     },
     Derived {
         lateral: bool,
@@ -725,8 +727,12 @@ impl fmt::Display for TableFactor {
                 args,
                 with_hints,
                 version,
+                partitions,
             } => {
                 write!(f, "{name}")?;
+                if !partitions.is_empty() {
+                    write!(f, "PARTITION ({})", display_comma_separated(partitions))?;
+                }
                 if let Some(args) = args {
                     write!(f, "({})", display_comma_separated(args))?;
                 }

src/keywords.rs

@@ -716,6 +716,8 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::QUALIFY,
     Keyword::WINDOW,
     Keyword::END,
+    // for MYSQL PARTITION SELECTION
+    Keyword::PARTITION,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`

src/parser/mod.rs

@@ -6210,6 +6210,20 @@ impl<'a> Parser<'a> {
         } else {
             let name = self.parse_object_name()?;
 
+            let partitions = if dialect_of!(self is MySqlDialect)
+                && self.parse_keyword(Keyword::PARTITION)
+            {
+                let mut partitions = self.parse_comma_separated(|p| p.parse_tuple(true, false))?;
+                if partitions.len() != 1 {
+                    return Err(ParserError::ParserError(format!(
+                        "Partition expect one tuple"
+                    )));
+                }
+                partitions.remove(0)
+            } else {
+                vec![]
+            };
+
             // Parse potential version qualifier
             let version = self.parse_table_version()?;
 
@@ -6244,6 +6258,7 @@ impl<'a> Parser<'a> {
                 args,
                 with_hints,
                 version,
+                partitions,
             })
         }
     }
@@ -8034,4 +8049,35 @@ mod tests {
             "sql parser error: Unexpected token following period in identifier: *",
         );
     }
+
+    #[test]
+    fn test_mysql_partition_selection() {
+        let sql = "SELECT * FROM employees PARTITION (p0, p2)";
+        let expected = vec!["p0", "p2"];
+
+        let ast: Vec<Statement> = Parser::parse_sql(&MySqlDialect {}, sql).unwrap();
+        assert_eq!(ast.len(), 1);
+        if let Statement::Query(v) = &ast[0] {
+            if let SetExpr::Select(select) = &*v.body {
+                assert_eq!(select.from.len(), 1);
+                let from: &TableWithJoins = &select.from[0];
+                let table_factor = &from.relation;
+                if let TableFactor::Table { partitions, .. } = table_factor {
+                    let actual: Vec<&str> = partitions
+                        .iter()
+                        .map(|expr| {
+                            if let Expr::Identifier(ident) = &expr {
+                                ident.value.as_str()
+                            } else {
+                                ""
+                            }
+                        })
+                        .collect();
+                    assert_eq!(expected, actual);
+                }
+            }
+        } else {
+            panic!("fail to parse mysql partition selection");
+        }
+    }
 }

src/test_utils.rs

@@ -222,6 +222,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         args: None,
         with_hints: vec![],
         version: None,
+        partitions: vec![],
     }
 }