common_function/
state.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_base::Plugins;
16
17use crate::handlers::{FlowServiceHandlerRef, ProcedureServiceHandlerRef, TableMutationHandlerRef};
18
19/// Shared state for SQL functions.
20/// The handlers in state may be `None` in cli command-line or test cases.
21#[derive(Clone, Default)]
22pub struct FunctionState {
23    /// Process plugins available to dynamically registered functions.
24    pub plugins: Plugins,
25    // The table mutation handler
26    pub table_mutation_handler: Option<TableMutationHandlerRef>,
27    // The procedure service handler
28    pub procedure_service_handler: Option<ProcedureServiceHandlerRef>,
29    // The flownode handler
30    pub flow_service_handler: Option<FlowServiceHandlerRef>,
31}
32
33impl FunctionState {
34    /// Create a mock [`FunctionState`] for test.
35    #[cfg(any(test, feature = "testing"))]
36    pub fn mock() -> Self {
37        Self::mock_with_procedure_state_counter(Default::default())
38    }
39
40    /// Creates a mock state that counts procedure state queries.
41    #[cfg(any(test, feature = "testing"))]
42    pub fn mock_with_procedure_state_counter(
43        calls: std::sync::Arc<std::sync::atomic::AtomicUsize>,
44    ) -> Self {
45        use std::sync::Arc;
46        use std::sync::atomic::{AtomicUsize, Ordering};
47
48        use api::v1::meta::{ProcedureStatus, ReconcileRequest};
49        use async_trait::async_trait;
50        use catalog::CatalogManagerRef;
51        use common_base::AffectedRows;
52        use common_meta::rpc::procedure::{
53            GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest,
54            MigrateRegionRequest, ProcedureStateResponse,
55        };
56        use common_query::Output;
57        use common_query::error::Result;
58        use session::context::QueryContextRef;
59        use store_api::storage::RegionId;
60        use table::requests::{
61            BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest,
62            InsertRequest,
63        };
64        use table::table_name::TableName;
65
66        use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler};
67        struct MockProcedureServiceHandler {
68            calls: Arc<AtomicUsize>,
69        }
70        struct MockTableMutationHandler;
71        struct MockFlowServiceHandler;
72        const ROWS: usize = 42;
73
74        #[async_trait]
75        impl ProcedureServiceHandler for MockProcedureServiceHandler {
76            async fn purge_table(
77                &self,
78                _query_ctx: QueryContextRef,
79                _table_name: table::table_name::TableName,
80            ) -> Result<()> {
81                Ok(())
82            }
83
84            async fn migrate_region(
85                &self,
86                _ctx: QueryContextRef,
87                _request: MigrateRegionRequest,
88            ) -> Result<Option<String>> {
89                Ok(Some("test_pid".to_string()))
90            }
91
92            async fn reconcile(&self, _request: ReconcileRequest) -> Result<Option<String>> {
93                Ok(Some("test_pid".to_string()))
94            }
95
96            async fn query_procedure_state(&self, _pid: &str) -> Result<ProcedureStateResponse> {
97                self.calls.fetch_add(1, Ordering::Relaxed);
98                Ok(ProcedureStateResponse {
99                    status: ProcedureStatus::Done.into(),
100                    error: "OK".to_string(),
101                    ..Default::default()
102                })
103            }
104
105            async fn manage_region_follower(
106                &self,
107                _query_ctx: QueryContextRef,
108                _request: ManageRegionFollowerRequest,
109            ) -> Result<()> {
110                Ok(())
111            }
112
113            async fn gc_regions(
114                &self,
115                _context: QueryContextRef,
116                _request: GcRegionsRequest,
117            ) -> Result<GcResponse> {
118                Ok(GcResponse {
119                    processed_regions: 1,
120                    need_retry_regions: vec![],
121                    deleted_files: 0,
122                    deleted_indexes: 0,
123                })
124            }
125
126            async fn gc_table(
127                &self,
128                _context: QueryContextRef,
129                _request: GcTableRequest,
130            ) -> Result<GcResponse> {
131                Ok(GcResponse {
132                    processed_regions: 1,
133                    need_retry_regions: vec![],
134                    deleted_files: 0,
135                    deleted_indexes: 0,
136                })
137            }
138
139            fn catalog_manager(&self) -> &CatalogManagerRef {
140                unimplemented!()
141            }
142        }
143
144        #[async_trait]
145        impl TableMutationHandler for MockTableMutationHandler {
146            async fn insert(
147                &self,
148                _request: InsertRequest,
149                _ctx: QueryContextRef,
150            ) -> Result<Output> {
151                Ok(Output::new_with_affected_rows(ROWS))
152            }
153
154            async fn delete(
155                &self,
156                _request: DeleteRequest,
157                _ctx: QueryContextRef,
158            ) -> Result<AffectedRows> {
159                Ok(ROWS)
160            }
161
162            async fn flush(
163                &self,
164                _request: FlushTableRequest,
165                _ctx: QueryContextRef,
166            ) -> Result<AffectedRows> {
167                Ok(ROWS)
168            }
169
170            async fn compact(
171                &self,
172                _request: CompactTableRequest,
173                _ctx: QueryContextRef,
174            ) -> Result<AffectedRows> {
175                Ok(ROWS)
176            }
177
178            async fn build_index(
179                &self,
180                _request: BuildIndexTableRequest,
181                _ctx: QueryContextRef,
182            ) -> Result<AffectedRows> {
183                Ok(ROWS)
184            }
185
186            async fn flush_region(
187                &self,
188                _region_id: RegionId,
189                _ctx: QueryContextRef,
190            ) -> Result<AffectedRows> {
191                Ok(ROWS)
192            }
193
194            async fn compact_region(
195                &self,
196                _region_id: RegionId,
197                _ctx: QueryContextRef,
198            ) -> Result<AffectedRows> {
199                Ok(ROWS)
200            }
201
202            async fn discard_unflushed_data(
203                &self,
204                _region_id: RegionId,
205                _ctx: QueryContextRef,
206            ) -> Result<AffectedRows> {
207                Ok(ROWS)
208            }
209
210            async fn discard_unflushed_data_by_table(
211                &self,
212                _table_name: TableName,
213                _ctx: QueryContextRef,
214            ) -> Result<AffectedRows> {
215                Ok(ROWS)
216            }
217        }
218
219        #[async_trait]
220        impl FlowServiceHandler for MockFlowServiceHandler {
221            async fn flush(
222                &self,
223                _catalog: &str,
224                _flow: &str,
225                _ctx: QueryContextRef,
226            ) -> Result<api::v1::flow::FlowResponse> {
227                todo!()
228            }
229        }
230
231        Self {
232            plugins: Plugins::default(),
233            table_mutation_handler: Some(Arc::new(MockTableMutationHandler)),
234            procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler { calls })),
235            flow_service_handler: Some(Arc::new(MockFlowServiceHandler)),
236        }
237    }
238}
common_function/state.rs

common_function/
state.rs
