iffyio commented on code in PR #1757: URL: https://github.com/apache/datafusion-sqlparser-rs/pull/1757#discussion_r1988400540
########## src/ast/mod.rs: ########## @@ -6148,12 +6218,12 @@ impl fmt::Display for GrantObjects { #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] -pub struct Assignment { +pub struct UpdateAssignment { Review Comment: is this rename necessary for the PR? ########## src/parser/mod.rs: ########## @@ -10955,134 +10955,226 @@ impl<'a> Parser<'a> { } else { Some(self.parse_identifier()?) }; - Ok(Statement::SetRole { + Ok(Statement::Set(Set::SetRole { context_modifier, role_name, - }) + })) } - pub fn parse_set(&mut self) -> Result<Statement, ParserError> { - 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 let Some(set_role_stmt) = - self.maybe_parse(|parser| parser.parse_set_role(modifier))? - { - return Ok(set_role_stmt); + fn parse_set_values( + &mut self, + parenthesized_assignment: bool, + ) -> Result<Vec<Expr>, ParserError> { + let mut values = vec![]; + + if parenthesized_assignment { + self.expect_token(&Token::LParen)?; + } + + loop { + let value = if let Some(expr) = self.try_parse_expr_sub_query()? { + expr + } else if let Ok(expr) = self.parse_expr() { + expr + } else { + self.expected("variable value", self.peek_token())? + }; + + values.push(value); + if self.consume_token(&Token::Comma) { + continue; + } + + if parenthesized_assignment { + self.expect_token(&Token::RParen)?; + } + return Ok(values); } + } - let variables = if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) { - OneOrManyWithParens::One(ObjectName::from(vec!["TIMEZONE".into()])) - } else if self.dialect.supports_parenthesized_set_variables() + fn parse_set_assignment( + &mut self, + ) -> Result<(OneOrManyWithParens<ObjectName>, Expr), ParserError> { + let variables = if self.dialect.supports_parenthesized_set_variables() && self.consume_token(&Token::LParen) { - let variables = OneOrManyWithParens::Many( + let vars = OneOrManyWithParens::Many( self.parse_comma_separated(|parser: &mut Parser<'a>| parser.parse_identifier())? .into_iter() .map(|ident| ObjectName::from(vec![ident])) .collect(), ); self.expect_token(&Token::RParen)?; - variables + vars } else { OneOrManyWithParens::One(self.parse_object_name(false)?) }; - let names = matches!(&variables, OneOrManyWithParens::One(variable) if variable.to_string().eq_ignore_ascii_case("NAMES")); + if !(self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO)) { + return self.expected("assignment operator", self.peek_token()); + } + + let values = self.parse_expr()?; + + Ok((variables, values)) + } + + fn parse_set(&mut self) -> Result<Statement, ParserError> { + let modifier = + self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]); - if names && self.dialect.supports_set_names() { + if let Some(Keyword::HIVEVAR) = modifier { + self.expect_token(&Token::Colon)?; + } + + if let Some(set_role_stmt) = self.maybe_parse(|parser| parser.parse_set_role(modifier))? { + return Ok(set_role_stmt); + } + + // Handle special cases first + if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) + || self.parse_keyword(Keyword::TIMEZONE) + { + if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) { + return Ok(Set::SingleAssignment { + local: modifier == Some(Keyword::LOCAL), + hivevar: modifier == Some(Keyword::HIVEVAR), + variable: ObjectName::from(vec!["TIMEZONE".into()]), + values: self.parse_set_values(false)?, + } + .into()); + } else if self.dialect.is::<PostgreSqlDialect>() { Review Comment: was there a requirement to gate this to the pg dialect? seems this is tied to the comment in the `parse_set_time_zone_alias` test, thinking we can probably allow the parser continue to work as before and reenable that test case? ########## src/parser/mod.rs: ########## @@ -10955,134 +10955,226 @@ impl<'a> Parser<'a> { } else { Some(self.parse_identifier()?) }; - Ok(Statement::SetRole { + Ok(Statement::Set(Set::SetRole { context_modifier, role_name, - }) + })) } - pub fn parse_set(&mut self) -> Result<Statement, ParserError> { - 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 let Some(set_role_stmt) = - self.maybe_parse(|parser| parser.parse_set_role(modifier))? - { - return Ok(set_role_stmt); + fn parse_set_values( + &mut self, + parenthesized_assignment: bool, + ) -> Result<Vec<Expr>, ParserError> { + let mut values = vec![]; + + if parenthesized_assignment { + self.expect_token(&Token::LParen)?; + } + + loop { + let value = if let Some(expr) = self.try_parse_expr_sub_query()? { + expr + } else if let Ok(expr) = self.parse_expr() { + expr + } else { + self.expected("variable value", self.peek_token())? + }; + + values.push(value); + if self.consume_token(&Token::Comma) { + continue; + } + + if parenthesized_assignment { + self.expect_token(&Token::RParen)?; + } + return Ok(values); } + } - let variables = if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) { - OneOrManyWithParens::One(ObjectName::from(vec!["TIMEZONE".into()])) - } else if self.dialect.supports_parenthesized_set_variables() + fn parse_set_assignment( + &mut self, + ) -> Result<(OneOrManyWithParens<ObjectName>, Expr), ParserError> { + let variables = if self.dialect.supports_parenthesized_set_variables() && self.consume_token(&Token::LParen) { - let variables = OneOrManyWithParens::Many( + let vars = OneOrManyWithParens::Many( self.parse_comma_separated(|parser: &mut Parser<'a>| parser.parse_identifier())? .into_iter() .map(|ident| ObjectName::from(vec![ident])) .collect(), ); self.expect_token(&Token::RParen)?; - variables + vars } else { OneOrManyWithParens::One(self.parse_object_name(false)?) }; - let names = matches!(&variables, OneOrManyWithParens::One(variable) if variable.to_string().eq_ignore_ascii_case("NAMES")); + if !(self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO)) { + return self.expected("assignment operator", self.peek_token()); + } + + let values = self.parse_expr()?; + + Ok((variables, values)) + } + + fn parse_set(&mut self) -> Result<Statement, ParserError> { + let modifier = + self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]); - if names && self.dialect.supports_set_names() { + if let Some(Keyword::HIVEVAR) = modifier { + self.expect_token(&Token::Colon)?; + } + + if let Some(set_role_stmt) = self.maybe_parse(|parser| parser.parse_set_role(modifier))? { + return Ok(set_role_stmt); + } + + // Handle special cases first + if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) + || self.parse_keyword(Keyword::TIMEZONE) + { + if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) { + return Ok(Set::SingleAssignment { + local: modifier == Some(Keyword::LOCAL), + hivevar: modifier == Some(Keyword::HIVEVAR), + variable: ObjectName::from(vec!["TIMEZONE".into()]), + values: self.parse_set_values(false)?, + } + .into()); + } else if self.dialect.is::<PostgreSqlDialect>() { + // Special case for Postgres + return Ok(Set::SetTimeZone { + local: modifier == Some(Keyword::LOCAL), + value: self.parse_expr()?, + } + .into()); + } else { + return self.expected("assignment operator", self.peek_token()); + } + } else if self.dialect.supports_set_names() && self.parse_keyword(Keyword::NAMES) { if self.parse_keyword(Keyword::DEFAULT) { - return Ok(Statement::SetNamesDefault {}); + return Ok(Set::SetNamesDefault {}.into()); } - let charset_name = self.parse_identifier()?; let collation_name = if self.parse_one_of_keywords(&[Keyword::COLLATE]).is_some() { Some(self.parse_literal_string()?) } else { None }; - return Ok(Statement::SetNames { + return Ok(Set::SetNames { charset_name, collation_name, - }); - } - - let parenthesized_assignment = matches!(&variables, OneOrManyWithParens::Many(_)); - - if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) { - if parenthesized_assignment { - self.expect_token(&Token::LParen)?; } - - let mut values = vec![]; - loop { - let value = if let Some(expr) = self.try_parse_expr_sub_query()? { - expr - } else if let Ok(expr) = self.parse_expr() { - expr - } else { - self.expected("variable value", self.peek_token())? - }; - - values.push(value); - if self.consume_token(&Token::Comma) { - continue; - } - - if parenthesized_assignment { - self.expect_token(&Token::RParen)?; - } - return Ok(Statement::SetVariable { - local: modifier == Some(Keyword::LOCAL), - hivevar: Some(Keyword::HIVEVAR) == modifier, - variables, - value: values, - }); - } - } - - let OneOrManyWithParens::One(variable) = variables else { - return self.expected("set variable", self.peek_token()); - }; - - if variable.to_string().eq_ignore_ascii_case("TIMEZONE") { - // for some db (e.g. postgresql), SET TIME ZONE <value> is an alias for SET TIMEZONE [TO|=] <value> - match self.parse_expr() { - Ok(expr) => Ok(Statement::SetTimeZone { - local: modifier == Some(Keyword::LOCAL), - value: expr, - }), - _ => self.expected("timezone value", self.peek_token())?, - } - } else if variable.to_string() == "CHARACTERISTICS" { + .into()); + } else if self.parse_keyword(Keyword::CHARACTERISTICS) { self.expect_keywords(&[Keyword::AS, Keyword::TRANSACTION])?; - Ok(Statement::SetTransaction { + return Ok(Set::SetTransaction { modes: self.parse_transaction_modes()?, snapshot: None, session: true, - }) - } else if variable.to_string() == "TRANSACTION" && modifier.is_none() { + } + .into()); + } else if self.parse_keyword(Keyword::TRANSACTION) { if self.parse_keyword(Keyword::SNAPSHOT) { let snapshot_id = self.parse_value()?.value; - return Ok(Statement::SetTransaction { + return Ok(Set::SetTransaction { modes: vec![], snapshot: Some(snapshot_id), session: false, - }); + } + .into()); } - Ok(Statement::SetTransaction { + return Ok(Set::SetTransaction { modes: self.parse_transaction_modes()?, snapshot: None, session: false, - }) - } else if self.dialect.supports_set_stmt_without_operator() { - self.prev_token(); - self.parse_set_session_params() + } + .into()); + } + + if self.dialect.supports_comma_separated_set_assignments() { + if let Ok(assignments) = + self.try_parse(|parser| parser.parse_comma_separated(Parser::parse_set_assignment)) Review Comment: we seem to be dropping the returned error here? should this use `self.maybe_parse` instead? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org