feat: impl migrate_region and procedure_state SQL function (#3325)

* fix: logical region can't find region routes

* feat: fetch partitions info in batch

* refactor: rename batch functions

* refactor: rename DdlTaskExecutor to ProcedureExecutor

* feat: impl migrate_region and query_procedure_state for ProcedureExecutor

* feat: adds SQL function procedure_state and finish migrate_region impl

* fix: constant vector

* feat: unit tests for migrate_region and procedure_state

* test: test region migration by SQL

* fix: compile error after rebeasing

* fix: clippy warnings

* feat: ensure procedure_state and migrate_region can be only called under greptime catalog

* fix: license header

src/common/function/Cargo.toml

@@ -12,8 +12,10 @@ api.workspace = true
 arc-swap = "1.0"
 async-trait.workspace = true
 chrono-tz = "0.6"
+common-catalog.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
+common-meta.workspace = true
 common-query.workspace = true
 common-runtime.workspace = true
 common-telemetry.workspace = true
@@ -26,6 +28,8 @@ num = "0.4"
 num-traits = "0.2"
 once_cell.workspace = true
 paste = "1.0"
+serde.workspace = true
+serde_json.workspace = true
 session.workspace = true
 snafu.workspace = true
 statrs = "0.16"

src/common/function/src/function.rs

@@ -30,6 +30,17 @@ pub struct FunctionContext {
     pub state: Arc<FunctionState>,
 }
 
+impl FunctionContext {
+    /// Create a mock [`FunctionContext`] for test.
+    #[cfg(any(test, feature = "testing"))]
+    pub fn mock() -> Self {
+        Self {
+            query_ctx: QueryContextBuilder::default().build(),
+            state: Arc::new(FunctionState::mock()),
+        }
+    }
+}
+
 impl Default for FunctionContext {
     fn default() -> Self {
         Self {

src/common/function/src/handlers.rs

@@ -13,10 +13,9 @@
 // limitations under the License.
 
 use std::sync::Arc;
-use std::time::Duration;
 
-use api::v1::meta::ProcedureStateResponse;
 use async_trait::async_trait;
+use common_meta::rpc::procedure::{MigrateRegionRequest, ProcedureStateResponse};
 use common_query::error::Result;
 use session::context::QueryContextRef;
 use table::requests::{DeleteRequest, InsertRequest};
@@ -31,24 +30,18 @@ pub trait TableMutationHandler: Send + Sync {
 
     /// 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`.
+/// A trait for handling procedure service requests in `QueryEngine`.
 #[async_trait]
-pub trait MetaServiceHandler: Send + Sync {
+pub trait ProcedureServiceHandler: Send + Sync {
+    /// Migrate a region from source peer to target peer, returns the procedure id if success.
+    async fn migrate_region(&self, request: MigrateRegionRequest) -> Result<Option<String>>;
+
     /// 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>;
+pub type ProcedureServiceHandlerRef = Arc<dyn ProcedureServiceHandler>;

src/common/function/src/lib.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod macros;
 pub mod scalars;
 mod system;
 mod table;

src/common/function/src/macros.rs

@@ -0,0 +1,27 @@
+// 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.
+
+/// Ensure current function is invokded under `greptime` catalog.
+#[macro_export]
+macro_rules! ensure_greptime {
+    ($func_ctx: expr) => {{
+        use common_catalog::consts::DEFAULT_CATALOG_NAME;
+        snafu::ensure!(
+            $func_ctx.query_ctx.current_catalog() == DEFAULT_CATALOG_NAME,
+            common_query::error::PermissionDeniedSnafu {
+                err_msg: format!("current catalog is not {DEFAULT_CATALOG_NAME}")
+            }
+        );
+    }};
+}

src/common/function/src/state.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::handlers::{MetaServiceHandlerRef, TableMutationHandlerRef};
+use crate::handlers::{ProcedureServiceHandlerRef, TableMutationHandlerRef};
 
 /// Shared state for SQL functions.
 /// The handlers in state may be `None` in cli command-line or test cases.
@@ -20,6 +20,45 @@ use crate::handlers::{MetaServiceHandlerRef, TableMutationHandlerRef};
 pub struct FunctionState {
     // The table mutation handler
     pub table_mutation_handler: Option<TableMutationHandlerRef>,
-    // The meta service handler
-    pub meta_service_handler: Option<MetaServiceHandlerRef>,
+    // The procedure service handler
+    pub procedure_service_handler: Option<ProcedureServiceHandlerRef>,
+}
+
+impl FunctionState {
+    /// Create a mock [`FunctionState`] for test.
+    #[cfg(any(test, feature = "testing"))]
+    pub fn mock() -> Self {
+        use std::sync::Arc;
+
+        use api::v1::meta::ProcedureStatus;
+        use async_trait::async_trait;
+        use common_meta::rpc::procedure::{MigrateRegionRequest, ProcedureStateResponse};
+        use common_query::error::Result;
+
+        use crate::handlers::ProcedureServiceHandler;
+        struct MockProcedureServiceHandler;
+
+        #[async_trait]
+        impl ProcedureServiceHandler for MockProcedureServiceHandler {
+            async fn migrate_region(
+                &self,
+                _request: MigrateRegionRequest,
+            ) -> Result<Option<String>> {
+                Ok(Some("test_pid".to_string()))
+            }
+
+            async fn query_procedure_state(&self, _pid: &str) -> Result<ProcedureStateResponse> {
+                Ok(ProcedureStateResponse {
+                    status: ProcedureStatus::Done.into(),
+                    error: "OK".to_string(),
+                    ..Default::default()
+                })
+            }
+        }
+
+        Self {
+            table_mutation_handler: None,
+            procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)),
+        }
+    }
 }

src/common/function/src/system.rs

@@ -14,6 +14,7 @@
 
 mod build;
 mod database;
+mod procedure_state;
 mod timezone;
 mod version;
 
@@ -21,6 +22,7 @@ use std::sync::Arc;
 
 use build::BuildFunction;
 use database::DatabaseFunction;
+use procedure_state::ProcedureStateFunction;
 use timezone::TimezoneFunction;
 use version::VersionFunction;
 
@@ -34,5 +36,6 @@ impl SystemFunction {
         registry.register(Arc::new(VersionFunction));
         registry.register(Arc::new(DatabaseFunction));
         registry.register(Arc::new(TimezoneFunction));
+        registry.register(Arc::new(ProcedureStateFunction));
     }
 }

src/common/function/src/system/build.rs

@@ -22,7 +22,7 @@ use datatypes::vectors::{StringVector, VectorRef};
 
 use crate::function::{Function, FunctionContext};
 
-/// Generates build information  
+/// Generates build information
 #[derive(Clone, Debug, Default)]
 pub struct BuildFunction;
 
@@ -42,11 +42,7 @@ impl Function for BuildFunction {
     }
 
     fn signature(&self) -> Signature {
-        Signature::uniform(
-            0,
-            vec![ConcreteDataType::string_datatype()],
-            Volatility::Immutable,
-        )
+        Signature::uniform(0, vec![], Volatility::Immutable)
     }
 
     fn eval(&self, _func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
@@ -75,7 +71,7 @@ mod tests {
                          Signature {
                              type_signature: TypeSignature::Uniform(0, valid_types),
                              volatility: Volatility::Immutable
-                         } if  valid_types == vec![ConcreteDataType::string_datatype()]
+                         } if  valid_types.is_empty()
         ));
         let build_info = common_version::build_info().to_string();
         let vector = build.eval(FunctionContext::default(), &[]).unwrap();

src/common/function/src/system/procedure_state.rs

@@ -0,0 +1,216 @@
+// 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;
+use std::sync::Arc;
+
+use api::v1::meta::ProcedureStatus;
+use common_meta::rpc::procedure::ProcedureStateResponse;
+use common_query::error::Error::ThreadJoin;
+use common_query::error::{
+    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result,
+    UnsupportedInputDataTypeSnafu,
+};
+use common_query::prelude::{Signature, Volatility};
+use common_telemetry::error;
+use datatypes::prelude::*;
+use datatypes::vectors::{ConstantVector, Helper, StringVector, VectorRef};
+use serde::Serialize;
+use snafu::{ensure, Location, OptionExt};
+
+use crate::function::{Function, FunctionContext};
+
+const NAME: &str = "procedure_state";
+
+/// A function to query procedure state by its id.
+/// Such as `procedure_state(pid)`.
+#[derive(Clone, Debug, Default)]
+pub struct ProcedureStateFunction;
+
+impl fmt::Display for ProcedureStateFunction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "PROCEDURE_STATE")
+    }
+}
+
+#[derive(Serialize)]
+struct ProcedureStateJson {
+    status: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
+}
+
+impl Function for ProcedureStateFunction {
+    fn name(&self) -> &str {
+        NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::uniform(
+            1,
+            vec![ConcreteDataType::string_datatype()],
+            Volatility::Immutable,
+        )
+    }
+
+    fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        crate::ensure_greptime!(func_ctx);
+
+        ensure!(
+            columns.len() == 1,
+            InvalidFuncArgsSnafu {
+                err_msg: format!(
+                    "The length of the args is not correct, expect 1, have: {}",
+                    columns.len()
+                ),
+            }
+        );
+
+        let pids = columns[0].clone();
+        let expect_len = pids.len();
+        let is_const = pids.is_const();
+
+        match pids.data_type() {
+            ConcreteDataType::String(_) => {
+                // TODO(dennis): datafusion UDF doesn't support async function currently
+                std::thread::spawn(move || {
+                    let pids: &StringVector = if is_const {
+                        let pids: &ConstantVector = unsafe { Helper::static_cast(&pids) };
+                        unsafe { Helper::static_cast(pids.inner()) }
+                    } else {
+                        unsafe { Helper::static_cast(&pids) }
+                    };
+
+                    let procedure_service_handler = func_ctx
+                        .state
+                        .procedure_service_handler
+                        .as_ref()
+                        .context(MissingProcedureServiceHandlerSnafu)?;
+
+                    let states = pids
+                        .iter_data()
+                        .map(|pid| {
+                            if let Some(pid) = pid {
+                                let ProcedureStateResponse { status, error, .. } =
+                                    common_runtime::block_on_read(async move {
+                                        procedure_service_handler.query_procedure_state(pid).await
+                                    })?;
+
+                                let status = ProcedureStatus::try_from(status)
+                                    .map(|v| v.as_str_name())
+                                    .unwrap_or("Unknown");
+
+                                let state = ProcedureStateJson {
+                                    status: status.to_string(),
+                                    error: if error.is_empty() { None } else { Some(error) },
+                                };
+
+                                Ok(Some(serde_json::to_string(&state).unwrap_or_default()))
+                            } else {
+                                Ok(None)
+                            }
+                        })
+                        .collect::<Result<Vec<_>>>()?;
+
+                    let results: VectorRef = Arc::new(StringVector::from(states));
+
+                    if is_const {
+                        Ok(Arc::new(ConstantVector::new(results, expect_len)) as _)
+                    } else {
+                        Ok(results)
+                    }
+                })
+                .join()
+                .map_err(|e| {
+                    error!(e; "Join thread error");
+                    ThreadJoin {
+                        location: Location::default(),
+                    }
+                })?
+            }
+            _ => UnsupportedInputDataTypeSnafu {
+                function: NAME,
+                datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
+            }
+            .fail(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use common_query::prelude::TypeSignature;
+    use datatypes::vectors::StringVector;
+
+    use super::*;
+
+    #[test]
+    fn test_procedure_state_misc() {
+        let f = ProcedureStateFunction;
+        assert_eq!("procedure_state", f.name());
+        assert_eq!(
+            ConcreteDataType::string_datatype(),
+            f.return_type(&[]).unwrap()
+        );
+        assert!(matches!(f.signature(),
+                         Signature {
+                             type_signature: TypeSignature::Uniform(1, valid_types),
+                             volatility: Volatility::Immutable
+                         } if valid_types == vec![ConcreteDataType::string_datatype()]
+        ));
+    }
+
+    #[test]
+    fn test_missing_procedure_service() {
+        let f = ProcedureStateFunction;
+
+        let args = vec!["pid"];
+
+        let args = args
+            .into_iter()
+            .map(|arg| Arc::new(StringVector::from_slice(&[arg])) as _)
+            .collect::<Vec<_>>();
+
+        let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+        assert_eq!(
+            "Missing ProcedureServiceHandler, not expected",
+            result.to_string()
+        );
+    }
+
+    #[test]
+    fn test_procedure_state() {
+        let f = ProcedureStateFunction;
+
+        let args = vec!["pid"];
+
+        let args = args
+            .into_iter()
+            .map(|arg| Arc::new(StringVector::from_slice(&[arg])) as _)
+            .collect::<Vec<_>>();
+
+        let result = f.eval(FunctionContext::mock(), &args).unwrap();
+
+        let expect: VectorRef = Arc::new(StringVector::from(vec![
+            "{\"status\":\"Done\",\"error\":\"OK\"}",
+        ]));
+        assert_eq!(expect, result);
+    }
+}

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

@@ -15,9 +15,10 @@
 use std::fmt::{self};
 use std::time::Duration;
 
+use common_meta::rpc::procedure::MigrateRegionRequest;
 use common_query::error::Error::ThreadJoin;
 use common_query::error::{
-    InvalidFuncArgsSnafu, InvalidInputTypeSnafu, MissingTableMutationHandlerSnafu, Result,
+    InvalidFuncArgsSnafu, InvalidInputTypeSnafu, MissingProcedureServiceHandlerSnafu, Result,
 };
 use common_query::prelude::{Signature, TypeSignature, Volatility};
 use common_telemetry::logging::error;
@@ -77,6 +78,8 @@ impl Function for MigrateRegionFunction {
     }
 
     fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        crate::ensure_greptime!(func_ctx);
+
         let (region_ids, from_peers, to_peers, replay_timeouts) = match columns.len() {
             3 => {
                 let region_ids = cast_u64_vector(&columns[0])?;
@@ -106,9 +109,15 @@ impl Function for MigrateRegionFunction {
             }
         };
 
+        // TODO(dennis): datafusion UDF doesn't support async function currently
         std::thread::spawn(move || {
             let len = region_ids.len();
             let mut results = StringVectorBuilder::with_capacity(len);
+            let procedure_service_handler = func_ctx
+                .state
+                .procedure_service_handler
+                .as_ref()
+                .context(MissingProcedureServiceHandlerSnafu)?;
 
             for index in 0..len {
                 let region_id = region_ids.get(index);
@@ -126,24 +135,18 @@ impl Function for MigrateRegionFunction {
                         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(
+                            procedure_service_handler
+                                .migrate_region(MigrateRegionRequest {
                                     region_id,
                                     from_peer,
                                     to_peer,
-                                    Duration::from_secs(replay_timeout),
-                                )
+                                    replay_timeout: Duration::from_secs(replay_timeout),
+                                })
                                 .await
                         })?;
 
-                        results.push(Some(&pid));
+                        results.push(pid.as_deref())
                     }
                     _ => {
                         results.push(None);
@@ -171,5 +174,60 @@ impl fmt::Display for MigrateRegionFunction {
 
 #[cfg(test)]
 mod tests {
-    // FIXME(dennis): test in the following PR.
+    use std::sync::Arc;
+
+    use common_query::prelude::TypeSignature;
+    use datatypes::vectors::{StringVector, UInt64Vector};
+
+    use super::*;
+
+    #[test]
+    fn test_migrate_region_misc() {
+        let f = MigrateRegionFunction;
+        assert_eq!("migrate_region", f.name());
+        assert_eq!(
+            ConcreteDataType::string_datatype(),
+            f.return_type(&[]).unwrap()
+        );
+        assert!(matches!(f.signature(),
+                         Signature {
+                             type_signature: TypeSignature::OneOf(sigs),
+                             volatility: Volatility::Immutable
+                         } if sigs.len() == 2));
+    }
+
+    #[test]
+    fn test_missing_procedure_service() {
+        let f = MigrateRegionFunction;
+
+        let args = vec![1, 1, 1];
+
+        let args = args
+            .into_iter()
+            .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
+            .collect::<Vec<_>>();
+
+        let result = f.eval(FunctionContext::default(), &args).unwrap_err();
+        assert_eq!(
+            "Missing ProcedureServiceHandler, not expected",
+            result.to_string()
+        );
+    }
+
+    #[test]
+    fn test_migrate_region() {
+        let f = MigrateRegionFunction;
+
+        let args = vec![1, 1, 1];
+
+        let args = args
+            .into_iter()
+            .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
+            .collect::<Vec<_>>();
+
+        let result = f.eval(FunctionContext::mock(), &args).unwrap();
+
+        let expect: VectorRef = Arc::new(StringVector::from(vec!["test_pid"]));
+        assert_eq!(expect, result);
+    }
 }