From f10849249ba21953a68159708066d56c58894c20 Mon Sep 17 00:00:00 2001 From: git-hulk Date: Wed, 10 Jul 2024 21:05:00 +0800 Subject: [PATCH 1/2] Enable PARTITON BY feature for Postgre while parsing the create table statement --- src/parser/mod.rs | 3 ++- tests/sqlparser_postgres.rs | 44 +++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index e89eba9b1..a2da73894 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -5416,7 +5416,8 @@ impl<'a> Parser<'a> { None }; - let big_query_config = if dialect_of!(self is BigQueryDialect | GenericDialect) { + let big_query_config = if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect) + { self.parse_optional_big_query_create_table_config()? } else { Default::default() diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 164bb72c7..ed17e9d8f 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -4039,6 +4039,50 @@ fn parse_create_table_with_alias() { } } +#[test] +fn parse_create_table_with_partition_by() { + let sql = "CREATE TABLE t1 (a INT, b TEXT) PARTITION BY RANGE(a)"; + match pg_and_generic().verified_stmt(sql) { + Statement::CreateTable(create_table) => { + assert_eq!("t1", create_table.name.to_string()); + assert_eq!( + vec![ + ColumnDef { + name: "a".into(), + data_type: DataType::Int(None), + collation: None, + options: vec![] + }, + ColumnDef { + name: "b".into(), + data_type: DataType::Text, + collation: None, + options: vec![] + } + ], + create_table.columns + ); + match *create_table.partition_by.unwrap() { + Expr::Function(f) => { + assert_eq!("RANGE", f.name.to_string()); + assert_eq!( + FunctionArguments::List(FunctionArgumentList { + duplicate_treatment: None, + clauses: vec![], + args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr( + Expr::Identifier(Ident::new("a")) + ))], + }), + f.args + ); + } + _ => unreachable!(), + } + } + _ => unreachable!(), + } +} + #[test] fn parse_join_constraint_unnest_alias() { assert_eq!( From dc38aa029cf92b23d37b0defad3f6ca4bec59bbc Mon Sep 17 00:00:00 2001 From: git-hulk Date: Fri, 12 Jul 2024 13:17:44 +0800 Subject: [PATCH 2/2] Fix review comments --- src/ast/helpers/stmt_create_table.rs | 4 +- src/parser/mod.rs | 60 ++++++++++++++-------------- 2 files changed, 33 insertions(+), 31 deletions(-) diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs index d862a36ae..92c75e6a4 100644 --- a/src/ast/helpers/stmt_create_table.rs +++ b/src/ast/helpers/stmt_create_table.rs @@ -496,9 +496,9 @@ impl TryFrom for CreateTableBuilder { } } -/// Helper return type when parsing configuration for a BigQuery `CREATE TABLE` statement. +/// Helper return type when parsing configuration for a `CREATE TABLE` statement. #[derive(Default)] -pub(crate) struct BigQueryTableConfiguration { +pub(crate) struct CreateTableConfiguration { pub partition_by: Option>, pub cluster_by: Option>>, pub options: Option>, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index a2da73894..4d2319a08 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -31,7 +31,7 @@ use recursion::RecursionCounter; use IsLateral::*; use IsOptional::*; -use crate::ast::helpers::stmt_create_table::{BigQueryTableConfiguration, CreateTableBuilder}; +use crate::ast::helpers::stmt_create_table::{CreateTableBuilder, CreateTableConfiguration}; use crate::ast::*; use crate::dialect::*; use crate::keywords::{Keyword, ALL_KEYWORDS}; @@ -5416,12 +5416,7 @@ impl<'a> Parser<'a> { None }; - let big_query_config = if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect) - { - self.parse_optional_big_query_create_table_config()? - } else { - Default::default() - }; + let create_table_config = self.parse_optional_create_table_config()?; // Parse optional `AS ( query )` let query = if self.parse_keyword(Keyword::AS) { @@ -5506,39 +5501,46 @@ impl<'a> Parser<'a> { .collation(collation) .on_commit(on_commit) .on_cluster(on_cluster) - .partition_by(big_query_config.partition_by) - .cluster_by(big_query_config.cluster_by) - .options(big_query_config.options) + .partition_by(create_table_config.partition_by) + .cluster_by(create_table_config.cluster_by) + .options(create_table_config.options) .primary_key(primary_key) .strict(strict) .build()) } - /// Parse configuration like partitioning, clustering information during big-query table creation. - /// - fn parse_optional_big_query_create_table_config( + /// Parse configuration like partitioning, clustering information during the table creation. + /// + /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_2) + /// [PostgreSQL](https://www.postgresql.org/docs/current/ddl-partitioning.html) + fn parse_optional_create_table_config( &mut self, - ) -> Result { - let mut partition_by = None; - if self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) { - partition_by = Some(Box::new(self.parse_expr()?)); + ) -> Result { + let partition_by = if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect) + && self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) + { + Some(Box::new(self.parse_expr()?)) + } else { + None }; let mut cluster_by = None; - if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) { - cluster_by = Some(WrappedCollection::NoWrapping( - self.parse_comma_separated(|p| p.parse_identifier(false))?, - )); - }; - let mut options = None; - if let Token::Word(word) = self.peek_token().token { - if word.keyword == Keyword::OPTIONS { - options = Some(self.parse_options(Keyword::OPTIONS)?); - } - }; + 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(false))?, + )); + }; + + if let Token::Word(word) = self.peek_token().token { + if word.keyword == Keyword::OPTIONS { + options = Some(self.parse_options(Keyword::OPTIONS)?); + } + }; + } - Ok(BigQueryTableConfiguration { + Ok(CreateTableConfiguration { partition_by, cluster_by, options,