fix: set local or session time_zone not work (#4064)

* fix: set local or session time_zone not work * chore: supports PostgreSQL-specific setting time zone
GreptimeTeam · May 29, 2024 · aafb468 · aafb468
1 parent 4aa756c
commit aafb468
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 17 deletions.
diff --git a/src/servers/src/mysql/federated.rs b/src/servers/src/mysql/federated.rs
@@ -72,7 +72,6 @@ static OTHER_NOT_SUPPORTED_STMT: Lazy<RegexSet> = Lazy::new(|| {
         "(?i)^(SELECT \\$\\$)",
 
         // mysqldump.
-        "(?i)^(SET SESSION(.*))",
         "(?i)^(SET SQL_QUOTE_SHOW_CREATE(.*))",
         "(?i)^(LOCK TABLES(.*))",
         "(?i)^(UNLOCK TABLES(.*))",

diff --git a/src/sql/src/parsers/set_var_parser.rs b/src/sql/src/parsers/set_var_parser.rs
@@ -15,6 +15,7 @@
 use snafu::ResultExt;
 use sqlparser::ast::Statement as SpStatement;
 
+use crate::ast::{Ident, ObjectName};
 use crate::error::{self, Result};
 use crate::parser::ParserContext;
 use crate::statements::set_variables::SetVariables;
@@ -29,11 +30,18 @@ impl<'a> ParserContext<'a> {
             SpStatement::SetVariable {
                 variable,
                 value,
-                local,
                 hivevar,
-            } if !local && !hivevar => {
-                Ok(Statement::SetVariables(SetVariables { variable, value }))
-            }
+                ..
+            } if !hivevar => Ok(Statement::SetVariables(SetVariables { variable, value })),
+
+            SpStatement::SetTimeZone { value, .. } => Ok(Statement::SetVariables(SetVariables {
+                variable: ObjectName(vec![Ident {
+                    value: "TIMEZONE".to_string(),
+                    quote_style: None,
+                }]),
+                value: vec![value],
+            })),
+
             unexp => error::UnsupportedSnafu {
                 sql: self.sql.to_string(),
                 keyword: unexp.to_string(),
@@ -51,10 +59,7 @@ mod tests {
     use crate::dialect::GreptimeDbDialect;
     use crate::parser::ParseOptions;
 
-    #[test]
-    pub fn test_set_timezone() {
-        // mysql style
-        let sql = "SET time_zone = 'UTC'";
+    fn assert_mysql_parse_result(sql: &str) {
         let result =
             ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
         let mut stmts = result.unwrap();
@@ -65,8 +70,9 @@ mod tests {
                 value: vec![Expr::Value(Value::SingleQuotedString("UTC".to_string()))]
             })
         );
-        // postgresql style
-        let sql = "SET TIMEZONE TO 'UTC'";
+    }
+
+    fn assert_pg_parse_result(sql: &str) {
         let result =
             ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
         let mut stmts = result.unwrap();
@@ -78,4 +84,22 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    pub fn test_set_timezone() {
+        // mysql style
+        let sql = "SET time_zone = 'UTC'";
+        assert_mysql_parse_result(sql);
+        // session or local style
+        let sql = "SET LOCAL time_zone = 'UTC'";
+        assert_mysql_parse_result(sql);
+        let sql = "SET SESSION time_zone = 'UTC'";
+        assert_mysql_parse_result(sql);
+
+        // postgresql style
+        let sql = "SET TIMEZONE TO 'UTC'";
+        assert_pg_parse_result(sql);
+        let sql = "SET TIMEZONE 'UTC'";
+        assert_pg_parse_result(sql);
+    }
 }
diff --git a/tests/cases/standalone/common/system/timezone.result b/tests/cases/standalone/common/system/timezone.result
@@ -187,7 +187,7 @@ select to_unixtime('2024-01-02 00:00:00+08:00');
 +------------------------------------------------+
 
 --- UTC-8 ---
-SET TIME_ZONE = '-8:00';
+SET SESSION TIME_ZONE = '-8:00';
 
 Affected Rows: 0
 
@@ -281,7 +281,7 @@ drop table test;
 Affected Rows: 0
 
 -- revert timezone to UTC
-SET TIME_ZONE = 'UTC';
+SET LOCAL TIME_ZONE = 'UTC';
 
 Affected Rows: 0
 

diff --git a/tests/cases/standalone/common/system/timezone.sql b/tests/cases/standalone/common/system/timezone.sql
@@ -48,7 +48,7 @@ select to_unixtime('2024-01-02 00:00:00');
 select to_unixtime('2024-01-02 00:00:00+08:00');
 
 --- UTC-8 ---
-SET TIME_ZONE = '-8:00';
+SET SESSION TIME_ZONE = '-8:00';
 
 SHOW VARIABLES time_zone;
 
@@ -71,7 +71,7 @@ select to_unixtime('2024-01-02 00:00:00+08:00');
 drop table test;
 
 -- revert timezone to UTC
-SET TIME_ZONE = 'UTC';
+SET LOCAL TIME_ZONE = 'UTC';
 
 SHOW VARIABLES time_zone;
 

diff --git a/tests/runner/src/env.rs b/tests/runner/src/env.rs
@@ -436,7 +436,9 @@ impl Database for GreptimeDB {
 
         let mut client = self.client.lock().await;
 
-        if query.trim().to_lowercase().starts_with("use ") {
+        let query_str = query.trim().to_lowercase();
+
+        if query_str.starts_with("use ") {
             // use [db]
             let database = query
                 .split_ascii_whitespace()
@@ -447,7 +449,10 @@ impl Database for GreptimeDB {
             Box::new(ResultDisplayer {
                 result: Ok(Output::new_with_affected_rows(0)),
             }) as _
-        } else if query.trim().to_lowercase().starts_with("set time_zone") {
+        } else if query_str.starts_with("set time_zone")
+            || query_str.starts_with("set session time_zone")
+            || query_str.starts_with("set local time_zone")
+        {
             // set time_zone='xxx'
             let timezone = query
                 .split('=')