From d46398f1019625d3eaa531554f93d11097daf000 Mon Sep 17 00:00:00 2001 From: Maciej Obuchowski Date: Mon, 7 Mar 2022 13:52:34 +0100 Subject: [PATCH 1/2] add support for MERGE statement Signed-off-by: Maciej Obuchowski --- src/ast/mod.rs | 89 ++++++++++++++++++++++++ src/keywords.rs | 1 + src/parser.rs | 85 +++++++++++++++++++++++ tests/sqlparser_common.rs | 138 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 313 insertions(+) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 232a506f4..62d4cc935 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -956,6 +956,19 @@ pub enum Statement { /// A SQL query that specifies what to explain statement: Box, }, + // MERGE INTO statement, based on Snowflake. See + Merge { + // Specifies the table to merge + table: TableFactor, + // Specifies the table or subquery to join with the target table + source: Box, + // Specifies alias to the table that is joined with target table + alias: Option, + // Specifies the expression on which to join the target table and source + on: Box, + // Specifies the actions to perform when values match or do not match. + clauses: Vec, + }, } impl fmt::Display for Statement { @@ -1606,6 +1619,20 @@ impl fmt::Display for Statement { write!(f, "NULL") } } + Statement::Merge { + table, + source, + alias, + on, + clauses, + } => { + write!(f, "MERGE INTO {} USING {} ", table, source)?; + if let Some(a) = alias { + write!(f, "as {} ", a)?; + }; + write!(f, "ON {} ", on)?; + write!(f, "{}", display_separated(clauses, " ")) + } } } } @@ -2137,6 +2164,68 @@ impl fmt::Display for SqliteOnConflict { } } +/// +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum MergeClause { + MatchedUpdate { + predicate: Option, + assignments: Vec, + }, + MatchedDelete(Option), + NotMatched { + predicate: Option, + columns: Vec, + values: Values, + }, +} + +impl fmt::Display for MergeClause { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + use MergeClause::*; + write!(f, "WHEN")?; + match self { + MatchedUpdate { + predicate, + assignments, + } => { + write!(f, " MATCHED")?; + if let Some(pred) = predicate { + write!(f, " AND {}", pred)?; + } + write!( + f, + " THEN UPDATE SET {}", + display_comma_separated(assignments) + ) + } + MatchedDelete(predicate) => { + write!(f, " MATCHED")?; + if let Some(pred) = predicate { + write!(f, " AND {}", pred)?; + } + write!(f, " THEN DELETE") + } + NotMatched { + predicate, + columns, + values, + } => { + write!(f, " NOT MATCHED")?; + if let Some(pred) = predicate { + write!(f, " AND {}", pred)?; + } + write!( + f, + " THEN INSERT ({}) {}", + display_comma_separated(columns), + values + ) + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/keywords.rs b/src/keywords.rs index b3c77eb53..9c2ff4c5c 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -295,6 +295,7 @@ define_keywords!( LOWER, MANAGEDLOCATION, MATCH, + MATCHED, MATERIALIZED, MAX, MEMBER, diff --git a/src/parser.rs b/src/parser.rs index 6ccd4980b..5fb43d66a 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -187,6 +187,7 @@ impl<'a> Parser<'a> { Keyword::DEALLOCATE => Ok(self.parse_deallocate()?), Keyword::EXECUTE => Ok(self.parse_execute()?), Keyword::PREPARE => Ok(self.parse_prepare()?), + Keyword::MERGE => Ok(self.parse_merge()?), Keyword::REPLACE if dialect_of!(self is SQLiteDialect ) => { self.prev_token(); Ok(self.parse_insert()?) @@ -3890,6 +3891,90 @@ impl<'a> Parser<'a> { comment, }) } + + pub fn parse_merge_clauses(&mut self) -> Result, ParserError> { + let mut clauses: Vec = vec![]; + loop { + if self.peek_token() == Token::EOF { + break; + } + self.expect_keyword(Keyword::WHEN)?; + + let is_not_matched = self.parse_keyword(Keyword::NOT); + self.expect_keyword(Keyword::MATCHED)?; + + let predicate = if self.parse_keyword(Keyword::AND) { + Some(self.parse_expr()?) + } else { + None + }; + + self.expect_keyword(Keyword::THEN)?; + + clauses.push( + match self.parse_one_of_keywords(&[ + Keyword::UPDATE, + Keyword::INSERT, + Keyword::DELETE, + ]) { + Some(Keyword::UPDATE) => { + if is_not_matched { + parser_err!("UPDATE in NOT MATCHED merge clause")?; + } + self.expect_keyword(Keyword::SET)?; + let assignments = self.parse_comma_separated(Parser::parse_assignment)?; + MergeClause::MatchedUpdate { + predicate, + assignments, + } + } + Some(Keyword::DELETE) => { + if is_not_matched { + parser_err!("DELETE in NOT MATCHED merge clause")?; + } + MergeClause::MatchedDelete(predicate) + } + Some(Keyword::INSERT) => { + if !is_not_matched { + parser_err!("INSERT in MATCHED merge clause")?; + } + let columns = self.parse_parenthesized_column_list(Optional)?; + self.expect_keyword(Keyword::VALUES)?; + let values = self.parse_values()?; + MergeClause::NotMatched { + predicate, + columns, + values, + } + } + Some(_) => parser_err!("expected UPDATE, DELETE or INSERT in merge clause")?, + None => parser_err!("expected UPDATE, DELETE or INSERT in merge clause")?, + }, + ); + } + Ok(clauses) + } + + pub fn parse_merge(&mut self) -> Result { + self.expect_keyword(Keyword::INTO)?; + + let table = self.parse_table_factor()?; + + self.expect_keyword(Keyword::USING)?; + let source = self.parse_query_body(0)?; + let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?; + self.expect_keyword(Keyword::ON)?; + let on = self.parse_expr()?; + let clauses = self.parse_merge_clauses()?; + + Ok(Statement::Merge { + table, + source: Box::new(source), + alias, + on: Box::new(on), + clauses, + }) + } } impl Word { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index c56cb37a6..788d4a854 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -4203,6 +4203,144 @@ fn test_revoke() { } } +#[test] +fn parse_merge() { + let sql = "MERGE INTO s.bar AS dest USING (SELECT * FROM s.foo) as stg ON dest.D = stg.D AND dest.E = stg.E WHEN NOT MATCHED THEN INSERT (A, B, C) VALUES (stg.A, stg.B, stg.C) WHEN MATCHED AND dest.A = 'a' THEN UPDATE SET dest.F = stg.F, dest.G = stg.G WHEN MATCHED THEN DELETE"; + match verified_stmt(sql) { + Statement::Merge { + table, + source, + alias, + on, + clauses, + } => { + assert_eq!( + table, + TableFactor::Table { + name: ObjectName(vec![Ident::new("s"), Ident::new("bar")]), + alias: Some(TableAlias { + name: Ident::new("dest"), + columns: vec![] + }), + args: vec![], + with_hints: vec![] + } + ); + assert_eq!( + source, + Box::new(SetExpr::Query(Box::new(Query { + with: None, + body: SetExpr::Select(Box::new(Select { + distinct: false, + top: None, + projection: vec![SelectItem::Wildcard], + from: vec![TableWithJoins { + relation: TableFactor::Table { + name: ObjectName(vec![Ident::new("s"), Ident::new("foo")]), + alias: None, + args: vec![], + with_hints: vec![], + }, + joins: vec![] + }], + lateral_views: vec![], + selection: None, + group_by: vec![], + cluster_by: vec![], + distribute_by: vec![], + sort_by: vec![], + having: None + })), + order_by: vec![], + limit: None, + offset: None, + fetch: None, + lock: None + }))) + ); + assert_eq!( + alias, + Some(TableAlias { + name: Ident::new("stg"), + columns: vec![] + }) + ); + assert_eq!( + on, + Box::new(Expr::BinaryOp { + left: Box::new(Expr::BinaryOp { + left: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("dest"), + Ident::new("D") + ])), + op: BinaryOperator::Eq, + right: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("stg"), + Ident::new("D") + ])) + }), + op: BinaryOperator::And, + right: Box::new(Expr::BinaryOp { + left: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("dest"), + Ident::new("E") + ])), + op: BinaryOperator::Eq, + right: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("stg"), + Ident::new("E") + ])) + }) + }) + ); + assert_eq!( + clauses, + vec![ + MergeClause::NotMatched { + predicate: None, + columns: vec![Ident::new("A"), Ident::new("B"), Ident::new("C")], + values: Values(vec![vec![ + Expr::CompoundIdentifier(vec![Ident::new("stg"), Ident::new("A")]), + Expr::CompoundIdentifier(vec![Ident::new("stg"), Ident::new("B")]), + Expr::CompoundIdentifier(vec![Ident::new("stg"), Ident::new("C")]), + ]]) + }, + MergeClause::MatchedUpdate { + predicate: Some(Expr::BinaryOp { + left: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("dest"), + Ident::new("A") + ])), + op: BinaryOperator::Eq, + right: Box::new(Expr::Value(Value::SingleQuotedString( + "a".to_string() + ))) + }), + assignments: vec![ + Assignment { + id: vec![Ident::new("dest"), Ident::new("F")], + value: Expr::CompoundIdentifier(vec![ + Ident::new("stg"), + Ident::new("F") + ]) + }, + Assignment { + id: vec![Ident::new("dest"), Ident::new("G")], + value: Expr::CompoundIdentifier(vec![ + Ident::new("stg"), + Ident::new("G") + ]) + } + ] + }, + MergeClause::MatchedDelete(None) + ] + ) + } + _ => unreachable!(), + } +} + #[test] fn test_lock() { let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE"; From 134fd6b7516636d3376854a5606785b6c332cbaf Mon Sep 17 00:00:00 2001 From: Maciej Obuchowski Date: Tue, 8 Mar 2022 13:28:54 +0100 Subject: [PATCH 2/2] fix lint errors Signed-off-by: Maciej Obuchowski --- src/parser.rs | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/parser.rs b/src/parser.rs index 5fb43d66a..cb87c99fd 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -3919,7 +3919,9 @@ impl<'a> Parser<'a> { ]) { Some(Keyword::UPDATE) => { if is_not_matched { - parser_err!("UPDATE in NOT MATCHED merge clause")?; + return Err(ParserError::ParserError( + "UPDATE in NOT MATCHED merge clause".to_string(), + )); } self.expect_keyword(Keyword::SET)?; let assignments = self.parse_comma_separated(Parser::parse_assignment)?; @@ -3930,13 +3932,17 @@ impl<'a> Parser<'a> { } Some(Keyword::DELETE) => { if is_not_matched { - parser_err!("DELETE in NOT MATCHED merge clause")?; + return Err(ParserError::ParserError( + "DELETE in NOT MATCHED merge clause".to_string(), + )); } MergeClause::MatchedDelete(predicate) } Some(Keyword::INSERT) => { if !is_not_matched { - parser_err!("INSERT in MATCHED merge clause")?; + return Err(ParserError::ParserError( + "INSERT in MATCHED merge clause".to_string(), + )); } let columns = self.parse_parenthesized_column_list(Optional)?; self.expect_keyword(Keyword::VALUES)?; @@ -3947,8 +3953,16 @@ impl<'a> Parser<'a> { values, } } - Some(_) => parser_err!("expected UPDATE, DELETE or INSERT in merge clause")?, - None => parser_err!("expected UPDATE, DELETE or INSERT in merge clause")?, + Some(_) => { + return Err(ParserError::ParserError( + "expected UPDATE, DELETE or INSERT in merge clause".to_string(), + )) + } + None => { + return Err(ParserError::ParserError( + "expected UPDATE, DELETE or INSERT in merge clause".to_string(), + )) + } }, ); }