laminar_connectors/kafka/
avro.rs

1//! Avro deserialization using `arrow-avro` with Confluent Schema Registry.
2//!
3//! [`AvroDeserializer`] implements [`RecordDeserializer`] by wrapping the
4//! `arrow-avro` push-based [`Decoder`], which
5//! natively supports the Confluent wire format (`0x00` + 4-byte BE schema ID
6//! + Avro payload).
7
8use std::collections::HashSet;
9use std::sync::Arc;
10
11use arrow_array::RecordBatch;
12use arrow_avro::reader::{Decoder, ReaderBuilder};
13use arrow_avro::schema::{AvroSchema, Fingerprint, FingerprintAlgorithm, SchemaStore};
14use arrow_schema::SchemaRef;
15use parking_lot::Mutex;
16
17use crate::error::{ConnectorError, SerdeError};
18use crate::kafka::schema_registry::SchemaRegistryClient;
19use crate::serde::{Format, RecordDeserializer};
20
21const DECODER_BATCH_CAPACITY: usize = 8192;
22
23/// Confluent wire format magic byte.
24const CONFLUENT_MAGIC: u8 = 0x00;
25
26/// Size of the Confluent wire format header (1 magic + 4 schema ID).
27const CONFLUENT_HEADER_SIZE: usize = 5;
28
29/// Avro deserializer backed by `arrow-avro` with optional Schema Registry.
30///
31/// Supports both raw Avro and the Confluent wire format. When a Schema
32/// Registry client is provided, unknown schema IDs are fetched and
33/// registered automatically.
34pub struct AvroDeserializer {
35    /// Schema store shared with the Decoder.
36    schema_store: SchemaStore,
37    /// Optional Schema Registry client for resolving unknown IDs.
38    schema_registry: Option<Arc<SchemaRegistryClient>>,
39    /// Set of schema IDs already registered in the store.
40    known_ids: HashSet<i32>,
41    /// Reused across batches; rebuilt when `register_schema` runs.
42    decoder: Mutex<Option<Decoder>>,
43}
44
45impl AvroDeserializer {
46    /// Creates a new Avro deserializer without Schema Registry integration.
47    ///
48    /// The caller must register schemas manually via [`register_schema`](Self::register_schema).
49    #[must_use]
50    pub fn new() -> Self {
51        Self {
52            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
53            schema_registry: None,
54            known_ids: HashSet::new(),
55            decoder: Mutex::new(None),
56        }
57    }
58
59    /// Creates a new Avro deserializer with Schema Registry integration.
60    ///
61    /// Unknown schema IDs encountered in the Confluent wire format will
62    /// be fetched from the registry automatically.
63    #[must_use]
64    pub fn with_schema_registry(registry: Arc<SchemaRegistryClient>) -> Self {
65        Self {
66            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
67            schema_registry: Some(registry),
68            known_ids: HashSet::new(),
69            decoder: Mutex::new(None),
70        }
71    }
72
73    /// Registers an Avro schema with a Confluent schema ID.
74    ///
75    /// # Errors
76    ///
77    /// Returns `SerdeError` if the fingerprint cannot be set.
78    #[allow(clippy::cast_sign_loss)]
79    pub fn register_schema(
80        &mut self,
81        schema_id: i32,
82        avro_schema_json: &str,
83    ) -> Result<(), SerdeError> {
84        let avro_schema = AvroSchema::new(avro_schema_json.to_string());
85        // Use Fingerprint::Id directly — NOT load_fingerprint_id which
86        // applies from_be byte-swap meant for raw wire bytes.
87        let fp = Fingerprint::Id(schema_id as u32);
88        self.schema_store
89            .set(fp, avro_schema)
90            .map_err(|e| SerdeError::MalformedInput(format!("failed to register schema: {e}")))?;
91        self.known_ids.insert(schema_id);
92        // Schema store changed — drop the cached decoder so it rebuilds
93        // against the new store on the next deserialize_batch call.
94        *self.decoder.lock() = None;
95        Ok(())
96    }
97
98    /// Ensures a schema ID is registered, fetching from SR if needed.
99    ///
100    /// Called by the Kafka source connector when an unknown schema ID is
101    /// encountered in the Confluent wire format during poll.
102    ///
103    /// # Errors
104    ///
105    /// Preserves Schema Registry `ConnectorError` classification for remote
106    /// resolution failures and returns `ConnectorError::Serde` for local
107    /// decoder registration failures.
108    /// Returns `Ok(true)` if this was a newly registered schema ID,
109    /// `Ok(false)` if already known.
110    pub async fn ensure_schema_registered(
111        &mut self,
112        schema_id: i32,
113    ) -> Result<bool, ConnectorError> {
114        if self.known_ids.contains(&schema_id) {
115            return Ok(false);
116        }
117
118        let registry = self.schema_registry.as_ref().ok_or(ConnectorError::Serde(
119            SerdeError::SchemaNotFound { schema_id },
120        ))?;
121
122        let cached = registry.resolve_confluent_id(schema_id).await?;
123
124        self.register_schema(schema_id, &cached.schema_str)
125            .map_err(ConnectorError::Serde)?;
126        Ok(true)
127    }
128
129    /// Extracts the Confluent schema ID from a wire-format message.
130    ///
131    /// Returns `None` if the message is not in Confluent wire format.
132    #[must_use]
133    pub fn extract_confluent_id(data: &[u8]) -> Option<i32> {
134        if data.len() < CONFLUENT_HEADER_SIZE || data[0] != CONFLUENT_MAGIC {
135            return None;
136        }
137        let id = i32::from_be_bytes([data[1], data[2], data[3], data[4]]);
138        Some(id)
139    }
140}
141
142impl Default for AvroDeserializer {
143    fn default() -> Self {
144        Self::new()
145    }
146}
147
148impl RecordDeserializer for AvroDeserializer {
149    fn deserialize(&self, data: &[u8], schema: &SchemaRef) -> Result<RecordBatch, SerdeError> {
150        self.deserialize_batch(&[data], schema)
151    }
152
153    fn deserialize_batch(
154        &self,
155        records: &[&[u8]],
156        schema: &SchemaRef,
157    ) -> Result<RecordBatch, SerdeError> {
158        if records.is_empty() {
159            return Ok(RecordBatch::new_empty(schema.clone()));
160        }
161
162        let mut guard = self.decoder.lock();
163        let decoder = if let Some(d) = guard.as_mut() {
164            d
165        } else {
166            let d = ReaderBuilder::new()
167                .with_batch_size(DECODER_BATCH_CAPACITY)
168                .with_writer_schema_store(self.schema_store.clone())
169                .build_decoder()
170                .map_err(|e| SerdeError::MalformedInput(format!("failed to build decoder: {e}")))?;
171            guard.insert(d)
172        };
173
174        let mut partials: Vec<RecordBatch> = Vec::new();
175        for record in records {
176            let mut offset = 0;
177            while offset < record.len() {
178                let consumed = decoder
179                    .decode(&record[offset..])
180                    .map_err(|e| SerdeError::MalformedInput(format!("Avro decode error: {e}")))?;
181                if consumed == 0 {
182                    break;
183                }
184                offset += consumed;
185                if decoder.batch_is_full() {
186                    if let Some(b) = decoder
187                        .flush()
188                        .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
189                    {
190                        partials.push(b);
191                    }
192                }
193            }
194        }
195        if let Some(b) = decoder
196            .flush()
197            .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
198        {
199            partials.push(b);
200        }
201
202        match partials.len() {
203            0 => Err(SerdeError::MalformedInput("no records decoded".into())),
204            1 => Ok(partials.pop().unwrap()),
205            _ => arrow_select::concat::concat_batches(schema, &partials)
206                .map_err(|e| SerdeError::MalformedInput(format!("concat: {e}"))),
207        }
208    }
209
210    fn format(&self) -> Format {
211        Format::Avro
212    }
213
214    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
215        Some(self)
216    }
217}
218
219impl std::fmt::Debug for AvroDeserializer {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        f.debug_struct("AvroDeserializer")
222            .field("known_ids", &self.known_ids)
223            .field("has_registry", &self.schema_registry.is_some())
224            .finish_non_exhaustive()
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use arrow_schema::{DataType, Field, Schema};
232
233    const TEST_AVRO_SCHEMA: &str = r#"{
234        "type": "record",
235        "name": "test",
236        "fields": [
237            {"name": "id", "type": "long"},
238            {"name": "name", "type": "string"}
239        ]
240    }"#;
241
242    fn test_arrow_schema() -> SchemaRef {
243        Arc::new(Schema::new(vec![
244            Field::new("id", DataType::Int64, false),
245            Field::new("name", DataType::Utf8, false),
246        ]))
247    }
248
249    /// Avro `union<null, map<string, union<null, double>>>` must produce
250    /// identical Arrow schemas from both `avro_to_arrow_schema` (SR path)
251    /// and `AvroDeserializer` (wire decode path).
252    #[test]
253    fn test_nullable_map_nullable_double_full_path() {
254        fn zigzag(val: i64) -> Vec<u8> {
255            let mut z = ((val << 1) ^ (val >> 63)) as u64;
256            let mut buf = Vec::new();
257            loop {
258                if z & !0x7F == 0 {
259                    buf.push(z as u8);
260                    break;
261                }
262                buf.push((z as u8 & 0x7F) | 0x80);
263                z >>= 7;
264            }
265            buf
266        }
267        fn avro_string(s: &str) -> Vec<u8> {
268            let mut b = zigzag(s.len() as i64);
269            b.extend_from_slice(s.as_bytes());
270            b
271        }
272
273        let avro_json = r#"{
274            "type": "record",
275            "name": "Metrics",
276            "fields": [
277                {"name": "sensor_id", "type": "string"},
278                {
279                    "name": "data",
280                    "type": ["null", {"type": "map", "values": ["null", "double"]}]
281                }
282            ]
283        }"#;
284
285        // Path 1: what the schema registry infers
286        let sr_schema = crate::kafka::schema_registry::avro_to_arrow_schema(avro_json)
287            .expect("avro_to_arrow_schema should handle nullable map");
288
289        // Path 2: what the decoder actually produces from wire bytes
290        let mut deser = AvroDeserializer::new();
291        deser.register_schema(42, avro_json).unwrap();
292
293        // Encode { sensor_id: "s1", data: {"temp": 23.5} } in Avro binary.
294        // Union branches are prefixed with a zigzag-encoded index.
295        let mut payload = Vec::new();
296        payload.extend_from_slice(&avro_string("s1"));
297        payload.extend_from_slice(&zigzag(1)); // data: branch 1 (map, not null)
298        payload.extend_from_slice(&zigzag(1)); // map block: 1 entry
299        payload.extend_from_slice(&avro_string("temp"));
300        payload.extend_from_slice(&zigzag(1)); // value: branch 1 (double, not null)
301        payload.extend_from_slice(&23.5_f64.to_le_bytes());
302        payload.extend_from_slice(&zigzag(0)); // end of map
303
304        // Confluent wire frame: magic + schema_id + payload
305        let mut msg = vec![0x00u8];
306        msg.extend_from_slice(&42i32.to_be_bytes());
307        msg.extend_from_slice(&payload);
308
309        let batch = deser
310            .deserialize_batch(&[msg.as_slice()], &sr_schema)
311            .expect("decode should succeed");
312        assert_eq!(batch.num_rows(), 1);
313
314        let sr_data = sr_schema.field_with_name("data").unwrap();
315        let decoded_schema = batch.schema();
316        let dec_data = decoded_schema.field_with_name("data").unwrap();
317        assert_eq!(
318            sr_data.data_type(),
319            dec_data.data_type(),
320            "schema registry and arrow-avro must produce identical Map types"
321        );
322    }
323
324    #[test]
325    fn test_extract_confluent_id() {
326        // Valid: 0x00 + 4-byte BE schema ID
327        let data = [0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03];
328        assert_eq!(AvroDeserializer::extract_confluent_id(&data), Some(1));
329
330        let data = [0x00, 0x00, 0x00, 0x01, 0x00, 0x02, 0x03];
331        assert_eq!(AvroDeserializer::extract_confluent_id(&data), Some(256));
332    }
333
334    #[test]
335    fn test_extract_confluent_id_not_confluent() {
336        let data = [0x01, 0x00, 0x00, 0x00, 0x01];
337        assert_eq!(AvroDeserializer::extract_confluent_id(&data), None);
338    }
339
340    #[test]
341    fn test_extract_confluent_id_too_short() {
342        let data = [0x00, 0x00];
343        assert_eq!(AvroDeserializer::extract_confluent_id(&data), None);
344    }
345
346    #[test]
347    fn test_new_deserializer() {
348        let deser = AvroDeserializer::new();
349        assert!(deser.schema_registry.is_none());
350        assert!(deser.known_ids.is_empty());
351    }
352
353    #[test]
354    fn test_register_schema() {
355        let mut deser = AvroDeserializer::new();
356        let result = deser.register_schema(1, TEST_AVRO_SCHEMA);
357        assert!(result.is_ok());
358        assert!(deser.known_ids.contains(&1));
359    }
360
361    #[tokio::test]
362    async fn hot_schema_resolution_preserves_registry_error_type_and_context() {
363        use wiremock::matchers::{method, path};
364        use wiremock::{Mock, MockServer, ResponseTemplate};
365
366        let server = MockServer::start().await;
367        Mock::given(method("GET"))
368            .and(path("/schemas/ids/42"))
369            .respond_with(ResponseTemplate::new(401).set_body_string("invalid credentials"))
370            .mount(&server)
371            .await;
372        let registry = Arc::new(SchemaRegistryClient::new(server.uri(), None).unwrap());
373        let mut deser = AvroDeserializer::with_schema_registry(registry);
374
375        let error = deser.ensure_schema_registered(42).await.unwrap_err();
376        assert!(matches!(error, ConnectorError::ConfigurationError(_)));
377        assert!(!error.is_transient());
378        assert!(error.to_string().contains("schema ID 42"));
379    }
380
381    #[tokio::test]
382    async fn missing_hot_path_registry_remains_a_serde_error() {
383        let mut deser = AvroDeserializer::new();
384        let error = deser.ensure_schema_registered(42).await.unwrap_err();
385        assert!(matches!(
386            error,
387            ConnectorError::Serde(SerdeError::SchemaNotFound { schema_id: 42 })
388        ));
389    }
390
391    #[test]
392    fn test_format() {
393        let deser = AvroDeserializer::new();
394        assert_eq!(deser.format(), Format::Avro);
395    }
396
397    #[test]
398    fn test_deserialize_empty_batch() {
399        let deser = AvroDeserializer::new();
400        let schema = test_arrow_schema();
401        let result = deser.deserialize_batch(&[], &schema);
402        assert!(result.is_ok());
403        assert_eq!(result.unwrap().num_rows(), 0);
404    }
405
406    #[test]
407    fn test_extract_confluent_id_large() {
408        // Schema ID 42
409        let mut data = vec![0x00u8];
410        data.extend_from_slice(&42i32.to_be_bytes());
411        data.push(0xFF);
412        assert_eq!(AvroDeserializer::extract_confluent_id(&data), Some(42));
413    }
414
415    #[test]
416    fn test_extract_confluent_id_edge_cases() {
417        // Empty slice
418        assert_eq!(AvroDeserializer::extract_confluent_id(&[]), None);
419
420        // Magic byte only (too short)
421        assert_eq!(AvroDeserializer::extract_confluent_id(&[0x00]), None);
422
423        // 4 bytes with magic (too short)
424        assert_eq!(
425            AvroDeserializer::extract_confluent_id(&[0x00, 0x00, 0x00, 0x00]),
426            None
427        );
428
429        // Exactly 5 bytes (boundary of CONFLUENT_HEADER_SIZE)
430        let data = [0x00, 0x00, 0x00, 0x00, 0x05];
431        assert_eq!(AvroDeserializer::extract_confluent_id(&data), Some(5));
432
433        // i32::MAX ID
434        let mut data = vec![0x00];
435        data.extend_from_slice(&i32::MAX.to_be_bytes());
436        assert_eq!(
437            AvroDeserializer::extract_confluent_id(&data),
438            Some(i32::MAX)
439        );
440
441        // i32::MIN ID
442        let mut data = vec![0x00];
443        data.extend_from_slice(&i32::MIN.to_be_bytes());
444        assert_eq!(
445            AvroDeserializer::extract_confluent_id(&data),
446            Some(i32::MIN)
447        );
448    }
449}
laminar_connectors/kafka/avro.rs

laminar_connectors/kafka/
avro.rs
