feat(query): support unset session setting (#15510)

TCeason · BohuTANG · web-flow · commit 2cbff7957f3c · 2024-05-14T17:09:32.000+08:00
* feat(query): support unset session setting

Databend now supports the `unset &lt;settings&gt;` command, which is designed to remove global settings from the metasrv. We plan to expand this functionality with additional unset commands:

```sql
unset [session] &lt;setting&gt;; -- Only unset a setting at the session level.
```

For example:
```
-- load_file_metadata_expire_hours default is 248

set global load_file_metadata_expire_hours=12;
-- show settings like '%load_file_metadata_expire_hours%';
-- 12 global-level

set load_file_metadata_expire_hours=13;
-- show settings like '%load_file_metadata_expire_hours%';
-- 13 session-level

unset session load_file_metadata_expire_hours;
-- show settings like '%load_file_metadata_expire_hours%';
-- 12 global-level
```

* optimize code

* fix test

---------

Co-authored-by: Bohu &lt;overred.shuttler@gmail.com&gt;
diff --git a/src/query/ast/src/ast/format/ast_format.rs b/src/query/ast/src/ast/format/ast_format.rs
@@ -1048,7 +1048,11 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
     }
 
     fn visit_unset_variable(&mut self, stmt: &'ast UnSetStmt) {
-        let name = format!("UnSet {}", stmt);
+        let name = if stmt.session_level {
+            format!("UnSet SESSION {}", stmt)
+        } else {
+            format!("UnSet {}", stmt)
+        };
         let format_ctx = AstFormatContext::new(name);
         let node = FormatTreeNode::new(format_ctx);
         self.children.push(node);
diff --git a/src/query/ast/src/ast/statements/unset.rs b/src/query/ast/src/ast/statements/unset.rs
@@ -23,12 +23,17 @@ use crate::ast::Identifier;
 
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub struct UnSetStmt {
+    #[drive(skip)]
+    pub session_level: bool,
     pub source: UnSetSource,
 }
 
 impl Display for UnSetStmt {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         write!(f, "UNSET ")?;
+        if self.session_level {
+            write!(f, "SESSION ")?;
+        }
         write!(f, "{}", self.source)
     }
 }
diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -362,10 +362,11 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let unset_variable = map(
         rule! {
-            UNSET ~ #unset_source
+            UNSET ~ SESSION? ~ #unset_source
         },
-        |(_, unset_source)| {
+        |(_, opt_session_level, unset_source)| {
             Statement::UnSetVariable(UnSetStmt {
+                session_level: opt_session_level.is_some(),
                 source: unset_source,
             })
         },
diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -980,6 +980,8 @@ pub enum TokenKind {
     SET,
     #[token("UNSET", ignore(ascii_case))]
     UNSET,
+    #[token("SESSION", ignore(ascii_case))]
+    SESSION,
     #[token("SETTINGS", ignore(ascii_case))]
     SETTINGS,
     #[token("STAGES", ignore(ascii_case))]
diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs
@@ -569,7 +569,9 @@ fn test_statement() {
         r#"SET max_threads = 10;"#,
         r#"SET max_threads = 10*2;"#,
         r#"UNSET max_threads;"#,
+        r#"UNSET session max_threads;"#,
         r#"UNSET (max_threads, sql_dialect);"#,
+        r#"UNSET session (max_threads, sql_dialect);"#,
         r#"select $1 FROM '@my_stage/my data/'"#,
         r#"
             SELECT t.c1 FROM @stage1/dir/file
diff --git a/src/query/ast/tests/it/testdata/statement.txt b/src/query/ast/tests/it/testdata/statement.txt
@@ -14988,6 +14988,7 @@ UNSET max_threads
 ---------- AST ------------
 UnSetVariable(
     UnSetStmt {
+        session_level: false,
         source: Var {
             variable: Identifier {
                 span: Some(
@@ -15002,13 +15003,36 @@ UnSetVariable(
 )
 
 
+---------- Input ----------
+UNSET session max_threads;
+---------- Output ---------
+UNSET SESSION max_threads
+---------- AST ------------
+UnSetVariable(
+    UnSetStmt {
+        session_level: true,
+        source: Var {
+            variable: Identifier {
+                span: Some(
+                    14..25,
+                ),
+                name: "max_threads",
+                quote: None,
+                is_hole: false,
+            },
+        },
+    },
+)
+
+
 ---------- Input ----------
 UNSET (max_threads, sql_dialect);
 ---------- Output ---------
 UNSET (max_threads, sql_dialect)
 ---------- AST ------------
 UnSetVariable(
     UnSetStmt {
+        session_level: false,
         source: Vars {
             variables: [
                 Identifier {
@@ -15033,6 +15057,38 @@ UnSetVariable(
 )
 
 
+---------- Input ----------
+UNSET session (max_threads, sql_dialect);
+---------- Output ---------
+UNSET SESSION (max_threads, sql_dialect)
+---------- AST ------------
+UnSetVariable(
+    UnSetStmt {
+        session_level: true,
+        source: Vars {
+            variables: [
+                Identifier {
+                    span: Some(
+                        15..26,
+                    ),
+                    name: "max_threads",
+                    quote: None,
+                    is_hole: false,
+                },
+                Identifier {
+                    span: Some(
+                        28..39,
+                    ),
+                    name: "sql_dialect",
+                    quote: None,
+                    is_hole: false,
+                },
+            ],
+        },
+    },
+)
+
+
 ---------- Input ----------
 select $1 FROM '@my_stage/my data/'
 ---------- Output ---------
diff --git a/src/query/service/src/interpreters/interpreter_unsetting.rs b/src/query/service/src/interpreters/interpreter_unsetting.rs
@@ -17,6 +17,7 @@ use std::sync::Arc;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
 use databend_common_sql::plans::UnSettingPlan;
+use databend_common_users::UserApiProvider;
 
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
@@ -51,18 +52,29 @@ impl Interpreter for UnSettingInterpreter {
         let mut keys: Vec<String> = vec![];
         let mut values: Vec<String> = vec![];
         let mut is_globals: Vec<bool> = vec![];
-
         let settings = self.ctx.get_shared_settings();
+        let session_level = self.set.session_level;
+        settings.load_changes().await?;
+
+        // Fetch global settings asynchronously if necessary
+        let global_settings = if !session_level {
+            UserApiProvider::instance()
+                .setting_api(&self.ctx.get_tenant())
+                .get_settings()
+                .await?
+        } else {
+            Vec::new()
+        };
+
         for var in plan.vars {
             let (ok, value) = match var.to_lowercase().as_str() {
                 // To be compatible with some drivers
                 "sql_mode" | "autocommit" => (false, String::from("")),
                 setting_key => {
                     // TODO(liyz): why drop the global setting without checking the variable is global or not?
-                    self.ctx
-                        .get_shared_settings()
-                        .try_drop_global_setting(setting_key)
-                        .await?;
+                    if !session_level {
+                        settings.try_drop_global_setting(setting_key).await?;
+                    }
 
                     let default_val = match setting_key {
                         "max_memory_usage" => {
@@ -90,7 +102,23 @@ impl Interpreter for UnSettingInterpreter {
                             .map(|v| v.to_string()),
                     };
                     match default_val {
-                        Some(val) => (true, val),
+                        Some(val) => {
+                            let final_val = if global_settings.is_empty() {
+                                val.to_string()
+                            } else {
+                                global_settings
+                                    .iter()
+                                    .find(|setting| setting.name.to_lowercase() == setting_key)
+                                    .map_or(
+                                        {
+                                            self.ctx.get_shared_settings().unset_setting(&var);
+                                            val.to_string()
+                                        },
+                                        |setting| setting.value.to_string(),
+                                    )
+                            };
+                            (true, final_val)
+                        }
                         None => {
                             self.ctx
                                 .push_warning(format!("Unknown setting: '{}'", setting_key));
@@ -100,8 +128,6 @@ impl Interpreter for UnSettingInterpreter {
                 }
             };
             if ok {
-                // reset the current ctx settings, just remove it.
-                self.ctx.get_shared_settings().unset_setting(&var);
                 // set effect, this can be considered to be removed in the future.
                 keys.push(var);
                 values.push(value);
diff --git a/src/query/sql/src/planner/binder/setting.rs b/src/query/sql/src/planner/binder/setting.rs
@@ -75,17 +75,24 @@ impl Binder {
         _bind_context: &BindContext,
         stmt: &UnSetStmt,
     ) -> Result<Plan> {
+        let session_level = stmt.session_level;
         match &stmt.source {
             UnSetSource::Var { variable } => {
                 let vars = vec![variable.name.to_lowercase()];
-                Ok(Plan::UnSetVariable(Box::new(UnSettingPlan { vars })))
+                Ok(Plan::UnSetVariable(Box::new(UnSettingPlan {
+                    session_level,
+                    vars,
+                })))
             }
             UnSetSource::Vars { variables } => {
                 let vars: Vec<String> = variables
                     .iter()
                     .map(|var| var.name.to_lowercase())
                     .collect();
-                Ok(Plan::UnSetVariable(Box::new(UnSettingPlan { vars })))
+                Ok(Plan::UnSetVariable(Box::new(UnSettingPlan {
+                    session_level,
+                    vars,
+                })))
             }
         }
     }
diff --git a/src/query/sql/src/planner/plans/setting.rs b/src/query/sql/src/planner/plans/setting.rs
@@ -26,5 +26,6 @@ pub struct SettingPlan {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct UnSettingPlan {
+    pub session_level: bool,
     pub vars: Vec<String>,
 }
diff --git a/tests/sqllogictests/suites/base/05_ddl/05_0026_ddl_unset_settings.test b/tests/sqllogictests/suites/base/05_ddl/05_0026_ddl_unset_settings.test
@@ -1,41 +1,33 @@
-onlyif mysql
 query TTTT
 SELECT name, value, default, level from system.settings where name in ('sql_dialect', 'timezone')
 ----
 sql_dialect PostgreSQL PostgreSQL DEFAULT
 timezone UTC UTC DEFAULT
 
-onlyif mysql
 statement ok
 SET GLOBAL sql_dialect='MySQL'
 
-onlyif mysql
 statement ok
 SET GLOBAL timezone='Asia/Shanghai'
 
-onlyif mysql
 query TTTT
 SELECT name, value, default, level from system.settings where name in ('sql_dialect', 'timezone')
 ----
 sql_dialect  MySQL  PostgreSQL  GLOBAL
 timezone  Asia/Shanghai  UTC  GLOBAL
 
-onlyif mysql
 statement ok
 UNSET (timezone)
 
-onlyif mysql
 statement error 2801
 SET stl_dialect='MySQL'
 
 statement ok
 UNSET stl_dialect
 
-onlyif mysql
 statement ok
 UNSET sql_dialect
 
-onlyif mysql
 query TTTT
 SELECT name, value, default, level from system.settings where name in ('sql_dialect', 'timezone')
 ----
@@ -58,3 +50,35 @@ show settings like 'DATA_RETENTION_TIME_IN_DAYS';
 ----
 data_retention_time_in_days 1 1 [0, 90] DEFAULT Sets the data retention time in days. UInt64
 
+statement ok
+set global load_file_metadata_expire_hours=12;
+
+query TT
+show settings like 'load_file_metadata_expire_hours';
+----
+load_file_metadata_expire_hours 12 24 [0, 18446744073709551615] GLOBAL Sets the hours that the metadata of files you load data from with COPY INTO will expire in. UInt64
+
+statement ok
+set load_file_metadata_expire_hours=13;
+
+query TTTTTTT
+show settings like 'load_file_metadata_expire_hours';
+----
+load_file_metadata_expire_hours 13 24 [0, 18446744073709551615] SESSION Sets the hours that the metadata of files you load data from with COPY INTO will expire in. UInt64
+
+statement ok
+unset session load_file_metadata_expire_hours;
+
+query TTTTTTT
+show settings like 'load_file_metadata_expire_hours';
+----
+load_file_metadata_expire_hours 12 24 [0, 18446744073709551615] GLOBAL Sets the hours that the metadata of files you load data from with COPY INTO will expire in. UInt64
+
+statement ok
+unset load_file_metadata_expire_hours
+
+query TTTTTTT
+show settings like 'load_file_metadata_expire_hours';
+----
+load_file_metadata_expire_hours 24 24 [0, 18446744073709551615] DEFAULT Sets the hours that the metadata of files you load data from with COPY INTO will expire in. UInt64
+

Original file line number	Diff line number	Diff line change
`@@ -23,12 +23,17 @@ use crate::ast::Identifier;`
`23`	`23`
`24`	`24`	`#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]`
`25`	`25`	`pub struct UnSetStmt {`
	`26`	`+ #[drive(skip)]`
	`27`	`+ pub session_level: bool,`
`26`	`28`	`pub source: UnSetSource,`
`27`	`29`	`}`
`28`	`30`
`29`	`31`	`impl Display for UnSetStmt {`
`30`	`32`	`fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {`
`31`	`33`	`write!(f, "UNSET ")?;`
	`34`	`+ if self.session_level {`
	`35`	`+ write!(f, "SESSION ")?;`
	`36`	`+ }`
`32`	`37`	`write!(f, "{}", self.source)`
`33`	`38`	`}`
`34`	`39`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,5 +26,6 @@ pub struct SettingPlan {`
`26`	`26`
`27`	`27`	`#[derive(Clone, Debug, PartialEq, Eq)]`
`28`	`28`	`pub struct UnSettingPlan {`
	`29`	`+ pub session_level: bool,`
`29`	`30`	`pub vars: Vec<String>,`
`30`	`31`	`}`