Fix: Add support for enums defined in message

Proto allows for an enum to be defined within a Message. Register nested enums in the type repo so they can be referenced but not created as types.
flipp-oss · May 1, 2024 · 08c1411 · 08c1411
1 parent f9e45cf
commit 08c1411
Show file tree

Hide file tree

Showing 11 changed files with 133 additions and 8 deletions.
diff --git a/avro/record.go b/avro/record.go
@@ -8,9 +8,9 @@ import (
 )
 
 type Record struct {
-	Name string
+	Name      string
 	Namespace string
-	Fields []Field
+	Fields    []Field
 }
 
 func (t Record) GetName() string {
@@ -61,13 +61,13 @@ func RecordFromProto(proto *descriptorpb.DescriptorProto, namespace string, type
 	var fields []Field
 	oneofs := make([]Field, len(proto.OneofDecl))
 	nested := make([]NamedType, len(proto.NestedType))
-//	enums := make([]Enum, len(proto.EnumType))
+	enums := make([]Enum, len(proto.EnumType))
 	for i, field := range proto.NestedType {
 		nested[i] = RecordFromProto(field, fmt.Sprintf("%s.%s", namespace, proto.GetName()), typeRepo)[0]
 	}
-	//for i, field := range proto.EnumType {
-	//	enums[i] = EnumFromProto(field)
-	//}
+	for i, field := range proto.EnumType {
+		enums[i] = EnumFromProto(field, fmt.Sprintf("%s.%s", namespace, proto.GetName()))
+	}
 	for i, oneof := range proto.OneofDecl {
 		oneofs[i] = Field{
 			Name: oneof.GetName(),
@@ -102,10 +102,13 @@ func RecordFromProto(proto *descriptorpb.DescriptorProto, namespace string, type
 	for _, subRecord := range nested {
 		types = append(types, subRecord)
 	}
+	for _, enum := range enums {
+		types = append(types, enum)
+	}
 	types = append(types, Record{
-		Name: proto.GetName(),
+		Name:      proto.GetName(),
 		Namespace: namespace,
-		Fields: fields,
+		Fields:    fields,
 	})
 	return types
 }
diff --git a/testdata/base/Foobar.avsc b/testdata/base/Foobar.avsc
@@ -143,6 +143,19 @@
         "items": "testdata.StringList"
       },
       "default": []
+    },
+    {
+      "name": "nested_enum",
+      "type": {
+        "type": "enum",
+        "name": "NestedEnum",
+        "symbols": [
+          "A",
+          "B"
+        ],
+        "default": "A"
+      },
+      "default": "A"
     }
   ]
 }
diff --git a/testdata/base/Widget.avsc b/testdata/base/Widget.avsc
@@ -220,6 +220,19 @@
               "items": "testdata.StringList"
             },
             "default": []
+          },
+          {
+            "name": "nested_enum",
+            "type": {
+              "type": "enum",
+              "name": "NestedEnum",
+              "symbols": [
+                "A",
+                "B"
+              ],
+              "default": "A"
+            },
+            "default": "A"
           }
         ]
       }

diff --git a/testdata/collapse_fields/Foobar.avsc b/testdata/collapse_fields/Foobar.avsc
@@ -137,6 +137,19 @@
         }
       },
       "default": []
+    },
+    {
+      "name": "nested_enum",
+      "type": {
+        "type": "enum",
+        "name": "NestedEnum",
+        "symbols": [
+          "A",
+          "B"
+        ],
+        "default": "A"
+      },
+      "default": "A"
     }
   ]
 }
diff --git a/testdata/collapse_fields/Widget.avsc b/testdata/collapse_fields/Widget.avsc
@@ -220,6 +220,19 @@
               }
             },
             "default": []
+          },
+          {
+            "name": "nested_enum",
+            "type": {
+              "type": "enum",
+              "name": "NestedEnum",
+              "symbols": [
+                "A",
+                "B"
+              ],
+              "default": "A"
+            },
+            "default": "A"
           }
         ]
       }

diff --git a/testdata/emit_only/Widget.avsc b/testdata/emit_only/Widget.avsc
@@ -220,6 +220,19 @@
               "items": "testdata.StringList"
             },
             "default": []
+          },
+          {
+            "name": "nested_enum",
+            "type": {
+              "type": "enum",
+              "name": "NestedEnum",
+              "symbols": [
+                "A",
+                "B"
+              ],
+              "default": "A"
+            },
+            "default": "A"
           }
         ]
       }

diff --git a/testdata/foobar.proto b/testdata/foobar.proto
@@ -17,6 +17,10 @@ message Yowza {
 }
 
 message Foobar {
+  enum NestedEnum {
+   A = 0;
+   B = 1;
+  }
   string name = 1;
   Blarp blarp = 2;
   Yowza yowza = 3;
@@ -32,4 +36,5 @@ message Foobar {
   map<bool, Yowza> a_yowza_map = 13;
   map<string, StringList> string_list_map = 14;
   repeated StringList string_lists = 15;
+  NestedEnum nested_enum = 16;
 }
diff --git a/testdata/namespace_map/Foobar.avsc b/testdata/namespace_map/Foobar.avsc
@@ -143,6 +143,19 @@
         "items": "mynamespace.StringList"
       },
       "default": []
+    },
+    {
+      "name": "nested_enum",
+      "type": {
+        "type": "enum",
+        "name": "NestedEnum",
+        "symbols": [
+          "A",
+          "B"
+        ],
+        "default": "A"
+      },
+      "default": "A"
     }
   ]
 }
diff --git a/testdata/namespace_map/Widget.avsc b/testdata/namespace_map/Widget.avsc
@@ -220,6 +220,19 @@
               "items": "mynamespace.StringList"
             },
             "default": []
+          },
+          {
+            "name": "nested_enum",
+            "type": {
+              "type": "enum",
+              "name": "NestedEnum",
+              "symbols": [
+                "A",
+                "B"
+              ],
+              "default": "A"
+            },
+            "default": "A"
           }
         ]
       }

diff --git a/testdata/preserve_non_string_maps/Foobar.avsc b/testdata/preserve_non_string_maps/Foobar.avsc
@@ -159,6 +159,19 @@
         "items": "testdata.StringList"
       },
       "default": []
+    },
+    {
+      "name": "nested_enum",
+      "type": {
+        "type": "enum",
+        "name": "NestedEnum",
+        "symbols": [
+          "A",
+          "B"
+        ],
+        "default": "A"
+      },
+      "default": "A"
     }
   ]
 }
diff --git a/testdata/preserve_non_string_maps/Widget.avsc b/testdata/preserve_non_string_maps/Widget.avsc
@@ -236,6 +236,19 @@
               "items": "testdata.StringList"
             },
             "default": []
+          },
+          {
+            "name": "nested_enum",
+            "type": {
+              "type": "enum",
+              "name": "NestedEnum",
+              "symbols": [
+                "A",
+                "B"
+              ],
+              "default": "A"
+            },
+            "default": "A"
           }
         ]
       }