datum_mq/
config.rs

1use std::{collections::BTreeMap, time::Duration};
2
3/// Kafka assignment start offset.
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum Offset {
6    Beginning,
7    End,
8    Stored,
9    Offset(i64),
10    OffsetTail(i64),
11    Invalid,
12}
13
14use crate::{MqError, MqResult, TopicPartition};
15
16/// Raw Kafka property map with typed helpers for common Datum defaults.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct KafkaConfig {
19    properties: BTreeMap<String, String>,
20}
21
22impl KafkaConfig {
23    #[must_use]
24    pub fn new(bootstrap_servers: impl Into<String>) -> Self {
25        let mut properties = BTreeMap::new();
26        properties.insert("bootstrap.servers".to_owned(), bootstrap_servers.into());
27        Self { properties }
28    }
29
30    #[must_use]
31    pub fn empty() -> Self {
32        Self {
33            properties: BTreeMap::new(),
34        }
35    }
36
37    #[must_use]
38    pub fn get(&self, key: &str) -> Option<&str> {
39        self.properties.get(key).map(String::as_str)
40    }
41
42    #[must_use]
43    pub fn properties(&self) -> &BTreeMap<String, String> {
44        &self.properties
45    }
46
47    pub fn set(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
48        self.properties.insert(key.into(), value.into());
49        self
50    }
51
52    #[must_use]
53    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
54        self.set(key, value);
55        self
56    }
57
58    /// Sets the common TLS properties understood by both backends.
59    #[must_use]
60    pub fn with_tls(
61        mut self,
62        ca_location: impl Into<String>,
63        certificate_location: impl Into<String>,
64        key_location: impl Into<String>,
65    ) -> Self {
66        self.set("security.protocol", "ssl");
67        self.set("ssl.ca.location", ca_location);
68        self.set("ssl.certificate.location", certificate_location);
69        self.set("ssl.key.location", key_location);
70        self
71    }
72
73    /// Sets SASL/SSL PLAIN properties.
74    ///
75    /// This is a property helper only. The selected backend must support the
76    /// requested SASL mechanism.
77    #[must_use]
78    pub fn with_sasl_plain(
79        mut self,
80        username: impl Into<String>,
81        password: impl Into<String>,
82    ) -> Self {
83        self.set("security.protocol", "sasl_ssl");
84        self.set("sasl.mechanism", "PLAIN");
85        self.set("sasl.username", username);
86        self.set("sasl.password", password);
87        self
88    }
89}
90
91/// Offset commit mode for `KafkaSource`.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum CommitPolicy {
94    /// Datum-managed at-least-once commits. This is the default.
95    Manual,
96    /// Use Kafka auto-commit. This is an escape hatch and may lose records.
97    AutoKafka,
98    /// Downstream stores offsets externally; Kafka commits are optional.
99    External,
100}
101
102/// Consumer implementation used by `KafkaSource`.
103///
104/// Datum's native consumer is the only supported backend. The producer has a
105/// separate [`KafkaProducerBackend`] selector.
106#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
107pub enum KafkaConsumerBackend {
108    #[default]
109    Native,
110}
111
112/// Producer implementation used by [`crate::KafkaSink`].
113///
114/// Datum's native producer is the only supported backend. It enables
115/// idempotent retries by default, preserving the no-duplicate producer
116/// guarantee.
117#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
118pub enum KafkaProducerBackend {
119    #[default]
120    Native,
121}
122
123/// Topic subscription mode.
124#[derive(Debug, Clone, PartialEq, Eq)]
125pub enum Subscription {
126    Topics(Vec<String>),
127    Pattern(String),
128    Assignment(Vec<TopicPartitionOffset>),
129}
130
131impl Subscription {
132    #[must_use]
133    pub fn topics<I, S>(topics: I) -> Self
134    where
135        I: IntoIterator<Item = S>,
136        S: Into<String>,
137    {
138        Self::Topics(topics.into_iter().map(Into::into).collect())
139    }
140
141    #[must_use]
142    pub fn pattern(pattern: impl Into<String>) -> Self {
143        let pattern = pattern.into();
144        if pattern.starts_with('^') {
145            Self::Pattern(pattern)
146        } else {
147            Self::Pattern(format!("^{pattern}"))
148        }
149    }
150
151    #[must_use]
152    pub fn assignment<I>(partitions: I) -> Self
153    where
154        I: IntoIterator<Item = TopicPartitionOffset>,
155    {
156        Self::Assignment(partitions.into_iter().collect())
157    }
158
159    pub(crate) fn validate(&self) -> MqResult<()> {
160        match self {
161            Subscription::Topics(topics) if topics.is_empty() => Err(MqError::InvalidConfig(
162                "Kafka subscription requires at least one topic".to_owned(),
163            )),
164            Subscription::Pattern(pattern) if pattern.is_empty() || pattern == "^" => {
165                Err(MqError::InvalidConfig(
166                    "Kafka pattern subscription requires a non-empty pattern".to_owned(),
167                ))
168            }
169            Subscription::Assignment(partitions) if partitions.is_empty() => {
170                Err(MqError::InvalidConfig(
171                    "Kafka assignment requires at least one partition".to_owned(),
172                ))
173            }
174            _ => Ok(()),
175        }
176    }
177}
178
179/// A concrete topic/partition plus the starting offset for assignment subscriptions.
180#[derive(Debug, Clone, PartialEq, Eq)]
181pub struct TopicPartitionOffset {
182    pub topic: String,
183    pub partition: i32,
184    pub offset: Offset,
185}
186
187impl TopicPartitionOffset {
188    #[must_use]
189    pub fn new(topic: impl Into<String>, partition: i32, offset: Offset) -> Self {
190        Self {
191            topic: topic.into(),
192            partition,
193            offset,
194        }
195    }
196
197    #[must_use]
198    pub fn beginning(topic: impl Into<String>, partition: i32) -> Self {
199        Self::new(topic, partition, Offset::Beginning)
200    }
201
202    #[must_use]
203    pub fn end(topic: impl Into<String>, partition: i32) -> Self {
204        Self::new(topic, partition, Offset::End)
205    }
206
207    #[must_use]
208    pub fn absolute(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
209        Self::new(topic, partition, Offset::Offset(offset))
210    }
211
212    #[must_use]
213    pub fn topic_partition(&self) -> TopicPartition {
214        TopicPartition::new(self.topic.clone(), self.partition)
215    }
216}
217
218/// Consumer settings for `KafkaSource`.
219#[derive(Debug, Clone)]
220pub struct KafkaConsumerSettings {
221    pub config: KafkaConfig,
222    pub group_id: String,
223    pub consumer_backend: KafkaConsumerBackend,
224    pub commit_policy: CommitPolicy,
225    pub high_watermark: usize,
226    pub low_watermark: usize,
227    pub drain_timeout: Duration,
228    pub poll_timeout: Duration,
229    pub poll_batch_size: usize,
230    pub commit_batch_size: usize,
231    pub commit_interval: Duration,
232    pub low_latency: bool,
233    pub commit_sync: bool,
234    pub include_payload_timestamps: bool,
235}
236
237impl KafkaConsumerSettings {
238    #[must_use]
239    pub fn new(bootstrap_servers: impl Into<String>, group_id: impl Into<String>) -> Self {
240        Self {
241            config: KafkaConfig::new(bootstrap_servers),
242            group_id: group_id.into(),
243            consumer_backend: KafkaConsumerBackend::Native,
244            commit_policy: CommitPolicy::Manual,
245            high_watermark: 4096,
246            low_watermark: 2048,
247            drain_timeout: Duration::from_secs(30),
248            poll_timeout: Duration::from_millis(100),
249            poll_batch_size: 256,
250            commit_batch_size: 10_000,
251            commit_interval: Duration::from_millis(100),
252            low_latency: false,
253            commit_sync: false,
254            include_payload_timestamps: true,
255        }
256    }
257
258    #[must_use]
259    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
260        self.config.set(key, value);
261        self
262    }
263
264    #[must_use]
265    pub fn with_commit_policy(mut self, policy: CommitPolicy) -> Self {
266        self.commit_policy = policy;
267        self
268    }
269
270    #[must_use]
271    pub fn with_consumer_backend(mut self, backend: KafkaConsumerBackend) -> Self {
272        self.consumer_backend = backend;
273        self
274    }
275
276    #[must_use]
277    pub fn with_backpressure(mut self, low_watermark: usize, high_watermark: usize) -> Self {
278        assert!(
279            low_watermark < high_watermark,
280            "low watermark must be below high watermark"
281        );
282        self.low_watermark = low_watermark;
283        self.high_watermark = high_watermark;
284        self
285    }
286
287    #[must_use]
288    pub fn with_drain_timeout(mut self, timeout: Duration) -> Self {
289        self.drain_timeout = timeout;
290        self
291    }
292
293    #[must_use]
294    pub fn with_poll_batch_size(mut self, batch_size: usize) -> Self {
295        assert!(batch_size > 0, "poll batch size must be greater than zero");
296        self.poll_batch_size = batch_size;
297        self
298    }
299
300    #[must_use]
301    pub fn with_commit_batch_size(mut self, batch_size: usize) -> Self {
302        assert!(
303            batch_size > 0,
304            "commit batch size must be greater than zero"
305        );
306        self.commit_batch_size = batch_size;
307        self
308    }
309
310    #[must_use]
311    pub fn with_commit_interval(mut self, interval: Duration) -> Self {
312        self.commit_interval = interval;
313        self
314    }
315
316    #[must_use]
317    pub fn with_low_latency(mut self, enabled: bool) -> Self {
318        self.low_latency = enabled;
319        self
320    }
321
322    #[must_use]
323    pub fn with_payload_timestamps(mut self, enabled: bool) -> Self {
324        self.include_payload_timestamps = enabled;
325        self
326    }
327}
328
329/// Producer settings for `KafkaSink`.
330#[derive(Debug, Clone)]
331pub struct KafkaProducerSettings {
332    pub config: KafkaConfig,
333    pub producer_backend: KafkaProducerBackend,
334    pub in_flight_limit: usize,
335    pub queue_timeout: Duration,
336    pub flush_timeout: Duration,
337    pub drain_timeout: Duration,
338}
339
340impl KafkaProducerSettings {
341    #[must_use]
342    pub fn new(bootstrap_servers: impl Into<String>) -> Self {
343        let settings = Self {
344            config: KafkaConfig::new(bootstrap_servers)
345                .with("enable.idempotence", "true")
346                .with("acks", "all")
347                .with("retries", "2147483647")
348                .with("max.in.flight.requests.per.connection", "5")
349                .with("linger.ms", "5")
350                .with("batch.num.messages", "10000")
351                .with("batch.size", "131072")
352                .with("queue.buffering.max.messages", "1000000")
353                .with("queue.buffering.max.kbytes", "1048576"),
354            producer_backend: KafkaProducerBackend::default(),
355            in_flight_limit: 65_536,
356            queue_timeout: Duration::from_secs(30),
357            flush_timeout: Duration::from_secs(30),
358            drain_timeout: Duration::from_secs(30),
359        };
360        settings.with_producer_backend(KafkaProducerBackend::default())
361    }
362
363    #[must_use]
364    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
365        self.config.set(key, value);
366        self
367    }
368
369    /// Selects the producer backend.
370    ///
371    /// Selecting the native backend applies idempotent producer defaults
372    /// (`enable.idempotence=true`): `InitProducerId` negotiates a producer
373    /// id/epoch and per-partition sequences deduplicate retried batches. This
374    /// requires `acks=all` and at most five in-flight requests per connection;
375    /// an incompatible config is rejected at materialization with a typed error
376    /// naming the setting. Native Produce requests default to one in flight per
377    /// broker connection, and requests for the same topic-partition are always
378    /// serialized. Users may explicitly opt out with
379    /// `.with("enable.idempotence", "false")`.
380    #[must_use]
381    pub fn with_producer_backend(mut self, backend: KafkaProducerBackend) -> Self {
382        self.producer_backend = backend;
383        self.config.set("enable.idempotence", "true");
384        self.config.set("retries", "3");
385        self.config
386            .set("max.in.flight.requests.per.connection", "1");
387        self
388    }
389
390    /// Sets the maximum number of native Produce requests in flight on each
391    /// broker connection.
392    ///
393    /// The native backend preserves per-partition order by never overlapping
394    /// two requests that touch the same topic-partition. On a retriable
395    /// failure it halts that connection's pipeline, settles every issued
396    /// response, and retries failed requests in their original order. Without
397    /// idempotence a lost response can still duplicate an accepted batch;
398    /// enabling `enable.idempotence` eliminates those duplicates and caps this
399    /// value at five. The value must be non-zero.
400    #[must_use]
401    pub fn with_max_in_flight_requests_per_connection(mut self, limit: usize) -> Self {
402        assert!(
403            limit > 0,
404            "max in-flight requests must be greater than zero"
405        );
406        self.config
407            .set("max.in.flight.requests.per.connection", limit.to_string());
408        self
409    }
410
411    #[must_use]
412    pub fn with_in_flight_limit(mut self, in_flight_limit: usize) -> Self {
413        assert!(
414            in_flight_limit > 0,
415            "in-flight limit must be greater than zero"
416        );
417        self.in_flight_limit = in_flight_limit;
418        self
419    }
420
421    #[must_use]
422    pub fn with_queue_timeout(mut self, timeout: Duration) -> Self {
423        self.queue_timeout = timeout;
424        self
425    }
426
427    #[must_use]
428    pub fn with_drain_timeout(mut self, timeout: Duration) -> Self {
429        self.drain_timeout = timeout;
430        self
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn raw_properties_are_preserved() {
440        let config = KafkaConfig::new("127.0.0.1:9092")
441            .with("client.id", "datum-test")
442            .with("linger.ms", "5");
443
444        assert_eq!(config.get("bootstrap.servers"), Some("127.0.0.1:9092"));
445        assert_eq!(config.get("client.id"), Some("datum-test"));
446        assert_eq!(config.get("linger.ms"), Some("5"));
447    }
448
449    #[test]
450    fn consumer_backend_defaults_to_native() {
451        let settings = KafkaConsumerSettings::new("localhost:9092", "group");
452        assert_eq!(settings.consumer_backend, KafkaConsumerBackend::Native);
453    }
454
455    #[test]
456    fn consumer_backend_can_be_selected_explicitly() {
457        let settings = KafkaConsumerSettings::new("localhost:9092", "group")
458            .with_consumer_backend(KafkaConsumerBackend::Native);
459        assert_eq!(settings.consumer_backend, KafkaConsumerBackend::Native);
460    }
461
462    #[test]
463    fn native_producer_backend_uses_idempotent_defaults() {
464        let settings = KafkaProducerSettings::new("localhost:9092")
465            .with_producer_backend(KafkaProducerBackend::Native);
466        assert_eq!(settings.producer_backend, KafkaProducerBackend::Native);
467        assert_eq!(settings.config.get("enable.idempotence"), Some("true"));
468        assert_eq!(settings.config.get("retries"), Some("3"));
469        assert_eq!(
470            settings.config.get("max.in.flight.requests.per.connection"),
471            Some("1")
472        );
473    }
474
475    #[test]
476    fn producer_backend_defaults_to_native() {
477        let settings = KafkaProducerSettings::new("localhost:9092");
478        assert_eq!(settings.producer_backend, KafkaProducerBackend::Native);
479        assert_eq!(settings.config.get("enable.idempotence"), Some("true"));
480        assert_eq!(settings.config.get("retries"), Some("3"));
481        assert_eq!(
482            settings.config.get("max.in.flight.requests.per.connection"),
483            Some("1")
484        );
485    }
486
487    #[test]
488    fn native_producer_idempotence_can_be_disabled_explicitly() {
489        let settings = KafkaProducerSettings::new("localhost:9092")
490            .with_producer_backend(KafkaProducerBackend::Native)
491            .with("enable.idempotence", "false");
492        assert_eq!(settings.producer_backend, KafkaProducerBackend::Native);
493        assert_eq!(settings.config.get("enable.idempotence"), Some("false"));
494    }
495
496    #[test]
497    fn tls_and_sasl_helpers_set_kafka_properties() {
498        let tls = KafkaConfig::new("broker").with_tls("/ca.pem", "/cert.pem", "/key.pem");
499        assert_eq!(tls.get("security.protocol"), Some("ssl"));
500        assert_eq!(tls.get("ssl.ca.location"), Some("/ca.pem"));
501
502        let sasl = KafkaConfig::new("broker").with_sasl_plain("user", "pass");
503        assert_eq!(sasl.get("security.protocol"), Some("sasl_ssl"));
504        assert_eq!(sasl.get("sasl.mechanism"), Some("PLAIN"));
505        assert_eq!(sasl.get("sasl.username"), Some("user"));
506        assert_eq!(sasl.get("sasl.password"), Some("pass"));
507    }
508
509    #[test]
510    fn subscription_validation_rejects_empty_inputs() {
511        assert!(
512            Subscription::topics(Vec::<String>::new())
513                .validate()
514                .is_err()
515        );
516        assert!(Subscription::pattern("").validate().is_err());
517        assert!(Subscription::assignment(Vec::new()).validate().is_err());
518    }
519}
datum_mq/config.rs

datum_mq/
config.rs
