common_function/admin/
reconcile_database.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 api::v1::meta::reconcile_request::Target;
16use api::v1::meta::{ReconcileDatabase, ReconcileRequest};
17use arrow::datatypes::DataType as ArrowDataType;
18use common_macro::admin_fn;
19use common_query::error::{
20    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result,
21    UnsupportedInputDataTypeSnafu,
22};
23use common_telemetry::info;
24use datafusion_expr::{Signature, TypeSignature, Volatility};
25use datatypes::data_type::DataType;
26use datatypes::prelude::*;
27use session::context::QueryContextRef;
28
29use crate::handlers::ProcedureServiceHandlerRef;
30use crate::helper::{
31    cast_u32, default_parallelism, default_resolve_strategy, get_string_from_params,
32    parse_resolve_strategy,
33};
34
35const FN_NAME: &str = "reconcile_database";
36
37/// A function to reconcile a database.
38/// Returns the procedure id if success.
39///
40/// - `reconcile_database(database_name)`.
41/// - `reconcile_database(database_name, resolve_strategy)`.
42/// - `reconcile_database(database_name, resolve_strategy, parallelism)`.
43///
44/// The parameters:
45/// - `database_name`:  the database name
46#[admin_fn(
47    name = ReconcileDatabaseFunction,
48    display_name = reconcile_database,
49    sig_fn = signature,
50    ret = string
51)]
52pub(crate) async fn reconcile_database(
53    procedure_service_handler: &ProcedureServiceHandlerRef,
54    query_ctx: &QueryContextRef,
55    params: &[ValueRef<'_>],
56) -> Result<Value> {
57    let (database_name, resolve_strategy, parallelism) = match params.len() {
58        1 => (
59            get_string_from_params(params, 0, FN_NAME)?,
60            default_resolve_strategy(),
61            default_parallelism(),
62        ),
63        2 => (
64            get_string_from_params(params, 0, FN_NAME)?,
65            parse_resolve_strategy(get_string_from_params(params, 1, FN_NAME)?)?,
66            default_parallelism(),
67        ),
68        3 => {
69            let Some(parallelism) = cast_u32(&params[2])? else {
70                return UnsupportedInputDataTypeSnafu {
71                    function: FN_NAME,
72                    datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
73                }
74                .fail();
75            };
76            (
77                get_string_from_params(params, 0, FN_NAME)?,
78                parse_resolve_strategy(get_string_from_params(params, 1, FN_NAME)?)?,
79                parallelism,
80            )
81        }
82        size => {
83            return InvalidFuncArgsSnafu {
84                err_msg: format!(
85                    "The length of the args is not correct, expect 1, 2 or 3, have: {}",
86                    size
87                ),
88            }
89            .fail();
90        }
91    };
92    info!(
93        "Reconciling database: {}, resolve_strategy: {:?}, parallelism: {}",
94        database_name, resolve_strategy, parallelism
95    );
96    let pid = procedure_service_handler
97        .reconcile(ReconcileRequest {
98            target: Some(Target::ReconcileDatabase(ReconcileDatabase {
99                catalog_name: query_ctx.current_catalog().to_string(),
100                database_name: database_name.to_string(),
101                parallelism,
102                resolve_strategy: resolve_strategy as i32,
103            })),
104            ..Default::default()
105        })
106        .await?;
107    match pid {
108        Some(pid) => {
109            query_ctx.set_last_procedure_id(pid.clone());
110            Ok(Value::from(pid))
111        }
112        None => Ok(Value::Null),
113    }
114}
115
116fn signature() -> Signature {
117    let nums = ConcreteDataType::numerics();
118    let mut signs = Vec::with_capacity(2 + nums.len());
119    signs.extend([
120        // reconcile_database(datanode_name)
121        TypeSignature::Exact(vec![ArrowDataType::Utf8]),
122        // reconcile_database(database_name, resolve_strategy)
123        TypeSignature::Exact(vec![ArrowDataType::Utf8, ArrowDataType::Utf8]),
124    ]);
125    for sign in nums {
126        // reconcile_database(database_name, resolve_strategy, parallelism)
127        signs.push(TypeSignature::Exact(vec![
128            ArrowDataType::Utf8,
129            ArrowDataType::Utf8,
130            sign.as_arrow_type(),
131        ]));
132    }
133    Signature::one_of(signs, Volatility::Immutable)
134}
135
136#[cfg(test)]
137mod tests {
138    use std::sync::Arc;
139
140    use arrow::array::{StringArray, UInt32Array};
141    use arrow::datatypes::{DataType, Field};
142    use datafusion_expr::ColumnarValue;
143
144    use crate::admin::reconcile_database::ReconcileDatabaseFunction;
145    use crate::function::FunctionContext;
146    use crate::function_factory::ScalarFunctionFactory;
147
148    #[tokio::test]
149    async fn test_reconcile_catalog() {
150        common_telemetry::init_default_ut_logging();
151
152        // reconcile_database(database_name)
153        let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
154        let context = FunctionContext::mock();
155        let query_ctx = context.query_ctx.clone();
156        query_ctx.set_last_procedure_id("previous-id".to_string());
157        let provider = factory.provide(context);
158        let f = provider.as_async().unwrap();
159
160        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
161            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
162                "test",
163            ])))],
164            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))],
165            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
166            number_rows: 1,
167            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
168        };
169        let result = f.invoke_async_with_args(func_args).await.unwrap();
170        assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
171        match result {
172            ColumnarValue::Array(array) => {
173                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
174                assert_eq!(result_array.value(0), "test_pid");
175            }
176            ColumnarValue::Scalar(scalar) => {
177                assert_eq!(
178                    scalar,
179                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
180                );
181            }
182        }
183
184        // reconcile_database(database_name, resolve_strategy)
185        let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
186        let provider = factory.provide(FunctionContext::mock());
187        let f = provider.as_async().unwrap();
188
189        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
190            args: vec![
191                ColumnarValue::Array(Arc::new(StringArray::from(vec!["test"]))),
192                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseLatest"]))),
193            ],
194            arg_fields: vec![
195                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
196                Arc::new(Field::new("arg_1", DataType::Utf8, false)),
197            ],
198            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
199            number_rows: 1,
200            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
201        };
202        let result = f.invoke_async_with_args(func_args).await.unwrap();
203        match result {
204            ColumnarValue::Array(array) => {
205                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
206                assert_eq!(result_array.value(0), "test_pid");
207            }
208            ColumnarValue::Scalar(scalar) => {
209                assert_eq!(
210                    scalar,
211                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
212                );
213            }
214        }
215
216        // reconcile_database(database_name, resolve_strategy, parallelism)
217        let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
218        let provider = factory.provide(FunctionContext::mock());
219        let f = provider.as_async().unwrap();
220
221        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
222            args: vec![
223                ColumnarValue::Array(Arc::new(StringArray::from(vec!["test"]))),
224                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseLatest"]))),
225                ColumnarValue::Array(Arc::new(UInt32Array::from(vec![10]))),
226            ],
227            arg_fields: vec![
228                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
229                Arc::new(Field::new("arg_1", DataType::Utf8, false)),
230                Arc::new(Field::new("arg_2", DataType::UInt32, false)),
231            ],
232            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
233            number_rows: 1,
234            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
235        };
236        let result = f.invoke_async_with_args(func_args).await.unwrap();
237        match result {
238            ColumnarValue::Array(array) => {
239                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
240                assert_eq!(result_array.value(0), "test_pid");
241            }
242            ColumnarValue::Scalar(scalar) => {
243                assert_eq!(
244                    scalar,
245                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
246                );
247            }
248        }
249
250        // invalid function args
251        let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
252        let provider = factory.provide(FunctionContext::mock());
253        let f = provider.as_async().unwrap();
254
255        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
256            args: vec![
257                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseLatest"]))),
258                ColumnarValue::Array(Arc::new(UInt32Array::from(vec![10]))),
259                ColumnarValue::Array(Arc::new(StringArray::from(vec!["v1"]))),
260                ColumnarValue::Array(Arc::new(StringArray::from(vec!["v2"]))),
261            ],
262            arg_fields: vec![
263                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
264                Arc::new(Field::new("arg_1", DataType::UInt32, false)),
265                Arc::new(Field::new("arg_2", DataType::Utf8, false)),
266                Arc::new(Field::new("arg_3", DataType::Utf8, false)),
267            ],
268            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
269            number_rows: 1,
270            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
271        };
272        let _err = f.invoke_async_with_args(func_args).await.unwrap_err();
273        // Note: Error type is DataFusionError at this level, not common_query::Error
274
275        // unsupported input data type
276        let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
277        let provider = factory.provide(FunctionContext::mock());
278        let f = provider.as_async().unwrap();
279
280        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
281            args: vec![
282                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseLatest"]))),
283                ColumnarValue::Array(Arc::new(UInt32Array::from(vec![10]))),
284                ColumnarValue::Array(Arc::new(StringArray::from(vec!["v1"]))),
285            ],
286            arg_fields: vec![
287                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
288                Arc::new(Field::new("arg_1", DataType::UInt32, false)),
289                Arc::new(Field::new("arg_2", DataType::Utf8, false)),
290            ],
291            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
292            number_rows: 1,
293            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
294        };
295        let _err = f.invoke_async_with_args(func_args).await.unwrap_err();
296        // Note: Error type is DataFusionError at this level, not common_query::Error
297    }
298}
common_function/admin/reconcile_database.rs

common_function/admin/
reconcile_database.rs
