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        use std::sync::Arc;
38
39        use api::v1::meta::{ProcedureStatus, ReconcileRequest};
40        use async_trait::async_trait;
41        use catalog::CatalogManagerRef;
42        use common_base::AffectedRows;
43        use common_meta::rpc::procedure::{
44            GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest,
45            MigrateRegionRequest, ProcedureStateResponse,
46        };
47        use common_query::Output;
48        use common_query::error::Result;
49        use session::context::QueryContextRef;
50        use store_api::storage::RegionId;
51        use table::requests::{
52            BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest,
53            InsertRequest,
54        };
55        use table::table_name::TableName;
56
57        use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler};
58        struct MockProcedureServiceHandler;
59        struct MockTableMutationHandler;
60        struct MockFlowServiceHandler;
61        const ROWS: usize = 42;
62
63        #[async_trait]
64        impl ProcedureServiceHandler for MockProcedureServiceHandler {
65            async fn purge_table(
66                &self,
67                _query_ctx: QueryContextRef,
68                _table_name: table::table_name::TableName,
69            ) -> Result<()> {
70                Ok(())
71            }
72
73            async fn migrate_region(
74                &self,
75                _ctx: QueryContextRef,
76                _request: MigrateRegionRequest,
77            ) -> Result<Option<String>> {
78                Ok(Some("test_pid".to_string()))
79            }
80
81            async fn reconcile(&self, _request: ReconcileRequest) -> Result<Option<String>> {
82                Ok(Some("test_pid".to_string()))
83            }
84
85            async fn query_procedure_state(&self, _pid: &str) -> Result<ProcedureStateResponse> {
86                Ok(ProcedureStateResponse {
87                    status: ProcedureStatus::Done.into(),
88                    error: "OK".to_string(),
89                    ..Default::default()
90                })
91            }
92
93            async fn manage_region_follower(
94                &self,
95                _request: ManageRegionFollowerRequest,
96            ) -> Result<()> {
97                Ok(())
98            }
99
100            async fn gc_regions(
101                &self,
102                _context: QueryContextRef,
103                _request: GcRegionsRequest,
104            ) -> Result<GcResponse> {
105                Ok(GcResponse {
106                    processed_regions: 1,
107                    need_retry_regions: vec![],
108                    deleted_files: 0,
109                    deleted_indexes: 0,
110                })
111            }
112
113            async fn gc_table(
114                &self,
115                _context: QueryContextRef,
116                _request: GcTableRequest,
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            fn catalog_manager(&self) -> &CatalogManagerRef {
127                unimplemented!()
128            }
129        }
130
131        #[async_trait]
132        impl TableMutationHandler for MockTableMutationHandler {
133            async fn insert(
134                &self,
135                _request: InsertRequest,
136                _ctx: QueryContextRef,
137            ) -> Result<Output> {
138                Ok(Output::new_with_affected_rows(ROWS))
139            }
140
141            async fn delete(
142                &self,
143                _request: DeleteRequest,
144                _ctx: QueryContextRef,
145            ) -> Result<AffectedRows> {
146                Ok(ROWS)
147            }
148
149            async fn flush(
150                &self,
151                _request: FlushTableRequest,
152                _ctx: QueryContextRef,
153            ) -> Result<AffectedRows> {
154                Ok(ROWS)
155            }
156
157            async fn compact(
158                &self,
159                _request: CompactTableRequest,
160                _ctx: QueryContextRef,
161            ) -> Result<AffectedRows> {
162                Ok(ROWS)
163            }
164
165            async fn build_index(
166                &self,
167                _request: BuildIndexTableRequest,
168                _ctx: QueryContextRef,
169            ) -> Result<AffectedRows> {
170                Ok(ROWS)
171            }
172
173            async fn flush_region(
174                &self,
175                _region_id: RegionId,
176                _ctx: QueryContextRef,
177            ) -> Result<AffectedRows> {
178                Ok(ROWS)
179            }
180
181            async fn compact_region(
182                &self,
183                _region_id: RegionId,
184                _ctx: QueryContextRef,
185            ) -> Result<AffectedRows> {
186                Ok(ROWS)
187            }
188
189            async fn discard_unflushed_data(
190                &self,
191                _region_id: RegionId,
192                _ctx: QueryContextRef,
193            ) -> Result<AffectedRows> {
194                Ok(ROWS)
195            }
196
197            async fn discard_unflushed_data_by_table(
198                &self,
199                _table_name: TableName,
200                _ctx: QueryContextRef,
201            ) -> Result<AffectedRows> {
202                Ok(ROWS)
203            }
204        }
205
206        #[async_trait]
207        impl FlowServiceHandler for MockFlowServiceHandler {
208            async fn flush(
209                &self,
210                _catalog: &str,
211                _flow: &str,
212                _ctx: QueryContextRef,
213            ) -> Result<api::v1::flow::FlowResponse> {
214                todo!()
215            }
216        }
217
218        Self {
219            plugins: Plugins::default(),
220            table_mutation_handler: Some(Arc::new(MockTableMutationHandler)),
221            procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)),
222            flow_service_handler: Some(Arc::new(MockFlowServiceHandler)),
223        }
224    }
225}
common_function/state.rs

common_function/
state.rs
