sql/parsers/
set_var_parser.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use snafu::ResultExt;
16use sqlparser::ast::{Set, Statement as SpStatement};
17
18use crate::ast::{Ident, ObjectName};
19use crate::error::{self, Result};
20use crate::parser::ParserContext;
21use crate::statements::set_variables::SetVariables;
22use crate::statements::statement::Statement;
23
24/// SET variables statement parser implementation
25impl ParserContext<'_> {
26    pub(crate) fn parse_set_variables(&mut self) -> Result<Statement> {
27        let spstatement = self.parser.parse_statement().context(error::SyntaxSnafu)?;
28        match spstatement {
29            SpStatement::Set(set) => match set {
30                Set::SingleAssignment {
31                    scope: _,
32                    hivevar,
33                    variable,
34                    values,
35                } if !hivevar => Ok(Statement::SetVariables(SetVariables {
36                    variable,
37                    value: values,
38                })),
39
40                Set::SetTimeZone { local: _, value } => Ok(Statement::SetVariables(SetVariables {
41                    variable: ObjectName::from(vec![Ident::new("TIMEZONE")]),
42                    value: vec![value],
43                })),
44
45                set => error::UnsupportedSnafu {
46                    keyword: set.to_string(),
47                }
48                .fail(),
49            },
50            unexp => error::UnsupportedSnafu {
51                keyword: unexp.to_string(),
52            }
53            .fail(),
54        }
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use sqlparser::ast::{Expr, Ident, ObjectName, Value};
61
62    use super::*;
63    use crate::dialect::GreptimeDbDialect;
64    use crate::parser::ParseOptions;
65
66    fn assert_mysql_parse_result(sql: &str, indent_str: &str, expr: Expr) {
67        let result =
68            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
69        let mut stmts = result.unwrap();
70        assert_eq!(
71            stmts.pop().unwrap(),
72            Statement::SetVariables(SetVariables {
73                variable: ObjectName::from(vec![Ident::new(indent_str)]),
74                value: vec![expr]
75            })
76        );
77    }
78
79    fn assert_pg_parse_result(sql: &str, indent: &str, expr: Expr) {
80        let result =
81            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
82        let mut stmts = result.unwrap();
83        assert_eq!(
84            stmts.pop().unwrap(),
85            Statement::SetVariables(SetVariables {
86                variable: ObjectName::from(vec![Ident::new(indent)]),
87                value: vec![expr],
88            })
89        );
90    }
91
92    #[test]
93    pub fn test_set_timezone() {
94        let expected_utc_expr = Expr::Value(Value::SingleQuotedString("UTC".to_string()).into());
95        // mysql style
96        let sql = "SET time_zone = 'UTC'";
97        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
98        // session or local style
99        let sql = "SET LOCAL time_zone = 'UTC'";
100        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
101        let sql = "SET SESSION time_zone = 'UTC'";
102        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
103
104        // postgresql style
105        let sql = "SET TIMEZONE TO 'UTC'";
106        assert_pg_parse_result(sql, "TIMEZONE", expected_utc_expr.clone());
107        let sql = "SET TIMEZONE 'UTC'";
108        assert_pg_parse_result(sql, "TIMEZONE", expected_utc_expr);
109    }
110
111    #[test]
112    pub fn test_set_query_timeout() {
113        let expected_query_timeout_expr =
114            Expr::Value(Value::Number("5000".to_string(), false).into());
115        // mysql style
116        let sql = "SET MAX_EXECUTION_TIME = 5000";
117        assert_mysql_parse_result(
118            sql,
119            "MAX_EXECUTION_TIME",
120            expected_query_timeout_expr.clone(),
121        );
122        // session or local style
123        let sql = "SET LOCAL MAX_EXECUTION_TIME = 5000";
124        assert_mysql_parse_result(
125            sql,
126            "MAX_EXECUTION_TIME",
127            expected_query_timeout_expr.clone(),
128        );
129        let sql = "SET SESSION MAX_EXECUTION_TIME = 5000";
130        assert_mysql_parse_result(
131            sql,
132            "MAX_EXECUTION_TIME",
133            expected_query_timeout_expr.clone(),
134        );
135
136        // postgresql style
137        let sql = "SET STATEMENT_TIMEOUT = 5000";
138        assert_pg_parse_result(
139            sql,
140            "STATEMENT_TIMEOUT",
141            expected_query_timeout_expr.clone(),
142        );
143        let sql = "SET STATEMENT_TIMEOUT TO 5000";
144        assert_pg_parse_result(sql, "STATEMENT_TIMEOUT", expected_query_timeout_expr);
145    }
146
147    #[test]
148    fn test_unsupported_set_variant_remains_rejected() {
149        let result = ParserContext::create_with_dialect(
150            "SET ROLE admin",
151            &GreptimeDbDialect {},
152            ParseOptions::default(),
153        );
154        assert!(result.is_err());
155    }
156}
sql/parsers/set_var_parser.rs

sql/parsers/
set_var_parser.rs
