From f0ece388a3ab7f4de1531c6b16071d6c986b5351 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 12 Jun 2025 12:46:58 +0300 Subject: [PATCH] Add support for cluster by expressions --- src/ast/dml.rs | 4 ++- src/ast/helpers/stmt_create_table.rs | 6 ++--- src/dialect/snowflake.rs | 2 +- src/parser/mod.rs | 2 +- tests/sqlparser_bigquery.rs | 4 +-- tests/sqlparser_mysql.rs | 8 +++--- tests/sqlparser_postgres.rs | 6 ++--- tests/sqlparser_snowflake.rs | 38 +++++++++++++++++++++++++--- 8 files changed, 50 insertions(+), 20 deletions(-) diff --git a/src/ast/dml.rs b/src/ast/dml.rs index 8cfc67414..3d63371e6 100644 --- a/src/ast/dml.rs +++ b/src/ast/dml.rs @@ -156,7 +156,9 @@ pub struct CreateTable { pub partition_by: Option>, /// BigQuery: Table clustering column list. /// - pub cluster_by: Option>>, + /// Snowflake: Table clustering list which contains base column, expressions on base columns. + /// + pub cluster_by: Option>>, /// Hive: Table clustering column list. /// pub clustered_by: Option, diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs index 344e9dec9..311012633 100644 --- a/src/ast/helpers/stmt_create_table.rs +++ b/src/ast/helpers/stmt_create_table.rs @@ -94,7 +94,7 @@ pub struct CreateTableBuilder { pub primary_key: Option>, pub order_by: Option>, pub partition_by: Option>, - pub cluster_by: Option>>, + pub cluster_by: Option>>, pub clustered_by: Option, pub options: Option>, pub strict: bool, @@ -316,7 +316,7 @@ impl CreateTableBuilder { self } - pub fn cluster_by(mut self, cluster_by: Option>>) -> Self { + pub fn cluster_by(mut self, cluster_by: Option>>) -> Self { self.cluster_by = cluster_by; self } @@ -589,7 +589,7 @@ impl TryFrom for CreateTableBuilder { #[derive(Default)] pub(crate) struct CreateTableConfiguration { pub partition_by: Option>, - pub cluster_by: Option>>, + pub cluster_by: Option>>, pub options: Option>, } diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs index 72252b277..2954dce0a 100644 --- a/src/dialect/snowflake.rs +++ b/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)?; diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 0c9544f41..d400a8dbd 100644 --- a/src/parser/mod.rs +++ b/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())?, )); }; diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs index 3037d4ae5..0bfa6613d 100644 --- a/tests/sqlparser_bigquery.rs +++ b/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 { diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs index 15f79b4c2..7e99eb017 100644 --- a/tests/sqlparser_mysql.rs +++ b/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); diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 7508218f9..bec60c20f 100644 --- a/tests/sqlparser_postgres.rs +++ b/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!( diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index 634b203db..8141e54f6 100644 --- a/tests/sqlparser_snowflake.rs +++ b/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 );