refactor: Unify the set and unset structures (#16622)

In order to simplify the query setting implementation
databendlabs · Oct 16, 2024 · f147d3d · f147d3d
1 parent b46195b
commit f147d3d
Showing 9 changed files with 489 additions and 452 deletions.
diff --git a/src/query/ast/src/ast/statements/mod.rs b/src/query/ast/src/ast/statements/mod.rs
@@ -42,6 +42,7 @@ mod replace;
 mod script;
 mod sequence;
 mod set;
+mod settings;
 mod show;
 mod stage;
 mod statement;
@@ -85,6 +86,7 @@ pub use replace::*;
 pub use script::*;
 pub use sequence::*;
 pub use set::*;
+pub use settings::*;
 pub use show::*;
 pub use stage::*;
 pub use statement::*;

diff --git a/src/query/ast/src/ast/statements/set.rs b/src/query/ast/src/ast/statements/set.rs
@@ -32,4 +32,6 @@ pub enum SetType {
 pub enum SetValues {
     Expr(Vec<Box<Expr>>),
     Query(Box<Query>),
+    // None means Unset Stmt
+    None,
 }
diff --git a/src/query/ast/src/ast/statements/settings.rs b/src/query/ast/src/ast/statements/settings.rs
@@ -0,0 +1,77 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use derive_visitor::Drive;
+use derive_visitor::DriveMut;
+
+use crate::ast::Identifier;
+use crate::ast::SetType;
+use crate::ast::SetValues;
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct Settings {
+    pub set_type: SetType,
+    pub identifiers: Vec<Identifier>,
+    pub values: SetValues,
+}
+
+impl Display for Settings {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self.set_type {
+            SetType::SettingsGlobal => write!(f, "GLOBAL ")?,
+            SetType::SettingsSession => write!(f, "SESSION ")?,
+            SetType::Variable => write!(f, "VARIABLE ")?,
+        }
+
+        if self.identifiers.len() > 1 {
+            write!(f, "(")?;
+        }
+        for (idx, variable) in self.identifiers.iter().enumerate() {
+            if idx > 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{variable}")?;
+        }
+        if self.identifiers.len() > 1 {
+            write!(f, ")")?;
+        }
+
+        match &self.values {
+            SetValues::Expr(exprs) => {
+                write!(f, " = ")?;
+                if exprs.len() > 1 {
+                    write!(f, "(")?;
+                }
+
+                for (idx, value) in exprs.iter().enumerate() {
+                    if idx > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "{value}")?;
+                }
+                if exprs.len() > 1 {
+                    write!(f, ")")?;
+                }
+            }
+            SetValues::Query(query) => {
+                write!(f, " = {query}")?;
+            }
+            SetValues::None => {}
+        }
+        Ok(())
+    }
+}
diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs
@@ -27,8 +27,8 @@ use super::*;
 use crate::ast::quote::QuotedString;
 use crate::ast::statements::connection::CreateConnectionStmt;
 use crate::ast::statements::pipe::CreatePipeStmt;
+use crate::ast::statements::settings::Settings;
 use crate::ast::statements::task::CreateTaskStmt;
-use crate::ast::write_comma_separated_list;
 use crate::ast::CreateOption;
 use crate::ast::Identifier;
 use crate::ast::Query;
@@ -86,14 +86,11 @@ pub enum Statement {
     },
 
     SetStmt {
-        set_type: SetType,
-        identifiers: Vec<Identifier>,
-        values: SetValues,
+        settings: Settings,
     },
 
     UnSetStmt {
-        unset_type: SetType,
-        identifiers: Vec<Identifier>,
+        settings: Settings,
     },
 
     ShowVariables {
@@ -496,71 +493,8 @@ impl Display for Statement {
                 }
                 write!(f, " '{object_id}'")?;
             }
-            Statement::SetStmt {
-                set_type,
-                identifiers,
-                values,
-            } => {
-                write!(f, "SET ")?;
-                match *set_type {
-                    SetType::SettingsGlobal => write!(f, "GLOBAL ")?,
-                    SetType::SettingsSession => {}
-                    SetType::Variable => write!(f, "VARIABLE ")?,
-                }
-
-                if identifiers.len() > 1 {
-                    write!(f, "(")?;
-                }
-                for (idx, variable) in identifiers.iter().enumerate() {
-                    if idx > 0 {
-                        write!(f, ", ")?;
-                    }
-                    write!(f, "{variable}")?;
-                }
-                if identifiers.len() > 1 {
-                    write!(f, ")")?;
-                }
-
-                match values {
-                    SetValues::Expr(exprs) => {
-                        write!(f, " = ")?;
-                        if exprs.len() > 1 {
-                            write!(f, "(")?;
-                        }
-
-                        for (idx, value) in exprs.iter().enumerate() {
-                            if idx > 0 {
-                                write!(f, ", ")?;
-                            }
-                            write!(f, "{value}")?;
-                        }
-                        if exprs.len() > 1 {
-                            write!(f, ")")?;
-                        }
-                    }
-                    SetValues::Query(query) => {
-                        write!(f, " = {query}")?;
-                    }
-                }
-            }
-            Statement::UnSetStmt {
-                unset_type,
-                identifiers,
-            } => {
-                write!(f, "UNSET ")?;
-                match *unset_type {
-                    SetType::SettingsSession => write!(f, "SESSION ")?,
-                    SetType::SettingsGlobal => write!(f, "GLOBAL ")?,
-                    SetType::Variable => write!(f, "VARIABLE ")?,
-                }
-                if identifiers.len() == 1 {
-                    write!(f, "{}", identifiers[0])?;
-                } else {
-                    write!(f, "(")?;
-                    write_comma_separated_list(f, identifiers)?;
-                    write!(f, ")")?;
-                }
-            }
+            Statement::SetStmt { settings } => write!(f, "SET {}", settings)?,
+            Statement::UnSetStmt { settings } => write!(f, "UNSET {}", settings)?,
             Statement::SetRole {
                 is_default,
                 role_name,

diff --git a/src/query/ast/src/parser/common.rs b/src/query/ast/src/parser/common.rs
@@ -244,26 +244,9 @@ pub fn set_type(i: Input) -> IResult<SetType> {
         },
         |res| match res {
             Some(token) => match token.kind {
-                TokenKind::GLOBAL => SetType::SettingsGlobal,
-                TokenKind::SESSION => SetType::SettingsSession,
-                TokenKind::VARIABLE => SetType::Variable,
-                _ => unreachable!(),
-            },
-            None => SetType::SettingsSession,
-        },
-    )(i)
-}
-
-pub fn unset_type(i: Input) -> IResult<SetType> {
-    map(
-        rule! {
-           (GLOBAL | SESSION | VARIABLE)?
-        },
-        |res| match res {
-            Some(token) => match token.kind {
-                TokenKind::GLOBAL => SetType::SettingsGlobal,
-                TokenKind::SESSION => SetType::SettingsSession,
-                TokenKind::VARIABLE => SetType::Variable,
+                GLOBAL => SetType::SettingsGlobal,
+                SESSION => SetType::SettingsSession,
+                VARIABLE => SetType::Variable,
                 _ => unreachable!(),
             },
             None => SetType::SettingsSession,

diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -351,11 +351,14 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let unset_stmt = map(
         rule! {
-            UNSET ~ #unset_type ~ #unset_source
+            UNSET ~ #set_type ~ #unset_source
         },
         |(_, unset_type, identifiers)| Statement::UnSetStmt {
-            unset_type,
-            identifiers,
+            settings: Settings {
+                set_type: unset_type,
+                identifiers,
+                values: SetValues::None,
+            },
         },
     );
 
@@ -389,9 +392,11 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
                 SET ~ #set_type ~ #ident ~ "=" ~ #subexpr(0)
             },
             |(_, set_type, var, _, value)| Statement::SetStmt {
-                set_type,
-                identifiers: vec![var],
-                values: SetValues::Expr(vec![Box::new(value)]),
+                settings: Settings {
+                    set_type,
+                    identifiers: vec![var],
+                    values: SetValues::Expr(vec![Box::new(value)]),
+                },
             },
         ),
         map_res(
@@ -402,9 +407,11 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             |(_, set_type, _, ids, _, _, _, values, _)| {
                 if ids.len() == values.len() {
                     Ok(Statement::SetStmt {
-                        set_type,
-                        identifiers: ids,
-                        values: SetValues::Expr(values.into_iter().map(|x| x.into()).collect()),
+                        settings: Settings {
+                            set_type,
+                            identifiers: ids,
+                            values: SetValues::Expr(values.into_iter().map(|x| x.into()).collect()),
+                        },
                     })
                 } else {
                     Err(nom::Err::Failure(ErrorKind::Other(
@@ -418,19 +425,23 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
                 SET ~ #set_type ~ #ident ~ "=" ~ #query
             },
             |(_, set_type, var, _, query)| Statement::SetStmt {
-                set_type,
-                identifiers: vec![var],
-                values: SetValues::Query(Box::new(query)),
+                settings: Settings {
+                    set_type,
+                    identifiers: vec![var],
+                    values: SetValues::Query(Box::new(query)),
+                },
             },
         ),
         map(
             rule! {
                 SET ~ #set_type ~ "(" ~ #comma_separated_list0(ident) ~ ")" ~ "=" ~ #query
             },
             |(_, set_type, _, vars, _, _, query)| Statement::SetStmt {
-                set_type,
-                identifiers: vars,
-                values: SetValues::Query(Box::new(query)),
+                settings: Settings {
+                    set_type,
+                    identifiers: vars,
+                    values: SetValues::Query(Box::new(query)),
+                },
             },
         ),
     ));

diff --git a/src/query/ast/tests/it/testdata/stmt.txt b/src/query/ast/tests/it/testdata/stmt.txt
diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs
@@ -21,6 +21,7 @@ use std::time::Instant;
 use chrono_tz::Tz;
 use databend_common_ast::ast::Hint;
 use databend_common_ast::ast::Identifier;
+use databend_common_ast::ast::Settings;
 use databend_common_ast::ast::Statement;
 use databend_common_ast::parser::parse_sql;
 use databend_common_ast::parser::tokenize_sql;
@@ -470,13 +471,15 @@ impl<'a> Binder {
 
             Statement::Presign(stmt) => self.bind_presign(bind_context, stmt).await?,
 
-            Statement::SetStmt {set_type, identifiers, values } => {
+            Statement::SetStmt { settings } => {
+                let Settings { set_type, identifiers, values } = settings;
                 self.bind_set(bind_context, *set_type, identifiers, values)
                     .await?
             }
 
-            Statement::UnSetStmt{unset_type, identifiers } => {
-                self.bind_unset(bind_context, *unset_type, identifiers)
+            Statement::UnSetStmt{settings } => {
+                let Settings { set_type, identifiers, .. } = settings;
+                self.bind_unset(bind_context, *set_type, identifiers)
                     .await?
             }
 

diff --git a/src/query/sql/src/planner/binder/set.rs b/src/query/sql/src/planner/binder/set.rs
@@ -71,6 +71,7 @@ impl Binder {
                 let p = self.clone().bind(&Statement::Query(query.clone())).await?;
                 SetScalarsOrQuery::Query(Box::new(p))
             }
+            SetValues::None => return Err(ErrorCode::SemanticError("set value can not be None")),
         };
 
         Ok(Plan::Set(Box::new(SetPlan {