ClickHouse: support of create table query with primary key and parametrised table engine

Author: aleksei.p

src/ast/helpers/stmt_create_table.rs

@@ -9,7 +9,7 @@ use sqlparser_derive::{Visit, VisitMut};
 
 use crate::ast::{
     ColumnDef, Expr, FileFormat, HiveDistributionStyle, HiveFormat, Ident, ObjectName, OnCommit,
-    Query, SqlOption, Statement, TableConstraint,
+    OneOrManyWithParens, Query, SqlOption, Statement, TableConstraint, TableEngine,
 };
 use crate::parser::ParserError;
 
@@ -64,14 +64,15 @@ pub struct CreateTableBuilder {
     pub without_rowid: bool,
     pub like: Option<ObjectName>,
     pub clone: Option<ObjectName>,
-    pub engine: Option<String>,
+    pub engine: Option<TableEngine>,
     pub comment: Option<String>,
     pub auto_increment_offset: Option<u32>,
     pub default_charset: Option<String>,
     pub collation: Option<String>,
     pub on_commit: Option<OnCommit>,
     pub on_cluster: Option<String>,
-    pub order_by: Option<Vec<Ident>>,
+    pub primary_key: Option<Box<Expr>>,
+    pub order_by: Option<OneOrManyWithParens<Expr>>,
     pub partition_by: Option<Box<Expr>>,
     pub cluster_by: Option<Vec<Ident>>,
     pub options: Option<Vec<SqlOption>>,
@@ -107,6 +108,7 @@ impl CreateTableBuilder {
             collation: None,
             on_commit: None,
             on_cluster: None,
+            primary_key: None,
             order_by: None,
             partition_by: None,
             cluster_by: None,
@@ -202,7 +204,7 @@ impl CreateTableBuilder {
         self
     }
 
-    pub fn engine(mut self, engine: Option<String>) -> Self {
+    pub fn engine(mut self, engine: Option<TableEngine>) -> Self {
         self.engine = engine;
         self
     }
@@ -237,7 +239,12 @@ impl CreateTableBuilder {
         self
     }
 
-    pub fn order_by(mut self, order_by: Option<Vec<Ident>>) -> Self {
+    pub fn primary_key(mut self, primary_key: Option<Box<Expr>>) -> Self {
+        self.primary_key = primary_key;
+        self
+    }
+
+    pub fn order_by(mut self, order_by: Option<OneOrManyWithParens<Expr>>) -> Self {
         self.order_by = order_by;
         self
     }
@@ -290,6 +297,7 @@ impl CreateTableBuilder {
             collation: self.collation,
             on_commit: self.on_commit,
             on_cluster: self.on_cluster,
+            primary_key: self.primary_key,
             order_by: self.order_by,
             partition_by: self.partition_by,
             cluster_by: self.cluster_by,
@@ -333,6 +341,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 collation,
                 on_commit,
                 on_cluster,
+                primary_key,
                 order_by,
                 partition_by,
                 cluster_by,
@@ -365,6 +374,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 collation,
                 on_commit,
                 on_cluster,
+                primary_key,
                 order_by,
                 partition_by,
                 cluster_by,

src/ast/mod.rs

@@ -1995,7 +1995,7 @@ pub enum Statement {
         without_rowid: bool,
         like: Option<ObjectName>,
         clone: Option<ObjectName>,
-        engine: Option<String>,
+        engine: Option<TableEngine>,
         comment: Option<String>,
         auto_increment_offset: Option<u32>,
         default_charset: Option<String>,
@@ -2004,10 +2004,13 @@ pub enum Statement {
         /// ClickHouse "ON CLUSTER" clause:
         /// <https://clickhouse.com/docs/en/sql-reference/distributed-ddl/>
         on_cluster: Option<String>,
+        /// ClickHouse "PRIMARY KEY " clause.
+        /// <https://clickhouse.com/docs/en/sql-reference/statements/create/table/>
+        primary_key: Option<Box<Expr>>,
         /// ClickHouse "ORDER BY " clause. Note that omitted ORDER BY is different
         /// than empty (represented as ()), the latter meaning "no sorting".
         /// <https://clickhouse.com/docs/en/sql-reference/statements/create/table/>
-        order_by: Option<Vec<Ident>>,
+        order_by: Option<OneOrManyWithParens<Expr>>,
         /// BigQuery: A partition expression for the table.
         /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#partition_expression>
         partition_by: Option<Box<Expr>>,
@@ -3392,6 +3395,7 @@ impl fmt::Display for Statement {
                 collation,
                 on_commit,
                 on_cluster,
+                primary_key,
                 order_by,
                 partition_by,
                 cluster_by,
@@ -3560,8 +3564,11 @@ impl fmt::Display for Statement {
                 if let Some(auto_increment_offset) = auto_increment_offset {
                     write!(f, " AUTO_INCREMENT {auto_increment_offset}")?;
                 }
+                if let Some(primary_key) = primary_key {
+                    write!(f, " PRIMARY KEY {primary_key}")?;
+                }
                 if let Some(order_by) = order_by {
-                    write!(f, " ORDER BY ({})", display_comma_separated(order_by))?;
+                    write!(f, " ORDER BY {order_by}")?;
                 }
                 if let Some(partition_by) = partition_by.as_ref() {
                     write!(f, " PARTITION BY {partition_by}")?;
@@ -6554,6 +6561,29 @@ impl Display for MySQLColumnPosition {
     }
 }
 
+/// Engine of DB. Some warehouse has parameters of engine, e.g. [clickhouse]
+///
+/// [clickhouse]: https://clickhouse.com/docs/en/engines/table-engines
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableEngine {
+    pub name: String,
+    pub parameters: Option<Vec<Ident>>,
+}
+
+impl Display for TableEngine {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.name)?;
+
+        if let Some(parameters) = self.parameters.as_ref() {
+            write!(f, "({})", display_comma_separated(parameters))?;
+        }
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/parser/mod.rs

@@ -5139,7 +5139,15 @@ impl<'a> Parser<'a> {
             self.expect_token(&Token::Eq)?;
             let next_token = self.next_token();
             match next_token.token {
-                Token::Word(w) => Some(w.value),
+                Token::Word(w) => {
+                    let name = w.value;
+                    let parameters = if self.peek_token() == Token::LParen {
+                        Some(self.parse_parenthesized_identifiers()?)
+                    } else {
+                        None
+                    };
+                    Some(TableEngine { name, parameters })
+                }
                 _ => self.expected("identifier", next_token)?,
             }
         } else {
@@ -5157,17 +5165,27 @@ impl<'a> Parser<'a> {
             None
         };
 
+        // ClickHouse supports `PRIMARY KEY`, before `ORDER BY`
+        // https://clickhouse.com/docs/en/sql-reference/statements/create/table#primary-key
+        let primary_key = if dialect_of!(self is ClickHouseDialect | GenericDialect)
+            && self.parse_keywords(&[Keyword::PRIMARY, Keyword::KEY])
+        {
+            Some(Box::new(self.parse_expr()?))
+        } else {
+            None
+        };
+
         let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
             if self.consume_token(&Token::LParen) {
                 let columns = if self.peek_token() != Token::RParen {
-                    self.parse_comma_separated(|p| p.parse_identifier(false))?
+                    self.parse_comma_separated(|p| p.parse_expr())?
                 } else {
                     vec![]
                 };
                 self.expect_token(&Token::RParen)?;
-                Some(columns)
+                Some(OneOrManyWithParens::Many(columns))
             } else {
-                Some(vec![self.parse_identifier(false)?])
+                Some(OneOrManyWithParens::One(self.parse_expr()?))
             }
         } else {
             None
@@ -5265,6 +5283,7 @@ impl<'a> Parser<'a> {
             .partition_by(big_query_config.partition_by)
             .cluster_by(big_query_config.cluster_by)
             .options(big_query_config.options)
+            .primary_key(primary_key)
             .strict(strict)
             .build())
     }
@@ -8818,7 +8837,7 @@ impl<'a> Parser<'a> {
             let partitions: Vec<Ident> = if dialect_of!(self is MySqlDialect | GenericDialect)
                 && self.parse_keyword(Keyword::PARTITION)
             {
-                self.parse_partitions()?
+                self.parse_parenthesized_identifiers()?
             } else {
                 vec![]
             };
@@ -10737,7 +10756,7 @@ impl<'a> Parser<'a> {
         })
     }
 
-    fn parse_partitions(&mut self) -> Result<Vec<Ident>, ParserError> {
+    fn parse_parenthesized_identifiers(&mut self) -> Result<Vec<Ident>, ParserError> {
         self.expect_token(&Token::LParen)?;
         let partitions = self.parse_comma_separated(|p| p.parse_identifier(false))?;
         self.expect_token(&Token::RParen)?;

tests/sqlparser_clickhouse.rs

@@ -211,15 +211,102 @@ fn parse_delimited_identifiers() {
 #[test]
 fn parse_create_table() {
     clickhouse().verified_stmt(r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY ("x")"#);
-    clickhouse().one_statement_parses_to(
-        r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY "x""#,
-        r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY ("x")"#,
-    );
+    clickhouse().verified_stmt(r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY "x""#);
     clickhouse().verified_stmt(
-        r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY ("x") AS SELECT * FROM "t" WHERE true"#,
+        r#"CREATE TABLE "x" ("a" "int") ENGINE=MergeTree ORDER BY "x" AS SELECT * FROM "t" WHERE true"#,
     );
 }
 
+#[test]
+fn parse_create_table_with_primary_key() {
+    match clickhouse_and_generic().one_statement_parses_to(
+        concat!(
+            r#"CREATE TABLE db.table (`i` Int, `k` Int)"#,
+            " ENGINE=SharedMergeTree('/clickhouse/tables/{uuid}/{shard}', '{replica}')",
+            " PRIMARY KEY tuple(i)",
+            " ORDER BY tuple(i)",
+        ),
+        concat!(
+            r#"CREATE TABLE db.table (`i` INT, `k` INT)"#,
+            " ENGINE=SharedMergeTree('/clickhouse/tables/{uuid}/{shard}', '{replica}')",
+            " PRIMARY KEY tuple(i)",
+            " ORDER BY tuple(i)",
+        ),
+    ) {
+        Statement::CreateTable {
+            name,
+            columns,
+            engine,
+            primary_key,
+            order_by,
+            ..
+        } => {
+            assert_eq!(name.to_string(), "db.table");
+            assert_eq!(
+                vec![
+                    ColumnDef {
+                        name: Ident::with_quote('`', "i"),
+                        data_type: DataType::Int(None),
+                        collation: None,
+                        options: vec![],
+                    },
+                    ColumnDef {
+                        name: Ident::with_quote('`', "k"),
+                        data_type: DataType::Int(None),
+                        collation: None,
+                        options: vec![],
+                    },
+                ],
+                columns
+            );
+            assert_eq!(
+                engine,
+                Some(TableEngine {
+                    name: "SharedMergeTree".to_string(),
+                    parameters: Some(vec![
+                        Ident::with_quote('\'', "/clickhouse/tables/{uuid}/{shard}"),
+                        Ident::with_quote('\'', "{replica}"),
+                    ]),
+                })
+            );
+            fn assert_function(actual: &Function, name: &str, arg: &str) -> bool {
+                assert_eq!(actual.name, ObjectName(vec![Ident::new(name)]));
+                assert_eq!(
+                    actual.args,
+                    FunctionArguments::List(FunctionArgumentList {
+                        args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(Identifier(
+                            Ident::new(arg)
+                        )),)],
+                        duplicate_treatment: None,
+                        clauses: vec![],
+                    })
+                );
+                true
+            }
+            match primary_key.unwrap().as_ref() {
+                Expr::Function(primary_key) => {
+                    assert!(assert_function(primary_key, "tuple", "i"));
+                }
+                _ => panic!("unexpected primary key type"),
+            }
+            match order_by {
+                Some(OneOrManyWithParens::One(Expr::Function(order_by))) => {
+                    assert!(assert_function(&order_by, "tuple", "i"));
+                }
+                _ => panic!("unexpected order by type"),
+            };
+        }
+        _ => unreachable!(),
+    }
+
+    clickhouse_and_generic()
+        .parse_sql_statements(concat!(
+            r#"CREATE TABLE db.table (`i` Int, `k` Int)"#,
+            " ORDER BY tuple(i), tuple(k)",
+        ))
+        .expect_err("ORDER BY supports one expression with tuple");
+}
+
 #[test]
 fn parse_create_view_with_fields_data_types() {
     match clickhouse().verified_stmt(r#"CREATE VIEW v (i "int", f "String") AS SELECT * FROM t"#) {

tests/sqlparser_mysql.rs

@@ -773,7 +773,13 @@ fn parse_create_table_engine_default_charset() {
                 },],
                 columns
             );
-            assert_eq!(engine, Some("InnoDB".to_string()));
+            assert_eq!(
+                engine,
+                Some(TableEngine {
+                    name: "InnoDB".to_string(),
+                    parameters: None
+                })
+            );
             assert_eq!(default_charset, Some("utf8mb3".to_string()));
         }
         _ => unreachable!(),