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) => {
98                    query_ctx.set_last_procedure_id(pid.clone());
99                    Ok(Value::from(pid))
100                }
101                None => Ok(Value::Null),
102            }
103        }
104
105        _ => Ok(Value::Null),
106    }
107}
108
109fn signature() -> Signature {
110    Signature::one_of(
111        vec![
112            // migrate_region(region_id, from_peer, to_peer)
113            TypeSignature::Uniform(
114                3,
115                ConcreteDataType::numerics()
116                    .into_iter()
117                    .map(|dt| dt.as_arrow_type())
118                    .collect(),
119            ),
120            // migrate_region(region_id, from_peer, to_peer, timeout(secs))
121            TypeSignature::Uniform(
122                4,
123                ConcreteDataType::numerics()
124                    .into_iter()
125                    .map(|dt| dt.as_arrow_type())
126                    .collect(),
127            ),
128        ],
129        Volatility::Immutable,
130    )
131}
132
133#[cfg(test)]
134mod tests {
135    use std::sync::Arc;
136
137    use arrow::array::{StringArray, UInt64Array};
138    use arrow::datatypes::{DataType, Field};
139    use datafusion_expr::ColumnarValue;
140
141    use super::*;
142    use crate::function::FunctionContext;
143    use crate::function_factory::ScalarFunctionFactory;
144
145    #[test]
146    fn test_migrate_region_misc() {
147        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
148        let f = factory.provide(FunctionContext::mock());
149        assert_eq!("migrate_region", f.name());
150        assert_eq!(DataType::Utf8, f.return_type(&[]).unwrap());
151        assert!(matches!(f.signature(),
152                         datafusion_expr::Signature {
153                             type_signature: datafusion_expr::TypeSignature::OneOf(sigs),
154                             volatility: datafusion_expr::Volatility::Immutable,
155                             ..
156                         } if sigs.len() == 2));
157    }
158
159    #[tokio::test]
160    async fn test_missing_procedure_service() {
161        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
162        let provider = factory.provide(FunctionContext::default());
163        let f = provider.as_async().unwrap();
164
165        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
166            args: vec![
167                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
168                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
169                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
170            ],
171            arg_fields: vec![
172                Arc::new(Field::new("arg_0", DataType::UInt64, false)),
173                Arc::new(Field::new("arg_1", DataType::UInt64, false)),
174                Arc::new(Field::new("arg_2", DataType::UInt64, false)),
175            ],
176            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
177            number_rows: 1,
178            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
179        };
180        let result = f.invoke_async_with_args(func_args).await.unwrap_err();
181        assert_eq!(
182            "Execution error: Missing ProcedureServiceHandler, not expected",
183            result.to_string()
184        );
185    }
186
187    #[tokio::test]
188    async fn test_migrate_region() {
189        let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
190        let context = FunctionContext::mock();
191        let query_ctx = context.query_ctx.clone();
192        query_ctx.set_last_procedure_id("previous-id".to_string());
193        let provider = factory.provide(context);
194        let f = provider.as_async().unwrap();
195
196        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
197            args: vec![
198                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
199                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
200                ColumnarValue::Array(Arc::new(UInt64Array::from(vec![1]))),
201            ],
202            arg_fields: vec![
203                Arc::new(Field::new("arg_0", DataType::UInt64, false)),
204                Arc::new(Field::new("arg_1", DataType::UInt64, false)),
205                Arc::new(Field::new("arg_2", DataType::UInt64, false)),
206            ],
207            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
208            number_rows: 1,
209            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
210        };
211        let result = f.invoke_async_with_args(func_args).await.unwrap();
212        assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
213
214        match result {
215            ColumnarValue::Array(array) => {
216                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
217                assert_eq!(result_array.value(0), "test_pid");
218            }
219            ColumnarValue::Scalar(scalar) => {
220                assert_eq!(
221                    scalar,
222                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
223                );
224            }
225        }
226    }
227}
common_function/admin/migrate_region.rs

common_function/admin/
migrate_region.rs
