common_function/admin/
purge_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 arrow::datatypes::DataType as ArrowDataType;
16use common_error::ext::BoxedError;
17use common_macro::admin_fn;
18use common_query::error::{
19    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, ProcedureServiceSnafu, Result,
20    UnsupportedInputDataTypeSnafu,
21};
22use datafusion_expr::{Signature, Volatility};
23use datatypes::prelude::{Value, ValueRef};
24use session::context::QueryContextRef;
25use session::table_name::table_name_to_full_name;
26use snafu::{ResultExt, ensure};
27use table::table_name::TableName;
28
29use crate::handlers::ProcedureServiceHandlerRef;
30
31/// Purges a dropped table after the ADMIN statement layer authorizes the request.
32/// Name validation retains the session helper's existing catalog-only permission check.
33#[admin_fn(
34    name = PurgeTableFunction,
35    display_name = purge_table,
36    sig_fn = purge_table_signature,
37    ret = uint64,
38    single_row
39)]
40pub(crate) async fn purge_table(
41    procedure_service_handler: &ProcedureServiceHandlerRef,
42    query_ctx: &QueryContextRef,
43    params: &[ValueRef<'_>],
44) -> Result<Value> {
45    ensure!(
46        params.len() == 1,
47        InvalidFuncArgsSnafu {
48            err_msg: format!(
49                "The length of the args is not correct, expect 1, have: {}",
50                params.len()
51            ),
52        }
53    );
54    let ValueRef::String(table_name) = params[0] else {
55        return UnsupportedInputDataTypeSnafu {
56            function: "purge_table",
57            datatypes: params
58                .iter()
59                .map(|value| value.data_type())
60                .collect::<Vec<_>>(),
61        }
62        .fail();
63    };
64    let (catalog_name, schema_name, table_name) = table_name_to_full_name(table_name, query_ctx)
65        .map_err(BoxedError::new)
66        .context(ProcedureServiceSnafu)?;
67
68    procedure_service_handler
69        .purge_table(
70            TableName::new(catalog_name, schema_name, table_name),
71            query_ctx.clone(),
72        )
73        .await?;
74    Ok(Value::from(0_u64))
75}
76
77fn purge_table_signature() -> Signature {
78    Signature::uniform(1, vec![ArrowDataType::Utf8], Volatility::Immutable)
79}
80
81#[cfg(test)]
82mod tests {
83    use std::sync::{Arc, Mutex};
84
85    use api::v1::meta::ReconcileRequest;
86    use arrow::array::StringArray;
87    use arrow::datatypes::{DataType as ArrowDataType, Field};
88    use async_trait::async_trait;
89    use catalog::CatalogManagerRef;
90    use common_macro::admin_fn;
91    use common_meta::rpc::procedure::{
92        GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest,
93        MigrateRegionRequest, ProcedureStateResponse,
94    };
95    use common_query::error::{InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result};
96    use datafusion_expr::{ColumnarValue, Signature, TypeSignature, Volatility};
97    use datatypes::prelude::{Value, ValueRef};
98    use session::context::{QueryContext, QueryContextBuilder, QueryContextRef};
99    use table::table_name::TableName;
100
101    use super::{PurgeTableFunction, purge_table};
102    use crate::function_factory::ScalarFunctionFactory;
103    use crate::function_registry::{FUNCTION_REGISTRY, get_admin_function};
104    use crate::handlers::{ProcedureServiceHandler, ProcedureServiceHandlerRef};
105    use crate::state::FunctionState;
106
107    #[admin_fn(
108        name = ZeroArgSingleRowFunction,
109        display_name = zero_arg_single_row,
110        sig_fn = zero_arg_single_row_signature,
111        ret = uint64,
112        single_row
113    )]
114    async fn zero_arg_single_row(
115        _procedure_service_handler: &ProcedureServiceHandlerRef,
116        _query_ctx: &QueryContextRef,
117        params: &[ValueRef<'_>],
118    ) -> Result<Value> {
119        assert!(params.is_empty());
120        Ok(Value::from(0_u64))
121    }
122
123    fn zero_arg_single_row_signature() -> Signature {
124        Signature::nullary(Volatility::Immutable)
125    }
126
127    #[derive(Default)]
128    struct RecordingHandler {
129        calls: Mutex<Vec<(TableName, QueryContextRef)>>,
130        fail: bool,
131    }
132
133    #[async_trait]
134    impl ProcedureServiceHandler for RecordingHandler {
135        async fn purge_table(
136            &self,
137            table_name: TableName,
138            query_ctx: QueryContextRef,
139        ) -> Result<()> {
140            if self.fail {
141                return InvalidFuncArgsSnafu {
142                    err_msg: "purge failed",
143                }
144                .fail();
145            }
146            self.calls.lock().unwrap().push((table_name, query_ctx));
147            Ok(())
148        }
149
150        async fn migrate_region(&self, _: MigrateRegionRequest) -> Result<Option<String>> {
151            unreachable!()
152        }
153        async fn reconcile(&self, _: ReconcileRequest) -> Result<Option<String>> {
154            unreachable!()
155        }
156        async fn query_procedure_state(&self, _: &str) -> Result<ProcedureStateResponse> {
157            unreachable!()
158        }
159        async fn manage_region_follower(&self, _: ManageRegionFollowerRequest) -> Result<()> {
160            unreachable!()
161        }
162        fn catalog_manager(&self) -> &CatalogManagerRef {
163            unreachable!()
164        }
165        async fn gc_regions(&self, _: GcRegionsRequest) -> Result<GcResponse> {
166            unreachable!()
167        }
168        async fn gc_table(&self, _: GcTableRequest) -> Result<GcResponse> {
169            unreachable!()
170        }
171    }
172
173    #[test]
174    fn test_purge_table_is_admin_only() {
175        assert!(get_admin_function("purge_table").is_some());
176        assert!(FUNCTION_REGISTRY.get_function("purge_table").is_none());
177    }
178
179    #[tokio::test]
180    async fn test_admin_only_factory_executes_without_registry_leak() {
181        let handler = Arc::new(RecordingHandler::default());
182        let state = FunctionState {
183            procedure_service_handler: Some(handler.clone()),
184            ..Default::default()
185        };
186        let function =
187            get_admin_function("purge_table")
188                .unwrap()
189                .provide(crate::function::FunctionContext {
190                    query_ctx: QueryContext::arc(),
191                    state: Arc::new(state),
192                });
193        let args = datafusion_expr::ScalarFunctionArgs {
194            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
195                "foo",
196            ])))],
197            arg_fields: vec![Arc::new(Field::new("arg_0", ArrowDataType::Utf8, false))],
198            return_field: Arc::new(Field::new("result", ArrowDataType::UInt64, false)),
199            number_rows: 1,
200            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
201        };
202
203        function
204            .as_async()
205            .unwrap()
206            .invoke_async_with_args(args)
207            .await
208            .unwrap();
209
210        assert_eq!(1, handler.calls.lock().unwrap().len());
211        assert!(FUNCTION_REGISTRY.get_function("purge_table").is_none());
212    }
213
214    #[test]
215    fn test_purge_table_signature() {
216        let factory: ScalarFunctionFactory = PurgeTableFunction::factory().into();
217        let function = factory.provide(crate::function::FunctionContext::mock());
218        assert_eq!("purge_table", function.name());
219        assert_eq!(ArrowDataType::UInt64, function.return_type(&[]).unwrap());
220        assert!(matches!(
221            function.signature(),
222            Signature { type_signature: TypeSignature::Uniform(1, types), volatility: Volatility::Immutable, .. }
223                if types == &vec![ArrowDataType::Utf8]
224        ));
225    }
226
227    #[tokio::test]
228    async fn test_purge_table_rejects_multi_row_before_handler_call() {
229        let handler = Arc::new(RecordingHandler::default());
230        let state = FunctionState {
231            procedure_service_handler: Some(handler.clone()),
232            ..Default::default()
233        };
234        let factory: ScalarFunctionFactory = PurgeTableFunction::factory().into();
235        let function = factory.provide(crate::function::FunctionContext {
236            query_ctx: QueryContext::arc(),
237            state: Arc::new(state),
238        });
239        let args = datafusion_expr::ScalarFunctionArgs {
240            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
241                Some("foo"),
242                None,
243            ])))],
244            arg_fields: vec![Arc::new(Field::new("arg_0", ArrowDataType::Utf8, true))],
245            return_field: Arc::new(Field::new("result", ArrowDataType::UInt64, true)),
246            number_rows: 2,
247            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
248        };
249
250        assert!(
251            function
252                .as_async()
253                .unwrap()
254                .invoke_async_with_args(args)
255                .await
256                .is_err()
257        );
258        assert!(handler.calls.lock().unwrap().is_empty());
259    }
260
261    #[tokio::test]
262    async fn test_single_row_admin_fn_rejects_zero_arg_multi_row_invocation() {
263        let handler = Arc::new(RecordingHandler::default());
264        let state = FunctionState {
265            procedure_service_handler: Some(handler),
266            ..Default::default()
267        };
268        let factory: ScalarFunctionFactory = ZeroArgSingleRowFunction::factory().into();
269        let function = factory.provide(crate::function::FunctionContext {
270            query_ctx: QueryContext::arc(),
271            state: Arc::new(state),
272        });
273        let args = datafusion_expr::ScalarFunctionArgs {
274            args: vec![],
275            arg_fields: vec![],
276            return_field: Arc::new(Field::new("result", ArrowDataType::UInt64, false)),
277            number_rows: 2,
278            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
279        };
280
281        let err = function
282            .as_async()
283            .unwrap()
284            .invoke_async_with_args(args)
285            .await
286            .unwrap_err();
287        assert!(err.to_string().contains("received 2"));
288    }
289
290    #[tokio::test]
291    async fn test_purge_table_resolves_names_and_returns_zero() {
292        let handler = Arc::new(RecordingHandler::default());
293        let handler_ref: ProcedureServiceHandlerRef = handler.clone();
294        let query_ctx = QueryContextBuilder::default()
295            .current_catalog("catalog".to_string())
296            .current_schema("schema".to_string())
297            .build()
298            .into();
299
300        for (input, expected) in [
301            ("foo", TableName::new("catalog", "schema", "foo")),
302            ("other.foo", TableName::new("catalog", "other", "foo")),
303            (
304                "catalog.other.foo",
305                TableName::new("catalog", "other", "foo"),
306            ),
307        ] {
308            assert_eq!(
309                Value::from(0_u64),
310                purge_table(&handler_ref, &query_ctx, &[ValueRef::String(input)])
311                    .await
312                    .unwrap()
313            );
314            let calls = handler.calls.lock().unwrap();
315            assert_eq!(expected, calls.last().unwrap().0);
316            assert!(Arc::ptr_eq(&query_ctx, &calls.last().unwrap().1));
317        }
318    }
319
320    #[tokio::test]
321    async fn test_purge_table_rejects_invalid_arguments() {
322        let recording_handler = Arc::new(RecordingHandler::default());
323        let handler: ProcedureServiceHandlerRef = recording_handler.clone();
324        for params in [vec![], vec![ValueRef::String("a"), ValueRef::String("b")]] {
325            assert!(
326                purge_table(&handler, &QueryContext::arc(), &params)
327                    .await
328                    .is_err()
329            );
330        }
331        assert!(
332            purge_table(&handler, &QueryContext::arc(), &[ValueRef::Int64(1)])
333                .await
334                .is_err()
335        );
336        assert!(
337            purge_table(&handler, &QueryContext::arc(), &[ValueRef::Null])
338                .await
339                .is_err()
340        );
341        assert!(
342            purge_table(
343                &handler,
344                &QueryContext::arc(),
345                &[ValueRef::String("a.b.c.d")]
346            )
347            .await
348            .is_err()
349        );
350        assert!(recording_handler.calls.lock().unwrap().is_empty());
351    }
352
353    #[tokio::test]
354    async fn test_purge_table_propagates_handler_error() {
355        let handler: ProcedureServiceHandlerRef = Arc::new(RecordingHandler {
356            fail: true,
357            ..Default::default()
358        });
359        let error = purge_table(&handler, &QueryContext::arc(), &[ValueRef::String("foo")])
360            .await
361            .unwrap_err();
362        assert!(error.to_string().contains("purge failed"));
363    }
364}
common_function/admin/purge_table.rs

common_function/admin/
purge_table.rs
