Merge pull request #13 from Embucket/issues/1036_cluster_by_expr

Add support for cluster by expressions

Author: osipovartem

src/ast/dml.rs

@@ -156,7 +156,9 @@ pub struct CreateTable {
     pub partition_by: Option<Box<Expr>>,
     /// BigQuery: Table clustering column list.
     /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#table_option_list>
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    /// Snowflake: Table clustering list which contains base column, expressions on base columns.
+    /// <https://docs.snowflake.com/en/user-guide/tables-clustering-keys#defining-a-clustering-key-for-a-table>
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     /// Hive: Table clustering column list.
     /// <https://cwiki.apache.org/confluence/display/Hive/LanguageManual+DDL#LanguageManualDDL-CreateTable>
     pub clustered_by: Option<ClusteredBy>,

src/ast/helpers/stmt_create_table.rs

@@ -94,7 +94,7 @@ pub struct CreateTableBuilder {
     pub primary_key: Option<Box<Expr>>,
     pub order_by: Option<OneOrManyWithParens<Expr>>,
     pub partition_by: Option<Box<Expr>>,
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     pub clustered_by: Option<ClusteredBy>,
     pub options: Option<Vec<SqlOption>>,
     pub strict: bool,
@@ -316,7 +316,7 @@ impl CreateTableBuilder {
         self
     }
 
-    pub fn cluster_by(mut self, cluster_by: Option<WrappedCollection<Vec<Ident>>>) -> Self {
+    pub fn cluster_by(mut self, cluster_by: Option<WrappedCollection<Vec<Expr>>>) -> Self {
         self.cluster_by = cluster_by;
         self
     }
@@ -589,7 +589,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
 #[derive(Default)]
 pub(crate) struct CreateTableConfiguration {
     pub partition_by: Option<Box<Expr>>,
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     pub options: Option<Vec<SqlOption>>,
 }
 

src/dialect/snowflake.rs

@@ -443,7 +443,7 @@ pub fn parse_create_table(
                     parser.expect_keyword_is(Keyword::BY)?;
                     parser.expect_token(&Token::LParen)?;
                     let cluster_by = Some(WrappedCollection::Parentheses(
-                        parser.parse_comma_separated(|p| p.parse_identifier())?,
+                        parser.parse_comma_separated(|p| p.parse_expr())?,
                     ));
                     parser.expect_token(&Token::RParen)?;
 

src/parser/mod.rs

@@ -6904,7 +6904,7 @@ impl<'a> Parser<'a> {
         if dialect_of!(self is BigQueryDialect | GenericDialect) {
             if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
                 cluster_by = Some(WrappedCollection::NoWrapping(
-                    self.parse_comma_separated(|p| p.parse_identifier())?,
+                    self.parse_comma_separated(|p| p.parse_expr())?,
                 ));
             };
 

tests/sqlparser_bigquery.rs

@@ -536,8 +536,8 @@ fn parse_create_table_with_options() {
                 (
                     Some(Box::new(Expr::Identifier(Ident::new("_PARTITIONDATE")))),
                     Some(WrappedCollection::NoWrapping(vec![
-                        Ident::new("userid"),
-                        Ident::new("age"),
+                        Expr::Identifier(Ident::new("userid")),
+                        Expr::Identifier(Ident::new("age")),
                     ])),
                     Some(vec![
                         SqlOption::KeyValue {

tests/sqlparser_mysql.rs

@@ -2158,11 +2158,13 @@ fn parse_alter_table_add_column() {
             if_exists,
             only,
             operations,
+            iceberg,
             location: _,
             on_cluster: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);
+            assert!(!iceberg);
             assert!(!only);
             assert_eq!(
                 operations,
@@ -2187,8 +2189,7 @@ fn parse_alter_table_add_column() {
             if_exists,
             only,
             operations,
-            location: _,
-            on_cluster: _,
+            ..
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);
@@ -2225,8 +2226,7 @@ fn parse_alter_table_add_columns() {
             if_exists,
             only,
             operations,
-            location: _,
-            on_cluster: _,
+            ..
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);

tests/sqlparser_postgres.rs

@@ -828,8 +828,7 @@ fn parse_alter_table_add_columns() {
             if_exists,
             only,
             operations,
-            location: _,
-            on_cluster: _,
+            ..
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(if_exists);
@@ -909,8 +908,7 @@ fn parse_alter_table_owner_to() {
                 if_exists: _,
                 only: _,
                 operations,
-                location: _,
-                on_cluster: _,
+                ..
             } => {
                 assert_eq!(name.to_string(), "tab");
                 assert_eq!(

tests/sqlparser_snowflake.rs

@@ -457,8 +457,38 @@ fn test_snowflake_create_table_cluster_by() {
             assert_eq!("my_table", name.to_string());
             assert_eq!(
                 Some(WrappedCollection::Parentheses(vec![
-                    Ident::new("a"),
-                    Ident::new("b"),
+                    Expr::Identifier(Ident::new("a")),
+                    Expr::Identifier(Ident::new("b")),
+                ])),
+                cluster_by
+            )
+        }
+        _ => unreachable!(),
+    }
+    match snowflake().verified_stmt("CREATE TABLE my_table (ts DATE, a TEXT) CLUSTER BY (to_date(ts), a)") {
+        Statement::CreateTable(CreateTable {
+                                   name, cluster_by, ..
+                               }) => {
+            assert_eq!("my_table", name.to_string());
+            assert_eq!(
+                Some(WrappedCollection::Parentheses(vec![
+                    Expr::Function(Function {
+                        name: ObjectName::from(vec![Ident::new("to_date")]),
+                        uses_odbc_syntax: false,
+                        parameters: FunctionArguments::None,
+                        args: FunctionArguments::List(FunctionArgumentList {
+                            args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                                Expr::Identifier(Ident::new("ts"))
+                            ))],
+                            duplicate_treatment: None,
+                            clauses: vec![],
+                        }),
+                        filter: None,
+                        null_treatment: None,
+                        over: None,
+                        within_group: vec![],
+                    }),
+                    Expr::Identifier(Ident::new("a")),
                 ])),
                 cluster_by
             )
@@ -869,8 +899,8 @@ fn test_snowflake_create_iceberg_table_all_options() {
             assert_eq!("my_table", name.to_string());
             assert_eq!(
                 Some(WrappedCollection::Parentheses(vec![
-                    Ident::new("a"),
-                    Ident::new("b"),
+                    Expr::Identifier(Ident::new("a")),
+                    Expr::Identifier(Ident::new("b")),
                 ])),
                 cluster_by
             );