common_function/admin/
remove_region_follower.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 common_macro::admin_fn;
16use common_meta::rpc::procedure::RemoveRegionFollowerRequest;
17use common_query::error::{
18    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result,
19    UnsupportedInputDataTypeSnafu,
20};
21use common_query::prelude::{Signature, TypeSignature, Volatility};
22use datatypes::prelude::ConcreteDataType;
23use datatypes::value::{Value, ValueRef};
24use session::context::QueryContextRef;
25use snafu::ensure;
26
27use crate::handlers::ProcedureServiceHandlerRef;
28use crate::helper::cast_u64;
29
30/// A function to remove a follower from a region.
31//// Only available in cluster mode.
32///
33/// - `remove_region_follower(region_id, peer_id)`.
34///
35/// The parameters:
36/// - `region_id`:  the region id
37/// - `peer_id`:  the peer id
38#[admin_fn(
39    name = RemoveRegionFollowerFunction,
40    display_name = remove_region_follower,
41    sig_fn = signature,
42    ret = uint64
43)]
44pub(crate) async fn remove_region_follower(
45    procedure_service_handler: &ProcedureServiceHandlerRef,
46    _ctx: &QueryContextRef,
47    params: &[ValueRef<'_>],
48) -> Result<Value> {
49    ensure!(
50        params.len() == 2,
51        InvalidFuncArgsSnafu {
52            err_msg: format!(
53                "The length of the args is not correct, expect exactly 2, have: {}",
54                params.len()
55            ),
56        }
57    );
58
59    let Some(region_id) = cast_u64(&params[0])? else {
60        return UnsupportedInputDataTypeSnafu {
61            function: "add_region_follower",
62            datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
63        }
64        .fail();
65    };
66    let Some(peer_id) = cast_u64(&params[1])? else {
67        return UnsupportedInputDataTypeSnafu {
68            function: "add_region_follower",
69            datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
70        }
71        .fail();
72    };
73
74    procedure_service_handler
75        .remove_region_follower(RemoveRegionFollowerRequest { region_id, peer_id })
76        .await?;
77
78    Ok(Value::from(0u64))
79}
80
81fn signature() -> Signature {
82    Signature::one_of(
83        vec![
84            // remove_region_follower(region_id, peer_id)
85            TypeSignature::Uniform(2, ConcreteDataType::numerics()),
86        ],
87        Volatility::Immutable,
88    )
89}
90
91#[cfg(test)]
92mod tests {
93    use std::sync::Arc;
94
95    use common_query::prelude::TypeSignature;
96    use datatypes::vectors::{UInt64Vector, VectorRef};
97
98    use super::*;
99    use crate::function::{AsyncFunction, FunctionContext};
100
101    #[test]
102    fn test_remove_region_follower_misc() {
103        let f = RemoveRegionFollowerFunction;
104        assert_eq!("remove_region_follower", f.name());
105        assert_eq!(
106            ConcreteDataType::uint64_datatype(),
107            f.return_type(&[]).unwrap()
108        );
109        assert!(matches!(f.signature(),
110                         Signature {
111                             type_signature: TypeSignature::OneOf(sigs),
112                             volatility: Volatility::Immutable
113                         } if sigs.len() == 1));
114    }
115
116    #[tokio::test]
117    async fn test_remove_region_follower() {
118        let f = RemoveRegionFollowerFunction;
119        let args = vec![1, 1];
120        let args = args
121            .into_iter()
122            .map(|arg| Arc::new(UInt64Vector::from_slice([arg])) as _)
123            .collect::<Vec<_>>();
124
125        let result = f.eval(FunctionContext::mock(), &args).await.unwrap();
126        let expect: VectorRef = Arc::new(UInt64Vector::from_slice([0u64]));
127        assert_eq!(result, expect);
128    }
129}