risingwave_connector/source/kafka/
mod.rs

1// Copyright 2022 RisingWave Labs
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::collections::HashMap;
16
17use serde::Deserialize;
18use serde_with::{DisplayFromStr, serde_as};
19
20use crate::connector_common::{AwsAuthProps, KafkaConnectionProps, KafkaPrivateLinkCommon};
21use crate::enforce_secret::EnforceSecret;
22use crate::error::ConnectorResult;
23
24mod client_context;
25pub mod enumerator;
26pub mod private_link;
27pub mod source;
28pub mod split;
29pub mod stats;
30
31pub use client_context::*;
32pub use enumerator::*;
33use risingwave_common::id::FragmentId;
34pub use source::*;
35pub use split::*;
36use with_options::WithOptions;
37
38use crate::connector_common::{KafkaCommon, RdKafkaPropertiesCommon};
39use crate::source::SourceProperties;
40
41pub const KAFKA_CONNECTOR: &str = "kafka";
42pub const KAFKA_PROPS_BROKER_KEY: &str = "properties.bootstrap.server";
43pub const KAFKA_PROPS_BROKER_KEY_ALIAS: &str = "kafka.brokers";
44pub const PRIVATELINK_CONNECTION: &str = "privatelink";
45
46/// Properties for the rdkafka library. Leave a field as `None` to use the default value.
47/// These properties are not intended to be exposed to users in the majority of cases.
48///
49/// See also <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>
50#[serde_as]
51#[derive(Clone, Debug, Deserialize, Default, WithOptions)]
52pub struct RdKafkaPropertiesConsumer {
53    /// Minimum number of messages per topic+partition librdkafka tries to maintain in the local
54    /// consumer queue.
55    #[serde(rename = "properties.queued.min.messages")]
56    #[serde_as(as = "Option<DisplayFromStr>")]
57    #[with_option(allow_alter_on_fly)]
58    pub queued_min_messages: Option<usize>,
59
60    #[serde(rename = "properties.queued.max.messages.kbytes")]
61    #[serde_as(as = "Option<DisplayFromStr>")]
62    #[with_option(allow_alter_on_fly)]
63    pub queued_max_messages_kbytes: Option<usize>,
64
65    /// Maximum time the broker may wait to fill the Fetch response with `fetch.min.`bytes of
66    /// messages.
67    #[serde(rename = "properties.fetch.wait.max.ms")]
68    #[serde_as(as = "Option<DisplayFromStr>")]
69    #[with_option(allow_alter_on_fly)]
70    pub fetch_wait_max_ms: Option<usize>,
71
72    /// Minimum number of bytes the broker responds with. If fetch.wait.max.ms expires the accumulated data will be sent to the client regardless of this setting.
73    #[serde(rename = "properties.fetch.min.bytes")]
74    #[serde_as(as = "Option<DisplayFromStr>")]
75    pub fetch_min_bytes: Option<usize>,
76
77    /// Initial maximum number of bytes per topic+partition to request when fetching messages from the broker. If the client encounters a message larger than this value it will gradually try to increase it until the entire message can be fetched.
78    #[serde(rename = "properties.fetch.message.max.bytes")]
79    #[serde_as(as = "Option<DisplayFromStr>")]
80    pub fetch_message_max_bytes: Option<usize>,
81
82    /// How long to postpone the next fetch request for a topic+partition in case the current fetch
83    /// queue thresholds (`queued.min.messages` or `queued.max.messages.kbytes`) have been
84    /// exceeded. This property may need to be decreased if the queue thresholds are set low
85    /// and the application is experiencing long (~1s) delays between messages. Low values may
86    /// increase CPU utilization.
87    #[serde(rename = "properties.fetch.queue.backoff.ms")]
88    #[serde_as(as = "Option<DisplayFromStr>")]
89    #[with_option(allow_alter_on_fly)]
90    pub fetch_queue_backoff_ms: Option<usize>,
91
92    /// Maximum amount of data the broker shall return for a Fetch request. Messages are fetched in
93    /// batches by the consumer and if the first message batch in the first non-empty partition of
94    /// the Fetch request is larger than this value, then the message batch will still be returned
95    /// to ensure the consumer can make progress. The maximum message batch size accepted by the
96    /// broker is defined via `message.max.bytes` (broker config) or `max.message.bytes` (broker
97    /// topic config). `fetch.max.bytes` is automatically adjusted upwards to be at least
98    /// `message.max.bytes` (consumer config).
99    #[serde(rename = "properties.fetch.max.bytes")]
100    #[serde_as(as = "Option<DisplayFromStr>")]
101    #[with_option(allow_alter_on_fly)]
102    pub fetch_max_bytes: Option<usize>,
103
104    /// Whether to automatically and periodically commit offsets in the background.
105    ///
106    /// Note that RisingWave does NOT rely on committed offsets. Committing offset is only for exposing the
107    /// progress for monitoring. Setting this to false can avoid creating consumer groups.
108    ///
109    /// default: true
110    #[serde(rename = "properties.enable.auto.commit")]
111    #[serde_as(as = "Option<DisplayFromStr>")]
112    #[with_option(allow_alter_on_fly)]
113    pub enable_auto_commit: Option<bool>,
114
115    /// The frequency in milliseconds that the consumer offsets are committed.
116    #[serde(rename = "properties.auto.commit.interval.ms")]
117    #[serde_as(as = "Option<DisplayFromStr>")]
118    #[with_option(allow_alter_on_fly)]
119    pub auto_commit_interval_ms: Option<usize>,
120}
121
122#[derive(Clone, Debug, Deserialize, WithOptions)]
123pub struct KafkaProperties {
124    /// This parameter is not intended to be exposed to users.
125    /// This parameter specifies only for one parallelism. The parallelism of kafka source
126    /// is equal to the parallelism passed into compute nodes. So users need to calculate
127    /// how many bytes will be consumed in total across all the parallelism by themselves.
128    #[serde(rename = "bytes.per.second", alias = "kafka.bytes.per.second")]
129    pub bytes_per_second: Option<String>,
130
131    /// This parameter is not intended to be exposed to users.
132    /// This parameter specifies only for one parallelism. The parallelism of kafka source
133    /// is equal to the parallelism passed into compute nodes. So users need to calculate
134    /// how many messages will be consumed in total across all the parallelism by themselves.
135    #[serde(rename = "max.num.messages", alias = "kafka.max.num.messages")]
136    pub max_num_messages: Option<String>,
137
138    #[serde(rename = "scan.startup.mode", alias = "kafka.scan.startup.mode")]
139    pub scan_startup_mode: Option<String>,
140
141    #[serde(
142        rename = "scan.startup.timestamp.millis",
143        alias = "kafka.time.offset",
144        alias = "scan.startup.timestamp_millis" // keep for compatibility
145    )]
146    pub time_offset: Option<String>,
147
148    /// Specify a custom consumer group id prefix for the source.
149    /// Defaults to `rw-consumer`.
150    ///
151    /// Notes:
152    /// - Each job (materialized view) will have a separated consumer group and
153    ///   contains a generated suffix in the group id.
154    ///   The consumer group will be `{group_id_prefix}-{fragment_id}`.
155    /// - The consumer group is solely for monintoring progress in some external
156    ///   Kafka tools, and for authorization. RisingWave does not rely on committed
157    ///   offsets, and does not join the consumer group. It just reports offsets
158    ///   to the group.
159    #[serde(rename = "group.id.prefix")]
160    #[with_option(allow_alter_on_fly)]
161    pub group_id_prefix: Option<String>,
162
163    /// This parameter is used to tell `KafkaSplitReader` to produce `UpsertMessage`s, which
164    /// combine both key and value fields of the Kafka message.
165    /// TODO: Currently, `Option<bool>` can not be parsed here.
166    #[serde(rename = "upsert")]
167    pub upsert: Option<String>,
168
169    #[serde(flatten)]
170    pub common: KafkaCommon,
171
172    #[serde(flatten)]
173    pub connection: KafkaConnectionProps,
174
175    #[serde(flatten)]
176    pub rdkafka_properties_common: RdKafkaPropertiesCommon,
177
178    #[serde(flatten)]
179    pub rdkafka_properties_consumer: RdKafkaPropertiesConsumer,
180
181    #[serde(flatten)]
182    pub privatelink_common: KafkaPrivateLinkCommon,
183
184    #[serde(flatten)]
185    pub aws_auth_props: AwsAuthProps,
186
187    #[serde(flatten)]
188    pub unknown_fields: HashMap<String, String>,
189}
190
191impl EnforceSecret for KafkaProperties {
192    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
193        for prop in prop_iter {
194            KafkaConnectionProps::enforce_one(prop)?;
195            AwsAuthProps::enforce_one(prop)?;
196        }
197        Ok(())
198    }
199}
200
201impl SourceProperties for KafkaProperties {
202    type Split = KafkaSplit;
203    type SplitEnumerator = KafkaSplitEnumerator;
204    type SplitReader = KafkaSplitReader;
205
206    const SOURCE_NAME: &'static str = KAFKA_CONNECTOR;
207}
208
209impl crate::source::UnknownFields for KafkaProperties {
210    fn unknown_fields(&self) -> HashMap<String, String> {
211        self.unknown_fields.clone()
212    }
213}
214
215impl KafkaProperties {
216    pub fn set_client(&self, c: &mut rdkafka::ClientConfig) {
217        self.rdkafka_properties_common.set_client(c);
218        self.rdkafka_properties_consumer.set_client(c);
219    }
220
221    pub fn group_id(&self, fragment_id: FragmentId) -> String {
222        format!(
223            "{}-{}",
224            self.group_id_prefix.as_deref().unwrap_or("rw-consumer"),
225            fragment_id
226        )
227    }
228}
229
230const KAFKA_ISOLATION_LEVEL: &str = "read_committed";
231
232impl RdKafkaPropertiesConsumer {
233    pub fn set_client(&self, c: &mut rdkafka::ClientConfig) {
234        if let Some(v) = &self.queued_min_messages {
235            c.set("queued.min.messages", v.to_string());
236        }
237        if let Some(v) = &self.queued_max_messages_kbytes {
238            c.set("queued.max.messages.kbytes", v.to_string());
239        }
240        if let Some(v) = &self.fetch_wait_max_ms {
241            c.set("fetch.wait.max.ms", v.to_string());
242        }
243        if let Some(v) = &self.fetch_min_bytes {
244            c.set("fetch.min.bytes", v.to_string());
245        }
246        if let Some(v) = &self.fetch_message_max_bytes {
247            c.set("fetch.message.max.bytes", v.to_string());
248        }
249        if let Some(v) = &self.fetch_queue_backoff_ms {
250            c.set("fetch.queue.backoff.ms", v.to_string());
251        }
252        if let Some(v) = &self.fetch_max_bytes {
253            c.set("fetch.max.bytes", v.to_string());
254        }
255        if let Some(v) = &self.enable_auto_commit {
256            c.set("enable.auto.commit", v.to_string());
257        }
258        if let Some(v) = &self.auto_commit_interval_ms {
259            c.set("auto.commit.interval.ms", v.to_string());
260        }
261    }
262}
263
264#[cfg(test)]
265mod test {
266    use std::collections::BTreeMap;
267
268    use maplit::btreemap;
269
270    use super::*;
271
272    #[test]
273    fn test_parse_config_consumer_common() {
274        let config: BTreeMap<String, String> = btreemap! {
275            // common
276            "properties.bootstrap.server".to_owned() => "127.0.0.1:9092".to_owned(),
277            "topic".to_owned() => "test".to_owned(),
278            // kafka props
279            "scan.startup.mode".to_owned() => "earliest".to_owned(),
280            // RdKafkaPropertiesCommon
281            "properties.message.max.bytes".to_owned() => "12345".to_owned(),
282            "properties.receive.message.max.bytes".to_owned() => "54321".to_owned(),
283            "properties.reconnect.backoff.ms".to_owned() => "1000".to_owned(),
284            "properties.reconnect.backoff.max.ms".to_owned() => "30000".to_owned(),
285            "properties.socket.connection.setup.timeout.ms".to_owned() => "60000".to_owned(),
286            "properties.retry.backoff.ms".to_owned() => "200".to_owned(),
287            "properties.retry.backoff.max.ms".to_owned() => "2000".to_owned(),
288            // RdKafkaPropertiesConsumer
289            "properties.queued.min.messages".to_owned() => "114514".to_owned(),
290            "properties.queued.max.messages.kbytes".to_owned() => "114514".to_owned(),
291            "properties.fetch.wait.max.ms".to_owned() => "114514".to_owned(),
292            "properties.fetch.max.bytes".to_owned() => "114514".to_owned(),
293            "properties.enable.auto.commit".to_owned() => "true".to_owned(),
294            "properties.auto.commit.interval.ms".to_owned() => "10000".to_owned(),
295            "properties.fetch.queue.backoff.ms".to_owned() => "114514".to_owned(),
296            // PrivateLink
297            "broker.rewrite.endpoints".to_owned() => "{\"broker1\": \"10.0.0.1:8001\"}".to_owned(),
298        };
299
300        let props: KafkaProperties =
301            serde_json::from_value(serde_json::to_value(config).unwrap()).unwrap();
302
303        assert_eq!(props.scan_startup_mode, Some("earliest".to_owned()));
304        assert_eq!(
305            props.rdkafka_properties_common.receive_message_max_bytes,
306            Some(54321)
307        );
308        assert_eq!(
309            props.rdkafka_properties_common.message_max_bytes,
310            Some(12345)
311        );
312        assert_eq!(
313            props.rdkafka_properties_common.reconnect_backoff_ms,
314            Some(1000)
315        );
316        assert_eq!(
317            props.rdkafka_properties_common.reconnect_backoff_max_ms,
318            Some(30000)
319        );
320        assert_eq!(
321            props
322                .rdkafka_properties_common
323                .socket_connection_setup_timeout_ms,
324            Some(60000)
325        );
326        assert_eq!(props.rdkafka_properties_common.retry_backoff_ms, Some(200));
327        assert_eq!(
328            props.rdkafka_properties_common.retry_backoff_max_ms,
329            Some(2000)
330        );
331        assert_eq!(
332            props.rdkafka_properties_consumer.queued_min_messages,
333            Some(114514)
334        );
335        assert_eq!(
336            props.rdkafka_properties_consumer.queued_max_messages_kbytes,
337            Some(114514)
338        );
339        assert_eq!(
340            props.rdkafka_properties_consumer.fetch_wait_max_ms,
341            Some(114514)
342        );
343        assert_eq!(
344            props.rdkafka_properties_consumer.fetch_max_bytes,
345            Some(114514)
346        );
347        assert_eq!(
348            props.rdkafka_properties_consumer.enable_auto_commit,
349            Some(true)
350        );
351        assert_eq!(
352            props.rdkafka_properties_consumer.auto_commit_interval_ms,
353            Some(10000)
354        );
355        assert_eq!(
356            props.rdkafka_properties_consumer.fetch_queue_backoff_ms,
357            Some(114514)
358        );
359        let mut client_config = rdkafka::ClientConfig::new();
360        props.set_client(&mut client_config);
361        assert_eq!(client_config.get("reconnect.backoff.ms"), Some("1000"));
362        assert_eq!(client_config.get("reconnect.backoff.max.ms"), Some("30000"));
363        assert_eq!(
364            client_config.get("socket.connection.setup.timeout.ms"),
365            Some("60000")
366        );
367        assert_eq!(client_config.get("retry.backoff.ms"), Some("200"));
368        assert_eq!(client_config.get("retry.backoff.max.ms"), Some("2000"));
369        assert_eq!(client_config.get("auto.commit.interval.ms"), Some("10000"));
370        let hashmap: BTreeMap<String, String> = btreemap! {
371            "broker1".to_owned() => "10.0.0.1:8001".to_owned()
372        };
373        assert_eq!(props.privatelink_common.broker_rewrite_map, Some(hashmap));
374    }
375}
risingwave_connector/source/kafka/mod.rs

risingwave_connector/source/kafka/
mod.rs
