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                _query_ctx: QueryContextRef,
96                _request: ManageRegionFollowerRequest,
97            ) -> Result<()> {
98                Ok(())
99            }
100
101            async fn gc_regions(
102                &self,
103                _context: QueryContextRef,
104                _request: GcRegionsRequest,
105            ) -> Result<GcResponse> {
106                Ok(GcResponse {
107                    processed_regions: 1,
108                    need_retry_regions: vec![],
109                    deleted_files: 0,
110                    deleted_indexes: 0,
111                })
112            }
113
114            async fn gc_table(
115                &self,
116                _context: QueryContextRef,
117                _request: GcTableRequest,
118            ) -> Result<GcResponse> {
119                Ok(GcResponse {
120                    processed_regions: 1,
121                    need_retry_regions: vec![],
122                    deleted_files: 0,
123                    deleted_indexes: 0,
124                })
125            }
126
127            fn catalog_manager(&self) -> &CatalogManagerRef {
128                unimplemented!()
129            }
130        }
131
132        #[async_trait]
133        impl TableMutationHandler for MockTableMutationHandler {
134            async fn insert(
135                &self,
136                _request: InsertRequest,
137                _ctx: QueryContextRef,
138            ) -> Result<Output> {
139                Ok(Output::new_with_affected_rows(ROWS))
140            }
141
142            async fn delete(
143                &self,
144                _request: DeleteRequest,
145                _ctx: QueryContextRef,
146            ) -> Result<AffectedRows> {
147                Ok(ROWS)
148            }
149
150            async fn flush(
151                &self,
152                _request: FlushTableRequest,
153                _ctx: QueryContextRef,
154            ) -> Result<AffectedRows> {
155                Ok(ROWS)
156            }
157
158            async fn compact(
159                &self,
160                _request: CompactTableRequest,
161                _ctx: QueryContextRef,
162            ) -> Result<AffectedRows> {
163                Ok(ROWS)
164            }
165
166            async fn build_index(
167                &self,
168                _request: BuildIndexTableRequest,
169                _ctx: QueryContextRef,
170            ) -> Result<AffectedRows> {
171                Ok(ROWS)
172            }
173
174            async fn flush_region(
175                &self,
176                _region_id: RegionId,
177                _ctx: QueryContextRef,
178            ) -> Result<AffectedRows> {
179                Ok(ROWS)
180            }
181
182            async fn compact_region(
183                &self,
184                _region_id: RegionId,
185                _ctx: QueryContextRef,
186            ) -> Result<AffectedRows> {
187                Ok(ROWS)
188            }
189
190            async fn discard_unflushed_data(
191                &self,
192                _region_id: RegionId,
193                _ctx: QueryContextRef,
194            ) -> Result<AffectedRows> {
195                Ok(ROWS)
196            }
197
198            async fn discard_unflushed_data_by_table(
199                &self,
200                _table_name: TableName,
201                _ctx: QueryContextRef,
202            ) -> Result<AffectedRows> {
203                Ok(ROWS)
204            }
205        }
206
207        #[async_trait]
208        impl FlowServiceHandler for MockFlowServiceHandler {
209            async fn flush(
210                &self,
211                _catalog: &str,
212                _flow: &str,
213                _ctx: QueryContextRef,
214            ) -> Result<api::v1::flow::FlowResponse> {
215                todo!()
216            }
217        }
218
219        Self {
220            plugins: Plugins::default(),
221            table_mutation_handler: Some(Arc::new(MockTableMutationHandler)),
222            procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)),
223            flow_service_handler: Some(Arc::new(MockFlowServiceHandler)),
224        }
225    }
226}
common_function/state.rs

common_function/
state.rs
