Skip to content

Commit

Permalink
feat: support mysql partition selection
Browse files Browse the repository at this point in the history
  • Loading branch information
chunshao90 committed Aug 29, 2023
1 parent 9c2e8bc commit f2b7af9
Show file tree
Hide file tree
Showing 4 changed files with 55 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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))?;
}
Expand Down
2 changes: 2 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down
46 changes: 46 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;

Expand Down Expand Up @@ -6244,6 +6258,7 @@ impl<'a> Parser<'a> {
args,
with_hints,
version,
partitions,
})
}
}
Expand Down Expand Up @@ -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");
}
}
}
1 change: 1 addition & 0 deletions src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}
}

Expand Down

0 comments on commit f2b7af9

Please sign in to comment.