Added DISTSTYLE and DISTKEY keywords parsing for snowflake

Author: romanoff

src/ast/ddl.rs

@@ -3032,6 +3032,12 @@ pub struct CreateTable {
     /// Snowflake "REQUIRE USER" clause for dybamic tables
     /// <https://docs.snowflake.com/en/sql-reference/sql/create-dynamic-table>
     pub require_user: bool,
+    /// Redshift `DISTSTYLE` option
+    /// <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
+    pub diststyle: Option<DistStyle>,
+    /// Redshift `DISTKEY` option
+    /// <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
+    pub distkey: Option<Ident>,
 }
 
 impl fmt::Display for CreateTable {
@@ -3330,6 +3336,12 @@ impl fmt::Display for CreateTable {
         if self.strict {
             write!(f, " STRICT")?;
         }
+        if let Some(diststyle) = &self.diststyle {
+            write!(f, " DISTSTYLE {diststyle}")?;
+        }
+        if let Some(distkey) = &self.distkey {
+            write!(f, " DISTKEY({distkey})")?;
+        }
         if let Some(query) = &self.query {
             write!(f, " AS {query}")?;
         }
@@ -3417,6 +3429,34 @@ impl fmt::Display for PartitionBoundValue {
     }
 }
 
+/// Redshift distribution style for `CREATE TABLE`.
+///
+/// See [Redshift](https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum DistStyle {
+    /// `DISTSTYLE AUTO`
+    Auto,
+    /// `DISTSTYLE EVEN`
+    Even,
+    /// `DISTSTYLE KEY`
+    Key,
+    /// `DISTSTYLE ALL`
+    All,
+}
+
+impl fmt::Display for DistStyle {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            DistStyle::Auto => write!(f, "AUTO"),
+            DistStyle::Even => write!(f, "EVEN"),
+            DistStyle::Key => write!(f, "KEY"),
+            DistStyle::All => write!(f, "ALL"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/ast/helpers/stmt_create_table.rs

@@ -25,10 +25,11 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::ast::{
-    ClusteredBy, ColumnDef, CommentDef, CreateTable, CreateTableLikeKind, CreateTableOptions, Expr,
-    FileFormat, ForValues, HiveDistributionStyle, HiveFormat, Ident, InitializeKind, ObjectName,
-    OnCommit, OneOrManyWithParens, Query, RefreshModeKind, RowAccessPolicy, Statement,
-    StorageSerializationPolicy, TableConstraint, TableVersion, Tag, WrappedCollection,
+    ClusteredBy, ColumnDef, CommentDef, CreateTable, CreateTableLikeKind, CreateTableOptions,
+    DistStyle, Expr, FileFormat, ForValues, HiveDistributionStyle, HiveFormat, Ident,
+    InitializeKind, ObjectName, OnCommit, OneOrManyWithParens, Query, RefreshModeKind,
+    RowAccessPolicy, Statement, StorageSerializationPolicy, TableConstraint, TableVersion, Tag,
+    WrappedCollection,
 };
 
 use crate::parser::ParserError;
@@ -170,6 +171,10 @@ pub struct CreateTableBuilder {
     pub initialize: Option<InitializeKind>,
     /// Whether operations require a user identity.
     pub require_user: bool,
+    /// Redshift `DISTSTYLE` option.
+    pub diststyle: Option<DistStyle>,
+    /// Redshift `DISTKEY` option.
+    pub distkey: Option<Ident>,
 }
 
 impl CreateTableBuilder {
@@ -229,6 +234,8 @@ impl CreateTableBuilder {
             refresh_mode: None,
             initialize: None,
             require_user: false,
+            diststyle: None,
+            distkey: None,
         }
     }
     /// Set `OR REPLACE` for the CREATE TABLE statement.
@@ -504,6 +511,16 @@ impl CreateTableBuilder {
         self.require_user = require_user;
         self
     }
+    /// Set Redshift `DISTSTYLE` option.
+    pub fn diststyle(mut self, diststyle: Option<DistStyle>) -> Self {
+        self.diststyle = diststyle;
+        self
+    }
+    /// Set Redshift `DISTKEY` option.
+    pub fn distkey(mut self, distkey: Option<Ident>) -> Self {
+        self.distkey = distkey;
+        self
+    }
     /// Consume the builder and produce a `CreateTable`.
     pub fn build(self) -> CreateTable {
         CreateTable {
@@ -560,6 +577,8 @@ impl CreateTableBuilder {
             refresh_mode: self.refresh_mode,
             initialize: self.initialize,
             require_user: self.require_user,
+            diststyle: self.diststyle,
+            distkey: self.distkey,
         }
     }
 }
@@ -635,6 +654,8 @@ impl From<CreateTable> for CreateTableBuilder {
             refresh_mode: table.refresh_mode,
             initialize: table.initialize,
             require_user: table.require_user,
+            diststyle: table.diststyle,
+            distkey: table.distkey,
         }
     }
 }

src/ast/mod.rs

@@ -70,7 +70,8 @@ pub use self::ddl::{
     ConstraintCharacteristics, CreateConnector, CreateDomain, CreateExtension, CreateFunction,
     CreateIndex, CreateOperator, CreateOperatorClass, CreateOperatorFamily, CreatePolicy,
     CreatePolicyCommand, CreatePolicyType, CreateTable, CreateTrigger, CreateView, Deduplicate,
-    DeferrableInitial, DropBehavior, DropExtension, DropFunction, DropOperator, DropOperatorClass,
+    DeferrableInitial, DistStyle, DropBehavior, DropExtension, DropFunction, DropOperator,
+    DropOperatorClass,
     DropOperatorFamily, DropOperatorSignature, DropPolicy, DropTrigger, ForValues, GeneratedAs,
     GeneratedExpressionMode, IdentityParameters, IdentityProperty, IdentityPropertyFormatKind,
     IdentityPropertyKind, IdentityPropertyOrder, IndexColumn, IndexOption, IndexType,

src/ast/spans.rs

@@ -581,6 +581,8 @@ impl Spanned for CreateTable {
             refresh_mode: _,
             initialize: _,
             require_user: _,
+            diststyle: _,  // enum, no span
+            distkey: _,    // Ident, todo
         } = self;
 
         union_spans(

src/keywords.rs

@@ -334,7 +334,9 @@ define_keywords!(
     DISCONNECT,
     DISTINCT,
     DISTINCTROW,
+    DISTKEY,
     DISTRIBUTE,
+    DISTSTYLE,
     DIV,
     DO,
     DOMAIN,
@@ -378,6 +380,7 @@ define_keywords!(
     ESCAPE,
     ESCAPED,
     ESTIMATE,
+    EVEN,
     EVENT,
     EVERY,
     EVOLVE,

src/parser/mod.rs

@@ -8056,6 +8056,23 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse Redshift `DISTSTYLE { AUTO | EVEN | KEY | ALL }`.
+    ///
+    /// See <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
+    fn parse_dist_style(&mut self) -> Result<DistStyle, ParserError> {
+        let token = self.next_token();
+        match &token.token {
+            Token::Word(w) => match w.keyword {
+                Keyword::AUTO => Ok(DistStyle::Auto),
+                Keyword::EVEN => Ok(DistStyle::Even),
+                Keyword::KEY => Ok(DistStyle::Key),
+                Keyword::ALL => Ok(DistStyle::All),
+                _ => self.expected("AUTO, EVEN, KEY, or ALL", token),
+            },
+            _ => self.expected("AUTO, EVEN, KEY, or ALL", token),
+        }
+    }
+
     /// Parse Hive formats.
     pub fn parse_hive_formats(&mut self) -> Result<Option<HiveFormat>, ParserError> {
         let mut hive_format: Option<HiveFormat> = None;
@@ -8326,6 +8343,21 @@ impl<'a> Parser<'a> {
 
         let strict = self.parse_keyword(Keyword::STRICT);
 
+        // Redshift: DISTSTYLE, DISTKEY
+        let diststyle = if self.parse_keyword(Keyword::DISTSTYLE) {
+            Some(self.parse_dist_style()?)
+        } else {
+            None
+        };
+        let distkey = if self.parse_keyword(Keyword::DISTKEY) {
+            self.expect_token(&Token::LParen)?;
+            let column = self.parse_identifier()?;
+            self.expect_token(&Token::RParen)?;
+            Some(column)
+        } else {
+            None
+        };
+
         // Parse optional `AS ( query )`
         let query = if self.parse_keyword(Keyword::AS) {
             Some(self.parse_query()?)
@@ -8365,6 +8397,8 @@ impl<'a> Parser<'a> {
             .table_options(create_table_config.table_options)
             .primary_key(primary_key)
             .strict(strict)
+            .diststyle(diststyle)
+            .distkey(distkey)
             .build())
     }
 

tests/sqlparser_duckdb.rs

@@ -784,6 +784,8 @@ fn test_duckdb_union_datatype() {
             refresh_mode: None,
             initialize: None,
             require_user: Default::default(),
+            diststyle: Default::default(),
+            distkey: Default::default(),
         }),
         stmt
     );

tests/sqlparser_mssql.rs

@@ -1969,6 +1969,8 @@ fn parse_create_table_with_valid_options() {
                 refresh_mode: None,
                 initialize: None,
                 require_user: false,
+                diststyle: None,
+                distkey: None,
             })
         );
     }
@@ -2137,6 +2139,8 @@ fn parse_create_table_with_identity_column() {
                 refresh_mode: None,
                 initialize: None,
                 require_user: false,
+                diststyle: None,
+                distkey: None,
             }),
         );
     }

tests/sqlparser_postgres.rs

@@ -6322,6 +6322,8 @@ fn parse_trigger_related_functions() {
             refresh_mode: None,
             initialize: None,
             require_user: false,
+            diststyle: None,
+            distkey: None,
         }
     );
 

tests/sqlparser_redshift.rs

@@ -452,3 +452,19 @@ fn parse_vacuum() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn test_create_table_diststyle_distkey() {
+    redshift().verified_stmt(
+        "CREATE TEMPORARY TABLE tmp_sbk_summary_pp DISTSTYLE KEY DISTKEY(bet_id) AS SELECT 1 AS bet_id",
+    );
+}
+
+#[test]
+fn test_create_table_diststyle() {
+    redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE AUTO");
+    redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE EVEN");
+    redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE KEY DISTKEY(c1)");
+    redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE ALL");
+}
+