add with_order option

synnada-ai · Aug 6, 2024 · 1375a21 · 1375a21
1 parent 8f8c96f
commit 1375a21
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 1 deletion.
diff --git a/src/ast/dml.rs b/src/ast/dml.rs
@@ -124,6 +124,7 @@ pub struct CreateTable {
     pub default_charset: Option<String>,
     pub collation: Option<String>,
     pub on_commit: Option<OnCommit>,
+    pub with_order: Option<Vec<OrderByExpr>>,
     /// ClickHouse "ON CLUSTER" clause:
     /// <https://clickhouse.com/docs/en/sql-reference/distributed-ddl/>
     pub on_cluster: Option<Ident>,

diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs
@@ -10,7 +10,7 @@ use sqlparser_derive::{Visit, VisitMut};
 use super::super::dml::CreateTable;
 use crate::ast::{
     ColumnDef, CommentDef, Expr, FileFormat, HiveDistributionStyle, HiveFormat, Ident, ObjectName,
-    OnCommit, OneOrManyWithParens, Query, RowAccessPolicy, SqlOption, Statement, TableConstraint,
+    OnCommit, OneOrManyWithParens, OrderByExpr, Query, RowAccessPolicy, SqlOption, Statement, TableConstraint,
     TableEngine, Tag, WrappedCollection,
 };
 use crate::parser::ParserError;
@@ -75,6 +75,7 @@ pub struct CreateTableBuilder {
     pub on_commit: Option<OnCommit>,
     pub on_cluster: Option<Ident>,
     pub primary_key: Option<Box<Expr>>,
+    pub with_order: Option<Vec<OrderByExpr>>,
     pub order_by: Option<OneOrManyWithParens<Expr>>,
     pub partition_by: Option<Box<Expr>>,
     pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
@@ -134,6 +135,7 @@ impl CreateTableBuilder {
             max_data_extension_time_in_days: None,
             default_ddl_collation: None,
             with_aggregation_policy: None,
+            with_order: None,
             with_row_access_policy: None,
             with_tags: None,
         }
@@ -271,6 +273,11 @@ impl CreateTableBuilder {
         self
     }
 
+    pub fn with_order(mut self, with_order: Option<Vec<OrderByExpr>>) -> Self {
+        self.with_order = with_order;
+        self
+    }
+
     pub fn order_by(mut self, order_by: Option<OneOrManyWithParens<Expr>>) -> Self {
         self.order_by = order_by;
         self
@@ -389,6 +396,7 @@ impl CreateTableBuilder {
             max_data_extension_time_in_days: self.max_data_extension_time_in_days,
             default_ddl_collation: self.default_ddl_collation,
             with_aggregation_policy: self.with_aggregation_policy,
+            with_order: self.with_order,
             with_row_access_policy: self.with_row_access_policy,
             with_tags: self.with_tags,
         })
@@ -443,6 +451,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 max_data_extension_time_in_days,
                 default_ddl_collation,
                 with_aggregation_policy,
+                with_order,
                 with_row_access_policy,
                 with_tags,
             }) => Ok(Self {
@@ -486,6 +495,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 default_ddl_collation,
                 with_aggregation_policy,
                 with_row_access_policy,
+                with_order,
                 with_tags,
                 volatile,
             }),

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -5469,6 +5469,22 @@ impl<'a> Parser<'a> {
             None
         };
 
+        let with_order = if self.parse_keywords(&[Keyword::WITH, Keyword::ORDER])
+        {
+            let mut values = vec![];
+            self.expect_token(&Token::LParen)?;
+            loop {
+                values.push(self.parse_order_by_expr()?);
+                if !self.consume_token(&Token::Comma) {
+                    self.expect_token(&Token::RParen)?;
+                    break
+                }
+            }
+            Some(values)
+        } 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 {
@@ -5574,6 +5590,7 @@ impl<'a> Parser<'a> {
             .cluster_by(create_table_config.cluster_by)
             .options(create_table_config.options)
             .primary_key(primary_key)
+            .with_order(with_order)
             .strict(strict)
             .build())
     }

diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs
@@ -638,6 +638,7 @@ fn test_duckdb_union_datatype() {
             max_data_extension_time_in_days: Default::default(),
             default_ddl_collation: Default::default(),
             with_aggregation_policy: Default::default(),
+            with_order: Default::default(),
             with_row_access_policy: Default::default(),
             with_tags: Default::default()
         }),