mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-22 15:04:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Support mysql partition to table selection (#959)

Browse source 
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
This commit is contained in:
chunshao.rcs 2023-09-15 02:21:47 +08:00 • committed by GitHub
parent a16791d019
commit f6e4be4c15
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 108 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/ast/query.rs
View file

@ -670,6 +670,8 @@ pub enum TableFactor {
/// Optional version qualifier to facilitate table time-travel, as /// Optional version qualifier to facilitate table time-travel, as
/// supported by BigQuery and MSSQL. /// supported by BigQuery and MSSQL.
version: Option<TableVersion>, version: Option<TableVersion>,
/// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
partitions: Vec<Ident>,
}, },
Derived { Derived {
lateral: bool, lateral: bool,
@ -730,8 +732,12 @@ impl fmt::Display for TableFactor {
args, args,
with_hints, with_hints,
version, version,
partitions,
} => { } => {
write!(f, "{name}")?; write!(f, "{name}")?;
if !partitions.is_empty() {
write!(f, "PARTITION ({})", display_comma_separated(partitions))?;
}
if let Some(args) = args { if let Some(args) = args {
write!(f, "({})", display_comma_separated(args))?; write!(f, "({})", display_comma_separated(args))?;
} }

2
src/keywords.rs
View file

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

41
src/parser/mod.rs
View file

@ -6277,6 +6277,14 @@ impl<'a> Parser<'a> {
} else { } else {
let name = self.parse_object_name()?; 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 // Parse potential version qualifier
let version = self.parse_table_version()?; let version = self.parse_table_version()?;
@ -6311,6 +6319,7 @@ impl<'a> Parser<'a> {
args, args,
with_hints, with_hints,
version, version,
partitions,
}) })
} }
} }
@ -7483,6 +7492,13 @@ impl<'a> Parser<'a> {
representation: UserDefinedTypeRepresentation::Composite { attributes }, 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 { impl Word {
@ -8100,4 +8116,29 @@ mod tests {
"sql parser error: Unexpected token following period in identifier: *", "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");
}
}
} }

1
src/test_utils.rs
View file

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

2
tests/sqlparser_bigquery.rs
View file

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

16
tests/sqlparser_clickhouse.rs
View file

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

29
tests/sqlparser_common.rs
View file

@ -215,6 +215,7 @@ fn parse_update_set_from() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}, },
@ -242,6 +243,7 @@ fn parse_update_set_from() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -308,6 +310,7 @@ fn parse_update_with_table_alias() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}, },
@ -371,6 +374,7 @@ fn parse_select_with_table_alias() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}] }]
@ -402,6 +406,7 @@ fn parse_delete_statement() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].relation from[0].relation
); );
@ -430,6 +435,7 @@ fn parse_delete_statement_for_multi_tables() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].relation from[0].relation
); );
@ -440,6 +446,7 @@ fn parse_delete_statement_for_multi_tables() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].joins[0].relation from[0].joins[0].relation
); );
@ -464,6 +471,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].relation from[0].relation
); );
@ -474,6 +482,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[1].relation from[1].relation
); );
@ -484,6 +493,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
using[0].relation using[0].relation
); );
@ -494,6 +504,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
using[0].joins[0].relation using[0].joins[0].relation
); );
@ -522,6 +533,7 @@ fn parse_where_delete_statement() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].relation, from[0].relation,
); );
@ -564,6 +576,7 @@ fn parse_where_delete_with_alias_statement() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
from[0].relation, from[0].relation,
); );
@ -578,6 +591,7 @@ fn parse_where_delete_with_alias_statement() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}]), }]),
@ -3552,6 +3566,7 @@ fn test_parse_named_window() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -3891,6 +3906,7 @@ fn parse_interval_and_or_xor() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -4496,6 +4512,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}, },
@ -4506,6 +4523,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}, },
@ -4524,6 +4542,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![Join { joins: vec![Join {
relation: TableFactor::Table { relation: TableFactor::Table {
@ -4532,6 +4551,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: JoinOperator::Inner(JoinConstraint::Natural), join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}], }],
@ -4543,6 +4563,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![Join { joins: vec![Join {
relation: TableFactor::Table { relation: TableFactor::Table {
@ -4551,6 +4572,7 @@ fn parse_implicit_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: JoinOperator::Inner(JoinConstraint::Natural), join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}], }],
@ -4572,6 +4594,7 @@ fn parse_cross_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: JoinOperator::CrossJoin, join_operator: JoinOperator::CrossJoin,
}, },
@ -4593,6 +4616,7 @@ fn parse_joins_on() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: f(JoinConstraint::On(Expr::BinaryOp { join_operator: f(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::Identifier("c1".into())), left: Box::new(Expr::Identifier("c1".into())),
@ -4663,6 +4687,7 @@ fn parse_joins_using() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: f(JoinConstraint::Using(vec!["c1".into()])), join_operator: f(JoinConstraint::Using(vec!["c1".into()])),
} }
@ -4725,6 +4750,7 @@ fn parse_natural_join() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: f(JoinConstraint::Natural), join_operator: f(JoinConstraint::Natural),
} }
@ -4990,6 +5016,7 @@ fn parse_derived_tables() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: JoinOperator::Inner(JoinConstraint::Natural), join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}], }],
@ -6318,6 +6345,7 @@ fn parse_merge() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
} }
); );
assert_eq!(table, table_no_into); assert_eq!(table, table_no_into);
@ -6342,6 +6370,7 @@ fn parse_merge() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],

4
tests/sqlparser_duckdb.rs
View file

@ -156,6 +156,7 @@ fn test_select_union_by_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -189,6 +190,7 @@ fn test_select_union_by_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -231,6 +233,7 @@ fn test_select_union_by_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],
@ -264,6 +267,7 @@ fn test_select_union_by_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
}], }],

1
tests/sqlparser_hive.rs
View file

@ -323,6 +323,7 @@ fn parse_delimited_identifiers() {
args, args,
with_hints, with_hints,
version, version,
partitions: _,
} => { } => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

3
tests/sqlparser_mssql.rs
View file

@ -58,6 +58,7 @@ fn parse_table_time_travel() {
version: Some(TableVersion::ForSystemTimeAsOf(Expr::Value( version: Some(TableVersion::ForSystemTimeAsOf(Expr::Value(
Value::SingleQuotedString(version) Value::SingleQuotedString(version)
))), ))),
partitions: vec![],
}, },
joins: vec![] joins: vec![]
},] },]
@ -309,6 +310,7 @@ fn parse_delimited_identifiers() {
args, args,
with_hints, with_hints,
version, version,
partitions: _,
} => { } => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);
@ -514,6 +516,7 @@ fn parse_substring_in_select() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![] joins: vec![]
}], }],

5
tests/sqlparser_mysql.rs
View file

@ -1151,6 +1151,7 @@ fn parse_select_with_numeric_prefix_column_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![] joins: vec![]
}], }],
@ -1200,6 +1201,7 @@ fn parse_select_with_concatenation_of_exp_number_and_numeric_prefix_column() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![] joins: vec![]
}], }],
@ -1260,6 +1262,7 @@ fn parse_update_with_joins() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![Join { joins: vec![Join {
relation: TableFactor::Table { relation: TableFactor::Table {
@ -1271,6 +1274,7 @@ fn parse_update_with_joins() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp { join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::CompoundIdentifier(vec![ left: Box::new(Expr::CompoundIdentifier(vec![
@ -1376,6 +1380,7 @@ fn parse_substring_in_select() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![] joins: vec![]
}], }],

1
tests/sqlparser_postgres.rs
View file

@ -2893,6 +2893,7 @@ fn parse_delimited_identifiers() {
args, args,
with_hints, with_hints,
version, version,
partitions: _,
} => { } => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

3
tests/sqlparser_redshift.rs
View file

@ -46,6 +46,7 @@ fn test_square_brackets_over_db_schema_table_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
} }
@ -91,6 +92,7 @@ fn test_double_quotes_over_db_schema_table_name() {
args: None, args: None,
with_hints: vec![], with_hints: vec![],
version: None, version: None,
partitions: vec![],
}, },
joins: vec![], joins: vec![],
} }
@ -111,6 +113,7 @@ fn parse_delimited_identifiers() {
args, args,
with_hints, with_hints,
version, version,
partitions: _,
} => { } => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

1
tests/sqlparser_snowflake.rs
View file

@ -224,6 +224,7 @@ fn parse_delimited_identifiers() {
args, args,
with_hints, with_hints,
version, version,
partitions: _,
} => { } => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);