From 16af309c74c384755f2eb5d491e6328aedf85173 Mon Sep 17 00:00:00 2001
From: Andy Richardson <andy.john.richardson@gmail.com>
Date: Thu, 28 Jul 2022 19:44:05 +0100
Subject: [PATCH] Add support for ClickHouse DDL query syntax (`on cluster`)
 (#527)

* Add on_cluster to Statement::CreateTable

* Add printing of on_cluster

* Update src/ast/mod.rs

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

* Fix fmt + nostd

* Remove unintended diff

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
---
 src/ast/mod.rs            | 11 +++++++++++
 src/parser.rs             | 14 ++++++++++++++
 tests/sqlparser_common.rs | 21 +++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index aa8aa05f..c3dd6578 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -883,6 +883,9 @@ pub enum Statement {
         default_charset: Option<String>,
         collation: Option<String>,
         on_commit: Option<OnCommit>,
+        /// Click house "ON CLUSTER" clause:
+        /// <https://clickhouse.com/docs/en/sql-reference/distributed-ddl/>
+        on_cluster: Option<String>,
     },
     /// SQLite's `CREATE VIRTUAL TABLE .. USING <module_name> (<module_args>)`
     CreateVirtualTable {
@@ -1509,6 +1512,7 @@ impl fmt::Display for Statement {
                 engine,
                 collation,
                 on_commit,
+                on_cluster,
             } => {
                 // We want to allow the following options
                 // Empty column list, allowed by PostgreSQL:
@@ -1535,6 +1539,13 @@ impl fmt::Display for Statement {
                     temporary = if *temporary { "TEMPORARY " } else { "" },
                     name = name,
                 )?;
+                if let Some(on_cluster) = on_cluster {
+                    write!(
+                        f,
+                        " ON CLUSTER {}",
+                        on_cluster.replace('{', "'{").replace('}', "}'")
+                    )?;
+                }
                 if !columns.is_empty() || !constraints.is_empty() {
                     write!(f, " ({}", display_comma_separated(columns))?;
                     if !columns.is_empty() && !constraints.is_empty() {
diff --git a/src/parser.rs b/src/parser.rs
index 6dd95eab..d98fddc8 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -1826,6 +1826,7 @@ impl<'a> Parser<'a> {
             engine: None,
             collation: None,
             on_commit: None,
+            on_cluster: None,
         })
     }
 
@@ -2107,6 +2108,18 @@ impl<'a> Parser<'a> {
     ) -> Result<Statement, ParserError> {
         let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
         let table_name = self.parse_object_name()?;
+
+        // Clickhouse has `ON CLUSTER 'cluster'` syntax for DDLs
+        let on_cluster = if self.parse_keywords(&[Keyword::ON, Keyword::CLUSTER]) {
+            match self.next_token() {
+                Token::SingleQuotedString(s) => Some(s),
+                Token::Word(s) => Some(s.to_string()),
+                unexpected => self.expected("identifier or cluster literal", unexpected)?,
+            }
+        } else {
+            None
+        };
+
         let like = if self.parse_keyword(Keyword::LIKE) || self.parse_keyword(Keyword::ILIKE) {
             self.parse_object_name().ok()
         } else {
@@ -2207,6 +2220,7 @@ impl<'a> Parser<'a> {
             default_charset,
             collation,
             on_commit,
+            on_cluster,
         })
     }
 
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index eb3c0570..18641a94 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -1995,6 +1995,27 @@ fn parse_create_table_as() {
     }
 }
 
+#[test]
+fn parse_create_table_on_cluster() {
+    // Using single-quote literal to define current cluster
+    let sql = "CREATE TABLE t ON CLUSTER '{cluster}' (a INT, b INT)";
+    match verified_stmt(sql) {
+        Statement::CreateTable { on_cluster, .. } => {
+            assert_eq!(on_cluster.unwrap(), "{cluster}".to_string());
+        }
+        _ => unreachable!(),
+    }
+
+    // Using explicitly declared cluster name
+    let sql = "CREATE TABLE t ON CLUSTER my_cluster (a INT, b INT)";
+    match verified_stmt(sql) {
+        Statement::CreateTable { on_cluster, .. } => {
+            assert_eq!(on_cluster.unwrap(), "my_cluster".to_string());
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_create_or_replace_table() {
     let sql = "CREATE OR REPLACE TABLE t (a INT)";