fix(codecs): support Avro Date and TimeMillis logical types

changelog.d/avro_date_time_millis_encoding.enhancement.md

@@ -0,0 +1,3 @@
+The Avro codec now supports encoding and decoding Date and TimeMillis values.
+
+authors: omwbennett

lib/codecs/src/decoding/format/avro.rs

@@ -91,15 +91,13 @@ impl From<&AvroDeserializerOptions> for AvroSerializerOptions {
 pub struct AvroDeserializerOptions {
     /// The Avro schema definition.
     /// **Note**: The following [`apache_avro::types::Value`] variants are *not* supported:
-    /// * `Date`
     /// * `Decimal`
     /// * `Duration`
     /// * `Fixed`
-    /// * `TimeMillis`
     #[configurable(metadata(
         docs::examples = r#"{ "type": "record", "name": "log", "fields": [{ "name": "message", "type": "string" }] }"#,
         docs::additional_props_description = r#"Supports most avro data types, unsupported data types includes
-        ["decimal", "duration", "local-timestamp-millis", "local-timestamp-micros"]"#,
+        ["decimal", "duration", "fixed"]"#,
     ))]
     pub schema: String,
 
@@ -190,9 +188,7 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
         }
         AvroValue::Boolean(boolean) => Ok(VrlValue::from(boolean)),
         AvroValue::Bytes(bytes) => Ok(VrlValue::from(bytes)),
-        AvroValue::Date(_) => Err(vector_common::Error::from(
-            "AvroValue::Date is not supported",
-        )),
+        AvroValue::Date(days) => Ok(VrlValue::from(days)),
         AvroValue::Decimal(_) => Err(vector_common::Error::from(
             "AvroValue::Decimal is not supported",
         )),
@@ -220,9 +216,7 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
             .map(|v| VrlValue::Object(v.into_iter().collect())),
         AvroValue::String(string) => Ok(VrlValue::from(string)),
         AvroValue::TimeMicros(time_micros) => Ok(VrlValue::from(time_micros)),
-        AvroValue::TimeMillis(_) => Err(vector_common::Error::from(
-            "AvroValue::TimeMillis is not supported",
-        )),
+        AvroValue::TimeMillis(millis) => Ok(VrlValue::from(millis)),
         AvroValue::TimestampMicros(ts_micros) => Ok(VrlValue::from(ts_micros)),
         AvroValue::TimestampMillis(ts_millis) => Ok(VrlValue::from(ts_millis)),
         AvroValue::Union(_, v) => try_from(*v),

lib/codecs/src/encoding/format/avro.rs

@@ -1,3 +1,6 @@
+use std::collections::HashMap;
+
+use apache_avro::Schema;
 use bytes::{BufMut, BytesMut};
 use serde::{Deserialize, Serialize};
 use tokio_util::codec::Encoder;
@@ -6,6 +9,125 @@ use vector_core::{config::DataType, event::Event, schema};
 
 use crate::encoding::BuildError;
 
+type AvroValue = apache_avro::types::Value;
+type NamedSchemas = HashMap<apache_avro::schema::Name, apache_avro::Schema>;
+
+fn resolve_named_schemas(schema: &apache_avro::Schema) -> Result<NamedSchemas, apache_avro::Error> {
+    let resolved = apache_avro::schema::ResolvedSchema::try_from(schema)?;
+    Ok(resolved
+        .get_names()
+        .iter()
+        .map(|(name, schema)| (name.clone(), (*schema).clone()))
+        .collect())
+}
+
+/// `apache_avro::to_value` serializes VRL values into Avro types which may not
+/// resolve against certain logical type schemas
+/// (e.g. VRL integer (i64) -> Avro `Long` which cannot resolve to Avro `Date`).
+/// This does a recursive pre-pass to coerce such values before resolution.
+fn coerce_logical_types(
+    value: AvroValue,
+    schema: &apache_avro::Schema,
+    names: &NamedSchemas,
+) -> vector_common::Result<AvroValue> {
+    match (value, schema) {
+        (AvroValue::Long(days), Schema::Date) => {
+            i32::try_from(days).map(AvroValue::Date).map_err(|_| {
+                vector_common::Error::from(format!(
+                    "Avro date value {days} is out of range for i32"
+                ))
+            })
+        }
+        (AvroValue::Long(millis), Schema::TimeMillis) => i32::try_from(millis)
+            .map(AvroValue::TimeMillis)
+            .map_err(|_| {
+                vector_common::Error::from(format!(
+                    "Avro time-millis value {millis} is out of range for i32"
+                ))
+            }),
+        (AvroValue::Record(fields), Schema::Record(record_schema)) => {
+            let fields = fields
+                .into_iter()
+                .map(|(name, value)| {
+                    let value = match record_schema.lookup.get(&name) {
+                        Some(index) => {
+                            let field_schema = &record_schema.fields[*index].schema;
+                            coerce_logical_types(value, field_schema, names)?
+                        }
+                        None => value,
+                    };
+                    Ok((name, value))
+                })
+                .collect::<vector_common::Result<Vec<_>>>()?;
+            Ok(AvroValue::Record(fields))
+        }
+        (AvroValue::Map(entries), Schema::Record(record_schema)) => {
+            let entries = entries
+                .into_iter()
+                .map(|(name, value)| {
+                    let value = match record_schema.lookup.get(&name) {
+                        Some(index) => {
+                            let field_schema = &record_schema.fields[*index].schema;
+                            coerce_logical_types(value, field_schema, names)?
+                        }
+                        None => value,
+                    };
+                    Ok((name, value))
+                })
+                .collect::<vector_common::Result<_>>()?;
+            Ok(AvroValue::Map(entries))
+        }
+        (AvroValue::Array(items), Schema::Array(array_schema)) => items
+            .into_iter()
+            .map(|item| coerce_logical_types(item, &array_schema.items, names))
+            .collect::<Result<Vec<_>, _>>()
+            .map(AvroValue::Array),
+        (AvroValue::Map(entries), Schema::Map(map_schema)) => entries
+            .into_iter()
+            .map(|(key, value)| {
+                coerce_logical_types(value, &map_schema.types, names).map(|value| (key, value))
+            })
+            .collect::<vector_common::Result<_>>()
+            .map(AvroValue::Map),
+        (AvroValue::Union(index, value), Schema::Union(union_schema)) => {
+            let schema = union_schema
+                .variants()
+                .get(index as usize)
+                .unwrap_or(schema);
+            coerce_logical_types(*value, schema, names)
+                .map(|value| AvroValue::Union(index, Box::new(value)))
+        }
+        (value, Schema::Union(union_schema)) => {
+            if let Ok(resolved) = value.clone().resolve(schema) {
+                return Ok(resolved);
+            }
+
+            let mut last_err = None;
+            for (index, variant) in union_schema.variants().iter().enumerate() {
+                match coerce_logical_types(value.clone(), variant, names) {
+                    Ok(coerced) if coerced.clone().resolve(variant).is_ok() => {
+                        return Ok(AvroValue::Union(index as u32, Box::new(coerced)));
+                    }
+                    Ok(_) => {}
+                    Err(err) => last_err = Some(err),
+                }
+            }
+
+            match last_err {
+                Some(err) => Err(err),
+                None => Ok(value),
+            }
+        }
+        (value, Schema::Ref { name }) => {
+            let schema = names.get(name).ok_or_else(|| {
+                vector_common::Error::from(format!("Unknown schema ref: {}", name.fullname(None)))
+            })?;
+            coerce_logical_types(value, schema, names)
+        }
+        (value, _) => Ok(value),
+    }
+}
+
 /// Config used to build a `AvroSerializer`.
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct AvroSerializerConfig {
@@ -25,7 +147,12 @@ impl AvroSerializerConfig {
     pub fn build(&self) -> Result<AvroSerializer, BuildError> {
         let schema = apache_avro::Schema::parse_str(&self.avro.schema)
             .map_err(|error| format!("Failed building Avro serializer: {error}"))?;
-        Ok(AvroSerializer { schema })
+        let named_schemas = resolve_named_schemas(&schema)
+            .map_err(|error| format!("Failed resolving Avro schema: {error}"))?;
+        Ok(AvroSerializer {
+            schema,
+            named_schemas: Some(named_schemas),
+        })
     }
 
     /// The data type of events that are accepted by `AvroSerializer`.
@@ -56,12 +183,16 @@ pub struct AvroSerializerOptions {
 #[derive(Debug, Clone)]
 pub struct AvroSerializer {
     schema: apache_avro::Schema,
+    named_schemas: Option<NamedSchemas>,
 }
 
 impl AvroSerializer {
     /// Creates a new `AvroSerializer`.
     pub const fn new(schema: apache_avro::Schema) -> Self {
-        Self { schema }
+        Self {
+            schema,
+            named_schemas: None,
+        }
     }
 }
 
@@ -71,6 +202,18 @@ impl Encoder<Event> for AvroSerializer {
     fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
         let log = event.into_log();
         let value = apache_avro::to_value(log)?;
+        if self.named_schemas.is_none() {
+            self.named_schemas = Some(resolve_named_schemas(&self.schema).map_err(|error| {
+                vector_common::Error::from(format!("Failed resolving Avro schema: {error}"))
+            })?);
+        }
+        let value = coerce_logical_types(
+            value,
+            &self.schema,
+            self.named_schemas
+                .as_ref()
+                .expect("named schemas are initialized above"),
+        )?;
         let value = value.resolve(&self.schema)?;
         let bytes = apache_avro::to_avro_datum(&self.schema, value)?;
         buffer.put_slice(&bytes);
@@ -113,4 +256,171 @@ mod tests {
 
         assert_eq!(bytes.freeze(), b"\0\x06bar".as_slice());
     }
+
+    #[test]
+    fn coerce_date_fields_recursively() {
+        let schema = apache_avro::Schema::parse_str(indoc! {r#"
+            {
+                "type": "record",
+                "name": "Outer",
+                "fields": [
+                    {
+                        "name": "direct_date",
+                        "type": {"type": "int", "logicalType": "date"}
+                    },
+                    {
+                        "name": "inner",
+                        "type": {
+                            "type": "record",
+                            "name": "Inner",
+                            "fields": [
+                                {
+                                    "name": "date",
+                                    "type": {"type": "int", "logicalType": "date"}
+                                }
+                            ]
+                        }
+                    },
+                    {
+                        "name": "record_as_map",
+                        "type": {
+                            "type": "record",
+                            "name": "MapBackedInner",
+                            "fields": [
+                                {
+                                    "name": "date",
+                                    "type": {"type": "int", "logicalType": "date"}
+                                }
+                            ]
+                        }
+                    },
+                    {
+                        "name": "inner_reference",
+                        "type": "Inner"
+                    },
+                    {
+                        "name": "date_array",
+                        "type": {
+                            "type": "array",
+                            "items": {"type": "int", "logicalType": "date"}
+                        }
+                    },
+                    {
+                        "name": "date_map",
+                        "type": {
+                            "type": "map",
+                            "values": {"type": "int", "logicalType": "date"}
+                        }
+                    },
+                    {
+                        "name": "union_date",
+                        "type": ["null", {"type": "int", "logicalType": "date"}]
+                    },
+                    {
+                        "name": "fallback_union_date",
+                        "type": [
+                            "null",
+                            {"type": "int", "logicalType": "date"},
+                            "long"
+                        ]
+                    },
+                    {
+                        "name": "logical_only_union_date",
+                        "type": [
+                            "null",
+                            {"type": "int", "logicalType": "date"}
+                        ]
+                    }
+                ]
+            }
+        "#})
+        .unwrap();
+        let value = AvroValue::Record(vec![
+            ("direct_date".to_owned(), AvroValue::Long(20_000)),
+            (
+                "inner".to_owned(),
+                AvroValue::Record(vec![("date".to_owned(), AvroValue::Long(20_001))]),
+            ),
+            (
+                "record_as_map".to_owned(),
+                AvroValue::Map(
+                    [("date".to_owned(), AvroValue::Long(20_002))]
+                        .into_iter()
+                        .collect(),
+                ),
+            ),
+            (
+                "inner_reference".to_owned(),
+                AvroValue::Record(vec![("date".to_owned(), AvroValue::Long(20_003))]),
+            ),
+            (
+                "date_array".to_owned(),
+                AvroValue::Array(vec![AvroValue::Long(20_004), AvroValue::Long(20_005)]),
+            ),
+            (
+                "date_map".to_owned(),
+                AvroValue::Map(
+                    [
+                        ("first".to_owned(), AvroValue::Long(20_006)),
+                        ("second".to_owned(), AvroValue::Long(20_007)),
+                    ]
+                    .into_iter()
+                    .collect(),
+                ),
+            ),
+            ("union_date".to_owned(), AvroValue::Long(20_008)),
+            ("fallback_union_date".to_owned(), AvroValue::Long(20_010)),
+            (
+                "logical_only_union_date".to_owned(),
+                AvroValue::Long(20_009),
+            ),
+        ]);
+
+        let named_schemas = resolve_named_schemas(&schema).unwrap();
+        let value = coerce_logical_types(value, &schema, &named_schemas).unwrap();
+        let value = value.resolve(&schema).unwrap();
+
+        assert!(matches!(
+            value,
+            AvroValue::Record(fields) if {
+                matches!(fields[0].1, AvroValue::Date(20_000))
+                    && matches!(
+                        &fields[1].1,
+                        AvroValue::Record(inner) if matches!(inner[0].1, AvroValue::Date(20_001))
+                    )
+                    && matches!(
+                        &fields[2].1,
+                        AvroValue::Record(inner) if matches!(inner[0].1, AvroValue::Date(20_002))
+                    )
+                    && matches!(
+                        &fields[3].1,
+                        AvroValue::Record(inner) if matches!(inner[0].1, AvroValue::Date(20_003))
+                    )
+                    && matches!(
+                        &fields[4].1,
+                        AvroValue::Array(items)
+                            if matches!(items.as_slice(), [AvroValue::Date(20_004), AvroValue::Date(20_005)])
+                    )
+                    && matches!(
+                        &fields[5].1,
+                        AvroValue::Map(entries)
+                            if matches!(entries.get("first"), Some(AvroValue::Date(20_006)))
+                                && matches!(entries.get("second"), Some(AvroValue::Date(20_007)))
+                    )
+                    && matches!(
+                        &fields[6].1,
+                        AvroValue::Union(1, value) if matches!(value.as_ref(), AvroValue::Date(20_008))
+                    )
+                    && matches!(
+                        &fields[7].1,
+                        AvroValue::Union(2, value)
+                            if matches!(value.as_ref(), AvroValue::Long(20_010))
+                    )
+                    && matches!(
+                        &fields[8].1,
+                        AvroValue::Union(1, value) if matches!(value.as_ref(), AvroValue::Date(20_009))
+                    )
+            }
+        ));
+    }
 }