Skip to content

Commit 6f1f135

Browse files
romanoffromanoff
committed
Added SORTKEY keyword parsing for redshift queries
1 parent 1da2ff7 commit 6f1f135

9 files changed

Lines changed: 90 additions & 8 deletions

File tree

‎src/ast/ddl.rs‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,12 @@ pub enum AlterTableOperation {
457457
},
458458
/// Remove the clustering key from the table.
459459
DropClusteringKey,
460+
/// Redshift `ALTER SORTKEY (column_list)`
461+
/// <https://docs.aws.amazon.com/redshift/latest/dg/r_ALTER_TABLE.html>
462+
AlterSortKey {
463+
/// Column references in the sort key.
464+
columns: Vec<Expr>,
465+
},
460466
/// Suspend background reclustering operations.
461467
SuspendRecluster,
462468
/// Resume background reclustering operations.
@@ -993,6 +999,14 @@ impl fmt::Display for AlterTableOperation {
993999
write!(f, "DROP CLUSTERING KEY")?;
9941000
Ok(())
9951001
}
1002+
AlterTableOperation::AlterSortKey { columns } => {
1003+
write!(
1004+
f,
1005+
"ALTER SORTKEY({})",
1006+
display_comma_separated(columns)
1007+
)?;
1008+
Ok(())
1009+
}
9961010
AlterTableOperation::SuspendRecluster => {
9971011
write!(f, "SUSPEND RECLUSTER")?;
9981012
Ok(())
@@ -3037,7 +3051,10 @@ pub struct CreateTable {
30373051
pub diststyle: Option<DistStyle>,
30383052
/// Redshift `DISTKEY` option
30393053
/// <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
3040-
pub distkey: Option<Ident>,
3054+
pub distkey: Option<Expr>,
3055+
/// Redshift `SORTKEY` option
3056+
/// <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
3057+
pub sortkey: Option<Vec<Expr>>,
30413058
}
30423059

30433060
impl fmt::Display for CreateTable {
@@ -3342,6 +3359,9 @@ impl fmt::Display for CreateTable {
33423359
if let Some(distkey) = &self.distkey {
33433360
write!(f, " DISTKEY({distkey})")?;
33443361
}
3362+
if let Some(sortkey) = &self.sortkey {
3363+
write!(f, " SORTKEY({})", display_comma_separated(sortkey))?;
3364+
}
33453365
if let Some(query) = &self.query {
33463366
write!(f, " AS {query}")?;
33473367
}

‎src/ast/helpers/stmt_create_table.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,9 @@ pub struct CreateTableBuilder {
174174
/// Redshift `DISTSTYLE` option.
175175
pub diststyle: Option<DistStyle>,
176176
/// Redshift `DISTKEY` option.
177-
pub distkey: Option<Ident>,
177+
pub distkey: Option<Expr>,
178+
/// Redshift `SORTKEY` option.
179+
pub sortkey: Option<Vec<Expr>>,
178180
}
179181

180182
impl CreateTableBuilder {
@@ -236,6 +238,7 @@ impl CreateTableBuilder {
236238
require_user: false,
237239
diststyle: None,
238240
distkey: None,
241+
sortkey: None,
239242
}
240243
}
241244
/// Set `OR REPLACE` for the CREATE TABLE statement.
@@ -517,10 +520,15 @@ impl CreateTableBuilder {
517520
self
518521
}
519522
/// Set Redshift `DISTKEY` option.
520-
pub fn distkey(mut self, distkey: Option<Ident>) -> Self {
523+
pub fn distkey(mut self, distkey: Option<Expr>) -> Self {
521524
self.distkey = distkey;
522525
self
523526
}
527+
/// Set Redshift `SORTKEY` option.
528+
pub fn sortkey(mut self, sortkey: Option<Vec<Expr>>) -> Self {
529+
self.sortkey = sortkey;
530+
self
531+
}
524532
/// Consume the builder and produce a `CreateTable`.
525533
pub fn build(self) -> CreateTable {
526534
CreateTable {
@@ -579,6 +587,7 @@ impl CreateTableBuilder {
579587
require_user: self.require_user,
580588
diststyle: self.diststyle,
581589
distkey: self.distkey,
590+
sortkey: self.sortkey,
582591
}
583592
}
584593
}
@@ -656,6 +665,7 @@ impl From<CreateTable> for CreateTableBuilder {
656665
require_user: table.require_user,
657666
diststyle: table.diststyle,
658667
distkey: table.distkey,
668+
sortkey: table.sortkey,
659669
}
660670
}
661671
}

‎src/ast/spans.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -582,8 +582,9 @@ impl Spanned for CreateTable {
582582
refresh_mode: _,
583583
initialize: _,
584584
require_user: _,
585-
diststyle: _, // enum, no span
586-
distkey: _, // Ident, todo
585+
diststyle: _,
586+
distkey: _,
587+
sortkey: _,
587588
} = self;
588589

589590
union_spans(
@@ -1193,6 +1194,7 @@ impl Spanned for AlterTableOperation {
11931194
AlterTableOperation::OwnerTo { .. } => Span::empty(),
11941195
AlterTableOperation::ClusterBy { exprs } => union_spans(exprs.iter().map(|e| e.span())),
11951196
AlterTableOperation::DropClusteringKey => Span::empty(),
1197+
AlterTableOperation::AlterSortKey { .. } => Span::empty(),
11961198
AlterTableOperation::SuspendRecluster => Span::empty(),
11971199
AlterTableOperation::ResumeRecluster => Span::empty(),
11981200
AlterTableOperation::Refresh { .. } => Span::empty(),

‎src/keywords.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -953,6 +953,7 @@ define_keywords!(
953953
SOME,
954954
SORT,
955955
SORTED,
956+
SORTKEY,
956957
SOURCE,
957958
SPATIAL,
958959
SPECIFIC,

‎src/parser/mod.rs‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8384,17 +8384,25 @@ impl<'a> Parser<'a> {
83848384

83858385
let strict = self.parse_keyword(Keyword::STRICT);
83868386

8387-
// Redshift: DISTSTYLE, DISTKEY
8387+
// Redshift: DISTSTYLE, DISTKEY, SORTKEY
83888388
let diststyle = if self.parse_keyword(Keyword::DISTSTYLE) {
83898389
Some(self.parse_dist_style()?)
83908390
} else {
83918391
None
83928392
};
83938393
let distkey = if self.parse_keyword(Keyword::DISTKEY) {
83948394
self.expect_token(&Token::LParen)?;
8395-
let column = self.parse_identifier()?;
8395+
let expr = self.parse_expr()?;
83968396
self.expect_token(&Token::RParen)?;
8397-
Some(column)
8397+
Some(expr)
8398+
} else {
8399+
None
8400+
};
8401+
let sortkey = if self.parse_keyword(Keyword::SORTKEY) {
8402+
self.expect_token(&Token::LParen)?;
8403+
let columns = self.parse_comma_separated(|p| p.parse_expr())?;
8404+
self.expect_token(&Token::RParen)?;
8405+
Some(columns)
83988406
} else {
83998407
None
84008408
};
@@ -8440,6 +8448,7 @@ impl<'a> Parser<'a> {
84408448
.strict(strict)
84418449
.diststyle(diststyle)
84428450
.distkey(distkey)
8451+
.sortkey(sortkey)
84438452
.build())
84448453
}
84458454

@@ -10257,6 +10266,14 @@ impl<'a> Parser<'a> {
1025710266
column_position,
1025810267
}
1025910268
} else if self.parse_keyword(Keyword::ALTER) {
10269+
// Redshift: ALTER SORTKEY (column_list)
10270+
if self.parse_keyword(Keyword::SORTKEY) {
10271+
self.expect_token(&Token::LParen)?;
10272+
let columns = self.parse_comma_separated(|p| p.parse_expr())?;
10273+
self.expect_token(&Token::RParen)?;
10274+
return Ok(AlterTableOperation::AlterSortKey { columns });
10275+
}
10276+
1026010277
let _ = self.parse_keyword(Keyword::COLUMN); // [ COLUMN ]
1026110278
let column_name = self.parse_identifier()?;
1026210279
let is_postgresql = dialect_of!(self is PostgreSqlDialect);

‎tests/sqlparser_duckdb.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -790,6 +790,7 @@ fn test_duckdb_union_datatype() {
790790
require_user: Default::default(),
791791
diststyle: Default::default(),
792792
distkey: Default::default(),
793+
sortkey: Default::default(),
793794
}),
794795
stmt
795796
);

‎tests/sqlparser_mssql.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2008,6 +2008,7 @@ fn parse_create_table_with_valid_options() {
20082008
require_user: false,
20092009
diststyle: None,
20102010
distkey: None,
2011+
sortkey: None,
20112012
})
20122013
);
20132014
}
@@ -2178,6 +2179,7 @@ fn parse_create_table_with_identity_column() {
21782179
require_user: false,
21792180
diststyle: None,
21802181
distkey: None,
2182+
sortkey: None,
21812183
}),
21822184
);
21832185
}

‎tests/sqlparser_postgres.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6420,6 +6420,7 @@ fn parse_trigger_related_functions() {
64206420
require_user: false,
64216421
diststyle: None,
64226422
distkey: None,
6423+
sortkey: None,
64236424
}
64246425
);
64256426

‎tests/sqlparser_redshift.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -467,3 +467,31 @@ fn test_create_table_diststyle() {
467467
redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE KEY DISTKEY(c1)");
468468
redshift().verified_stmt("CREATE TABLE t1 (c1 INT) DISTSTYLE ALL");
469469
}
470+
471+
#[test]
472+
fn test_create_table_sortkey() {
473+
redshift()
474+
.verified_stmt("CREATE TABLE t1 (c1 INT, c2 INT, c3 TIMESTAMP) SORTKEY(c3)");
475+
redshift()
476+
.verified_stmt("CREATE TABLE t1 (c1 INT, c2 INT) SORTKEY(c1, c2)");
477+
}
478+
479+
#[test]
480+
fn test_create_table_distkey_sortkey_with_ctas() {
481+
redshift().verified_stmt(
482+
"CREATE TABLE t1 DISTKEY(1) SORTKEY(1, 3) AS SELECT eventid, venueid, dateid, eventname FROM event",
483+
);
484+
}
485+
486+
#[test]
487+
fn test_create_table_diststyle_distkey_sortkey() {
488+
redshift().verified_stmt(
489+
"CREATE TABLE t1 (c1 INT, c2 INT) DISTSTYLE KEY DISTKEY(c1) SORTKEY(c1, c2)",
490+
);
491+
}
492+
493+
#[test]
494+
fn test_alter_table_alter_sortkey() {
495+
redshift().verified_stmt("ALTER TABLE users ALTER SORTKEY(created_at)");
496+
redshift().verified_stmt("ALTER TABLE users ALTER SORTKEY(c1, c2)");
497+
}

0 commit comments

Comments
 (0)