common_function/admin/
migrate_region.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::time::Duration;
16
17use common_macro::admin_fn;
18use common_meta::rpc::procedure::MigrateRegionRequest;
19use common_query::error::{InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result};
20use datafusion_expr::{Signature, TypeSignature, Volatility};
21use datatypes::data_type::DataType;
22use datatypes::prelude::ConcreteDataType;
23use datatypes::value::{Value, ValueRef};
24use session::context::QueryContextRef;
25
26use crate::handlers::ProcedureServiceHandlerRef;
27use crate::helper::cast_u64;
28
29/// The default timeout for migrate region procedure.
30const DEFAULT_TIMEOUT_SECS: u64 = 300;
31
32/// A function to migrate a region from source peer to target peer.
33/// Returns the submitted procedure id if success. Only available in cluster mode.
34///
35/// - `migrate_region(region_id, from_peer, to_peer)`, with timeout(300 seconds).
36/// - `migrate_region(region_id, from_peer, to_peer, timeout(secs))`.
37///
38/// The parameters:
39/// - `region_id`:  the region id
40/// - `from_peer`:  the source peer id
41/// - `to_peer`:  the target peer id
42#[admin_fn(
43    name = MigrateRegionFunction,
44    display_name = migrate_region,
45    sig_fn = signature,
46    ret = string
47)]
48pub(crate) async fn migrate_region(
49    procedure_service_handler: &ProcedureServiceHandlerRef,
50    query_ctx: &QueryContextRef,
51    params: &[ValueRef<'_>],
52) -> Result<Value> {
53    let (region_id, from_peer, to_peer, timeout) = match params.len() {
54        3 => {
55            let region_id = cast_u64(&params[0])?;
56            let from_peer = cast_u64(&params[1])?;
57            let to_peer = cast_u64(&params[2])?;
58
59            (region_id, from_peer, to_peer, Some(DEFAULT_TIMEOUT_SECS))
60        }
61
62        4 => {
63            let region_id = cast_u64(&params[0])?;
64            let from_peer = cast_u64(&params[1])?;
65            let to_peer = cast_u64(&params[2])?;
66            let replay_timeout = cast_u64(&params[3])?;
67
68            (region_id, from_peer, to_peer, replay_timeout)
69        }
70
71        size => {
72            return InvalidFuncArgsSnafu {
73                err_msg: format!(
74                    "The length of the args is not correct, expect exactly 3 or 4, have: {}",
75                    size
76                ),
77            }
78            .fail();
79        }
80    };
81
82    match (region_id, from_peer, to_peer, timeout) {
83        (Some(region_id), Some(from_peer), Some(to_peer), Some(timeout)) => {
84            let pid = procedure_service_handler
85                .migrate_region(
86                    query_ctx.clone(),
87                    MigrateRegionRequest {
88                        region_id,
89                        from_peer,
90                        to_peer,
91                        timeout: Duration::from_secs(timeout),
92                    },
93                )
94                .await?;
95
96            match pid {
97                Some(pid) => Ok(Value::from(pid)),
98                None => Ok(Value::Null),
99            }
100        }
101
102        _ => Ok(Value::Null),
103    }
104}
105
106fn signature() -> Signature {
107    Signature::one_of(
108        vec![
109            // migrate_region(region_id, from_peer, to_peer)
110            TypeSignature::Uniform(
111                3,
112                ConcreteDataType::numerics()
113                    .into_iter()
114                    .map(|dt| dt.as_arrow_type())
115                    .collect(),
116            ),
117            // migrate_region(region_id, from_peer, to_peer, timeout(secs))
118            TypeSignature::Uniform(
119                4,
120                ConcreteDataType::numerics()
121                    .into_iter()
122                    .map(|dt| dt.as_arrow_type())
123                    .collect(),
124            ),
125        ],
126        Volatility::Immutable,
127    )
128}
129
130#[cfg(test)]
131mod tests {
132    use std::sync::Arc;
133
134    use arrow::array::{StringArray, UInt64Array};
135    use arrow::datatypes::{DataType, Field};
136    use datafusion_expr::ColumnarValue;
137
138    use super::*;
139    use crate::function::FunctionContext;
140    use crate::function_factory::ScalarFunctionFactory;
141
142    #[test]
143    fn test_migrate_region_misc() {
144        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
145        let f = factory.provide(FunctionContext::mock());
146        assert_eq!("migrate_region", f.name());
147        assert_eq!(DataType::Utf8, f.return_type(&[]).unwrap());
148        assert!(matches!(f.signature(),
149                         datafusion_expr::Signature {
150                             type_signature: datafusion_expr::TypeSignature::OneOf(sigs),
151                             volatility: datafusion_expr::Volatility::Immutable,
152                             ..
153                         } if sigs.len() == 2));
154    }
155
156    #[tokio::test]
157    async fn test_missing_procedure_service() {
158        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
159        let provider = factory.provide(FunctionContext::default());
160        let f = provider.as_async().unwrap();
161
162        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
163            args: vec![
164                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
165                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
166                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
167            ],
168            arg_fields: vec![
169                Arc::new(Field::new("arg_0", DataType::UInt64, false)),
170                Arc::new(Field::new("arg_1", DataType::UInt64, false)),
171                Arc::new(Field::new("arg_2", DataType::UInt64, false)),
172            ],
173            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
174            number_rows: 1,
175            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
176        };
177        let result = f.invoke_async_with_args(func_args).await.unwrap_err();
178        assert_eq!(
179            "Execution error: Missing ProcedureServiceHandler, not expected",
180            result.to_string()
181        );
182    }
183
184    #[tokio::test]
185    async fn test_migrate_region() {
186        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
187        let provider = factory.provide(FunctionContext::mock());
188        let f = provider.as_async().unwrap();
189
190        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
191            args: vec![
192                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
193                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
194                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
195            ],
196            arg_fields: vec![
197                Arc::new(Field::new("arg_0", DataType::UInt64, false)),
198                Arc::new(Field::new("arg_1", DataType::UInt64, false)),
199                Arc::new(Field::new("arg_2", DataType::UInt64, false)),
200            ],
201            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
202            number_rows: 1,
203            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
204        };
205        let result = f.invoke_async_with_args(func_args).await.unwrap();
206
207        match result {
208            ColumnarValue::Array(array) => {
209                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
210                assert_eq!(result_array.value(0), "test_pid");
211            }
212            ColumnarValue::Scalar(scalar) => {
213                assert_eq!(
214                    scalar,
215                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
216                );
217            }
218        }
219    }
220}
common_function/admin/migrate_region.rs

common_function/admin/
migrate_region.rs
