krabka_connect_worker/
kafka.rs

1//! Kafka sink and compacted-topic checkpoint store.
2
3use std::collections::HashSet;
4
5use async_trait::async_trait;
6use bytes::Bytes;
7use krabka_client_core::security::ClientSecurity;
8use krabka_client_producer::{Acks, DeliveryHandle, Header, Producer, ProducerRecord};
9use krabka_connect::{CheckpointStore, ConnectError, ConnectRecord, Sink, SourceOffset};
10use krabka_replicator::config::ReplicationFactor;
11use krabka_units::{ByteSize, convert::ByteSizeExt as _};
12
13use crate::metrics::WorkerMetrics;
14
15/// Compacted Kafka topic holding the most recent source offset per connector.
16pub const CHECKPOINT_TOPIC: &str = "__krabka_connect_offsets";
17
18#[derive(Clone)]
19pub(crate) struct KafkaClientConfig {
20    pub(crate) bootstrap: String,
21    pub(crate) security: Option<ClientSecurity>,
22    pub(crate) replication_factor: ReplicationFactor,
23    pub(crate) dispatch_queue_capacity: usize,
24    pub(crate) frame_max_bytes: u64,
25}
26
27async fn build_producer(
28    config: &KafkaClientConfig,
29    client_id: String,
30) -> Result<Producer, ConnectError> {
31    let builder = Producer::builder()
32        .bootstrap(config.bootstrap.clone())
33        .client_id(client_id)
34        .dispatch_queue_capacity(config.dispatch_queue_capacity)
35        .frame_max(ByteSize::from_bytes(config.frame_max_bytes))
36        .enable_idempotence(true)
37        .acks(Acks::All);
38    match config.security.clone() {
39        Some(security) => builder.security(security).build().await,
40        None => builder.build().await,
41    }
42    .map_err(|error| ConnectError::Backend(error.to_string()))
43}
44
45/// At-least-once Kafka sink backed by an idempotent `acks=all` producer.
46pub struct KafkaSink {
47    producer: Producer,
48    client: KafkaClientConfig,
49    topic_prefix: String,
50    ensured_topics: HashSet<String>,
51    pending: Vec<DeliveryHandle>,
52    metrics: WorkerMetrics,
53}
54
55impl KafkaSink {
56    /// Connect an idempotent Kafka producer.
57    ///
58    /// # Errors
59    ///
60    /// Returns [`ConnectError::Backend`] when the producer cannot connect.
61    pub async fn start(
62        bootstrap: impl Into<String>,
63        topic_prefix: impl Into<String>,
64    ) -> Result<Self, ConnectError> {
65        let client = KafkaClientConfig {
66            bootstrap: bootstrap.into(),
67            security: None,
68            replication_factor: ReplicationFactor::default(),
69            dispatch_queue_capacity: 64,
70            frame_max_bytes: 100 * 1024 * 1024,
71        };
72        Self::start_with_config(client, topic_prefix.into(), WorkerMetrics::new()).await
73    }
74
75    pub(crate) async fn start_with_config(
76        client: KafkaClientConfig,
77        topic_prefix: String,
78        metrics: WorkerMetrics,
79    ) -> Result<Self, ConnectError> {
80        let producer = build_producer(&client, "krabka-connect-worker-data".to_owned()).await?;
81        Ok(Self {
82            producer,
83            client,
84            topic_prefix,
85            ensured_topics: HashSet::new(),
86            pending: Vec::new(),
87            metrics,
88        })
89    }
90}
91
92#[async_trait]
93impl Sink<Bytes, Bytes> for KafkaSink {
94    async fn put(&mut self, records: Vec<ConnectRecord<Bytes, Bytes>>) -> Result<(), ConnectError> {
95        for record in records {
96            let output = to_producer_record(record, &self.topic_prefix)?;
97            if self.ensured_topics.insert(output.topic.clone()) {
98                krabka_replicator::admin_util::ensure_topic_with_replication_factor(
99                    &self.client.bootstrap,
100                    &output.topic,
101                    1,
102                    self.client.security.clone(),
103                    self.client.replication_factor,
104                )
105                .await
106                .map_err(ConnectError::Backend)?;
107            }
108            self.pending.push(
109                self.producer
110                    .enqueue(output)
111                    .await
112                    .map_err(|error| ConnectError::Backend(error.to_string()))?,
113            );
114        }
115        Ok(())
116    }
117
118    async fn flush(&mut self) -> Result<(), ConnectError> {
119        let mut first_error = None;
120        for acknowledgement in std::mem::take(&mut self.pending) {
121            match acknowledgement.await {
122                Ok(_) => self.metrics.record_produced(),
123                Err(error) => {
124                    first_error.get_or_insert_with(|| error.to_string());
125                }
126            }
127        }
128        if let Err(error) = self.producer.flush().await {
129            first_error.get_or_insert_with(|| error.to_string());
130        }
131        if let Some(error) = first_error {
132            self.metrics.record_error();
133            Err(ConnectError::Backend(error))
134        } else {
135            Ok(())
136        }
137    }
138
139    async fn close(&mut self) -> Result<(), ConnectError> {
140        self.flush().await
141    }
142}
143
144fn to_producer_record(
145    record: ConnectRecord<Bytes, Bytes>,
146    topic_prefix: &str,
147) -> Result<ProducerRecord, ConnectError> {
148    let source_topic = record.topic.ok_or_else(|| {
149        ConnectError::Backend("source record is missing destination topic".to_owned())
150    })?;
151    if matches!(record.partition, Some(partition) if partition < 0) {
152        return Err(ConnectError::Backend(
153            "source record has a negative destination partition".to_owned(),
154        ));
155    }
156    Ok(ProducerRecord {
157        topic: prefixed_topic(topic_prefix, &source_topic),
158        partition: record.partition,
159        key: record.key,
160        value: record.value,
161        headers: record
162            .headers
163            .into_iter()
164            .map(|header| Header {
165                key: header.key,
166                value: header.value,
167            })
168            .collect(),
169        timestamp_ms: record.timestamp,
170    })
171}
172
173fn prefixed_topic(prefix: &str, topic: &str) -> String {
174    let prefix = prefix.trim_end_matches('.');
175    if prefix.is_empty() {
176        topic.to_owned()
177    } else {
178        format!("{prefix}.{}", topic.trim_start_matches('.'))
179    }
180}
181
182/// Durable connector checkpoint store keyed by connector and source identity.
183pub struct KafkaCheckpointStore {
184    producer: Producer,
185    client: KafkaClientConfig,
186    checkpoint_key: String,
187    metrics: WorkerMetrics,
188}
189
190impl KafkaCheckpointStore {
191    /// Ensure the compacted checkpoint topic and connect an idempotent producer.
192    ///
193    /// # Errors
194    ///
195    /// Returns [`ConnectError::Offset`] when topic creation or producer startup fails.
196    pub async fn start(
197        bootstrap: impl Into<String>,
198        checkpoint_key: impl Into<String>,
199    ) -> Result<Self, ConnectError> {
200        let client = KafkaClientConfig {
201            bootstrap: bootstrap.into(),
202            security: None,
203            replication_factor: ReplicationFactor::default(),
204            dispatch_queue_capacity: 64,
205            frame_max_bytes: 100 * 1024 * 1024,
206        };
207        Self::start_with_config(client, checkpoint_key.into(), WorkerMetrics::new()).await
208    }
209
210    pub(crate) async fn start_with_config(
211        client: KafkaClientConfig,
212        checkpoint_key: String,
213        metrics: WorkerMetrics,
214    ) -> Result<Self, ConnectError> {
215        krabka_replicator::admin_util::ensure_compacted_topic_with_replication_factor(
216            &client.bootstrap,
217            CHECKPOINT_TOPIC,
218            client.security.clone(),
219            client.replication_factor,
220        )
221        .await
222        .map_err(ConnectError::Offset)?;
223        let producer = build_producer(&client, "krabka-connect-worker-checkpoint".to_owned())
224            .await
225            .map_err(|error| ConnectError::Offset(error.to_string()))?;
226        Ok(Self {
227            producer,
228            client,
229            checkpoint_key,
230            metrics,
231        })
232    }
233}
234
235#[async_trait]
236impl CheckpointStore for KafkaCheckpointStore {
237    async fn save(&self, offset: &SourceOffset) -> Result<(), ConnectError> {
238        let value =
239            serde_json::to_vec(offset).map_err(|error| ConnectError::Offset(error.to_string()))?;
240        self.producer
241            .send(ProducerRecord {
242                topic: CHECKPOINT_TOPIC.to_owned(),
243                partition: Some(0),
244                key: Some(Bytes::copy_from_slice(self.checkpoint_key.as_bytes())),
245                value: Some(Bytes::from(value)),
246                headers: Vec::new(),
247                timestamp_ms: None,
248            })
249            .await
250            .map_err(|error| ConnectError::Offset(error.to_string()))?;
251        self.producer
252            .flush()
253            .await
254            .map_err(|error| ConnectError::Offset(error.to_string()))?;
255        self.metrics.record_checkpoint();
256        Ok(())
257    }
258
259    async fn load(&self) -> Result<Option<SourceOffset>, ConnectError> {
260        let value = krabka_replicator::admin_util::read_last_value_for_key(
261            &self.client.bootstrap,
262            CHECKPOINT_TOPIC,
263            self.checkpoint_key.as_bytes(),
264            self.client.security.clone(),
265        )
266        .await
267        .map_err(ConnectError::Offset)?;
268        value
269            .map(|bytes| {
270                serde_json::from_slice(&bytes)
271                    .map_err(|error| ConnectError::Offset(error.to_string()))
272            })
273            .transpose()
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use assert2::assert;
280
281    use super::*;
282
283    #[test]
284    fn record_mapping_preserves_kafka_fields_and_applies_prefix() {
285        let source = ConnectRecord::new(
286            Some(Bytes::from_static(b"key")),
287            Some(Bytes::from_static(b"value")),
288        )
289        .with_topic("public.orders")
290        .with_partition(2)
291        .with_timestamp(42)
292        .with_header("trace", Some(Bytes::from_static(b"abc")));
293        let output = to_producer_record(source, "db.").expect("record maps");
294        assert!(
295            output
296                == ProducerRecord {
297                    topic: "db.public.orders".to_owned(),
298                    partition: Some(2),
299                    key: Some(Bytes::from_static(b"key")),
300                    value: Some(Bytes::from_static(b"value")),
301                    headers: vec![Header {
302                        key: "trace".to_owned(),
303                        value: Some(Bytes::from_static(b"abc")),
304                    }],
305                    timestamp_ms: Some(42),
306                }
307        );
308    }
309
310    #[test]
311    fn record_mapping_preserves_tombstone_and_rejects_missing_topic() {
312        let tombstone =
313            ConnectRecord::new(Some(Bytes::from_static(b"key")), None).with_topic("orders");
314        let output = to_producer_record(tombstone, "").expect("tombstone maps");
315        assert!(output.value.is_none());
316        assert!(output.topic == "orders");
317        assert!(to_producer_record(ConnectRecord::new(None, None), "db").is_err());
318    }
319}
krabka_connect_worker/kafka.rs

krabka_connect_worker/
kafka.rs
