Add support of parsing ON CLUSTER in ALTER TABLE for ClickHouse (apache#1342)

Author: git-hulk

src/ast/dml.rs

@@ -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>>,
@@ -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))?;

src/ast/helpers/stmt_create_table.rs

@@ -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>>,
@@ -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
     }

src/ast/mod.rs

@@ -2183,6 +2183,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
@@ -3653,6 +3657,7 @@ impl fmt::Display for Statement {
                 only,
                 operations,
                 location,
+                on_cluster,
             } => {
                 write!(f, "ALTER TABLE ")?;
                 if *if_exists {
@@ -3661,9 +3666,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 {

src/parser/mod.rs

@@ -5393,6 +5393,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,
@@ -5405,16 +5413,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()
@@ -6597,6 +6596,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;
@@ -6618,6 +6618,7 @@ impl<'a> Parser<'a> {
                     only,
                     operations,
                     location,
+                    on_cluster,
                 })
             }
             Keyword::INDEX => {

src/test_utils.rs

@@ -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);

tests/sqlparser_common.rs

@@ -3507,7 +3507,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!(),
     }
@@ -3516,7 +3516,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!(),
     }
@@ -3823,6 +3823,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";

tests/sqlparser_mysql.rs

@@ -1976,6 +1976,7 @@ fn parse_alter_table_add_column() {
             only,
             operations,
             location: _,
+            on_cluster: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);
@@ -2005,6 +2006,7 @@ fn parse_alter_table_add_column() {
             only,
             operations,
             location: _,
+            on_cluster: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);
@@ -2042,6 +2044,7 @@ fn parse_alter_table_add_columns() {
             only,
             operations,
             location: _,
+            on_cluster: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);

tests/sqlparser_postgres.rs

@@ -677,6 +677,7 @@ fn parse_alter_table_add_columns() {
             only,
             operations,
             location: _,
+            on_cluster: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(if_exists);
@@ -759,6 +760,7 @@ fn parse_alter_table_owner_to() {
                 only: _,
                 operations,
                 location: _,
+                on_cluster: _,
             } => {
                 assert_eq!(name.to_string(), "tab");
                 assert_eq!(