feat!: impl admin command (#4600)

* feat: impl admin statement parser

* feat: introduce AsyncFunction and implements it for admin functions

* feat: execute admin functions

* fix: license header

* fix: panic in test

* chore: fixed by code review

src/sql/src/parser.rs

@@ -159,8 +159,10 @@ impl<'a> ParserContext<'a> {
 
                     Keyword::SET => self.parse_set_variables(),
 
+                    Keyword::ADMIN => self.parse_admin_command(),
+
                     Keyword::NoKeyword
-                        if w.value.to_uppercase() == tql_parser::TQL && w.quote_style.is_none() =>
+                        if w.quote_style.is_none() && w.value.to_uppercase() == tql_parser::TQL =>
                     {
                         self.parse_tql()
                     }

src/sql/src/parsers.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub(crate) mod admin_parser;
 mod alter_parser;
 pub(crate) mod copy_parser;
 pub(crate) mod create_parser;

src/sql/src/parsers/admin_parser.rs

@@ -0,0 +1,124 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use snafu::ResultExt;
+
+use crate::ast::Expr;
+use crate::error::{Result, SyntaxSnafu};
+use crate::parser::ParserContext;
+use crate::statements::admin::Admin;
+use crate::statements::statement::Statement;
+
+/// `admin` extension parser: `admin function(arg1, arg2, ...)`
+/// or `admin function`
+impl<'a> ParserContext<'a> {
+    /// Parse `admin function(arg1, arg2, ...)` or `admin function` statement
+    pub(crate) fn parse_admin_command(&mut self) -> Result<Statement> {
+        let _token = self.parser.next_token();
+
+        let object_name = self.parser.parse_object_name(false).context(SyntaxSnafu)?;
+
+        let func = match self
+            .parser
+            .parse_function(object_name)
+            .context(SyntaxSnafu)?
+        {
+            Expr::Function(f) => f,
+            _ => {
+                return self.unsupported(self.peek_token_as_string());
+            }
+        };
+
+        Ok(Statement::Admin(Admin::Func(func)))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ast::{Expr, Function, FunctionArg, FunctionArgExpr, Value};
+    use crate::dialect::GreptimeDbDialect;
+    use crate::parser::ParseOptions;
+
+    #[test]
+    fn test_parse_admin_function() {
+        let sql = "ADMIN flush_table('test')";
+
+        let mut result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, result.len());
+        let stmt = result.remove(0);
+        match &stmt {
+            Statement::Admin(Admin::Func(Function { name, args, .. })) => {
+                assert_eq!("flush_table", name.to_string());
+                assert_eq!(args.len(), 1);
+                assert!(matches!(&args[0],
+                                 FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                                     Expr::Value(Value::SingleQuotedString(s))
+                                 )) if s == "test"));
+            }
+            _ => unreachable!(),
+        }
+
+        assert_eq!(sql, stmt.to_string());
+    }
+
+    #[test]
+    fn test_parse_admin_function_without_args() {
+        let sql = "ADMIN test()";
+
+        let mut result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, result.len());
+        let stmt = result.remove(0);
+        match &stmt {
+            Statement::Admin(Admin::Func(Function { name, args, .. })) => {
+                assert_eq!("test", name.to_string());
+                assert_eq!(args.len(), 0);
+            }
+            _ => unreachable!(),
+        }
+
+        assert_eq!("ADMIN test()", stmt.to_string());
+    }
+
+    #[test]
+    fn test_invalid_admin_statement() {
+        let sql = "ADMIN";
+        assert!(ParserContext::create_with_dialect(
+            sql,
+            &GreptimeDbDialect {},
+            ParseOptions::default()
+        )
+        .is_err());
+
+        let sql = "ADMIN test";
+        assert!(ParserContext::create_with_dialect(
+            sql,
+            &GreptimeDbDialect {},
+            ParseOptions::default()
+        )
+        .is_err());
+
+        let sql = "ADMIN test test";
+        assert!(ParserContext::create_with_dialect(
+            sql,
+            &GreptimeDbDialect {},
+            ParseOptions::default()
+        )
+        .is_err());
+    }
+}

src/sql/src/statements.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod admin;
 pub mod alter;
 pub mod copy;
 pub mod create;

src/sql/src/statements/admin.rs

@@ -0,0 +1,34 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Display;
+
+use sqlparser_derive::{Visit, VisitMut};
+
+use crate::ast::Function;
+
+/// `ADMIN` statement to execute some administration commands.
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub enum Admin {
+    /// Run a admin function.
+    Func(Function),
+}
+
+impl Display for Admin {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Admin::Func(func) => write!(f, "ADMIN {func}"),
+        }
+    }
+}

src/sql/src/statements/statement.rs

@@ -19,6 +19,7 @@ use sqlparser::ast::Statement as SpStatement;
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::error::{ConvertToDfStatementSnafu, Error};
+use crate::statements::admin::Admin;
 use crate::statements::alter::AlterTable;
 use crate::statements::create::{
     CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, CreateView,
@@ -110,6 +111,8 @@ pub enum Statement {
     ShowVariables(ShowVariables),
     // USE
     Use(String),
+    // Admin statement(extension)
+    Admin(Admin),
 }
 
 impl Display for Statement {
@@ -154,6 +157,7 @@ impl Display for Statement {
             }
             Statement::CreateView(s) => s.fmt(f),
             Statement::Use(s) => s.fmt(f),
+            Statement::Admin(admin) => admin.fmt(f),
         }
     }
 }