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::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 _ = self.parser.next_token();
28        let spstatement = self.parser.parse_set().context(error::SyntaxSnafu)?;
29        match spstatement {
30            SpStatement::SetVariable {
31                variables,
32                value,
33                hivevar,
34                ..
35            } if !hivevar => Ok(Statement::SetVariables(SetVariables {
36                variable: (*variables)[0].clone(),
37                value,
38            })),
39
40            SpStatement::SetTimeZone { value, .. } => Ok(Statement::SetVariables(SetVariables {
41                variable: ObjectName(vec![Ident::new("TIMEZONE")]),
42                value: vec![value],
43            })),
44
45            unexp => error::UnsupportedSnafu {
46                keyword: unexp.to_string(),
47            }
48            .fail(),
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use sqlparser::ast::{Expr, Ident, ObjectName, Value};
56
57    use super::*;
58    use crate::dialect::GreptimeDbDialect;
59    use crate::parser::ParseOptions;
60
61    fn assert_mysql_parse_result(sql: &str, indent_str: &str, expr: Expr) {
62        let result =
63            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
64        let mut stmts = result.unwrap();
65        assert_eq!(
66            stmts.pop().unwrap(),
67            Statement::SetVariables(SetVariables {
68                variable: ObjectName(vec![Ident::new(indent_str)]),
69                value: vec![expr]
70            })
71        );
72    }
73
74    fn assert_pg_parse_result(sql: &str, indent: &str, expr: Expr) {
75        let result =
76            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
77        let mut stmts = result.unwrap();
78        assert_eq!(
79            stmts.pop().unwrap(),
80            Statement::SetVariables(SetVariables {
81                variable: ObjectName(vec![Ident::new(indent)]),
82                value: vec![expr],
83            })
84        );
85    }
86
87    #[test]
88    pub fn test_set_timezone() {
89        let expected_utc_expr = Expr::Value(Value::SingleQuotedString("UTC".to_string()));
90        // mysql style
91        let sql = "SET time_zone = 'UTC'";
92        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
93        // session or local style
94        let sql = "SET LOCAL time_zone = 'UTC'";
95        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
96        let sql = "SET SESSION time_zone = 'UTC'";
97        assert_mysql_parse_result(sql, "time_zone", expected_utc_expr.clone());
98
99        // postgresql style
100        let sql = "SET TIMEZONE TO 'UTC'";
101        assert_pg_parse_result(sql, "TIMEZONE", expected_utc_expr.clone());
102        let sql = "SET TIMEZONE 'UTC'";
103        assert_pg_parse_result(sql, "TIMEZONE", expected_utc_expr);
104    }
105
106    #[test]
107    pub fn test_set_query_timeout() {
108        let expected_query_timeout_expr = Expr::Value(Value::Number("5000".to_string(), false));
109        // mysql style
110        let sql = "SET MAX_EXECUTION_TIME = 5000";
111        assert_mysql_parse_result(
112            sql,
113            "MAX_EXECUTION_TIME",
114            expected_query_timeout_expr.clone(),
115        );
116        // session or local style
117        let sql = "SET LOCAL MAX_EXECUTION_TIME = 5000";
118        assert_mysql_parse_result(
119            sql,
120            "MAX_EXECUTION_TIME",
121            expected_query_timeout_expr.clone(),
122        );
123        let sql = "SET SESSION MAX_EXECUTION_TIME = 5000";
124        assert_mysql_parse_result(
125            sql,
126            "MAX_EXECUTION_TIME",
127            expected_query_timeout_expr.clone(),
128        );
129
130        // postgresql style
131        let sql = "SET STATEMENT_TIMEOUT = 5000";
132        assert_pg_parse_result(
133            sql,
134            "STATEMENT_TIMEOUT",
135            expected_query_timeout_expr.clone(),
136        );
137        let sql = "SET STATEMENT_TIMEOUT TO 5000";
138        assert_pg_parse_result(sql, "STATEMENT_TIMEOUT", expected_query_timeout_expr);
139    }
140}
sql/parsers/set_var_parser.rs

sql/parsers/
set_var_parser.rs
