operator/
procedure.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::ReconcileRequest;
16use async_trait::async_trait;
17use catalog::CatalogManagerRef;
18use common_error::ext::BoxedError;
19use common_function::handlers::ProcedureServiceHandler;
20use common_meta::key::TableMetadataManagerRef;
21use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutorRef};
22use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, TriggerReason};
23use common_meta::rpc::procedure::{
24    GcRegionsRequest as MetaGcRegionsRequest, GcResponse as MetaGcResponse,
25    GcTableRequest as MetaGcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest,
26    ProcedureStateResponse,
27};
28use common_query::error as query_error;
29use common_query::error::Result as QueryResult;
30use session::context::QueryContextRef;
31use snafu::ResultExt;
32use table::table_name::TableName;
33
34use crate::error;
35use crate::utils::to_executor_context;
36
37/// The operator for procedures which implements [`ProcedureServiceHandler`].
38#[derive(Clone)]
39pub struct ProcedureServiceOperator {
40    procedure_executor: ProcedureExecutorRef,
41    catalog_manager: CatalogManagerRef,
42    table_metadata_manager: TableMetadataManagerRef,
43}
44
45impl ProcedureServiceOperator {
46    pub fn new(
47        procedure_executor: ProcedureExecutorRef,
48        catalog_manager: CatalogManagerRef,
49        table_metadata_manager: TableMetadataManagerRef,
50    ) -> Self {
51        Self {
52            procedure_executor,
53            catalog_manager,
54            table_metadata_manager,
55        }
56    }
57}
58
59#[async_trait]
60impl ProcedureServiceHandler for ProcedureServiceOperator {
61    async fn purge_table(
62        &self,
63        query_ctx: QueryContextRef,
64        table_name: TableName,
65    ) -> QueryResult<()> {
66        let dropped = self
67            .table_metadata_manager
68            .get_dropped_table(&table_name)
69            .await
70            .map_err(BoxedError::new)
71            .context(query_error::ProcedureServiceSnafu)?
72            .ok_or_else(|| {
73                error::TableNotFoundSnafu {
74                    table_name: table_name.to_string(),
75                }
76                .build()
77            })
78            .map_err(BoxedError::new)
79            .context(query_error::ProcedureServiceSnafu)?;
80        let executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
81        let request = SubmitDdlTaskRequest::new(DdlTask::new_purge_dropped_table(dropped.table_id));
82        self.procedure_executor
83            .submit_ddl_task(executor_context, request)
84            .await
85            .map_err(BoxedError::new)
86            .context(query_error::ProcedureServiceSnafu)?;
87        Ok(())
88    }
89
90    async fn migrate_region(
91        &self,
92        query_ctx: QueryContextRef,
93        request: MigrateRegionRequest,
94    ) -> QueryResult<Option<String>> {
95        let executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
96        Ok(self
97            .procedure_executor
98            .migrate_region(&executor_context, request)
99            .await
100            .map_err(BoxedError::new)
101            .context(query_error::ProcedureServiceSnafu)?
102            .pid
103            .map(|pid| String::from_utf8_lossy(&pid.key).to_string()))
104    }
105
106    async fn reconcile(&self, request: ReconcileRequest) -> QueryResult<Option<String>> {
107        Ok(self
108            .procedure_executor
109            .reconcile(&ExecutorContext::default(), request)
110            .await
111            .map_err(BoxedError::new)
112            .context(query_error::ProcedureServiceSnafu)?
113            .pid
114            .map(|pid| String::from_utf8_lossy(&pid.key).to_string()))
115    }
116
117    async fn query_procedure_state(&self, pid: &str) -> QueryResult<ProcedureStateResponse> {
118        self.procedure_executor
119            .query_procedure_state(&ExecutorContext::default(), pid)
120            .await
121            .map_err(BoxedError::new)
122            .context(query_error::ProcedureServiceSnafu)
123    }
124
125    async fn manage_region_follower(
126        &self,
127        query_ctx: QueryContextRef,
128        request: ManageRegionFollowerRequest,
129    ) -> QueryResult<()> {
130        let executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
131        self.procedure_executor
132            .manage_region_follower(&executor_context, request)
133            .await
134            .map_err(BoxedError::new)
135            .context(query_error::ProcedureServiceSnafu)
136    }
137
138    fn catalog_manager(&self) -> &CatalogManagerRef {
139        &self.catalog_manager
140    }
141
142    async fn gc_regions(
143        &self,
144        query_ctx: QueryContextRef,
145        request: MetaGcRegionsRequest,
146    ) -> QueryResult<MetaGcResponse> {
147        let executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
148        self.procedure_executor
149            .gc_regions(&executor_context, request)
150            .await
151            .map_err(BoxedError::new)
152            .context(query_error::ProcedureServiceSnafu)
153    }
154
155    async fn gc_table(
156        &self,
157        query_ctx: QueryContextRef,
158        request: MetaGcTableRequest,
159    ) -> QueryResult<MetaGcResponse> {
160        let executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
161        self.procedure_executor
162            .gc_table(&executor_context, request)
163            .await
164            .map_err(BoxedError::new)
165            .context(query_error::ProcedureServiceSnafu)
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use std::collections::HashMap;
172    use std::sync::{Arc, Mutex};
173
174    use api::v1::meta::{ProcedureDetailResponse, ReconcileResponse};
175    use common_meta::key::TableMetadataManager;
176    use common_meta::key::table_route::TableRouteValue;
177    use common_meta::key::test_utils::new_test_table_info_with_name;
178    use common_meta::kv_backend::memory::MemoryKvBackend;
179    use common_meta::procedure_executor::{
180        ExecutorContext, ProcedureExecutor, ProcedureExecutorRef,
181    };
182    use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
183    use common_meta::rpc::procedure::{
184        AddRegionFollowerRequest, ManageRegionFollowerRequest, MigrateRegionResponse,
185        ProcedureStateResponse,
186    };
187    use session::context::{QueryContextBuilder, QueryContextRef};
188    use table::table_name::TableName;
189
190    use super::*;
191
192    #[derive(Default)]
193    struct RecordingProcedureExecutor {
194        requests: Mutex<Vec<SubmitDdlTaskRequest>>,
195        contexts: Mutex<Vec<ExecutorContext>>,
196        fail: bool,
197    }
198
199    #[async_trait]
200    impl ProcedureExecutor for RecordingProcedureExecutor {
201        async fn submit_ddl_task(
202            &self,
203            context: ExecutorContext,
204            request: SubmitDdlTaskRequest,
205        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
206            self.contexts.lock().unwrap().push(context);
207            self.requests.lock().unwrap().push(request);
208            if self.fail {
209                return common_meta::error::UnsupportedSnafu {
210                    operation: "test purge failure",
211                }
212                .fail();
213            }
214            Ok(SubmitDdlTaskResponse::default())
215        }
216        async fn migrate_region(
217            &self,
218            _: &ExecutorContext,
219            _: MigrateRegionRequest,
220        ) -> common_meta::error::Result<MigrateRegionResponse> {
221            unimplemented!()
222        }
223        async fn reconcile(
224            &self,
225            _: &ExecutorContext,
226            _: ReconcileRequest,
227        ) -> common_meta::error::Result<ReconcileResponse> {
228            unimplemented!()
229        }
230        async fn query_procedure_state(
231            &self,
232            _: &ExecutorContext,
233            _: &str,
234        ) -> common_meta::error::Result<ProcedureStateResponse> {
235            unimplemented!()
236        }
237        async fn manage_region_follower(
238            &self,
239            context: &ExecutorContext,
240            _: ManageRegionFollowerRequest,
241        ) -> common_meta::error::Result<()> {
242            self.contexts.lock().unwrap().push(context.clone());
243            Ok(())
244        }
245        async fn list_procedures(
246            &self,
247            _: &ExecutorContext,
248        ) -> common_meta::error::Result<ProcedureDetailResponse> {
249            unimplemented!()
250        }
251    }
252
253    async fn operator_with_tombstone(
254        table_id: u32,
255        name: &TableName,
256        fail: bool,
257    ) -> (ProcedureServiceOperator, Arc<RecordingProcedureExecutor>) {
258        let manager = Arc::new(TableMetadataManager::new(Arc::new(
259            MemoryKvBackend::default(),
260        )));
261        let mut info = new_test_table_info_with_name(table_id, &name.table_name);
262        info.catalog_name = name.catalog_name.clone();
263        info.schema_name = name.schema_name.clone();
264        let route = TableRouteValue::physical(vec![]);
265        manager
266            .create_table_metadata(info, route.clone(), HashMap::new())
267            .await
268            .unwrap();
269        manager
270            .delete_table_metadata(table_id, name, &route, &HashMap::new(), None)
271            .await
272            .unwrap();
273        let executor = Arc::new(RecordingProcedureExecutor {
274            fail,
275            ..Default::default()
276        });
277        let operator = ProcedureServiceOperator::new(
278            executor.clone() as ProcedureExecutorRef,
279            catalog::memory::MemoryCatalogManager::new(),
280            manager,
281        );
282        (operator, executor)
283    }
284
285    #[tokio::test]
286    async fn test_purge_table_submits_tombstone_id_and_query_context() {
287        let name = TableName::new("catalog", "schema", "metrics");
288        let (operator, executor) = operator_with_tombstone(42, &name, false).await;
289        let manager = operator.table_metadata_manager.clone();
290        let mut live = new_test_table_info_with_name(99, "metrics");
291        live.catalog_name = "catalog".to_string();
292        live.schema_name = "schema".to_string();
293        manager
294            .create_table_metadata(live, TableRouteValue::physical(vec![]), HashMap::new())
295            .await
296            .unwrap();
297        let query_ctx = QueryContextBuilder::default()
298            .current_catalog("catalog".to_string())
299            .current_schema("schema".to_string())
300            .build()
301            .into();
302
303        operator.purge_table(query_ctx, name.clone()).await.unwrap();
304
305        let requests = executor.requests.lock().unwrap();
306        assert!(
307            matches!(&requests[0].task, DdlTask::PurgeDroppedTable(task) if task.table_id == 42)
308        );
309        let contexts = executor.contexts.lock().unwrap();
310        let query_context = contexts[0].query_context.as_ref().unwrap();
311        assert_eq!(query_context.current_catalog, "catalog");
312        assert_eq!(query_context.current_schema, "schema");
313        assert_eq!(
314            contexts[0].event_input.as_ref().map(|input| input.reason),
315            Some(TriggerReason::Manual)
316        );
317    }
318
319    #[tokio::test]
320    async fn test_manage_region_follower_forwards_query_context() {
321        let manager = Arc::new(TableMetadataManager::new(Arc::new(
322            MemoryKvBackend::default(),
323        )));
324        let executor = Arc::new(RecordingProcedureExecutor::default());
325        let operator = ProcedureServiceOperator::new(
326            executor.clone(),
327            catalog::memory::MemoryCatalogManager::new(),
328            manager,
329        );
330        let query_ctx: QueryContextRef = QueryContextBuilder::default().build().into();
331
332        operator
333            .manage_region_follower(
334                query_ctx,
335                ManageRegionFollowerRequest::AddRegionFollower(AddRegionFollowerRequest {
336                    region_id: 1,
337                    peer_id: 2,
338                }),
339            )
340            .await
341            .unwrap();
342
343        let contexts = executor.contexts.lock().unwrap();
344        assert_eq!(contexts.len(), 1);
345        // The default query context runs as the default user.
346        assert_eq!(contexts[0].actor.as_deref(), Some("greptime"));
347        assert_eq!(
348            contexts[0].event_input.as_ref().map(|input| input.reason),
349            Some(TriggerReason::Manual)
350        );
351    }
352
353    #[tokio::test]
354    async fn test_purge_table_missing_tombstone_is_table_not_found() {
355        let manager = Arc::new(TableMetadataManager::new(Arc::new(
356            MemoryKvBackend::default(),
357        )));
358        let executor = Arc::new(RecordingProcedureExecutor::default());
359        let operator = ProcedureServiceOperator::new(
360            executor.clone(),
361            catalog::memory::MemoryCatalogManager::new(),
362            manager,
363        );
364        let error = operator
365            .purge_table(
366                QueryContextBuilder::default()
367                    .current_catalog("catalog".to_string())
368                    .current_schema("schema".to_string())
369                    .build()
370                    .into(),
371                TableName::new("catalog", "schema", "missing"),
372            )
373            .await
374            .unwrap_err();
375        assert_eq!(
376            common_error::status_code::StatusCode::TableNotFound,
377            common_error::ext::ErrorExt::status_code(&error)
378        );
379        assert!(executor.requests.lock().unwrap().is_empty());
380    }
381
382    #[tokio::test]
383    async fn test_purge_table_propagates_submission_failure() {
384        let name = TableName::new("catalog", "schema", "metrics");
385        let (operator, executor) = operator_with_tombstone(42, &name, true).await;
386        assert!(
387            operator
388                .purge_table(
389                    QueryContextBuilder::default()
390                        .current_catalog("catalog".to_string())
391                        .current_schema("schema".to_string())
392                        .build()
393                        .into(),
394                    name,
395                )
396                .await
397                .is_err()
398        );
399        assert_eq!(1, executor.requests.lock().unwrap().len());
400    }
401}
operator/procedure.rs

operator/
procedure.rs
