Skip to main content

apache_avro/reader/
single_object.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::{io::Read, marker::PhantomData};
19
20use bon::bon;
21use serde::de::DeserializeOwned;
22
23use crate::{
24    AvroResult, AvroSchema, Schema,
25    decode::{DecodeContext, decode_internal},
26    error::Details,
27    headers::{HeaderBuilder, RabinFingerprintHeader},
28    schema::ResolvedOwnedSchema,
29    serde::deser_schema::{Config, SchemaAwareDeserializer},
30    types::Value,
31    util::is_human_readable,
32};
33
34pub struct GenericSingleObjectReader {
35    write_schema: ResolvedOwnedSchema,
36    expected_header: Vec<u8>,
37    human_readable: bool,
38}
39
40#[bon]
41impl GenericSingleObjectReader {
42    #[builder]
43    pub fn new(
44        schema: Schema,
45        /// The expected header.
46        #[builder(default = RabinFingerprintHeader::from_schema(&schema).build_header())]
47        header: Vec<u8>,
48        /// Was the data serialized with `human_readable`.
49        #[builder(default = is_human_readable())]
50        human_readable: bool,
51    ) -> AvroResult<GenericSingleObjectReader> {
52        Ok(Self {
53            write_schema: schema.try_into()?,
54            expected_header: header,
55            human_readable,
56        })
57    }
58}
59
60impl GenericSingleObjectReader {
61    pub fn read_value<R: Read>(&self, reader: &mut R) -> AvroResult<Value> {
62        self.read_header(reader)?;
63        decode_internal(
64            self.write_schema.get_root_schema(),
65            self.write_schema.get_names(),
66            None,
67            reader,
68            &mut DecodeContext::new(),
69        )
70    }
71
72    pub fn read_deser<T: DeserializeOwned>(&self, reader: &mut impl Read) -> AvroResult<T> {
73        self.read_header(reader)?;
74        let config = Config {
75            names: self.write_schema.get_names(),
76            human_readable: self.human_readable,
77            recursion_depth: 0,
78        };
79        T::deserialize(SchemaAwareDeserializer::new(
80            reader,
81            self.write_schema.get_root_schema(),
82            config,
83        )?)
84    }
85
86    fn read_header(&self, reader: &mut impl Read) -> AvroResult<()> {
87        let mut header = vec![0; self.expected_header.len()];
88        reader
89            .read_exact(&mut header)
90            .map_err(Details::ReadHeader)?;
91        if self.expected_header == header {
92            Ok(())
93        } else {
94            Err(Details::SingleObjectHeaderMismatch(self.expected_header.clone(), header).into())
95        }
96    }
97}
98
99pub struct SpecificSingleObjectReader<T>
100where
101    T: AvroSchema,
102{
103    inner: GenericSingleObjectReader,
104    _model: PhantomData<T>,
105}
106
107impl<T> SpecificSingleObjectReader<T>
108where
109    T: AvroSchema,
110{
111    pub fn new() -> AvroResult<SpecificSingleObjectReader<T>> {
112        Ok(SpecificSingleObjectReader {
113            inner: GenericSingleObjectReader::builder()
114                .schema(T::get_schema())
115                .build()?,
116            _model: PhantomData,
117        })
118    }
119}
120
121impl<T> SpecificSingleObjectReader<T>
122where
123    T: AvroSchema + From<Value>,
124{
125    pub fn read_from_value<R: Read>(&self, reader: &mut R) -> AvroResult<T> {
126        self.inner.read_value(reader).map(|v| v.into())
127    }
128}
129
130impl<T> SpecificSingleObjectReader<T>
131where
132    T: AvroSchema + DeserializeOwned,
133{
134    pub fn read<R: Read>(&self, reader: &mut R) -> AvroResult<T> {
135        self.inner.read_deser(reader)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use apache_avro_test_helper::TestResult;
142    use serde::Deserialize;
143    use uuid::Uuid;
144
145    use super::*;
146    use crate::{AvroSchema, Schema, encode::encode, headers::GlueSchemaUuidHeader, rabin::Rabin};
147
148    #[derive(Deserialize, Clone, PartialEq, Debug)]
149    struct TestSingleObjectReader {
150        a: i64,
151        b: f64,
152        c: Vec<String>,
153    }
154
155    impl AvroSchema for TestSingleObjectReader {
156        fn get_schema() -> Schema {
157            let schema = r#"
158            {
159                "type":"record",
160                "name":"TestSingleObjectReader",
161                "fields":[
162                    {
163                        "name":"a",
164                        "type":"long"
165                    },
166                    {
167                        "name":"b",
168                        "type":"double"
169                    },
170                    {
171                        "name":"c",
172                        "type":{
173                            "type":"array",
174                            "items":"string"
175                        }
176                    }
177                ]
178            }
179            "#;
180            Schema::parse_str(schema).unwrap()
181        }
182    }
183
184    impl From<Value> for TestSingleObjectReader {
185        fn from(obj: Value) -> TestSingleObjectReader {
186            if let Value::Record(fields) = obj {
187                let mut a = None;
188                let mut b = None;
189                let mut c = vec![];
190                for (field_name, v) in fields {
191                    match (field_name.as_str(), v) {
192                        ("a", Value::Long(i)) => a = Some(i),
193                        ("b", Value::Double(d)) => b = Some(d),
194                        ("c", Value::Array(v)) => {
195                            for inner_val in v {
196                                if let Value::String(s) = inner_val {
197                                    c.push(s);
198                                }
199                            }
200                        }
201                        (key, value) => panic!("Unexpected pair: {key:?} -> {value:?}"),
202                    }
203                }
204                TestSingleObjectReader {
205                    a: a.unwrap(),
206                    b: b.unwrap(),
207                    c,
208                }
209            } else {
210                panic!("Expected a Value::Record but was {obj:?}")
211            }
212        }
213    }
214
215    impl From<TestSingleObjectReader> for Value {
216        fn from(obj: TestSingleObjectReader) -> Value {
217            Value::Record(vec![
218                ("a".into(), obj.a.into()),
219                ("b".into(), obj.b.into()),
220                (
221                    "c".into(),
222                    Value::Array(obj.c.into_iter().map(|s| s.into()).collect()),
223                ),
224            ])
225        }
226    }
227
228    #[test]
229    fn test_avro_3507_single_object_reader() -> TestResult {
230        let obj = TestSingleObjectReader {
231            a: 42,
232            b: 3.33,
233            c: vec!["cat".into(), "dog".into()],
234        };
235        let mut to_read = Vec::<u8>::new();
236        to_read.extend_from_slice(&[0xC3, 0x01]);
237        to_read.extend_from_slice(
238            &TestSingleObjectReader::get_schema()
239                .fingerprint::<Rabin>()
240                .bytes[..],
241        );
242        encode(
243            &obj.clone().into(),
244            &TestSingleObjectReader::get_schema(),
245            &mut to_read,
246        )
247        .expect("Encode should succeed");
248        let mut to_read = &to_read[..];
249        let generic_reader = GenericSingleObjectReader::builder()
250            .schema(TestSingleObjectReader::get_schema())
251            .build()
252            .expect("Schema should resolve");
253        let val = generic_reader
254            .read_value(&mut to_read)
255            .expect("Should read");
256        let expected_value: Value = obj.into();
257        pretty_assertions::assert_eq!(expected_value, val);
258
259        Ok(())
260    }
261
262    #[test]
263    fn avro_3642_test_single_object_reader_incomplete_reads() -> TestResult {
264        let obj = TestSingleObjectReader {
265            a: 42,
266            b: 3.33,
267            c: vec!["cat".into(), "dog".into()],
268        };
269        // The two-byte marker, to show that the message uses this single-record format
270        let to_read_1 = [0xC3, 0x01];
271        let mut to_read_2 = Vec::<u8>::new();
272        to_read_2.extend_from_slice(
273            &TestSingleObjectReader::get_schema()
274                .fingerprint::<Rabin>()
275                .bytes[..],
276        );
277        let mut to_read_3 = Vec::<u8>::new();
278        encode(
279            &obj.clone().into(),
280            &TestSingleObjectReader::get_schema(),
281            &mut to_read_3,
282        )
283        .expect("Encode should succeed");
284        let mut to_read = (&to_read_1[..]).chain(&to_read_2[..]).chain(&to_read_3[..]);
285        let generic_reader = GenericSingleObjectReader::builder()
286            .schema(TestSingleObjectReader::get_schema())
287            .build()
288            .expect("Schema should resolve");
289        let val = generic_reader
290            .read_value(&mut to_read)
291            .expect("Should read");
292        let expected_value: Value = obj.into();
293        pretty_assertions::assert_eq!(expected_value, val);
294
295        Ok(())
296    }
297
298    #[test]
299    fn test_avro_3507_reader_parity() -> TestResult {
300        let obj = TestSingleObjectReader {
301            a: 42,
302            b: 3.33,
303            c: vec!["cat".into(), "dog".into()],
304        };
305
306        let mut to_read = Vec::<u8>::new();
307        to_read.extend_from_slice(&[0xC3, 0x01]);
308        to_read.extend_from_slice(
309            &TestSingleObjectReader::get_schema()
310                .fingerprint::<Rabin>()
311                .bytes[..],
312        );
313        encode(
314            &obj.clone().into(),
315            &TestSingleObjectReader::get_schema(),
316            &mut to_read,
317        )
318        .expect("Encode should succeed");
319        let generic_reader = GenericSingleObjectReader::builder()
320            .schema(TestSingleObjectReader::get_schema())
321            .build()
322            .expect("Schema should resolve");
323        let specific_reader = SpecificSingleObjectReader::<TestSingleObjectReader>::new()
324            .expect("schema should resolve");
325        let mut to_read1 = &to_read[..];
326        let mut to_read2 = &to_read[..];
327        let mut to_read3 = &to_read[..];
328
329        let val = generic_reader
330            .read_value(&mut to_read1)
331            .expect("Should read");
332        let read_obj1 = specific_reader
333            .read_from_value(&mut to_read2)
334            .expect("Should read from value");
335        let read_obj2 = specific_reader
336            .read(&mut to_read3)
337            .expect("Should read from deserilize");
338        let expected_value: Value = obj.clone().into();
339        pretty_assertions::assert_eq!(obj, read_obj1);
340        pretty_assertions::assert_eq!(obj, read_obj2);
341        pretty_assertions::assert_eq!(val, expected_value);
342
343        Ok(())
344    }
345
346    #[test]
347    fn avro_rs_164_generic_reader_alternate_header() -> TestResult {
348        let schema_uuid = Uuid::parse_str("b2f1cf00-0434-013e-439a-125eb8485a5f")?;
349        let header_builder = GlueSchemaUuidHeader::from_uuid(schema_uuid);
350        let generic_reader = GenericSingleObjectReader::builder()
351            .schema(TestSingleObjectReader::get_schema())
352            .header(header_builder.build_header())
353            .build()
354            .expect("failed to build reader");
355        // First 18 bytes are the header, then it's a varint 0, double 0 and an empty list (0)
356        let data_to_read: Vec<u8> = vec![
357            3, 0, 178, 241, 207, 0, 4, 52, 1, 62, 67, 154, 18, 94, 184, 72, 90, 95, 0, 0, 0, 0, 0,
358            0, 0, 0, 0, 0,
359        ];
360        let mut to_read = &data_to_read[..];
361        let read_result = generic_reader.read_value(&mut to_read)?;
362        assert_eq!(
363            read_result,
364            Value::Record(vec![
365                ("a".to_string(), Value::Long(0)),
366                ("b".to_string(), Value::Double(0.0)),
367                ("c".to_string(), Value::Array(vec![]))
368            ])
369        );
370        Ok(())
371    }
372}