From a1d92bdd9c4b90124e4901845556df3c537d8b53 Mon Sep 17 00:00:00 2001 From: jonathanlehto Date: Thu, 29 Feb 2024 08:50:55 -0500 Subject: [PATCH 1/4] adding unload --- src/ast/mod.rs | 17 +++++++++++ src/keywords.rs | 1 + src/parser/mod.rs | 18 ++++++++++++ tests/sqlparser_common.rs | 61 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 97 insertions(+) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index a2c28c810e..0df2e288bb 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -2301,6 +2301,14 @@ pub enum Statement { /// ``` /// Note: this is a MySQL-specific statement. See UnlockTables, + /// ```sql + /// UNLOAD(statement) TO [ WITH options ] + /// ``` + Unload { + query: Box, + to: Ident, + with: Vec + } } impl fmt::Display for Statement { @@ -3820,6 +3828,15 @@ impl fmt::Display for Statement { Statement::UnlockTables => { write!(f, "UNLOCK TABLES") } + Statement::Unload { query, to, with } => { + write!(f, "UNLOAD({query}) TO {to}")?; + + if with.len() > 0 { + write!(f, " WITH ({})", display_comma_separated(with))?; + } + + Ok(()) + } } } } diff --git a/src/keywords.rs b/src/keywords.rs index f14b92b76b..9186dbdb13 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -684,6 +684,7 @@ define_keywords!( UNCOMMITTED, UNION, UNIQUE, + UNLOAD, UNKNOWN, UNLOCK, UNLOGGED, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 36ac2fd28c..292e77933c 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -516,6 +516,7 @@ impl<'a> Parser<'a> { Keyword::MERGE => Ok(self.parse_merge()?), // `PRAGMA` is sqlite specific https://www.sqlite.org/pragma.html Keyword::PRAGMA => Ok(self.parse_pragma()?), + Keyword::UNLOAD => Ok(self.parse_unload()?), _ => self.expected("an SQL statement", next_token), }, Token::LParen => { @@ -8621,6 +8622,23 @@ impl<'a> Parser<'a> { }) } + pub fn parse_unload(&mut self) -> Result { + self.expect_token(&Token::LParen)?; + let query = self.parse_query()?; + self.expect_token(&Token::RParen)?; + + self.expect_keyword(Keyword::TO)?; + let to = self.parse_identifier(false)?; + + let with_options = self.parse_options(Keyword::WITH)?; + + Ok(Statement::Unload { + query: Box::new(query), + to: to, + with: with_options, + }) + } + pub fn parse_merge_clauses(&mut self) -> Result, ParserError> { let mut clauses: Vec = vec![]; loop { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index fcacb34596..bb05d05ba3 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -8351,6 +8351,67 @@ fn parse_binary_operators_without_whitespace() { ); } +#[test] +fn parse_unload() { + let unload = verified_stmt("UNLOAD(SELECT cola FROM tab) TO 's3://...' WITH (format = 'AVRO')"); + assert_eq!( + unload, + Statement::Unload { + query: Box::new(Query { + body: Box::new( + SetExpr::Select(Box::new(Select { + distinct: None, + top: None, + projection: vec![ + UnnamedExpr(Expr::Identifier(Ident::new("cola"))), + ], + into: None, + from: vec![TableWithJoins { + relation: TableFactor::Table { + name: ObjectName(vec![Ident::new("tab")]), + alias: None, + args: None, + with_hints: vec![], + version: None, + partitions: vec![], + }, + joins: vec![], + }], + lateral_views: vec![], + selection: None, + group_by: GroupByExpr::Expressions(vec![]), + cluster_by: vec![], + distribute_by: vec![], + sort_by: vec![], + having: None, + named_window: vec![], + qualify: None + })) + ), + with: None, + limit: None, + limit_by: vec![], + offset: None, + fetch: None, + locks: vec![], + for_clause: None, + order_by: vec![], + }), + to: Ident { + value: "s3://...".to_string(), + quote_style: Some('\'') + }, + with: vec![SqlOption { + name: Ident { + value: "format".to_string(), + quote_style: None + }, + value: Expr::Value(Value::SingleQuotedString("AVRO".to_string())) + }] + } + ); +} + #[test] fn test_savepoint() { match verified_stmt("SAVEPOINT test1") { From 4ae4c69bc4827927efd5fce4cdd802b1059eade6 Mon Sep 17 00:00:00 2001 From: jonathanlehto Date: Thu, 29 Feb 2024 09:00:45 -0500 Subject: [PATCH 2/4] moved keyword --- src/keywords.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/keywords.rs b/src/keywords.rs index 9186dbdb13..25acaa9a7f 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -684,8 +684,8 @@ define_keywords!( UNCOMMITTED, UNION, UNIQUE, - UNLOAD, UNKNOWN, + UNLOAD, UNLOCK, UNLOGGED, UNNEST, From 3860f992c9cc3329ef62cab49b2c29167a0aaf0e Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 29 Feb 2024 14:46:01 -0500 Subject: [PATCH 3/4] Fix clppy --- src/ast/mod.rs | 6 ++--- tests/sqlparser_common.rs | 57 +++++++++++++++++++-------------------- 2 files changed, 30 insertions(+), 33 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 277fa971fe..a53d83f151 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -2553,8 +2553,8 @@ pub enum Statement { Unload { query: Box, to: Ident, - with: Vec - } + with: Vec, + }, } impl fmt::Display for Statement { @@ -4071,7 +4071,7 @@ impl fmt::Display for Statement { Statement::Unload { query, to, with } => { write!(f, "UNLOAD({query}) TO {to}")?; - if with.len() > 0 { + if with.is_empty() { write!(f, " WITH ({})", display_comma_separated(with))?; } diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index b502ca7152..405395b7d3 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -8441,36 +8441,33 @@ fn parse_unload() { unload, Statement::Unload { query: Box::new(Query { - body: Box::new( - SetExpr::Select(Box::new(Select { - distinct: None, - top: None, - projection: vec![ - UnnamedExpr(Expr::Identifier(Ident::new("cola"))), - ], - into: None, - from: vec![TableWithJoins { - relation: TableFactor::Table { - name: ObjectName(vec![Ident::new("tab")]), - alias: None, - args: None, - with_hints: vec![], - version: None, - partitions: vec![], - }, - joins: vec![], - }], - lateral_views: vec![], - selection: None, - group_by: GroupByExpr::Expressions(vec![]), - cluster_by: vec![], - distribute_by: vec![], - sort_by: vec![], - having: None, - named_window: vec![], - qualify: None - })) - ), + body: Box::new(SetExpr::Select(Box::new(Select { + distinct: None, + top: None, + projection: vec![UnnamedExpr(Expr::Identifier(Ident::new("cola"))),], + into: None, + from: vec![TableWithJoins { + relation: TableFactor::Table { + name: ObjectName(vec![Ident::new("tab")]), + alias: None, + args: None, + with_hints: vec![], + version: None, + partitions: vec![], + }, + joins: vec![], + }], + lateral_views: vec![], + selection: None, + group_by: GroupByExpr::Expressions(vec![]), + cluster_by: vec![], + distribute_by: vec![], + sort_by: vec![], + having: None, + named_window: vec![], + qualify: None, + value_table_mode: None, + }))), with: None, limit: None, limit_by: vec![], From 11eb83d7b8ebe9d659986df71456f022cd045079 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 29 Feb 2024 14:49:19 -0500 Subject: [PATCH 4/4] fix bug and doc --- src/ast/mod.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index a53d83f151..d1c7d76ffe 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -2550,6 +2550,8 @@ pub enum Statement { /// ```sql /// UNLOAD(statement) TO [ WITH options ] /// ``` + /// See Redshift and + // Athena Unload { query: Box, to: Ident, @@ -4071,7 +4073,7 @@ impl fmt::Display for Statement { Statement::Unload { query, to, with } => { write!(f, "UNLOAD({query}) TO {to}")?; - if with.is_empty() { + if !with.is_empty() { write!(f, " WITH ({})", display_comma_separated(with))?; }