Enable PARTITION BY feature for PostgreSQL while parsing the create table statement (apache#1338)

Author: git-hulk

src/ast/helpers/stmt_create_table.rs

@@ -496,9 +496,9 @@ impl TryFrom<Statement> 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<Box<Expr>>,
     pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
     pub options: Option<Vec<SqlOption>>,

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};
@@ -5430,11 +5430,7 @@ impl<'a> Parser<'a> {
             None
         };
 
-        let big_query_config = if dialect_of!(self is BigQueryDialect | 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) {
@@ -5519,39 +5515,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.
-    /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_2>
-    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<BigQueryTableConfiguration, ParserError> {
-        let mut partition_by = None;
-        if self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) {
-            partition_by = Some(Box::new(self.parse_expr()?));
+    ) -> Result<CreateTableConfiguration, ParserError> {
+        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,

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!(