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

common_function/
state.rs
