common_function/admin/
gc.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_error::ext::BoxedError;
16use common_macro::admin_fn;
17use common_meta::rpc::procedure::{GcRegionsRequest, GcTableRequest};
18use common_query::error::{
19    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result, TableMutationSnafu,
20    UnsupportedInputDataTypeSnafu,
21};
22use datafusion_expr::{Signature, TypeSignature, Volatility};
23use datatypes::arrow::datatypes::DataType as ArrowDataType;
24use datatypes::prelude::*;
25use session::context::QueryContextRef;
26use snafu::{ResultExt, ensure};
27
28use crate::handlers::ProcedureServiceHandlerRef;
29use crate::helper::cast_u64;
30
31const DEFAULT_FULL_FILE_LISTING: bool = false;
32
33#[admin_fn(
34    name = GcRegionsFunction,
35    display_name = gc_regions,
36    sig_fn = gc_regions_signature,
37    ret = uint64
38)]
39pub(crate) async fn gc_regions(
40    procedure_service_handler: &ProcedureServiceHandlerRef,
41    _ctx: &QueryContextRef,
42    params: &[ValueRef<'_>],
43) -> Result<Value> {
44    let (region_ids, full_file_listing) = parse_gc_regions_params(params)?;
45
46    let resp = procedure_service_handler
47        .gc_regions(GcRegionsRequest {
48            region_ids,
49            full_file_listing,
50            timeout: None,
51        })
52        .await?;
53
54    Ok(Value::from(resp.processed_regions))
55}
56
57#[admin_fn(
58    name = GcTableFunction,
59    display_name = gc_table,
60    sig_fn = gc_table_signature,
61    ret = uint64
62)]
63pub(crate) async fn gc_table(
64    procedure_service_handler: &ProcedureServiceHandlerRef,
65    query_ctx: &QueryContextRef,
66    params: &[ValueRef<'_>],
67) -> Result<Value> {
68    let (catalog_name, schema_name, table_name, full_file_listing) =
69        parse_gc_table_params(params, query_ctx)?;
70
71    let resp = procedure_service_handler
72        .gc_table(GcTableRequest {
73            catalog_name,
74            schema_name,
75            table_name,
76            full_file_listing,
77            timeout: None,
78        })
79        .await?;
80
81    Ok(Value::from(resp.processed_regions))
82}
83
84fn parse_gc_regions_params(params: &[ValueRef<'_>]) -> Result<(Vec<u64>, bool)> {
85    ensure!(
86        !params.is_empty(),
87        InvalidFuncArgsSnafu {
88            err_msg: "The length of the args is not correct, expect at least 1 region id, have 0"
89                .to_string(),
90        }
91    );
92
93    let (full_file_listing, region_params) = match params.last() {
94        Some(ValueRef::Boolean(value)) => (*value, &params[..params.len() - 1]),
95        _ => (DEFAULT_FULL_FILE_LISTING, params),
96    };
97
98    ensure!(
99        !region_params.is_empty(),
100        InvalidFuncArgsSnafu {
101            err_msg: "The length of the args is not correct, expect at least 1 region id"
102                .to_string(),
103        }
104    );
105
106    let mut region_ids = Vec::with_capacity(region_params.len());
107    for param in region_params {
108        let Some(region_id) = cast_u64(param)? else {
109            return UnsupportedInputDataTypeSnafu {
110                function: "gc_regions",
111                datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
112            }
113            .fail();
114        };
115        region_ids.push(region_id);
116    }
117
118    Ok((region_ids, full_file_listing))
119}
120
121fn parse_gc_table_params(
122    params: &[ValueRef<'_>],
123    query_ctx: &QueryContextRef,
124) -> Result<(String, String, String, bool)> {
125    ensure!(
126        matches!(params.len(), 1 | 2),
127        InvalidFuncArgsSnafu {
128            err_msg: format!(
129                "The length of the args is not correct, expect 1 or 2, have: {}",
130                params.len()
131            ),
132        }
133    );
134
135    let ValueRef::String(table_name) = params[0] else {
136        return UnsupportedInputDataTypeSnafu {
137            function: "gc_table",
138            datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
139        }
140        .fail();
141    };
142
143    let full_file_listing = if params.len() == 2 {
144        let ValueRef::Boolean(value) = params[1] else {
145            return UnsupportedInputDataTypeSnafu {
146                function: "gc_table",
147                datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
148            }
149            .fail();
150        };
151        value
152    } else {
153        DEFAULT_FULL_FILE_LISTING
154    };
155
156    let (catalog_name, schema_name, table_name) =
157        session::table_name::table_name_to_full_name(table_name, query_ctx)
158            .map_err(BoxedError::new)
159            .context(TableMutationSnafu)?;
160
161    Ok((catalog_name, schema_name, table_name, full_file_listing))
162}
163
164fn gc_regions_signature() -> Signature {
165    Signature::variadic_any(Volatility::Immutable)
166}
167
168fn gc_table_signature() -> Signature {
169    Signature::one_of(
170        vec![
171            TypeSignature::Uniform(1, vec![ArrowDataType::Utf8]),
172            TypeSignature::Exact(vec![ArrowDataType::Utf8, ArrowDataType::Boolean]),
173        ],
174        Volatility::Immutable,
175    )
176}
177
178#[cfg(test)]
179mod tests {
180    use std::sync::{Arc, Mutex};
181
182    use api::v1::meta::ReconcileRequest;
183    use async_trait::async_trait;
184    use catalog::CatalogManagerRef;
185    use common_meta::rpc::procedure::{
186        GcResponse, ManageRegionFollowerRequest, MigrateRegionRequest, ProcedureStateResponse,
187    };
188    use session::context::QueryContext;
189
190    use super::*;
191    use crate::handlers::ProcedureServiceHandler;
192
193    #[test]
194    fn test_parse_gc_regions_params_with_full_file_listing() {
195        let params = vec![
196            ValueRef::UInt64(1),
197            ValueRef::UInt64(2),
198            ValueRef::Boolean(true),
199        ];
200        let (region_ids, full_file_listing) = parse_gc_regions_params(&params).unwrap();
201
202        assert_eq!(region_ids, vec![1, 2]);
203        assert!(full_file_listing);
204    }
205
206    #[test]
207    fn test_parse_gc_regions_params_default_full_file_listing() {
208        let params = vec![ValueRef::UInt64(1), ValueRef::UInt32(2)];
209        let (region_ids, full_file_listing) = parse_gc_regions_params(&params).unwrap();
210
211        assert_eq!(region_ids, vec![1, 2]);
212        assert!(!full_file_listing);
213    }
214
215    #[test]
216    fn test_parse_gc_table_params_with_full_file_listing() {
217        let params = vec![ValueRef::String("public.t"), ValueRef::Boolean(true)];
218        let (catalog, schema, table, full_file_listing) =
219            parse_gc_table_params(&params, &QueryContext::arc()).unwrap();
220
221        assert_eq!(catalog, "greptime");
222        assert_eq!(schema, "public");
223        assert_eq!(table, "t");
224        assert!(full_file_listing);
225    }
226
227    #[tokio::test]
228    async fn test_gc_regions_uses_meta_gc_timeout_config() {
229        let handler = Arc::new(MockProcedureServiceHandler::default());
230        let handler_ref: ProcedureServiceHandlerRef = handler.clone();
231        let params = vec![ValueRef::UInt64(1), ValueRef::Boolean(true)];
232
233        super::gc_regions(&handler_ref, &QueryContext::arc(), &params)
234            .await
235            .unwrap();
236
237        let request = handler.gc_regions_request.lock().unwrap().clone().unwrap();
238        assert_eq!(request.region_ids, vec![1]);
239        assert!(request.full_file_listing);
240        assert_eq!(request.timeout, None);
241    }
242
243    #[tokio::test]
244    async fn test_gc_table_uses_meta_gc_timeout_config() {
245        let handler = Arc::new(MockProcedureServiceHandler::default());
246        let handler_ref: ProcedureServiceHandlerRef = handler.clone();
247        let params = vec![ValueRef::String("public.t"), ValueRef::Boolean(true)];
248
249        super::gc_table(&handler_ref, &QueryContext::arc(), &params)
250            .await
251            .unwrap();
252
253        let request = handler.gc_table_request.lock().unwrap().clone().unwrap();
254        assert_eq!(request.catalog_name, "greptime");
255        assert_eq!(request.schema_name, "public");
256        assert_eq!(request.table_name, "t");
257        assert!(request.full_file_listing);
258        assert_eq!(request.timeout, None);
259    }
260
261    #[derive(Default)]
262    struct MockProcedureServiceHandler {
263        gc_regions_request: Mutex<Option<GcRegionsRequest>>,
264        gc_table_request: Mutex<Option<GcTableRequest>>,
265    }
266
267    #[async_trait]
268    impl ProcedureServiceHandler for MockProcedureServiceHandler {
269        async fn migrate_region(&self, _request: MigrateRegionRequest) -> Result<Option<String>> {
270            unreachable!()
271        }
272
273        async fn reconcile(&self, _request: ReconcileRequest) -> Result<Option<String>> {
274            unreachable!()
275        }
276
277        async fn query_procedure_state(&self, _pid: &str) -> Result<ProcedureStateResponse> {
278            unreachable!()
279        }
280
281        async fn manage_region_follower(
282            &self,
283            _request: ManageRegionFollowerRequest,
284        ) -> Result<()> {
285            unreachable!()
286        }
287
288        fn catalog_manager(&self) -> &CatalogManagerRef {
289            unreachable!()
290        }
291
292        async fn gc_regions(&self, request: GcRegionsRequest) -> Result<GcResponse> {
293            *self.gc_regions_request.lock().unwrap() = Some(request);
294            Ok(GcResponse::default())
295        }
296
297        async fn gc_table(&self, request: GcTableRequest) -> Result<GcResponse> {
298            *self.gc_table_request.lock().unwrap() = Some(request);
299            Ok(GcResponse::default())
300        }
301    }
302}
common_function/admin/gc.rs

common_function/admin/
gc.rs
