common_function/admin/
reconcile_table.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::{ReconcileRequest, ReconcileTable, ResolveStrategy};
17use arrow::datatypes::DataType as ArrowDataType;
18use common_catalog::format_full_table_name;
19use common_error::ext::BoxedError;
20use common_macro::admin_fn;
21use common_query::error::{
22    MissingProcedureServiceHandlerSnafu, Result, TableMutationSnafu, UnsupportedInputDataTypeSnafu,
23};
24use common_telemetry::info;
25use datafusion_expr::{Signature, TypeSignature, Volatility};
26use datatypes::prelude::*;
27use session::context::QueryContextRef;
28use session::table_name::table_name_to_full_name;
29use snafu::ResultExt;
30
31use crate::handlers::ProcedureServiceHandlerRef;
32use crate::helper::parse_resolve_strategy;
33
34const FN_NAME: &str = "reconcile_table";
35
36/// A function to reconcile a table.
37/// Returns the procedure id if success.
38///
39/// - `reconcile_table(table_name)`.
40/// - `reconcile_table(table_name, resolve_strategy)`.
41///
42/// The parameters:
43/// - `table_name`:  the table name
44#[admin_fn(
45    name = ReconcileTableFunction,
46    display_name = reconcile_table,
47    sig_fn = signature,
48    ret = string
49)]
50pub(crate) async fn reconcile_table(
51    procedure_service_handler: &ProcedureServiceHandlerRef,
52    query_ctx: &QueryContextRef,
53    params: &[ValueRef<'_>],
54) -> Result<Value> {
55    let (table_name, resolve_strategy) = match params {
56        [ValueRef::String(table_name)] => (table_name, ResolveStrategy::UseLatest),
57        [
58            ValueRef::String(table_name),
59            ValueRef::String(resolve_strategy),
60        ] => (table_name, parse_resolve_strategy(resolve_strategy)?),
61        _ => {
62            return UnsupportedInputDataTypeSnafu {
63                function: FN_NAME,
64                datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
65            }
66            .fail();
67        }
68    };
69    let (catalog_name, schema_name, table_name) = table_name_to_full_name(table_name, query_ctx)
70        .map_err(BoxedError::new)
71        .context(TableMutationSnafu)?;
72    info!(
73        "Reconciling table: {} with resolve_strategy: {:?}",
74        format_full_table_name(&catalog_name, &schema_name, &table_name),
75        resolve_strategy
76    );
77    let pid = procedure_service_handler
78        .reconcile(ReconcileRequest {
79            target: Some(Target::ReconcileTable(ReconcileTable {
80                catalog_name,
81                schema_name,
82                table_name,
83                resolve_strategy: resolve_strategy as i32,
84            })),
85            ..Default::default()
86        })
87        .await?;
88    match pid {
89        Some(pid) => {
90            query_ctx.set_last_procedure_id(pid.clone());
91            Ok(Value::from(pid))
92        }
93        None => Ok(Value::Null),
94    }
95}
96
97fn signature() -> Signature {
98    Signature::one_of(
99        vec![
100            // reconcile_table(table_name)
101            TypeSignature::Exact(vec![ArrowDataType::Utf8]),
102            // reconcile_table(table_name, resolve_strategy)
103            TypeSignature::Exact(vec![ArrowDataType::Utf8, ArrowDataType::Utf8]),
104        ],
105        Volatility::Immutable,
106    )
107}
108
109#[cfg(test)]
110mod tests {
111    use std::sync::Arc;
112
113    use arrow::array::StringArray;
114    use arrow::datatypes::{DataType, Field};
115    use datafusion_expr::ColumnarValue;
116
117    use crate::admin::reconcile_table::ReconcileTableFunction;
118    use crate::function::FunctionContext;
119    use crate::function_factory::ScalarFunctionFactory;
120
121    #[tokio::test]
122    async fn test_reconcile_table() {
123        common_telemetry::init_default_ut_logging();
124
125        // reconcile_table(table_name)
126        let factory: ScalarFunctionFactory = ReconcileTableFunction::factory().into();
127        let context = FunctionContext::mock();
128        let query_ctx = context.query_ctx.clone();
129        query_ctx.set_last_procedure_id("previous-id".to_string());
130        let provider = factory.provide(context);
131        let f = provider.as_async().unwrap();
132
133        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
134            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
135                "test",
136            ])))],
137            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))],
138            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
139            number_rows: 1,
140            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
141        };
142        let result = f.invoke_async_with_args(func_args).await.unwrap();
143        assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
144        match result {
145            ColumnarValue::Array(array) => {
146                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
147                assert_eq!(result_array.value(0), "test_pid");
148            }
149            ColumnarValue::Scalar(scalar) => {
150                assert_eq!(
151                    scalar,
152                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
153                );
154            }
155        }
156
157        // reconcile_table(table_name, resolve_strategy)
158        let factory: ScalarFunctionFactory = ReconcileTableFunction::factory().into();
159        let provider = factory.provide(FunctionContext::mock());
160        let f = provider.as_async().unwrap();
161
162        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
163            args: vec![
164                ColumnarValue::Array(Arc::new(StringArray::from(vec!["test"]))),
165                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseMetasrv"]))),
166            ],
167            arg_fields: vec![
168                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
169                Arc::new(Field::new("arg_1", DataType::Utf8, false)),
170            ],
171            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
172            number_rows: 1,
173            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
174        };
175        let result = f.invoke_async_with_args(func_args).await.unwrap();
176        match result {
177            ColumnarValue::Array(array) => {
178                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
179                assert_eq!(result_array.value(0), "test_pid");
180            }
181            ColumnarValue::Scalar(scalar) => {
182                assert_eq!(
183                    scalar,
184                    datafusion_common::ScalarValue::Utf8(Some("test_pid".to_string()))
185                );
186            }
187        }
188
189        // unsupported input data type
190        let factory: ScalarFunctionFactory = ReconcileTableFunction::factory().into();
191        let provider = factory.provide(FunctionContext::mock());
192        let f = provider.as_async().unwrap();
193
194        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
195            args: vec![
196                ColumnarValue::Array(Arc::new(StringArray::from(vec!["test"]))),
197                ColumnarValue::Array(Arc::new(StringArray::from(vec!["UseMetasrv"]))),
198                ColumnarValue::Array(Arc::new(StringArray::from(vec!["10"]))),
199            ],
200            arg_fields: vec![
201                Arc::new(Field::new("arg_0", DataType::Utf8, false)),
202                Arc::new(Field::new("arg_1", DataType::Utf8, false)),
203                Arc::new(Field::new("arg_2", DataType::Utf8, false)),
204            ],
205            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
206            number_rows: 1,
207            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
208        };
209        let _err = f.invoke_async_with_args(func_args).await.unwrap_err();
210        // Note: Error type is DataFusionError at this level, not common_query::Error
211    }
212}
common_function/admin/reconcile_table.rs

common_function/admin/
reconcile_table.rs
