Skip to content

SET statements: scope modifier for multiple assignments #1772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 22, 2025
Merged

SET statements: scope modifier for multiple assignments #1772

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
de590fa
SET statements: scope modifier for multiple assignments
MohamedAbdeen21 Mar 20, 2025
94f5266
fix unit test error msg
MohamedAbdeen21 Mar 20, 2025
8a4f530
replace ContextModifier with Option and remove None variant
MohamedAbdeen21 Mar 21, 2025
d131c79
fix clippy
MohamedAbdeen21 Mar 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/ast/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5727,13 +5727,14 @@ impl fmt::Display for SequenceOptions {
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct SetAssignment {
pub scope: ContextModifier,
pub name: ObjectName,
pub value: Expr,
}

impl fmt::Display for SetAssignment {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{} = {}", self.name, self.value)
write!(f, "{}{} = {}", self.scope, self.name, self.value)
}
}

Expand Down
4 changes: 4 additions & 0 deletions src/dialect/generic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,4 +159,8 @@ impl Dialect for GenericDialect {
fn supports_set_names(&self) -> bool {
true
}

fn supports_comma_separated_set_assignments(&self) -> bool {
true
}
}
108 changes: 49 additions & 59 deletions src/parser/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11145,17 +11145,16 @@ impl<'a> Parser<'a> {
}

/// Parse a `SET ROLE` statement. Expects SET to be consumed already.
fn parse_set_role(&mut self, modifier: Option<Keyword>) -> Result<Statement, ParserError> {
fn parse_set_role(&mut self, modifier: ContextModifier) -> Result<Statement, ParserError> {
self.expect_keyword_is(Keyword::ROLE)?;
let context_modifier = Self::keyword_to_modifier(modifier);

let role_name = if self.parse_keyword(Keyword::NONE) {
None
} else {
Some(self.parse_identifier()?)
};
Ok(Statement::Set(Set::SetRole {
context_modifier,
context_modifier: modifier,
role_name,
}))
}
Expand Down Expand Up @@ -11191,46 +11190,52 @@ impl<'a> Parser<'a> {
}
}

fn parse_set_assignment(
&mut self,
) -> Result<(OneOrManyWithParens<ObjectName>, Expr), ParserError> {
let variables = if self.dialect.supports_parenthesized_set_variables()
fn parse_context_modifier(&mut self) -> ContextModifier {
let modifier =
self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::GLOBAL]);

Self::keyword_to_modifier(modifier)
}

/// Parse a single SET statement assignment `var = expr`.
fn parse_set_assignment(&mut self) -> Result<SetAssignment, ParserError> {
let scope = self.parse_context_modifier();

let name = if self.dialect.supports_parenthesized_set_variables()
&& self.consume_token(&Token::LParen)
{
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)?;
vars
// Parenthesized assignments are handled in the `parse_set` function after
// trying to parse list of assignments using this function.
// If a dialect supports both, and we find a LParen, we early exit from this function.
self.expected("Unparenthesized assignment", self.peek_token())?
} else {
OneOrManyWithParens::One(self.parse_object_name(false)?)
self.parse_object_name(false)?
};

if !(self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO)) {
return self.expected("assignment operator", self.peek_token());
}

let values = self.parse_expr()?;
let value = self.parse_expr()?;

Ok((variables, values))
Ok(SetAssignment { scope, name, value })
}

fn parse_set(&mut self) -> Result<Statement, ParserError> {
let modifier = self.parse_one_of_keywords(&[
Keyword::SESSION,
Keyword::LOCAL,
Keyword::HIVEVAR,
Keyword::GLOBAL,
]);

if let Some(Keyword::HIVEVAR) = modifier {
let hivevar = self.parse_keyword(Keyword::HIVEVAR);

// Modifier is either HIVEVAR: or a ContextModifier (LOCAL, SESSION, etc), not both
let scope = if !hivevar {
self.parse_context_modifier()
} else {
ContextModifier::None
};

if hivevar {
self.expect_token(&Token::Colon)?;
}

if let Some(set_role_stmt) = self.maybe_parse(|parser| parser.parse_set_role(modifier))? {
if let Some(set_role_stmt) = self.maybe_parse(|parser| parser.parse_set_role(scope))? {
return Ok(set_role_stmt);
}

Expand All @@ -11240,8 +11245,8 @@ impl<'a> Parser<'a> {
{
if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
return Ok(Set::SingleAssignment {
scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR),
scope,
hivevar,
variable: ObjectName::from(vec!["TIMEZONE".into()]),
values: self.parse_set_values(false)?,
}
Expand All @@ -11251,7 +11256,7 @@ impl<'a> Parser<'a> {
// the assignment operator. It's originally PostgreSQL specific,
// but we allow it for all the dialects
return Ok(Set::SetTimeZone {
local: modifier == Some(Keyword::LOCAL),
local: scope == ContextModifier::Local,
value: self.parse_expr()?,
}
.into());
Expand Down Expand Up @@ -11299,41 +11304,26 @@ impl<'a> Parser<'a> {
}

if self.dialect.supports_comma_separated_set_assignments() {
if scope != ContextModifier::None {
self.prev_token();
}

if let Some(assignments) = self
.maybe_parse(|parser| parser.parse_comma_separated(Parser::parse_set_assignment))?
{
return if assignments.len() > 1 {
let assignments = assignments
.into_iter()
.map(|(var, val)| match var {
OneOrManyWithParens::One(v) => Ok(SetAssignment {
name: v,
value: val,
}),
OneOrManyWithParens::Many(_) => {
self.expected("List of single identifiers", self.peek_token())
}
})
.collect::<Result<_, _>>()?;

Ok(Set::MultipleAssignments { assignments }.into())
} else {
let (vars, values): (Vec<_>, Vec<_>) = assignments.into_iter().unzip();

let variable = match vars.into_iter().next() {
Some(OneOrManyWithParens::One(v)) => Ok(v),
Some(OneOrManyWithParens::Many(_)) => self.expected(
"Single assignment or list of assignments",
self.peek_token(),
),
None => self.expected("At least one identifier", self.peek_token()),
}?;
let SetAssignment { scope, name, value } =
assignments.into_iter().next().ok_or_else(|| {
ParserError::ParserError("Expected at least one assignment".to_string())
})?;

Ok(Set::SingleAssignment {
scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR),
variable,
values,
scope,
hivevar,
variable: name,
values: vec![value],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I might be misunderstanding something, but I thought the reason this was a vec was to support some kind of SET var = a, b, c syntax, but I don't see that tested anywhere and don't remember where I saw it mentioned. If that syntax doesn't actually exist, could SingleAssignment be reworked to contain a SetAssignment or something?

Copy link
Contributor Author

@MohamedAbdeen21 MohamedAbdeen21 Mar 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe you're referring to the test called parse_set_hivevar:

#[test]
fn parse_set_hivevar() {
    let set = "SET HIVEVAR:name = a, b, c_d";
    hive().verified_stmt(set);
}

The highlighted line is inside the block that parses a list of assignments and should fail for stmts like name = a, b since b by itself is not a valid assignment. Therefore, for this specific line, we always receive a single value and have to wrap it in a vec.

However, the code blocks (i.e. parsing rules) after this one produce a vec of values.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For places where we produce a list of values, search for uses of parse_set_values.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, thank you!

}
.into())
};
Expand All @@ -11358,8 +11348,8 @@ impl<'a> Parser<'a> {
if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
let stmt = match variables {
OneOrManyWithParens::One(var) => Set::SingleAssignment {
scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR),
scope,
hivevar,
variable: var,
values: self.parse_set_values(false)?,
},
Expand Down
35 changes: 35 additions & 0 deletions tests/sqlparser_common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14856,10 +14856,12 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
assignments,
vec![
SetAssignment {
scope: ContextModifier::None,
name: ObjectName::from(vec!["@a".into()]),
value: Expr::value(number("1"))
},
SetAssignment {
scope: ContextModifier::None,
name: ObjectName::from(vec!["b".into()]),
value: Expr::value(number("2"))
}
Expand All @@ -14869,6 +14871,39 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
_ => panic!("Expected SetVariable with 2 variables and 2 values"),
};

let stmt = dialects.verified_stmt("SET GLOBAL @a = 1, SESSION b = 2, LOCAL c = 3, d = 4");

match stmt {
Statement::Set(Set::MultipleAssignments { assignments }) => {
assert_eq!(
assignments,
vec![
SetAssignment {
scope: ContextModifier::Global,
name: ObjectName::from(vec!["@a".into()]),
value: Expr::value(number("1"))
},
SetAssignment {
scope: ContextModifier::Session,
name: ObjectName::from(vec!["b".into()]),
value: Expr::value(number("2"))
},
SetAssignment {
scope: ContextModifier::Local,
name: ObjectName::from(vec!["c".into()]),
value: Expr::value(number("3"))
},
SetAssignment {
scope: ContextModifier::None,
name: ObjectName::from(vec!["d".into()]),
value: Expr::value(number("4"))
}
]
);
}
_ => panic!("Expected MultipleAssignments with 4 scoped variables and 4 values"),
};

Ok(())
}

Expand Down