feat!: impl admin command (GreptimeTeam#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
CookiePieWw · Sep 17, 2024 · f15db85 · f15db85
1 parent 2ac6d25
commit f15db85
Show file tree

Hide file tree

Showing 39 changed files with 777 additions and 322 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml
@@ -39,3 +39,4 @@ table.workspace = true
 [dev-dependencies]
 ron = "0.7"
 serde = { version = "1.0", features = ["derive"] }
+tokio.workspace = true
diff --git a/src/common/function/src/flush_flow.rs b/src/common/function/src/flush_flow.rs
@@ -110,7 +110,7 @@ mod test {
     use session::context::QueryContext;
 
     use super::*;
-    use crate::function::{Function, FunctionContext};
+    use crate::function::{AsyncFunction, FunctionContext};
 
     #[test]
     fn test_flush_flow_metadata() {
@@ -130,8 +130,8 @@ mod test {
         );
     }
 
-    #[test]
-    fn test_missing_flow_service() {
+    #[tokio::test]
+    async fn test_missing_flow_service() {
         let f = FlushFlowFunction;
 
         let args = vec!["flow_name"];
@@ -140,7 +140,7 @@ mod test {
             .map(|arg| Arc::new(StringVector::from_slice(&[arg])) as _)
             .collect::<Vec<_>>();
 
-        let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+        let result = f.eval(FunctionContext::default(), &args).await.unwrap_err();
         assert_eq!(
             "Missing FlowServiceHandler, not expected",
             result.to_string()

diff --git a/src/common/function/src/function.rs b/src/common/function/src/function.rs
@@ -56,12 +56,33 @@ pub trait Function: fmt::Display + Sync + Send {
     /// Returns the name of the function, should be unique.
     fn name(&self) -> &str;
 
+    /// The returned data type of function execution.
     fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType>;
 
+    /// The signature of function.
     fn signature(&self) -> Signature;
 
     /// Evaluate the function, e.g. run/execute the function.
     fn eval(&self, _func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef>;
 }
 
 pub type FunctionRef = Arc<dyn Function>;
+
+/// Async Scalar function trait
+#[async_trait::async_trait]
+pub trait AsyncFunction: fmt::Display + Sync + Send {
+    /// Returns the name of the function, should be unique.
+    fn name(&self) -> &str;
+
+    /// The returned data type of function execution.
+    fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType>;
+
+    /// The signature of function.
+    fn signature(&self) -> Signature;
+
+    /// Evaluate the function, e.g. run/execute the function.
+    /// TODO(dennis): simplify the signature and refactor all the admin functions.
+    async fn eval(&self, _func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef>;
+}
+
+pub type AsyncFunctionRef = Arc<dyn AsyncFunction>;
diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs
@@ -18,7 +18,7 @@ use std::sync::{Arc, RwLock};
 
 use once_cell::sync::Lazy;
 
-use crate::function::FunctionRef;
+use crate::function::{AsyncFunctionRef, FunctionRef};
 use crate::scalars::aggregate::{AggregateFunctionMetaRef, AggregateFunctions};
 use crate::scalars::date::DateFunction;
 use crate::scalars::expression::ExpressionFunction;
@@ -32,6 +32,7 @@ use crate::table::TableFunction;
 #[derive(Default)]
 pub struct FunctionRegistry {
     functions: RwLock<HashMap<String, FunctionRef>>,
+    async_functions: RwLock<HashMap<String, AsyncFunctionRef>>,
     aggregate_functions: RwLock<HashMap<String, AggregateFunctionMetaRef>>,
 }
 
@@ -44,6 +45,27 @@ impl FunctionRegistry {
             .insert(func.name().to_string(), func);
     }
 
+    pub fn register_async(&self, func: AsyncFunctionRef) {
+        let _ = self
+            .async_functions
+            .write()
+            .unwrap()
+            .insert(func.name().to_string(), func);
+    }
+
+    pub fn get_async_function(&self, name: &str) -> Option<AsyncFunctionRef> {
+        self.async_functions.read().unwrap().get(name).cloned()
+    }
+
+    pub fn async_functions(&self) -> Vec<AsyncFunctionRef> {
+        self.async_functions
+            .read()
+            .unwrap()
+            .values()
+            .cloned()
+            .collect()
+    }
+
     pub fn register_aggregate_function(&self, func: AggregateFunctionMetaRef) {
         let _ = self
             .aggregate_functions

diff --git a/src/common/function/src/system.rs b/src/common/function/src/system.rs
@@ -38,7 +38,7 @@ impl SystemFunction {
         registry.register(Arc::new(VersionFunction));
         registry.register(Arc::new(DatabaseFunction));
         registry.register(Arc::new(TimezoneFunction));
-        registry.register(Arc::new(ProcedureStateFunction));
+        registry.register_async(Arc::new(ProcedureStateFunction));
         PGCatalogFunction::register(registry);
     }
 }
diff --git a/src/common/function/src/system/procedure_state.rs b/src/common/function/src/system/procedure_state.rs
@@ -96,7 +96,7 @@ mod tests {
     use datatypes::vectors::StringVector;
 
     use super::*;
-    use crate::function::{Function, FunctionContext};
+    use crate::function::{AsyncFunction, FunctionContext};
 
     #[test]
     fn test_procedure_state_misc() {
@@ -114,8 +114,8 @@ mod tests {
         ));
     }
 
-    #[test]
-    fn test_missing_procedure_service() {
+    #[tokio::test]
+    async fn test_missing_procedure_service() {
         let f = ProcedureStateFunction;
 
         let args = vec!["pid"];
@@ -125,15 +125,15 @@ mod tests {
             .map(|arg| Arc::new(StringVector::from_slice(&[arg])) as _)
             .collect::<Vec<_>>();
 
-        let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+        let result = f.eval(FunctionContext::default(), &args).await.unwrap_err();
         assert_eq!(
             "Missing ProcedureServiceHandler, not expected",
             result.to_string()
         );
     }
 
-    #[test]
-    fn test_procedure_state() {
+    #[tokio::test]
+    async fn test_procedure_state() {
         let f = ProcedureStateFunction;
 
         let args = vec!["pid"];
@@ -143,7 +143,7 @@ mod tests {
             .map(|arg| Arc::new(StringVector::from_slice(&[arg])) as _)
             .collect::<Vec<_>>();
 
-        let result = f.eval(FunctionContext::mock(), &args).unwrap();
+        let result = f.eval(FunctionContext::mock(), &args).await.unwrap();
 
         let expect: VectorRef = Arc::new(StringVector::from(vec![
             "{\"status\":\"Done\",\"error\":\"OK\"}",

diff --git a/src/common/function/src/table.rs b/src/common/function/src/table.rs
@@ -31,11 +31,11 @@ pub(crate) struct TableFunction;
 impl TableFunction {
     /// Register all table functions to [`FunctionRegistry`].
     pub fn register(registry: &FunctionRegistry) {
-        registry.register(Arc::new(MigrateRegionFunction));
-        registry.register(Arc::new(FlushRegionFunction));
-        registry.register(Arc::new(CompactRegionFunction));
-        registry.register(Arc::new(FlushTableFunction));
-        registry.register(Arc::new(CompactTableFunction));
-        registry.register(Arc::new(FlushFlowFunction));
+        registry.register_async(Arc::new(MigrateRegionFunction));
+        registry.register_async(Arc::new(FlushRegionFunction));
+        registry.register_async(Arc::new(CompactRegionFunction));
+        registry.register_async(Arc::new(FlushTableFunction));
+        registry.register_async(Arc::new(CompactTableFunction));
+        registry.register_async(Arc::new(FlushFlowFunction));
     }
 }
diff --git a/src/common/function/src/table/flush_compact_region.rs b/src/common/function/src/table/flush_compact_region.rs
@@ -77,7 +77,7 @@ mod tests {
     use datatypes::vectors::UInt64Vector;
 
     use super::*;
-    use crate::function::{Function, FunctionContext};
+    use crate::function::{AsyncFunction, FunctionContext};
 
     macro_rules! define_region_function_test {
         ($name: ident, $func: ident) => {
@@ -97,8 +97,8 @@ mod tests {
                                      } if valid_types == ConcreteDataType::numerics()));
                 }
 
-                #[test]
-                fn [<test_ $name _missing_table_mutation>]() {
+                #[tokio::test]
+                async fn [<test_ $name _missing_table_mutation>]() {
                     let f = $func;
 
                     let args = vec![99];
@@ -108,15 +108,15 @@ mod tests {
                         .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
                         .collect::<Vec<_>>();
 
-                    let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+                    let result = f.eval(FunctionContext::default(), &args).await.unwrap_err();
                     assert_eq!(
                         "Missing TableMutationHandler, not expected",
                         result.to_string()
                     );
                 }
 
-                #[test]
-                fn [<test_ $name>]() {
+                #[tokio::test]
+                async fn [<test_ $name>]() {
                     let f = $func;
 
 
@@ -127,7 +127,7 @@ mod tests {
                         .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
                         .collect::<Vec<_>>();
 
-                    let result = f.eval(FunctionContext::mock(), &args).unwrap();
+                    let result = f.eval(FunctionContext::mock(), &args).await.unwrap();
 
                     let expect: VectorRef = Arc::new(UInt64Vector::from_slice([42]));
                     assert_eq!(expect, result);

diff --git a/src/common/function/src/table/flush_compact_table.rs b/src/common/function/src/table/flush_compact_table.rs
@@ -210,7 +210,7 @@ mod tests {
     use session::context::QueryContext;
 
     use super::*;
-    use crate::function::{Function, FunctionContext};
+    use crate::function::{AsyncFunction, FunctionContext};
 
     macro_rules! define_table_function_test {
         ($name: ident, $func: ident) => {
@@ -230,8 +230,8 @@ mod tests {
                                      } if valid_types == vec![ConcreteDataType::string_datatype()]));
                 }
 
-                #[test]
-                fn [<test_ $name _missing_table_mutation>]() {
+                #[tokio::test]
+                async fn [<test_ $name _missing_table_mutation>]() {
                     let f = $func;
 
                     let args = vec!["test"];
@@ -241,15 +241,15 @@ mod tests {
                         .map(|arg| Arc::new(StringVector::from(vec![arg])) as _)
                         .collect::<Vec<_>>();
 
-                    let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+                    let result = f.eval(FunctionContext::default(), &args).await.unwrap_err();
                     assert_eq!(
                         "Missing TableMutationHandler, not expected",
                         result.to_string()
                     );
                 }
 
-                #[test]
-                fn [<test_ $name>]() {
+                #[tokio::test]
+                async fn [<test_ $name>]() {
                     let f = $func;
 
 
@@ -260,7 +260,7 @@ mod tests {
                         .map(|arg| Arc::new(StringVector::from(vec![arg])) as _)
                         .collect::<Vec<_>>();
 
-                    let result = f.eval(FunctionContext::mock(), &args).unwrap();
+                    let result = f.eval(FunctionContext::mock(), &args).await.unwrap();
 
                     let expect: VectorRef = Arc::new(UInt64Vector::from_slice([42]));
                     assert_eq!(expect, result);

diff --git a/src/common/function/src/table/migrate_region.rs b/src/common/function/src/table/migrate_region.rs
@@ -123,7 +123,7 @@ mod tests {
     use datatypes::vectors::{StringVector, UInt64Vector, VectorRef};
 
     use super::*;
-    use crate::function::{Function, FunctionContext};
+    use crate::function::{AsyncFunction, FunctionContext};
 
     #[test]
     fn test_migrate_region_misc() {
@@ -140,8 +140,8 @@ mod tests {
                          } if sigs.len() == 2));
     }
 
-    #[test]
-    fn test_missing_procedure_service() {
+    #[tokio::test]
+    async fn test_missing_procedure_service() {
         let f = MigrateRegionFunction;
 
         let args = vec![1, 1, 1];
@@ -151,15 +151,15 @@ mod tests {
             .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
             .collect::<Vec<_>>();
 
-        let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+        let result = f.eval(FunctionContext::default(), &args).await.unwrap_err();
         assert_eq!(
             "Missing ProcedureServiceHandler, not expected",
             result.to_string()
         );
     }
 
-    #[test]
-    fn test_migrate_region() {
+    #[tokio::test]
+    async fn test_migrate_region() {
         let f = MigrateRegionFunction;
 
         let args = vec![1, 1, 1];
@@ -169,7 +169,7 @@ mod tests {
             .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
             .collect::<Vec<_>>();
 
-        let result = f.eval(FunctionContext::mock(), &args).unwrap();
+        let result = f.eval(FunctionContext::mock(), &args).await.unwrap();
 
         let expect: VectorRef = Arc::new(StringVector::from(vec!["test_pid"]));
         assert_eq!(expect, result);