AVRO-3687 [Rust - avro_derive]: Add support for default enum values f…

…or rust derive macros (#2954) * Add support for default enum values for rust derive macros * AVRO-3687: Better error messages when multiple enum variants are marked as default Signed-off-by: Martin Tzvetanov Grigorov <[email protected]> --------- Signed-off-by: Martin Tzvetanov Grigorov <[email protected]> Co-authored-by: Martin Tzvetanov Grigorov <[email protected]>
apache · Jul 4, 2024 · 3413ac5 · 3413ac5
1 parent f3b6ee2
commit 3413ac5
Showing 1 changed file with 128 additions and 2 deletions.
diff --git a/lang/rust/avro_derive/src/lib.rs b/lang/rust/avro_derive/src/lib.rs
@@ -23,7 +23,10 @@ use darling::FromAttributes;
 use proc_macro2::{Span, TokenStream};
 use quote::quote;
 
-use syn::{parse_macro_input, spanned::Spanned, AttrStyle, Attribute, DeriveInput, Type, TypePath};
+use syn::{
+    parse_macro_input, spanned::Spanned, AttrStyle, Attribute, DeriveInput, Ident, Meta, Type,
+    TypePath,
+};
 
 #[derive(darling::FromAttributes)]
 #[darling(attributes(avro))]
@@ -214,6 +217,8 @@ fn get_data_enum_schema_def(
     let doc = preserve_optional(doc);
     let enum_aliases = preserve_vec(aliases);
     if e.variants.iter().all(|v| syn::Fields::Unit == v.fields) {
+        let default_value = default_enum_variant(e, error_span)?;
+        let default = preserve_optional(default_value);
         let symbols: Vec<String> = e
             .variants
             .iter()
@@ -225,7 +230,7 @@ fn get_data_enum_schema_def(
                 aliases: #enum_aliases,
                 doc: #doc,
                 symbols: vec![#(#symbols.to_owned()),*],
-                default: None,
+                default: #default,
                 attributes: Default::default(),
             })
         })
@@ -281,6 +286,35 @@ fn type_to_schema_expr(ty: &Type) -> Result<TokenStream, Vec<syn::Error>> {
     }
 }
 
+fn default_enum_variant(
+    data_enum: &syn::DataEnum,
+    error_span: Span,
+) -> Result<Option<String>, Vec<syn::Error>> {
+    match data_enum
+        .variants
+        .iter()
+        .filter(|v| v.attrs.iter().any(is_default_attr))
+        .collect::<Vec<_>>()
+    {
+        variants if variants.is_empty() => Ok(None),
+        single if single.len() == 1 => Ok(Some(single[0].ident.to_string())),
+        multiple => Err(vec![syn::Error::new(
+            error_span,
+            format!(
+                "Multiple defaults defined: {:?}",
+                multiple
+                    .iter()
+                    .map(|v| v.ident.to_string())
+                    .collect::<Vec<String>>()
+            ),
+        )]),
+    }
+}
+
+fn is_default_attr(attr: &Attribute) -> bool {
+    matches!(attr, Attribute { meta: Meta::Path(path), .. } if path.get_ident().map(Ident::to_string).as_deref() == Some("default"))
+}
+
 /// Generates the schema def expression for fully qualified type paths using the associated function
 /// - `A -> <A as apache_avro::schema::derive::AvroSchemaComponent>::get_schema_in_ctxt()`
 /// - `A<T> -> <A<T> as apache_avro::schema::derive::AvroSchemaComponent>::get_schema_in_ctxt()`
@@ -433,6 +467,98 @@ mod tests {
         };
     }
 
+    #[test]
+    fn avro_3687_basic_enum_with_default() {
+        let basic_enum = quote! {
+            enum Basic {
+                #[default]
+                A,
+                B,
+                C,
+                D
+            }
+        };
+        match syn::parse2::<DeriveInput>(basic_enum) {
+            Ok(mut input) => {
+                let derived = derive_avro_schema(&mut input);
+                assert!(derived.is_ok());
+                assert_eq!(derived.unwrap().to_string(), quote! {
+                    impl apache_avro::schema::derive::AvroSchemaComponent for Basic {
+                        fn get_schema_in_ctxt(
+                            named_schemas: &mut std::collections::HashMap<
+                                apache_avro::schema::Name,
+                                apache_avro::schema::Schema
+                            >,
+                            enclosing_namespace: &Option<String>
+                        ) -> apache_avro::schema::Schema {
+                            let name = apache_avro::schema::Name::new("Basic")
+                                .expect(&format!("Unable to parse schema name {}", "Basic")[..])
+                                .fully_qualified_name(enclosing_namespace);
+                            let enclosing_namespace = &name.namespace;
+                            if named_schemas.contains_key(&name) {
+                                apache_avro::schema::Schema::Ref { name: name.clone() }
+                            } else {
+                                named_schemas.insert(
+                                    name.clone(),
+                                    apache_avro::schema::Schema::Ref { name: name.clone() }
+                                );
+                                apache_avro::schema::Schema::Enum(apache_avro::schema::EnumSchema {
+                                    name: apache_avro::schema::Name::new("Basic").expect(
+                                        &format!("Unable to parse enum name for schema {}", "Basic")[..]
+                                    ),
+                                    aliases: None,
+                                    doc: None,
+                                    symbols: vec![
+                                        "A".to_owned(),
+                                        "B".to_owned(),
+                                        "C".to_owned(),
+                                        "D".to_owned()
+                                    ],
+                                    default: Some("A".into()),
+                                    attributes: Default::default(),
+                                })
+                            }
+                        }
+                    }
+                }.to_string());
+            }
+            Err(error) => panic!(
+                "Failed to parse as derive input when it should be able to. Error: {error:?}"
+            ),
+        };
+    }
+
+    #[test]
+    fn avro_3687_basic_enum_with_default_twice() {
+        let non_basic_enum = quote! {
+            enum Basic {
+                #[default]
+                A,
+                B,
+                #[default]
+                C,
+                D
+            }
+        };
+        match syn::parse2::<DeriveInput>(non_basic_enum) {
+            Ok(mut input) => match derive_avro_schema(&mut input) {
+                Ok(_) => {
+                    panic!("Should not be able to derive schema for enum with multiple defaults")
+                }
+                Err(errors) => {
+                    assert_eq!(errors.len(), 1);
+                    assert_eq!(
+                        errors[0].to_string(),
+                        r#"Multiple defaults defined: ["A", "C"]"#
+                    );
+                }
+            },
+            Err(error) => panic!(
+                "Failed to parse as derive input when it should be able to. Error: {error:?}"
+            ),
+        };
+    }
+
     #[test]
     fn test_non_basic_enum() {
         let non_basic_enum = quote! {