common_wal/
lib.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::net::SocketAddr;
16
17use error::{EndpointIPV4NotFoundSnafu, ResolveEndpointSnafu, Result};
18use serde::{Deserialize, Serialize};
19use snafu::{OptionExt, ResultExt};
20use tokio::net;
21
22pub mod config;
23pub mod error;
24pub mod kafka;
25pub mod options;
26#[cfg(any(test, feature = "testing"))]
27pub mod test_util;
28
29pub const BROKER_ENDPOINT: &str = "127.0.0.1:9092";
30pub const TOPIC_NAME_PREFIX: &str = "greptimedb_wal_topic";
31
32/// The type of the topic selector, i.e. with which strategy to select a topic.
33// The enum is defined here to work around cyclic dependency issues.
34#[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
35#[serde(rename_all = "snake_case")]
36pub enum TopicSelectorType {
37    #[default]
38    RoundRobin,
39}
40
41pub async fn resolve_to_ipv4<T: AsRef<str>>(endpoints: &[T]) -> Result<Vec<String>> {
42    futures_util::future::try_join_all(endpoints.iter().map(resolve_to_ipv4_one)).await
43}
44
45async fn resolve_to_ipv4_one<T: AsRef<str>>(endpoint: T) -> Result<String> {
46    let endpoint = endpoint.as_ref();
47    net::lookup_host(endpoint)
48        .await
49        .context(ResolveEndpointSnafu {
50            broker_endpoint: endpoint,
51        })?
52        .find(SocketAddr::is_ipv4)
53        .map(|addr| addr.to_string())
54        .context(EndpointIPV4NotFoundSnafu {
55            broker_endpoint: endpoint,
56        })
57}
58
59#[cfg(test)]
60mod tests {
61    use std::assert_matches;
62
63    use common_telemetry::warn;
64    use rskafka::client::{Credentials, SaslConfig};
65
66    use super::*;
67    use crate::error::Error;
68
69    // test for resolve_broker_endpoint
70    #[tokio::test]
71    async fn test_valid_host() {
72        let host = "localhost:9092";
73        let got = resolve_to_ipv4_one(host).await;
74        assert_eq!(got.unwrap(), "127.0.0.1:9092");
75    }
76
77    #[tokio::test]
78    async fn test_valid_host_ipv6() {
79        // the host is valid, it is an IPv6 address, but we only accept IPv4 addresses
80        let host = "::1:9092";
81        let got = resolve_to_ipv4_one(host).await;
82        assert_matches!(got.unwrap_err(), Error::EndpointIPV4NotFound { .. });
83    }
84
85    #[tokio::test]
86    async fn test_invalid_host() {
87        let host = "non-exist-host:9092";
88        let got = resolve_to_ipv4_one(host).await;
89        assert_matches!(got.unwrap_err(), Error::ResolveEndpoint { .. });
90    }
91
92    #[tokio::test]
93    async fn test_sasl() {
94        common_telemetry::init_default_ut_logging();
95        let Ok(broker_endpoints) = std::env::var("GT_KAFKA_SASL_ENDPOINTS") else {
96            warn!("The endpoints is empty, skipping the test 'test_sasl'");
97            return;
98        };
99        let broker_endpoints = broker_endpoints
100            .split(',')
101            .map(|s| s.trim().to_string())
102            .collect::<Vec<_>>();
103
104        let username = "user_kafka";
105        let password = "secret";
106        let _ = rskafka::client::ClientBuilder::new(broker_endpoints.clone())
107            .sasl_config(SaslConfig::Plain(Credentials::new(
108                username.to_string(),
109                password.to_string(),
110            )))
111            .build()
112            .await
113            .unwrap();
114        let _ = rskafka::client::ClientBuilder::new(broker_endpoints.clone())
115            .sasl_config(SaslConfig::ScramSha256(Credentials::new(
116                username.to_string(),
117                password.to_string(),
118            )))
119            .build()
120            .await
121            .unwrap();
122        let _ = rskafka::client::ClientBuilder::new(broker_endpoints)
123            .sasl_config(SaslConfig::ScramSha512(Credentials::new(
124                username.to_string(),
125                password.to_string(),
126            )))
127            .build()
128            .await
129            .unwrap();
130    }
131}
common_wal/lib.rs

common_wal/
lib.rs
