common_meta/
test_util.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 std::path::PathBuf;
16use std::sync::Arc;
17
18use api::region::RegionResponse;
19use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse};
20use api::v1::region::{InsertRequests, RegionRequest};
21pub use common_base::AffectedRows;
22use common_query::request::QueryRequest;
23use common_recordbatch::SendableRecordBatchStream;
24use common_wal::config::kafka::MetasrvKafkaConfig;
25use common_wal::config::kafka::common::{KafkaConnectionConfig, KafkaTopicConfig};
26
27use crate::cache_invalidator::DummyCacheInvalidator;
28use crate::ddl::flow_meta::FlowMetadataAllocator;
29use crate::ddl::table_meta::TableMetadataAllocator;
30use crate::ddl::{DdlContext, NoopRegionFailureDetectorControl};
31use crate::error::Result;
32use crate::key::TableMetadataManager;
33use crate::key::flow::FlowMetadataManager;
34use crate::kv_backend::KvBackendRef;
35use crate::kv_backend::memory::MemoryKvBackend;
36use crate::node_manager::{
37    Datanode, DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, NodeManagerRef,
38};
39use crate::peer::{Peer, PeerResolver};
40use crate::region_keeper::MemoryRegionKeeper;
41use crate::region_registry::LeaderRegionRegistry;
42use crate::sequence::SequenceBuilder;
43use crate::wal_provider::topic_pool::KafkaTopicPool;
44use crate::wal_provider::{WalProvider, build_kafka_topic_creator};
45use crate::{DatanodeId, FlownodeId};
46
47#[async_trait::async_trait]
48pub trait MockDatanodeHandler: Sync + Send + Clone {
49    async fn handle(&self, peer: &Peer, request: RegionRequest) -> Result<RegionResponse>;
50
51    async fn handle_query(
52        &self,
53        peer: &Peer,
54        request: QueryRequest,
55    ) -> Result<SendableRecordBatchStream>;
56}
57
58#[async_trait::async_trait]
59pub trait MockFlownodeHandler: Sync + Send + Clone {
60    async fn handle(&self, _peer: &Peer, _request: FlowRequest) -> Result<FlowResponse> {
61        unimplemented!()
62    }
63
64    async fn handle_inserts(
65        &self,
66        _peer: &Peer,
67        _requests: InsertRequests,
68    ) -> Result<FlowResponse> {
69        unimplemented!()
70    }
71
72    async fn handle_mark_window_dirty(
73        &self,
74        _peer: &Peer,
75        _req: DirtyWindowRequests,
76    ) -> Result<FlowResponse> {
77        unimplemented!()
78    }
79}
80
81/// A mock struct implements [NodeManager] only implement the `datanode` method.
82#[derive(Clone)]
83pub struct MockDatanodeManager<T> {
84    handler: T,
85}
86
87impl<T> MockDatanodeManager<T> {
88    pub fn new(handler: T) -> Self {
89        Self { handler }
90    }
91}
92
93/// A mock struct implements [NodeManager] only implement the `flownode` method.
94#[derive(Clone)]
95pub struct MockFlownodeManager<T> {
96    handler: T,
97}
98
99impl<T> MockFlownodeManager<T> {
100    pub fn new(handler: T) -> Self {
101        Self { handler }
102    }
103}
104
105/// A mock struct implements [Datanode].
106#[derive(Clone)]
107struct MockNode<T> {
108    peer: Peer,
109    handler: T,
110}
111
112#[async_trait::async_trait]
113impl<T: MockDatanodeHandler> Datanode for MockNode<T> {
114    async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
115        self.handler.handle(&self.peer, request).await
116    }
117
118    async fn handle_query(&self, request: QueryRequest) -> Result<SendableRecordBatchStream> {
119        self.handler.handle_query(&self.peer, request).await
120    }
121}
122
123#[async_trait::async_trait]
124impl<T: MockDatanodeHandler + 'static> DatanodeManager for MockDatanodeManager<T> {
125    async fn datanode(&self, peer: &Peer) -> DatanodeRef {
126        Arc::new(MockNode {
127            peer: peer.clone(),
128            handler: self.handler.clone(),
129        })
130    }
131}
132
133#[async_trait::async_trait]
134impl<T: 'static + Send + Sync> FlownodeManager for MockDatanodeManager<T> {
135    async fn flownode(&self, _peer: &Peer) -> FlownodeRef {
136        unimplemented!()
137    }
138}
139
140#[async_trait::async_trait]
141impl<T: MockFlownodeHandler> Flownode for MockNode<T> {
142    async fn handle(&self, request: FlowRequest) -> Result<FlowResponse> {
143        self.handler.handle(&self.peer, request).await
144    }
145
146    async fn handle_inserts(&self, requests: InsertRequests) -> Result<FlowResponse> {
147        self.handler.handle_inserts(&self.peer, requests).await
148    }
149
150    async fn handle_mark_window_dirty(&self, req: DirtyWindowRequests) -> Result<FlowResponse> {
151        self.handler.handle_mark_window_dirty(&self.peer, req).await
152    }
153}
154
155#[async_trait::async_trait]
156impl<T: MockFlownodeHandler + 'static> FlownodeManager for MockFlownodeManager<T> {
157    async fn flownode(&self, peer: &Peer) -> FlownodeRef {
158        Arc::new(MockNode {
159            peer: peer.clone(),
160            handler: self.handler.clone(),
161        })
162    }
163}
164
165#[async_trait::async_trait]
166impl<T: 'static + Send + Sync> DatanodeManager for MockFlownodeManager<T> {
167    async fn datanode(&self, _peer: &Peer) -> DatanodeRef {
168        unimplemented!()
169    }
170}
171
172/// Returns a test purpose [DdlContext].
173pub fn new_ddl_context(node_manager: NodeManagerRef) -> DdlContext {
174    let kv_backend = Arc::new(MemoryKvBackend::new());
175    new_ddl_context_with_kv_backend(node_manager, kv_backend)
176}
177
178/// Returns a test purpose [DdlContext] with a specified [KvBackendRef].
179pub fn new_ddl_context_with_kv_backend(
180    node_manager: NodeManagerRef,
181    kv_backend: KvBackendRef,
182) -> DdlContext {
183    let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone()));
184    let table_metadata_allocator = Arc::new(TableMetadataAllocator::new(
185        Arc::new(
186            SequenceBuilder::new("test", kv_backend.clone())
187                .initial(1024)
188                .build(),
189        ),
190        Arc::new(WalProvider::default()),
191    ));
192    let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone()));
193    let flow_metadata_allocator =
194        Arc::new(FlowMetadataAllocator::with_noop_peer_allocator(Arc::new(
195            SequenceBuilder::new("flow-test", kv_backend)
196                .initial(1024)
197                .build(),
198        )));
199    DdlContext {
200        node_manager,
201        cache_invalidator: Arc::new(DummyCacheInvalidator),
202        memory_region_keeper: Arc::new(MemoryRegionKeeper::new()),
203        leader_region_registry: Arc::new(LeaderRegionRegistry::default()),
204        table_metadata_allocator,
205        table_metadata_manager,
206        flow_metadata_allocator,
207        flow_metadata_manager,
208        region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl),
209        soft_drop_enabled: false,
210        soft_drop_retention: Some(std::time::Duration::from_secs(7 * 24 * 60 * 60)),
211        create_database_metadata_committer: None,
212    }
213}
214
215pub struct NoopPeerResolver;
216
217#[async_trait::async_trait]
218impl PeerResolver for NoopPeerResolver {
219    async fn datanode(&self, id: DatanodeId) -> Result<Option<Peer>> {
220        Ok(Some(Peer::empty(id)))
221    }
222
223    async fn flownode(&self, id: FlownodeId) -> Result<Option<Peer>> {
224        Ok(Some(Peer::empty(id)))
225    }
226}
227
228/// Create a kafka topic pool for testing.
229pub async fn test_kafka_topic_pool(
230    broker_endpoints: Vec<String>,
231    num_topics: usize,
232    auto_create_topics: bool,
233    topic_name_prefix: Option<&str>,
234) -> KafkaTopicPool {
235    let mut config = MetasrvKafkaConfig {
236        connection: KafkaConnectionConfig {
237            broker_endpoints,
238            ..Default::default()
239        },
240        kafka_topic: KafkaTopicConfig {
241            num_topics,
242
243            ..Default::default()
244        },
245        auto_create_topics,
246        ..Default::default()
247    };
248    if let Some(prefix) = topic_name_prefix {
249        config.kafka_topic.topic_name_prefix = prefix.to_string();
250    }
251    let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
252    let topic_creator = build_kafka_topic_creator(&config.connection, &config.kafka_topic)
253        .await
254        .unwrap();
255
256    KafkaTopicPool::new(&config, kv_backend, topic_creator)
257}
258
259#[macro_export]
260/// Skip the test if the environment variable `GT_POSTGRES_ENDPOINTS` is not set.
261///
262/// The format of the environment variable is:
263/// ```text
264/// GT_POSTGRES_ENDPOINTS=localhost:9092,localhost:9093
265/// ```
266macro_rules! maybe_skip_postgres_integration_test {
267    () => {
268        if std::env::var("GT_POSTGRES_ENDPOINTS").is_err() {
269            common_telemetry::warn!("The endpoints is empty, skipping the test");
270            return;
271        }
272    };
273}
274
275#[macro_export]
276/// Skip the test if the environment variable `GT_MYSQL_ENDPOINTS` is not set.
277///
278/// The format of the environment variable is:
279/// ```text
280/// GT_MYSQL_ENDPOINTS=localhost:9092,localhost:9093
281/// ```
282macro_rules! maybe_skip_mysql_integration_test {
283    () => {
284        if std::env::var("GT_MYSQL_ENDPOINTS").is_err() {
285            common_telemetry::warn!("The endpoints is empty, skipping the test");
286            return;
287        }
288    };
289}
290
291#[macro_export]
292/// Skip the test if the environment variable `GT_POSTGRES15_ENDPOINTS` is not set.
293///
294/// The format of the environment variable is:
295/// ```text
296/// GT_POSTGRES15_ENDPOINTS=postgres://user:password@127.0.0.1:5433/postgres
297/// ```
298macro_rules! maybe_skip_postgres15_integration_test {
299    () => {
300        if std::env::var("GT_POSTGRES15_ENDPOINTS").is_err() {
301            common_telemetry::warn!("The PG15 endpoints is empty, skipping the test");
302            return;
303        }
304    };
305}
306
307#[macro_export]
308/// Skip the test if the environment variable `GT_ETCD_TLS_ENDPOINTS` is not set.
309///
310/// The format of the environment variable is:
311/// ```text
312/// GT_ETCD_TLS_ENDPOINTS=localhost:9092,localhost:9093
313/// ```
314macro_rules! maybe_skip_etcd_tls_integration_test {
315    () => {
316        if std::env::var("GT_ETCD_TLS_ENDPOINTS").is_err() {
317            common_telemetry::warn!("The etcd with tls endpoints is empty, skipping the test");
318            return;
319        }
320    };
321}
322
323/// Returns the directory of the etcd TLS certs.
324pub fn etcd_certs_dir() -> PathBuf {
325    let project_path = env!("CARGO_MANIFEST_DIR");
326    let project_path = PathBuf::from(project_path);
327    let base = project_path.ancestors().nth(3).unwrap();
328    base.join("tests-integration")
329        .join("fixtures")
330        .join("etcd-tls-certs")
331}
332
333/// Returns the directory of the test certs.
334pub fn test_certs_dir() -> PathBuf {
335    let project_path = env!("CARGO_MANIFEST_DIR");
336    let project_path = PathBuf::from(project_path);
337    let base = project_path.ancestors().nth(3).unwrap();
338    base.join("tests-integration")
339        .join("fixtures")
340        .join("certs")
341}
common_meta/test_util.rs

common_meta/
test_util.rs
