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

common_function/admin/
reconcile_catalog.rs
