Support mysql partition to table selection (#959)

Co-authored-by: Andrew Lamb <[email protected]>

src/ast/query.rs

@@ -670,6 +670,8 @@ pub enum TableFactor {
         /// Optional version qualifier to facilitate table time-travel, as
         /// supported by BigQuery and MSSQL.
         version: Option<TableVersion>,
+        /// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
+        partitions: Vec<Ident>,
     },
     Derived {
         lateral: bool,
@@ -730,8 +732,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

@@ -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");
+        }
+    }
 }

src/test_utils.rs

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

tests/sqlparser_bigquery.rs

@@ -96,6 +96,7 @@ fn parse_table_identifiers() {
                     args: None,
                     with_hints: vec![],
                     version: None,
+                    partitions: vec![],
                 },
                 joins: vec![]
             },]
@@ -160,6 +161,7 @@ fn parse_table_time_travel() {
                 version: Some(TableVersion::ForSystemTimeAsOf(Expr::Value(
                     Value::SingleQuotedString(version)
                 ))),
+                partitions: vec![],
             },
             joins: vec![]
         },]

tests/sqlparser_clickhouse.rs

@@ -63,15 +63,16 @@ fn parse_map_access_expr() {
                     args: None,
                     with_hints: vec![],
                     version: None,
+                    partitions: vec![],
                 },
-                joins: vec![]
+                joins: vec![],
             }],
             lateral_views: vec![],
             selection: Some(BinaryOp {
                 left: Box::new(BinaryOp {
                     left: Box::new(Identifier(Ident::new("id"))),
                     op: BinaryOperator::Eq,
-                    right: Box::new(Expr::Value(Value::SingleQuotedString("test".to_string())))
+                    right: Box::new(Expr::Value(Value::SingleQuotedString("test".to_string()))),
                 }),
                 op: BinaryOperator::And,
                 right: Box::new(BinaryOp {
@@ -91,19 +92,19 @@ fn parse_map_access_expr() {
                             distinct: false,
                             special: false,
                             order_by: vec![],
-                        })]
+                        })],
                     }),
                     op: BinaryOperator::NotEq,
-                    right: Box::new(Expr::Value(Value::SingleQuotedString("foo".to_string())))
-                })
+                    right: Box::new(Expr::Value(Value::SingleQuotedString("foo".to_string()))),
+                }),
             }),
             group_by: GroupByExpr::Expressions(vec![]),
             cluster_by: vec![],
             distribute_by: vec![],
             sort_by: vec![],
             having: None,
             named_window: vec![],
-            qualify: None
+            qualify: None,
         },
         select
     );
@@ -117,7 +118,7 @@ fn parse_array_expr() {
         &Expr::Array(Array {
             elem: vec![
                 Expr::Value(Value::SingleQuotedString("1".to_string())),
-                Expr::Value(Value::SingleQuotedString("2".to_string()))
+                Expr::Value(Value::SingleQuotedString("2".to_string())),
             ],
             named: false,
         }),
@@ -171,6 +172,7 @@ fn parse_delimited_identifiers() {
             args,
             with_hints,
             version,
+            partitions: _,
         } => {
             assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
             assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);