mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-23 15:34:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add GLOBAL context/modifier to SET statements (#1767)

Browse source
This commit is contained in:
Mohamed Abdeen 2025-03-20 07:52:56 +02:00 • committed by GitHub
parent e3e88290cd
commit f487cbe004
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 72 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/ast/mod.rs
View file

@ -2629,7 +2629,7 @@ pub enum Set {
/// SQL Standard-style /// SQL Standard-style
/// SET a = 1; /// SET a = 1;
SingleAssignment { SingleAssignment {
local: bool, scope: ContextModifier,
hivevar: bool, hivevar: bool,
variable: ObjectName, variable: ObjectName,
values: Vec<Expr>, values: Vec<Expr>,
@ -2711,7 +2711,7 @@ impl Display for Set {
role_name, role_name,
} => { } => {
let role_name = role_name.clone().unwrap_or_else(|| Ident::new("NONE")); let role_name = role_name.clone().unwrap_or_else(|| Ident::new("NONE"));
write!(f, "SET{context_modifier} ROLE {role_name}") write!(f, "SET {context_modifier}ROLE {role_name}")
} }
Self::SetSessionParam(kind) => write!(f, "SET {kind}"), Self::SetSessionParam(kind) => write!(f, "SET {kind}"),
Self::SetTransaction { Self::SetTransaction {
@ -2758,7 +2758,7 @@ impl Display for Set {
Ok(()) Ok(())
} }
Set::SingleAssignment { Set::SingleAssignment {
local, scope,
hivevar, hivevar,
variable, variable,
values, values,
@ -2766,7 +2766,7 @@ impl Display for Set {
write!( write!(
f, f,
"SET {}{}{} = {}", "SET {}{}{} = {}",
if *local { "LOCAL " } else { "" }, scope,
if *hivevar { "HIVEVAR:" } else { "" }, if *hivevar { "HIVEVAR:" } else { "" },
variable, variable,
display_comma_separated(values) display_comma_separated(values)
@ -7955,7 +7955,7 @@ impl fmt::Display for FlushLocation {
} }
} }
/// Optional context modifier for statements that can be or `LOCAL`, or `SESSION`. /// Optional context modifier for statements that can be or `LOCAL`, `GLOBAL`, or `SESSION`.
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@ -7966,6 +7966,8 @@ pub enum ContextModifier {
Local, Local,
/// `SESSION` identifier /// `SESSION` identifier
Session, Session,
/// `GLOBAL` identifier
Global,
} }
impl fmt::Display for ContextModifier { impl fmt::Display for ContextModifier {
@ -7975,10 +7977,13 @@ impl fmt::Display for ContextModifier {
write!(f, "") write!(f, "")
} }
Self::Local => { Self::Local => {
write!(f, " LOCAL") write!(f, "LOCAL ")
} }
Self::Session => { Self::Session => {
write!(f, " SESSION") write!(f, "SESSION ")
}
Self::Global => {
write!(f, "GLOBAL ")
} }
} }
} }

29
src/parser/mod.rs
View file

@ -1819,6 +1819,15 @@ impl<'a> Parser<'a> {
}) })
} }
fn keyword_to_modifier(k: Option<Keyword>) -> ContextModifier {
match k {
Some(Keyword::LOCAL) => ContextModifier::Local,
Some(Keyword::GLOBAL) => ContextModifier::Global,
Some(Keyword::SESSION) => ContextModifier::Session,
_ => ContextModifier::None,
}
}
/// Check if the root is an identifier and all fields are identifiers. /// Check if the root is an identifier and all fields are identifiers.
fn is_all_ident(root: &Expr, fields: &[AccessExpr]) -> bool { fn is_all_ident(root: &Expr, fields: &[AccessExpr]) -> bool {
if !matches!(root, Expr::Identifier(_)) { if !matches!(root, Expr::Identifier(_)) {
@ -11138,11 +11147,7 @@ impl<'a> Parser<'a> {
/// Parse a `SET ROLE` statement. Expects SET to be consumed already. /// 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: Option<Keyword>) -> Result<Statement, ParserError> {
self.expect_keyword_is(Keyword::ROLE)?; self.expect_keyword_is(Keyword::ROLE)?;
let context_modifier = match modifier { let context_modifier = Self::keyword_to_modifier(modifier);
Some(Keyword::LOCAL) => ContextModifier::Local,
Some(Keyword::SESSION) => ContextModifier::Session,
_ => ContextModifier::None,
};
let role_name = if self.parse_keyword(Keyword::NONE) { let role_name = if self.parse_keyword(Keyword::NONE) {
None None
@ -11214,8 +11219,12 @@ impl<'a> Parser<'a> {
} }
fn parse_set(&mut self) -> Result<Statement, ParserError> { fn parse_set(&mut self) -> Result<Statement, ParserError> {
let modifier = let modifier = self.parse_one_of_keywords(&[
self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]); Keyword::SESSION,
Keyword::LOCAL,
Keyword::HIVEVAR,
Keyword::GLOBAL,
]);
if let Some(Keyword::HIVEVAR) = modifier { if let Some(Keyword::HIVEVAR) = modifier {
self.expect_token(&Token::Colon)?; self.expect_token(&Token::Colon)?;
@ -11231,7 +11240,7 @@ impl<'a> Parser<'a> {
{ {
if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) { if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
return Ok(Set::SingleAssignment { return Ok(Set::SingleAssignment {
local: modifier == Some(Keyword::LOCAL), scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR), hivevar: modifier == Some(Keyword::HIVEVAR),
variable: ObjectName::from(vec!["TIMEZONE".into()]), variable: ObjectName::from(vec!["TIMEZONE".into()]),
values: self.parse_set_values(false)?, values: self.parse_set_values(false)?,
@ -11321,7 +11330,7 @@ impl<'a> Parser<'a> {
}?; }?;
Ok(Set::SingleAssignment { Ok(Set::SingleAssignment {
local: modifier == Some(Keyword::LOCAL), scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR), hivevar: modifier == Some(Keyword::HIVEVAR),
variable, variable,
values, values,
@ -11349,7 +11358,7 @@ impl<'a> Parser<'a> {
if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) { if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
let stmt = match variables { let stmt = match variables {
OneOrManyWithParens::One(var) => Set::SingleAssignment { OneOrManyWithParens::One(var) => Set::SingleAssignment {
local: modifier == Some(Keyword::LOCAL), scope: Self::keyword_to_modifier(modifier),
hivevar: modifier == Some(Keyword::HIVEVAR), hivevar: modifier == Some(Keyword::HIVEVAR),
variable: var, variable: var,
values: self.parse_set_values(false)?, values: self.parse_set_values(false)?,

32
tests/sqlparser_common.rs
View file

@ -8627,12 +8627,12 @@ fn parse_set_transaction() {
fn parse_set_variable() { fn parse_set_variable() {
match verified_stmt("SET SOMETHING = '1'") { match verified_stmt("SET SOMETHING = '1'") {
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local, scope,
hivevar, hivevar,
variable, variable,
values, values,
}) => { }) => {
assert!(!local); assert_eq!(scope, ContextModifier::None);
assert!(!hivevar); assert!(!hivevar);
assert_eq!(variable, ObjectName::from(vec!["SOMETHING".into()])); assert_eq!(variable, ObjectName::from(vec!["SOMETHING".into()]));
assert_eq!( assert_eq!(
@ -8645,6 +8645,26 @@ fn parse_set_variable() {
_ => unreachable!(), _ => unreachable!(),
} }
match verified_stmt("SET GLOBAL VARIABLE = 'Value'") {
Statement::Set(Set::SingleAssignment {
scope,
hivevar,
variable,
values,
}) => {
assert_eq!(scope, ContextModifier::Global);
assert!(!hivevar);
assert_eq!(variable, ObjectName::from(vec!["VARIABLE".into()]));
assert_eq!(
values,
vec![Expr::Value(
(Value::SingleQuotedString("Value".into())).with_empty_span()
)]
);
}
_ => unreachable!(),
}
let multi_variable_dialects = all_dialects_where(|d| d.supports_parenthesized_set_variables()); let multi_variable_dialects = all_dialects_where(|d| d.supports_parenthesized_set_variables());
let sql = r#"SET (a, b, c) = (1, 2, 3)"#; let sql = r#"SET (a, b, c) = (1, 2, 3)"#;
match multi_variable_dialects.verified_stmt(sql) { match multi_variable_dialects.verified_stmt(sql) {
@ -8719,12 +8739,12 @@ fn parse_set_variable() {
fn parse_set_role_as_variable() { fn parse_set_role_as_variable() {
match verified_stmt("SET role = 'foobar'") { match verified_stmt("SET role = 'foobar'") {
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local, scope,
hivevar, hivevar,
variable, variable,
values, values,
}) => { }) => {
assert!(!local); assert_eq!(scope, ContextModifier::None);
assert!(!hivevar); assert!(!hivevar);
assert_eq!(variable, ObjectName::from(vec!["role".into()])); assert_eq!(variable, ObjectName::from(vec!["role".into()]));
assert_eq!( assert_eq!(
@ -8766,12 +8786,12 @@ fn parse_double_colon_cast_at_timezone() {
fn parse_set_time_zone() { fn parse_set_time_zone() {
match verified_stmt("SET TIMEZONE = 'UTC'") { match verified_stmt("SET TIMEZONE = 'UTC'") {
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local, scope,
hivevar, hivevar,
variable, variable,
values, values,
}) => { }) => {
assert!(!local); assert_eq!(scope, ContextModifier::None);
assert!(!hivevar); assert!(!hivevar);
assert_eq!(variable, ObjectName::from(vec!["TIMEZONE".into()])); assert_eq!(variable, ObjectName::from(vec!["TIMEZONE".into()]));
assert_eq!( assert_eq!(

9
tests/sqlparser_hive.rs
View file

@ -21,9 +21,10 @@
//! is also tested (on the inputs it can handle). //! is also tested (on the inputs it can handle).
use sqlparser::ast::{ use sqlparser::ast::{
ClusteredBy, CommentDef, CreateFunction, CreateFunctionBody, CreateFunctionUsing, CreateTable, ClusteredBy, CommentDef, ContextModifier, CreateFunction, CreateFunctionBody,
Expr, Function, FunctionArgumentList, FunctionArguments, Ident, ObjectName, OrderByExpr, CreateFunctionUsing, CreateTable, Expr, Function, FunctionArgumentList, FunctionArguments,
OrderByOptions, SelectItem, Set, Statement, TableFactor, UnaryOperator, Use, Value, Ident, ObjectName, OrderByExpr, OrderByOptions, SelectItem, Set, Statement, TableFactor,
UnaryOperator, Use, Value,
}; };
use sqlparser::dialect::{GenericDialect, HiveDialect, MsSqlDialect}; use sqlparser::dialect::{GenericDialect, HiveDialect, MsSqlDialect};
use sqlparser::parser::ParserError; use sqlparser::parser::ParserError;
@ -369,7 +370,7 @@ fn set_statement_with_minus() {
assert_eq!( assert_eq!(
hive().verified_stmt("SET hive.tez.java.opts = -Xmx4g"), hive().verified_stmt("SET hive.tez.java.opts = -Xmx4g"),
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![ variable: ObjectName::from(vec![
Ident::new("hive"), Ident::new("hive"),

2
tests/sqlparser_mssql.rs
View file

@ -1251,7 +1251,7 @@ fn parse_mssql_declare() {
}] }]
}, },
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("@bar")]), variable: ObjectName::from(vec![Ident::new("@bar")]),
values: vec![Expr::Value( values: vec![Expr::Value(

2
tests/sqlparser_mysql.rs
View file

@ -618,7 +618,7 @@ fn parse_set_variables() {
assert_eq!( assert_eq!(
mysql_and_generic().verified_stmt("SET LOCAL autocommit = 1"), mysql_and_generic().verified_stmt("SET LOCAL autocommit = 1"),
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: true, scope: ContextModifier::Local,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec!["autocommit".into()]), variable: ObjectName::from(vec!["autocommit".into()]),
values: vec![Expr::value(number("1"))], values: vec![Expr::value(number("1"))],

18
tests/sqlparser_postgres.rs
View file

@ -988,8 +988,7 @@ fn parse_create_schema_if_not_exists() {
Statement::CreateSchema { Statement::CreateSchema {
if_not_exists: true, if_not_exists: true,
schema_name, schema_name,
options: _, ..
default_collate_spec: _,
} => assert_eq!("schema_name", schema_name.to_string()), } => assert_eq!("schema_name", schema_name.to_string()),
_ => unreachable!(), _ => unreachable!(),
} }
@ -1433,7 +1432,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a")]), variable: ObjectName::from(vec![Ident::new("a")]),
values: vec![Expr::Identifier(Ident { values: vec![Expr::Identifier(Ident {
@ -1448,7 +1447,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a")]), variable: ObjectName::from(vec![Ident::new("a")]),
values: vec![Expr::Value( values: vec![Expr::Value(
@ -1461,7 +1460,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a")]), variable: ObjectName::from(vec![Ident::new("a")]),
values: vec![Expr::value(number("0"))], values: vec![Expr::value(number("0"))],
@ -1472,7 +1471,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a")]), variable: ObjectName::from(vec![Ident::new("a")]),
values: vec![Expr::Identifier(Ident::new("DEFAULT"))], values: vec![Expr::Identifier(Ident::new("DEFAULT"))],
@ -1483,7 +1482,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: true, scope: ContextModifier::Local,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a")]), variable: ObjectName::from(vec![Ident::new("a")]),
values: vec![Expr::Identifier("b".into())], values: vec![Expr::Identifier("b".into())],
@ -1494,7 +1493,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![Ident::new("a"), Ident::new("b"), Ident::new("c")]), variable: ObjectName::from(vec![Ident::new("a"), Ident::new("b"), Ident::new("c")]),
values: vec![Expr::Identifier(Ident { values: vec![Expr::Identifier(Ident {
@ -1512,7 +1511,7 @@ fn parse_set() {
assert_eq!( assert_eq!(
stmt, stmt,
Statement::Set(Set::SingleAssignment { Statement::Set(Set::SingleAssignment {
local: false, scope: ContextModifier::None,
hivevar: false, hivevar: false,
variable: ObjectName::from(vec![ variable: ObjectName::from(vec![
Ident::new("hive"), Ident::new("hive"),
@ -1526,7 +1525,6 @@ fn parse_set() {
); );
pg_and_generic().one_statement_parses_to("SET a TO b", "SET a = b"); pg_and_generic().one_statement_parses_to("SET a TO b", "SET a = b");
pg_and_generic().one_statement_parses_to("SET SESSION a = b", "SET a = b");
assert_eq!( assert_eq!(
pg_and_generic().parse_sql_statements("SET"), pg_and_generic().parse_sql_statements("SET"),