From 75a5014ab49838f0bc4703d2d1e2acaf82e9dbed Mon Sep 17 00:00:00 2001 From: lovasoa Date: Sat, 12 Oct 2024 13:06:37 +0200 Subject: [PATCH 1/2] this fixes a regression that was introduced with the new SetRole statement --- src/parser/mod.rs | 43 ++++++++++++++++++++++++++------------- tests/sqlparser_common.rs | 24 ++++++++++++++++++++++ 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index cd9be1d8f..83935e0ee 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9416,27 +9416,42 @@ impl<'a> Parser<'a> { } } + fn parse_set_role(&mut self, modifier: Option) -> Option { + let context_modifier = match modifier { + Some(Keyword::LOCAL) => ContextModifier::Local, + Some(Keyword::SESSION) => ContextModifier::Session, + _ => ContextModifier::None, + }; + + if self.parse_keyword(Keyword::NONE) { + Some(Statement::SetRole { + context_modifier, + role_name: None, + }) + } else if matches!( + self.peek_token().token, + Token::Word(_) | Token::DoubleQuotedString(_) | Token::SingleQuotedString(_) + ) { + Some(Statement::SetRole { + context_modifier, + role_name: Some(self.parse_identifier(false).ok()?), + }) + } else { + None + } + } + pub fn parse_set(&mut self) -> Result { let modifier = self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]); if let Some(Keyword::HIVEVAR) = modifier { self.expect_token(&Token::Colon)?; } else if self.parse_keyword(Keyword::ROLE) { - let context_modifier = match modifier { - Some(Keyword::LOCAL) => ContextModifier::Local, - Some(Keyword::SESSION) => ContextModifier::Session, - _ => ContextModifier::None, - }; - - let role_name = if self.parse_keyword(Keyword::NONE) { - None + if let Some(set_role_stmt) = self.parse_set_role(modifier) { + return Ok(set_role_stmt); } else { - Some(self.parse_identifier(false)?) - }; - return Ok(Statement::SetRole { - context_modifier, - role_name, - }); + self.prev_token(); // replace the "role" token which was actually a variable name + } } let variables = if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 5327880a4..447fea318 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -7665,6 +7665,30 @@ fn parse_set_variable() { one_statement_parses_to("SET SOMETHING TO '1'", "SET SOMETHING = '1'"); } +#[test] +fn parse_set_role_as_variable() { + match verified_stmt("SET role = 'foobar'") { + Statement::SetVariable { + local, + hivevar, + variables, + value, + } => { + assert!(!local); + assert!(!hivevar); + assert_eq!( + variables, + OneOrManyWithParens::One(ObjectName(vec!["role".into()])) + ); + assert_eq!( + value, + vec![Expr::Value(Value::SingleQuotedString("foobar".into()))] + ); + } + _ => unreachable!(), + } +} + #[test] fn parse_double_colon_cast_at_timezone() { let sql = "SELECT '2001-01-01T00:00:00.000Z'::TIMESTAMP AT TIME ZONE 'Europe/Brussels' FROM t"; From 88a1d61f347b490cff79b0cb174ec10a258c9130 Mon Sep 17 00:00:00 2001 From: lovasoa Date: Sun, 13 Oct 2024 22:37:11 +0200 Subject: [PATCH 2/2] use maybe_parse Thanks @iffyio --- src/parser/mod.rs | 37 +++++++++++++++---------------------- 1 file changed, 15 insertions(+), 22 deletions(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 83935e0ee..b4c0487b4 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9416,29 +9416,24 @@ impl<'a> Parser<'a> { } } - fn parse_set_role(&mut self, modifier: Option) -> Option { + /// Parse a `SET ROLE` statement. Expects SET to be consumed already. + fn parse_set_role(&mut self, modifier: Option) -> Result { + self.expect_keyword(Keyword::ROLE)?; let context_modifier = match modifier { Some(Keyword::LOCAL) => ContextModifier::Local, Some(Keyword::SESSION) => ContextModifier::Session, _ => ContextModifier::None, }; - if self.parse_keyword(Keyword::NONE) { - Some(Statement::SetRole { - context_modifier, - role_name: None, - }) - } else if matches!( - self.peek_token().token, - Token::Word(_) | Token::DoubleQuotedString(_) | Token::SingleQuotedString(_) - ) { - Some(Statement::SetRole { - context_modifier, - role_name: Some(self.parse_identifier(false).ok()?), - }) - } else { + let role_name = if self.parse_keyword(Keyword::NONE) { None - } + } else { + Some(self.parse_identifier(false)?) + }; + Ok(Statement::SetRole { + context_modifier, + role_name, + }) } pub fn parse_set(&mut self) -> Result { @@ -9446,12 +9441,10 @@ impl<'a> Parser<'a> { self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]); if let Some(Keyword::HIVEVAR) = modifier { self.expect_token(&Token::Colon)?; - } else if self.parse_keyword(Keyword::ROLE) { - if let Some(set_role_stmt) = self.parse_set_role(modifier) { - return Ok(set_role_stmt); - } else { - self.prev_token(); // replace the "role" token which was actually a variable name - } + } else if let Some(set_role_stmt) = + self.maybe_parse(|parser| parser.parse_set_role(modifier)) + { + return Ok(set_role_stmt); } let variables = if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) {