From cb2cc91546b9cb3b317fc7ffd7f68c7966689a62 Mon Sep 17 00:00:00 2001 From: Zdenko Nevrala Date: Thu, 25 Apr 2024 15:14:00 +0200 Subject: [PATCH 1/2] Support distribution styles for redshift. --- src/keywords.rs | 3 +++ tests/sqlparser_snowflake.rs | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/src/keywords.rs b/src/keywords.rs index fe0ab0f4c..3f256e3f2 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -224,6 +224,8 @@ define_keywords!( DISCONNECT, DISTINCT, DISTRIBUTE, + DISTSTYLE, + DISTKEY, DIV, DO, DOUBLE, @@ -250,6 +252,7 @@ define_keywords!( EQUALS, ERROR, ESCAPE, + EVEN, EVENT, EVERY, EXCEPT, diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index a275c34a5..186fbde8d 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -1392,3 +1392,9 @@ fn parse_constraints() { fn test_sf_trailing_commas() { snowflake().verified_only_select_with_canonical("SELECT 1, 2, FROM t", "SELECT 1, 2 FROM t"); } + +#[test] +fn test_distribution_styles() { + let sql = "CREATE TABLE foo (id VARCHAR(32)) DISTSTYLE KEY DISTKEY (id) COMPOUND SORTKEY(id)"; + snowflake().verified_stmt(sql); +} From 0429ff2f178d5143659cdb70b3de11fbf82e744d Mon Sep 17 00:00:00 2001 From: Zdenko Nevrala Date: Thu, 2 May 2024 13:04:58 +0200 Subject: [PATCH 2/2] Working. --- src/ast/helpers/stmt_create_table.rs | 36 ++++++++++++++++++-- src/ast/mod.rs | 51 ++++++++++++++++++++++++++++ src/keywords.rs | 4 ++- src/parser/mod.rs | 45 ++++++++++++++++++++++++ tests/sqlparser_redshift.rs | 6 ++++ tests/sqlparser_snowflake.rs | 6 ---- 6 files changed, 138 insertions(+), 10 deletions(-) diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs index e0df72fd2..7c5827122 100644 --- a/src/ast/helpers/stmt_create_table.rs +++ b/src/ast/helpers/stmt_create_table.rs @@ -1,7 +1,7 @@ #[cfg(not(feature = "std"))] use alloc::{boxed::Box, format, string::String, vec, vec::Vec}; -use crate::ast::WithSpan; +use crate::ast::{SortKey, WithSpan}; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; @@ -10,8 +10,8 @@ use serde::{Deserialize, Serialize}; use sqlparser_derive::{Visit, VisitMut}; use crate::ast::{ - ColumnDef, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat, Ident, ObjectName, - OnCommit, Query, SqlOption, Statement, TableConstraint, + ColumnDef, DistributionStyle, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat, + Ident, ObjectName, OnCommit, Query, SqlOption, Statement, TableConstraint, }; use crate::parser::ParserError; @@ -58,6 +58,9 @@ pub struct CreateTableBuilder { pub constraints: Vec, pub hive_distribution: HiveDistributionStyle, pub hive_formats: Option, + pub dist_style: Option, + pub dist_key: Option>, + pub sort_key: Option, pub table_options: Vec, pub table_properties: Vec, pub with_options: Vec, @@ -98,6 +101,9 @@ impl CreateTableBuilder { constraints: vec![], hive_distribution: HiveDistributionStyle::NONE, hive_formats: None, + dist_style: None, + dist_key: None, + sort_key: None, table_options: vec![], table_properties: vec![], with_options: vec![], @@ -174,6 +180,21 @@ impl CreateTableBuilder { self } + pub fn dist_style(mut self, dist_style: Option) -> Self { + self.dist_style = dist_style; + self + } + + pub fn dist_key(mut self, dist_key: Option>) -> Self { + self.dist_key = dist_key; + self + } + + pub fn sort_key(mut self, sort_key: Option) -> Self { + self.sort_key = sort_key; + self + } + pub fn table_options(mut self, table_options: Vec) -> Self { self.table_options = table_options; self @@ -305,6 +326,9 @@ impl CreateTableBuilder { constraints: self.constraints, hive_distribution: self.hive_distribution, hive_formats: self.hive_formats, + dist_style: self.dist_style, + dist_key: self.dist_key, + sort_key: self.sort_key, table_options: self.table_options, table_properties: self.table_properties, with_options: self.with_options, @@ -352,6 +376,9 @@ impl TryFrom for CreateTableBuilder { constraints, hive_distribution, hive_formats, + dist_style, + dist_key, + sort_key, table_options, table_properties, with_options, @@ -388,6 +415,9 @@ impl TryFrom for CreateTableBuilder { constraints, hive_distribution, hive_formats, + dist_style, + dist_key, + sort_key, table_options, table_properties, with_options, diff --git a/src/ast/mod.rs b/src/ast/mod.rs index a05f10fda..2e3ed0eed 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1643,6 +1643,9 @@ pub enum Statement { constraints: Vec, hive_distribution: HiveDistributionStyle, hive_formats: Option, + dist_style: Option, + dist_key: Option>, + sort_key: Option, table_properties: Vec, table_options: Vec, with_options: Vec, @@ -2752,6 +2755,9 @@ impl fmt::Display for Statement { transient, hive_distribution, hive_formats, + dist_style, + dist_key, + sort_key, external, global, temporary, @@ -2821,6 +2827,22 @@ impl fmt::Display for Statement { write!(f, " ()")?; } + if let Some(dist_style) = dist_style { + write!(f, " DISTSTYLE {style}", style = dist_style)?; + } + + if let Some(dist_key) = dist_key { + write!(f, " DISTKEY({dist_key})")?; + } + + if let Some(SortKey { compound, columns }) = sort_key { + if *compound { + write!(f, " COMPOUND SORTKEY({})", display_comma_separated(columns))?; + } else { + write!(f, " SORTKEY({})", display_comma_separated(columns))?; + } + } + if let Some(using) = using { write!(f, " USING {using}")?; } @@ -4356,6 +4378,35 @@ impl fmt::Display for KillType { } } +#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum DistributionStyle { + Auto, + Even, + Key, + All, +} + +impl fmt::Display for DistributionStyle { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str(match self { + DistributionStyle::Auto => "AUTO", + DistributionStyle::Even => "EVEN", + DistributionStyle::Key => "KEY", + DistributionStyle::All => "ALL", + }) + } +} + +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub struct SortKey { + pub compound: bool, + pub columns: Vec>, +} + #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] diff --git a/src/keywords.rs b/src/keywords.rs index 3f256e3f2..92ab5952d 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -157,6 +157,7 @@ define_keywords!( COMMENT, COMMIT, COMMITTED, + COMPOUND, COMPRESSION, COMPUTE, CONCURRENTLY, @@ -223,9 +224,9 @@ define_keywords!( DISCARD, DISCONNECT, DISTINCT, + DISTKEY, DISTRIBUTE, DISTSTYLE, - DISTKEY, DIV, DO, DOUBLE, @@ -583,6 +584,7 @@ define_keywords!( SNAPSHOT, SOME, SORT, + SORTKEY, SPATIAL, SPECIFIC, SPECIFICTYPE, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 84dca46c6..81fed9fb7 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -4082,6 +4082,48 @@ impl<'a> Parser<'a> { None }; + // Redshift allows specifying DISTSTYLE after column definitions + let dist_style = if self.parse_keywords(&[Keyword::DISTSTYLE]) { + match self.parse_one_of_keywords(&[ + Keyword::EVEN, + Keyword::ALL, + Keyword::AUTO, + Keyword::KEY, + ]) { + Some(Keyword::EVEN) => Some(DistributionStyle::Even), + Some(Keyword::ALL) => Some(DistributionStyle::All), + Some(Keyword::AUTO) => Some(DistributionStyle::Auto), + Some(Keyword::KEY) => Some(DistributionStyle::Key), + _ => self.expected("KEY, EVEN, ALL or AUTO", self.peek_token())?, + } + } else { + None + }; + + // Redshift allows specifying DISTKEY after column definitions + let dist_key = if self.parse_keywords(&[Keyword::DISTKEY]) { + self.expect_token(&Token::LParen)?; + let key = self.parse_identifier(false)?; + self.expect_token(&Token::RParen)?; + Some(key) + } else { + None + }; + + // Redshift allows specifying SORTKEY after column definitions + let compound_sort_key = self.parse_keywords(&[Keyword::COMPOUND]); + let sort_key = if self.parse_keywords(&[Keyword::SORTKEY]) { + self.expect_token(&Token::LParen)?; + let columns = self.parse_comma_separated(|p| p.parse_identifier(false))?; + self.expect_token(&Token::RParen)?; + Some(SortKey { + compound: compound_sort_key, + columns, + }) + } else { + None + }; + // SQLite supports `WITHOUT ROWID` at the end of `CREATE TABLE` let without_rowid = self.parse_keywords(&[Keyword::WITHOUT, Keyword::ROWID]); @@ -4269,6 +4311,9 @@ impl<'a> Parser<'a> { .transient(transient) .hive_distribution(hive_distribution) .hive_formats(Some(hive_formats)) + .dist_style(dist_style) + .dist_key(dist_key) + .sort_key(sort_key) .global(global) .query(query) .without_rowid(without_rowid) diff --git a/tests/sqlparser_redshift.rs b/tests/sqlparser_redshift.rs index 1206b63ff..ee8a7826b 100644 --- a/tests/sqlparser_redshift.rs +++ b/tests/sqlparser_redshift.rs @@ -363,3 +363,9 @@ fn test_escape_string() { redshift_unescaped().verified_stmt(r#"SELECT 'I\\\'m fine'"#); redshift_unescaped().verified_stmt(r#"SELECT '[\'\\[\\]]'"#); } + +#[test] +fn test_distribution_styles() { + let sql = "CREATE TABLE foo (id VARCHAR(32)) DISTSTYLE KEY DISTKEY(id) COMPOUND SORTKEY(id)"; + redshift().verified_stmt(sql); +} diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index 186fbde8d..a275c34a5 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -1392,9 +1392,3 @@ fn parse_constraints() { fn test_sf_trailing_commas() { snowflake().verified_only_select_with_canonical("SELECT 1, 2, FROM t", "SELECT 1, 2 FROM t"); } - -#[test] -fn test_distribution_styles() { - let sql = "CREATE TABLE foo (id VARCHAR(32)) DISTSTYLE KEY DISTKEY (id) COMPOUND SORTKEY(id)"; - snowflake().verified_stmt(sql); -}