Skip to content

Commit

Permalink
Add support of parsing ON CLUSTER in ALTER TABLE for ClickHouse (#1342)
Browse files Browse the repository at this point in the history
  • Loading branch information
git-hulk authored Jul 30, 2024
1 parent f966580 commit cc13841
Show file tree
Hide file tree
Showing 8 changed files with 67 additions and 21 deletions.
8 changes: 2 additions & 6 deletions src/ast/dml.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ pub struct CreateTable {
pub on_commit: Option<OnCommit>,
/// ClickHouse "ON CLUSTER" clause:
/// <https://clickhouse.com/docs/en/sql-reference/distributed-ddl/>
pub on_cluster: Option<String>,
pub on_cluster: Option<Ident>,
/// ClickHouse "PRIMARY KEY " clause.
/// <https://clickhouse.com/docs/en/sql-reference/statements/create/table/>
pub primary_key: Option<Box<Expr>>,
Expand Down Expand Up @@ -206,11 +206,7 @@ impl Display for CreateTable {
name = self.name,
)?;
if let Some(on_cluster) = &self.on_cluster {
write!(
f,
" ON CLUSTER {}",
on_cluster.replace('{', "'{").replace('}', "}'")
)?;
write!(f, " ON CLUSTER {}", on_cluster)?;
}
if !self.columns.is_empty() || !self.constraints.is_empty() {
write!(f, " ({}", display_comma_separated(&self.columns))?;
Expand Down
4 changes: 2 additions & 2 deletions src/ast/helpers/stmt_create_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ pub struct CreateTableBuilder {
pub default_charset: Option<String>,
pub collation: Option<String>,
pub on_commit: Option<OnCommit>,
pub on_cluster: Option<String>,
pub on_cluster: Option<Ident>,
pub primary_key: Option<Box<Expr>>,
pub order_by: Option<OneOrManyWithParens<Expr>>,
pub partition_by: Option<Box<Expr>>,
Expand Down Expand Up @@ -261,7 +261,7 @@ impl CreateTableBuilder {
self
}

pub fn on_cluster(mut self, on_cluster: Option<String>) -> Self {
pub fn on_cluster(mut self, on_cluster: Option<Ident>) -> Self {
self.on_cluster = on_cluster;
self
}
Expand Down
11 changes: 10 additions & 1 deletion src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2162,6 +2162,10 @@ pub enum Statement {
only: bool,
operations: Vec<AlterTableOperation>,
location: Option<HiveSetLocation>,
/// ClickHouse dialect supports `ON CLUSTER` clause for ALTER TABLE
/// For example: `ALTER TABLE table_name ON CLUSTER cluster_name ADD COLUMN c UInt32`
/// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/alter/update)
on_cluster: Option<Ident>,
},
/// ```sql
/// ALTER INDEX
Expand Down Expand Up @@ -3632,6 +3636,7 @@ impl fmt::Display for Statement {
only,
operations,
location,
on_cluster,
} => {
write!(f, "ALTER TABLE ")?;
if *if_exists {
Expand All @@ -3640,9 +3645,13 @@ impl fmt::Display for Statement {
if *only {
write!(f, "ONLY ")?;
}
write!(f, "{name} ", name = name)?;
if let Some(cluster) = on_cluster {
write!(f, "ON CLUSTER {cluster} ")?;
}
write!(
f,
"{name} {operations}",
"{operations}",
operations = display_comma_separated(operations)
)?;
if let Some(loc) = location {
Expand Down
21 changes: 11 additions & 10 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5379,6 +5379,14 @@ impl<'a> Parser<'a> {
}
}

fn parse_optional_on_cluster(&mut self) -> Result<Option<Ident>, ParserError> {
if self.parse_keywords(&[Keyword::ON, Keyword::CLUSTER]) {
Ok(Some(self.parse_identifier(false)?))
} else {
Ok(None)
}
}

pub fn parse_create_table(
&mut self,
or_replace: bool,
Expand All @@ -5391,16 +5399,7 @@ impl<'a> Parser<'a> {
let table_name = self.parse_object_name(allow_unquoted_hyphen)?;

// Clickhouse has `ON CLUSTER 'cluster'` syntax for DDLs
let on_cluster = if self.parse_keywords(&[Keyword::ON, Keyword::CLUSTER]) {
let next_token = self.next_token();
match next_token.token {
Token::SingleQuotedString(s) => Some(s),
Token::Word(s) => Some(s.to_string()),
_ => self.expected("identifier or cluster literal", next_token)?,
}
} else {
None
};
let on_cluster = self.parse_optional_on_cluster()?;

let like = if self.parse_keyword(Keyword::LIKE) || self.parse_keyword(Keyword::ILIKE) {
self.parse_object_name(allow_unquoted_hyphen).ok()
Expand Down Expand Up @@ -6583,6 +6582,7 @@ impl<'a> Parser<'a> {
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
let only = self.parse_keyword(Keyword::ONLY); // [ ONLY ]
let table_name = self.parse_object_name(false)?;
let on_cluster = self.parse_optional_on_cluster()?;
let operations = self.parse_comma_separated(Parser::parse_alter_table_operation)?;

let mut location = None;
Expand All @@ -6604,6 +6604,7 @@ impl<'a> Parser<'a> {
only,
operations,
location,
on_cluster,
})
}
Keyword::INDEX => {
Expand Down
1 change: 1 addition & 0 deletions src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,7 @@ pub fn alter_table_op_with_name(stmt: Statement, expected_name: &str) -> AlterTa
if_exists,
only: is_only,
operations,
on_cluster: _,
location: _,
} => {
assert_eq!(name.to_string(), expected_name);
Expand Down
38 changes: 36 additions & 2 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3506,7 +3506,7 @@ fn parse_create_table_on_cluster() {
let sql = "CREATE TABLE t ON CLUSTER '{cluster}' (a INT, b INT)";
match generic.verified_stmt(sql) {
Statement::CreateTable(CreateTable { on_cluster, .. }) => {
assert_eq!(on_cluster.unwrap(), "{cluster}".to_string());
assert_eq!(on_cluster.unwrap().to_string(), "'{cluster}'".to_string());
}
_ => unreachable!(),
}
Expand All @@ -3515,7 +3515,7 @@ fn parse_create_table_on_cluster() {
let sql = "CREATE TABLE t ON CLUSTER my_cluster (a INT, b INT)";
match generic.verified_stmt(sql) {
Statement::CreateTable(CreateTable { on_cluster, .. }) => {
assert_eq!(on_cluster.unwrap(), "my_cluster".to_string());
assert_eq!(on_cluster.unwrap().to_string(), "my_cluster".to_string());
}
_ => unreachable!(),
}
Expand Down Expand Up @@ -3822,6 +3822,40 @@ fn parse_alter_table() {
}
}

#[test]
fn test_alter_table_with_on_cluster() {
match all_dialects()
.verified_stmt("ALTER TABLE t ON CLUSTER 'cluster' ADD CONSTRAINT bar PRIMARY KEY (baz)")
{
Statement::AlterTable {
name, on_cluster, ..
} => {
std::assert_eq!(name.to_string(), "t");
std::assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
}
_ => unreachable!(),
}

match all_dialects()
.verified_stmt("ALTER TABLE t ON CLUSTER cluster_name ADD CONSTRAINT bar PRIMARY KEY (baz)")
{
Statement::AlterTable {
name, on_cluster, ..
} => {
std::assert_eq!(name.to_string(), "t");
std::assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
}
_ => unreachable!(),
}

let res = all_dialects()
.parse_sql_statements("ALTER TABLE t ON CLUSTER 123 ADD CONSTRAINT bar PRIMARY KEY (baz)");
std::assert_eq!(
res.unwrap_err(),
ParserError::ParserError("Expected: identifier, found: 123".to_string())
)
}

#[test]
fn parse_alter_index() {
let rename_index = "ALTER INDEX idx RENAME TO new_idx";
Expand Down
3 changes: 3 additions & 0 deletions tests/sqlparser_mysql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1976,6 +1976,7 @@ fn parse_alter_table_add_column() {
only,
operations,
location: _,
on_cluster: _,
} => {
assert_eq!(name.to_string(), "tab");
assert!(!if_exists);
Expand Down Expand Up @@ -2005,6 +2006,7 @@ fn parse_alter_table_add_column() {
only,
operations,
location: _,
on_cluster: _,
} => {
assert_eq!(name.to_string(), "tab");
assert!(!if_exists);
Expand Down Expand Up @@ -2042,6 +2044,7 @@ fn parse_alter_table_add_columns() {
only,
operations,
location: _,
on_cluster: _,
} => {
assert_eq!(name.to_string(), "tab");
assert!(!if_exists);
Expand Down
2 changes: 2 additions & 0 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,7 @@ fn parse_alter_table_add_columns() {
only,
operations,
location: _,
on_cluster: _,
} => {
assert_eq!(name.to_string(), "tab");
assert!(if_exists);
Expand Down Expand Up @@ -759,6 +760,7 @@ fn parse_alter_table_owner_to() {
only: _,
operations,
location: _,
on_cluster: _,
} => {
assert_eq!(name.to_string(), "tab");
assert_eq!(
Expand Down

0 comments on commit cc13841

Please sign in to comment.