feat: administration functions (#3236)

* feat: adds database() function to return current db

* refactor: refactor meta src and client with new protos

* feat: impl migrate_region and query_procedure_state for procedure service/client

* fix: format

* temp commit

* feat: impl migrate_region SQL function

* chore: clean code for review

* fix: license header

* fix: toml format

* chore: update proto dependency

* chore: apply suggestion

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* chore: apply suggestion

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* chore: apply suggestion

Co-authored-by: JeremyHi <jiachun_feng@proton.me>

* chore: apply suggestion

Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com>

* chore: print key when parsing procedure id fails

* chore: comment

* chore: comment for MigrateRegionFunction

---------

Co-authored-by: Weny Xu <wenymedia@gmail.com>
Co-authored-by: JeremyHi <jiachun_feng@proton.me>
Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com>

src/common/function/Cargo.toml

@@ -5,12 +5,16 @@ version.workspace = true
 license.workspace = true
 
 [dependencies]
+api.workspace = true
 arc-swap = "1.0"
+async-trait.workspace = true
 build-data = "0.1"
 chrono-tz = "0.6"
 common-error.workspace = true
 common-macro.workspace = true
 common-query.workspace = true
+common-runtime.workspace = true
+common-telemetry.workspace = true
 common-time.workspace = true
 datafusion.workspace = true
 datatypes.workspace = true
@@ -22,6 +26,7 @@ paste = "1.0"
 session.workspace = true
 snafu.workspace = true
 statrs = "0.16"
+table.workspace = true
 
 [dev-dependencies]
 ron = "0.7"

src/common/function/src/function.rs

@@ -21,16 +21,20 @@ use datatypes::data_type::ConcreteDataType;
 use datatypes::vectors::VectorRef;
 use session::context::{QueryContextBuilder, QueryContextRef};
 
+use crate::state::FunctionState;
+
 /// The function execution context
 #[derive(Clone)]
 pub struct FunctionContext {
     pub query_ctx: QueryContextRef,
+    pub state: Arc<FunctionState>,
 }
 
 impl Default for FunctionContext {
     fn default() -> Self {
         Self {
             query_ctx: QueryContextBuilder::default().build(),
+            state: Arc::new(FunctionState::default()),
         }
     }
 }

src/common/function/src/function_registry.rs

@@ -25,6 +25,7 @@ use crate::scalars::math::MathFunction;
 use crate::scalars::numpy::NumpyFunction;
 use crate::scalars::timestamp::TimestampFunction;
 use crate::system::SystemFunction;
+use crate::table::TableFunction;
 
 #[derive(Default)]
 pub struct FunctionRegistry {
@@ -74,13 +75,19 @@ impl FunctionRegistry {
 pub static FUNCTION_REGISTRY: Lazy<Arc<FunctionRegistry>> = Lazy::new(|| {
     let function_registry = FunctionRegistry::default();
 
+    // Utility functions
     MathFunction::register(&function_registry);
     NumpyFunction::register(&function_registry);
     TimestampFunction::register(&function_registry);
     DateFunction::register(&function_registry);
 
+    // Aggregate functions
     AggregateFunctions::register(&function_registry);
+
+    // System and administration functions
     SystemFunction::register(&function_registry);
+    TableFunction::register(&function_registry);
+
     Arc::new(function_registry)
 });
 

src/common/function/src/handlers.rs

@@ -0,0 +1,54 @@
+// 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::sync::Arc;
+use std::time::Duration;
+
+use api::v1::meta::ProcedureStateResponse;
+use async_trait::async_trait;
+use common_query::error::Result;
+use session::context::QueryContextRef;
+use table::requests::{DeleteRequest, InsertRequest};
+
+pub type AffectedRows = usize;
+
+/// A trait for handling table mutations in `QueryEngine`.
+#[async_trait]
+pub trait TableMutationHandler: Send + Sync {
+    /// Inserts rows into the table.
+    async fn insert(&self, request: InsertRequest, ctx: QueryContextRef) -> Result<AffectedRows>;
+
+    /// Delete rows from the table.
+    async fn delete(&self, request: DeleteRequest, ctx: QueryContextRef) -> Result<AffectedRows>;
+
+    /// Migrate a region from source peer to target peer, returns the procedure id if success.
+    async fn migrate_region(
+        &self,
+        region_id: u64,
+        from_peer: u64,
+        to_peer: u64,
+        replay_timeout: Duration,
+    ) -> Result<String>;
+}
+
+/// A trait for handling meta service requests in `QueryEngine`.
+#[async_trait]
+pub trait MetaServiceHandler: Send + Sync {
+    /// Query the procedure' state by its id
+    async fn query_procedure_state(&self, pid: &str) -> Result<ProcedureStateResponse>;
+}
+
+pub type TableMutationHandlerRef = Arc<dyn TableMutationHandler>;
+
+pub type MetaServiceHandlerRef = Arc<dyn MetaServiceHandler>;

src/common/function/src/lib.rs

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 pub mod scalars;
-pub mod system;
+mod system;
+mod table;
 
 pub mod function;
 pub mod function_registry;
+pub mod handlers;
 pub mod helper;
+pub mod state;

src/common/function/src/scalars/udf.rs

@@ -25,9 +25,14 @@ use session::context::QueryContextRef;
 use snafu::ResultExt;
 
 use crate::function::{FunctionContext, FunctionRef};
+use crate::state::FunctionState;
 
-/// Create a ScalarUdf from function and query context.
-pub fn create_udf(func: FunctionRef, query_ctx: QueryContextRef) -> ScalarUdf {
+/// Create a ScalarUdf from function, query context and state.
+pub fn create_udf(
+    func: FunctionRef,
+    query_ctx: QueryContextRef,
+    state: Arc<FunctionState>,
+) -> ScalarUdf {
     let func_cloned = func.clone();
     let return_type: ReturnTypeFunction = Arc::new(move |input_types: &[ConcreteDataType]| {
         Ok(Arc::new(func_cloned.return_type(input_types)?))
@@ -38,6 +43,7 @@ pub fn create_udf(func: FunctionRef, query_ctx: QueryContextRef) -> ScalarUdf {
     let fun: ScalarFunctionImplementation = Arc::new(move |args: &[ColumnarValue]| {
         let func_ctx = FunctionContext {
             query_ctx: query_ctx.clone(),
+            state: state.clone(),
         };
 
         let len = args
@@ -101,7 +107,7 @@ mod tests {
         }
 
         // create a udf and test it again
-        let udf = create_udf(f.clone(), query_ctx);
+        let udf = create_udf(f.clone(), query_ctx, Arc::new(FunctionState::default()));
 
         assert_eq!("test_and", udf.name);
         assert_eq!(f.signature(), udf.signature);

src/common/function/src/state.rs

@@ -0,0 +1,25 @@
+// 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 crate::handlers::{MetaServiceHandlerRef, TableMutationHandlerRef};
+
+/// Shared state for SQL functions.
+/// The handlers in state may be `None` in cli command-line or test cases.
+#[derive(Clone, Default)]
+pub struct FunctionState {
+    // The table mutation handler
+    pub table_mutation_handler: Option<TableMutationHandlerRef>,
+    // The meta service handler
+    pub meta_service_handler: Option<MetaServiceHandlerRef>,
+}

src/common/function/src/system.rs

@@ -12,12 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod build;
-pub mod version;
+mod build;
+mod database;
+mod timezone;
+mod version;
 
 use std::sync::Arc;
 
 use build::BuildFunction;
+use database::DatabaseFunction;
+use timezone::TimezoneFunction;
 use version::VersionFunction;
 
 use crate::function_registry::FunctionRegistry;
@@ -28,5 +32,7 @@ impl SystemFunction {
     pub fn register(registry: &FunctionRegistry) {
         registry.register(Arc::new(BuildFunction));
         registry.register(Arc::new(VersionFunction));
+        registry.register(Arc::new(DatabaseFunction));
+        registry.register(Arc::new(TimezoneFunction));
     }
 }

src/common/function/src/system/database.rs

@@ -0,0 +1,92 @@
+// 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::{self};
+use std::sync::Arc;
+
+use common_query::error::Result;
+use common_query::prelude::{Signature, Volatility};
+use datatypes::prelude::{ConcreteDataType, ScalarVector};
+use datatypes::vectors::{StringVector, VectorRef};
+
+use crate::function::{Function, FunctionContext};
+
+/// A function to return current schema name.
+#[derive(Clone, Debug, Default)]
+pub struct DatabaseFunction;
+
+const NAME: &str = "database";
+
+impl Function for DatabaseFunction {
+    fn name(&self) -> &str {
+        NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::uniform(0, vec![], Volatility::Immutable)
+    }
+
+    fn eval(&self, func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
+        let db = func_ctx.query_ctx.current_schema();
+
+        Ok(Arc::new(StringVector::from_slice(&[db])) as _)
+    }
+}
+
+impl fmt::Display for DatabaseFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "DATABASE")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use common_query::prelude::TypeSignature;
+    use session::context::QueryContextBuilder;
+
+    use super::*;
+    #[test]
+    fn test_build_function() {
+        let build = DatabaseFunction;
+        assert_eq!("database", build.name());
+        assert_eq!(
+            ConcreteDataType::string_datatype(),
+            build.return_type(&[]).unwrap()
+        );
+        assert!(matches!(build.signature(),
+                         Signature {
+                             type_signature: TypeSignature::Uniform(0, valid_types),
+                             volatility: Volatility::Immutable
+                         } if  valid_types == vec![]
+        ));
+
+        let query_ctx = QueryContextBuilder::default()
+            .current_schema("test_db".to_string())
+            .build();
+
+        let func_ctx = FunctionContext {
+            query_ctx,
+            ..Default::default()
+        };
+        let vector = build.eval(func_ctx, &[]).unwrap();
+        let expect: VectorRef = Arc::new(StringVector::from(vec!["test_db"]));
+        assert_eq!(expect, vector);
+    }
+}

src/common/function/src/system/timezone.rs

@@ -0,0 +1,90 @@
+// 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::{self};
+use std::sync::Arc;
+
+use common_query::error::Result;
+use common_query::prelude::{Signature, Volatility};
+use datatypes::prelude::{ConcreteDataType, ScalarVector};
+use datatypes::vectors::{StringVector, VectorRef};
+
+use crate::function::{Function, FunctionContext};
+
+/// A function to return current session timezone.
+#[derive(Clone, Debug, Default)]
+pub struct TimezoneFunction;
+
+const NAME: &str = "timezone";
+
+impl Function for TimezoneFunction {
+    fn name(&self) -> &str {
+        NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::uniform(0, vec![], Volatility::Immutable)
+    }
+
+    fn eval(&self, func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
+        let tz = func_ctx.query_ctx.timezone().to_string();
+
+        Ok(Arc::new(StringVector::from_slice(&[&tz])) as _)
+    }
+}
+
+impl fmt::Display for TimezoneFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "TIMEZONE")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use common_query::prelude::TypeSignature;
+    use session::context::QueryContextBuilder;
+
+    use super::*;
+    #[test]
+    fn test_build_function() {
+        let build = TimezoneFunction;
+        assert_eq!("timezone", build.name());
+        assert_eq!(
+            ConcreteDataType::string_datatype(),
+            build.return_type(&[]).unwrap()
+        );
+        assert!(matches!(build.signature(),
+                         Signature {
+                             type_signature: TypeSignature::Uniform(0, valid_types),
+                             volatility: Volatility::Immutable
+                         } if  valid_types == vec![]
+        ));
+
+        let query_ctx = QueryContextBuilder::default().build();
+
+        let func_ctx = FunctionContext {
+            query_ctx,
+            ..Default::default()
+        };
+        let vector = build.eval(func_ctx, &[]).unwrap();
+        let expect: VectorRef = Arc::new(StringVector::from(vec!["UTC"]));
+        assert_eq!(expect, vector);
+    }
+}

src/common/function/src/table.rs

@@ -0,0 +1,31 @@
+// 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.
+
+mod migrate_region;
+
+use std::sync::Arc;
+
+use migrate_region::MigrateRegionFunction;
+
+use crate::function_registry::FunctionRegistry;
+
+/// Table functions
+pub(crate) struct TableFunction;
+
+impl TableFunction {
+    /// Register all table functions to [`FunctionRegistry`].
+    pub fn register(registry: &FunctionRegistry) {
+        registry.register(Arc::new(MigrateRegionFunction));
+    }
+}

src/common/function/src/table/migrate_region.rs

@@ -0,0 +1,175 @@
+// 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::{self};
+use std::time::Duration;
+
+use common_query::error::Error::ThreadJoin;
+use common_query::error::{
+    InvalidFuncArgsSnafu, InvalidInputTypeSnafu, MissingTableMutationHandlerSnafu, Result,
+};
+use common_query::prelude::{Signature, TypeSignature, Volatility};
+use common_telemetry::logging::error;
+use datatypes::prelude::{ConcreteDataType, MutableVector, ScalarVectorBuilder};
+use datatypes::value::Value;
+use datatypes::vectors::{StringVectorBuilder, VectorRef};
+use snafu::{Location, OptionExt, ResultExt};
+
+use crate::function::{Function, FunctionContext};
+
+/// A function to migrate a region from source peer to target peer.
+/// Returns the submitted procedure id if success. Only available in cluster mode.
+///
+/// - `migrate_region(region_id, from_peer, to_peer)`, with default replay WAL timeout(10 seconds).
+/// - `migrate_region(region_id, from_peer, to_peer, timeout(secs))`
+///
+/// The parameters:
+/// - `region_id`:  the region id
+/// - `from_peer`:  the source peer id
+/// - `to_peer`:  the target peer id
+#[derive(Clone, Debug, Default)]
+pub struct MigrateRegionFunction;
+
+const NAME: &str = "migrate_region";
+const DEFAULT_REPLAY_TIMEOUT_SECS: u64 = 10;
+
+fn cast_u64_vector(vector: &VectorRef) -> Result<VectorRef> {
+    vector
+        .cast(&ConcreteDataType::uint64_datatype())
+        .context(InvalidInputTypeSnafu {
+            err_msg: format!(
+                "Failed to cast input into uint64, actual type: {:#?}",
+                vector.data_type(),
+            ),
+        })
+}
+
+impl Function for MigrateRegionFunction {
+    fn name(&self) -> &str {
+        NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::one_of(
+            vec![
+                // migrate_region(region_id, from_peer, to_peer)
+                TypeSignature::Uniform(3, ConcreteDataType::numerics()),
+                // migrate_region(region_id, from_peer, to_peer, timeout(secs))
+                TypeSignature::Uniform(4, ConcreteDataType::numerics()),
+            ],
+            Volatility::Immutable,
+        )
+    }
+
+    fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        let (region_ids, from_peers, to_peers, replay_timeouts) = match columns.len() {
+            3 => {
+                let region_ids = cast_u64_vector(&columns[0])?;
+                let from_peers = cast_u64_vector(&columns[1])?;
+                let to_peers = cast_u64_vector(&columns[2])?;
+
+                (region_ids, from_peers, to_peers, None)
+            }
+
+            4 => {
+                let region_ids = cast_u64_vector(&columns[0])?;
+                let from_peers = cast_u64_vector(&columns[1])?;
+                let to_peers = cast_u64_vector(&columns[2])?;
+                let replay_timeouts = cast_u64_vector(&columns[3])?;
+
+                (region_ids, from_peers, to_peers, Some(replay_timeouts))
+            }
+
+            size => {
+                return InvalidFuncArgsSnafu {
+                    err_msg: format!(
+                        "The length of the args is not correct, expect exactly 3 or 4, have: {}",
+                        size
+                    ),
+                }
+                .fail();
+            }
+        };
+
+        std::thread::spawn(move || {
+            let len = region_ids.len();
+            let mut results = StringVectorBuilder::with_capacity(len);
+
+            for index in 0..len {
+                let region_id = region_ids.get(index);
+                let from_peer = from_peers.get(index);
+                let to_peer = to_peers.get(index);
+                let replay_timeout = match &replay_timeouts {
+                    Some(replay_timeouts) => replay_timeouts.get(index),
+                    None => Value::UInt64(DEFAULT_REPLAY_TIMEOUT_SECS),
+                };
+
+                match (region_id, from_peer, to_peer, replay_timeout) {
+                    (
+                        Value::UInt64(region_id),
+                        Value::UInt64(from_peer),
+                        Value::UInt64(to_peer),
+                        Value::UInt64(replay_timeout),
+                    ) => {
+                        let func_ctx = func_ctx.clone();
+
+                        let pid = common_runtime::block_on_read(async move {
+                            func_ctx
+                                .state
+                                .table_mutation_handler
+                                .as_ref()
+                                .context(MissingTableMutationHandlerSnafu)?
+                                .migrate_region(
+                                    region_id,
+                                    from_peer,
+                                    to_peer,
+                                    Duration::from_secs(replay_timeout),
+                                )
+                                .await
+                        })?;
+
+                        results.push(Some(&pid));
+                    }
+                    _ => {
+                        results.push(None);
+                    }
+                }
+            }
+
+            Ok(results.to_vector())
+        })
+        .join()
+        .map_err(|e| {
+            error!(e; "Join thread error");
+            ThreadJoin {
+                location: Location::default(),
+            }
+        })?
+    }
+}
+
+impl fmt::Display for MigrateRegionFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "MIGRATE_REGION")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    // FIXME(dennis): test in the following PR.
+}