rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-24 08:50:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: impl admin command (#4600)

Browse Source 
* 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
This commit is contained in:
dennis zhuang
2024-08-26 15:53:40 +08:00
committed by GitHub
parent da337a9635
commit 63f2463273
39 changed files with 777 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/common/function/Cargo.toml
View File

@@ -39,3 +39,4 @@ table.workspace = true
[dev-dependencies]
ron = "0.7"
serde = { version = "1.0", features = ["derive"] }
tokio.workspace = true

8
src/common/function/src/flush_flow.rs
View File

@@ -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()

21
src/common/function/src/function.rs
View File

@@ -56,8 +56,10 @@ 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.
@@ -65,3 +67,22 @@ pub trait Function: fmt::Display + Sync + Send {
}
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>;

24
src/common/function/src/function_registry.rs
View File

@@ -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

2
src/common/function/src/system.rs
View File

@@ -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);
}
}

14
src/common/function/src/system/procedure_state.rs
View File

@@ -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\"}",

12
src/common/function/src/table.rs
View File

@@ -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));
}
}

14
src/common/function/src/table/flush_compact_region.rs
View File

@@ -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);

14
src/common/function/src/table/flush_compact_table.rs
View File

@@ -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);

14
src/common/function/src/table/migrate_region.rs
View File

@@ -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);