Support mysql partition to table selection (#959)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

src/parser/mod.rs

@@ -6277,6 +6277,14 @@ impl<'a> Parser<'a> {
         } else {
             let name = self.parse_object_name()?;
 
+            let partitions: Vec<Ident> = if dialect_of!(self is MySqlDialect | GenericDialect)
+                && self.parse_keyword(Keyword::PARTITION)
+            {
+                self.parse_partitions()?
+            } else {
+                vec![]
+            };
+
             // Parse potential version qualifier
             let version = self.parse_table_version()?;
 
@@ -6311,6 +6319,7 @@ impl<'a> Parser<'a> {
                 args,
                 with_hints,
                 version,
+                partitions,
             })
         }
     }
@@ -7483,6 +7492,13 @@ impl<'a> Parser<'a> {
             representation: UserDefinedTypeRepresentation::Composite { attributes },
         })
     }
+
+    fn parse_partitions(&mut self) -> Result<Vec<Ident>, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let partitions = self.parse_comma_separated(Parser::parse_identifier)?;
+        self.expect_token(&Token::RParen)?;
+        Ok(partitions)
+    }
 }
 
 impl Word {
@@ -8100,4 +8116,29 @@ 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(|ident| ident.value.as_str())
+                        .collect();
+                    assert_eq!(expected, actual);
+                }
+            }
+        } else {
+            panic!("fail to parse mysql partition selection");
+        }
+    }
 }