sql/parsers/
admin_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;
16
17use crate::ast::Expr;
18use crate::error::{Result, SyntaxSnafu};
19use crate::parser::ParserContext;
20use crate::statements::admin::Admin;
21use crate::statements::statement::Statement;
22
23/// `admin` extension parser: `admin function(arg1, arg2, ...)`
24/// or `admin function`
25impl ParserContext<'_> {
26    /// Parse `admin function(arg1, arg2, ...)` or `admin function` statement
27    pub(crate) fn parse_admin_command(&mut self) -> Result<Statement> {
28        let _token = self.parser.next_token();
29
30        let object_name = self.parser.parse_object_name(false).context(SyntaxSnafu)?;
31
32        let func = match self
33            .parser
34            .parse_function(object_name)
35            .context(SyntaxSnafu)?
36        {
37            Expr::Function(f) => f,
38            _ => {
39                return self.unsupported(self.peek_token_as_string());
40            }
41        };
42
43        Ok(Statement::Admin(Admin::Func(func)))
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use sqlparser::ast::FunctionArguments;
50
51    use super::*;
52    use crate::ast::{Expr, Function, FunctionArg, FunctionArgExpr, Value};
53    use crate::dialect::GreptimeDbDialect;
54    use crate::parser::ParseOptions;
55
56    #[test]
57    fn test_parse_admin_function() {
58        let sql = "ADMIN flush_table('test')";
59
60        let mut result =
61            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
62                .unwrap();
63        assert_eq!(1, result.len());
64        let stmt = result.remove(0);
65        match &stmt {
66            Statement::Admin(Admin::Func(Function { name, args, .. })) => {
67                let FunctionArguments::List(arg_list) = args else {
68                    unreachable!()
69                };
70                assert_eq!("flush_table", name.to_string());
71                assert_eq!(arg_list.args.len(), 1);
72                assert!(matches!(&arg_list.args[0],
73                                 FunctionArg::Unnamed(FunctionArgExpr::Expr(
74                                     Expr::Value(Value::SingleQuotedString(s))
75                                 )) if s == "test"));
76            }
77            _ => unreachable!(),
78        }
79
80        assert_eq!(sql, stmt.to_string());
81    }
82
83    #[test]
84    fn test_parse_admin_function_without_args() {
85        let sql = "ADMIN test()";
86
87        let mut result =
88            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
89                .unwrap();
90        assert_eq!(1, result.len());
91        let stmt = result.remove(0);
92        match &stmt {
93            Statement::Admin(Admin::Func(Function { name, args, .. })) => {
94                let FunctionArguments::List(arg_list) = args else {
95                    unreachable!()
96                };
97                assert_eq!("test", name.to_string());
98                assert_eq!(arg_list.args.len(), 0);
99            }
100            _ => unreachable!(),
101        }
102
103        assert_eq!("ADMIN test()", stmt.to_string());
104    }
105
106    #[test]
107    fn test_invalid_admin_statement() {
108        let sql = "ADMIN";
109        assert!(ParserContext::create_with_dialect(
110            sql,
111            &GreptimeDbDialect {},
112            ParseOptions::default()
113        )
114        .is_err());
115
116        let sql = "ADMIN test";
117        assert!(ParserContext::create_with_dialect(
118            sql,
119            &GreptimeDbDialect {},
120            ParseOptions::default()
121        )
122        .is_err());
123
124        let sql = "ADMIN test test";
125        assert!(ParserContext::create_with_dialect(
126            sql,
127            &GreptimeDbDialect {},
128            ParseOptions::default()
129        )
130        .is_err());
131    }
132}