flipp-oss · dorner · Feb 27, 2024 · Feb 27, 2024 · Feb 27, 2024
diff --git a/CHANGELOG b/CHANGELOG
@@ -2,6 +2,8 @@
 
 # UNRELEASED
 
+- Feature / breaking change: add `preserve_non_string_maps` and change default behavior to turn all maps into string-keyed ones.
+
 # 0.6.1 - 2024-02-27
 
 - Fix: `collapse_fields` was only working on the first instance of the field, not when it was found inside a map etc. later on.

diff --git a/README.md b/README.md
@@ -94,6 +94,58 @@ enum Category {
 }
 ```
 
+* `preserve_non_string_maps` - if set to true, will replace maps with non-string keys with records. E.g. if you have a map like:
+```protobuf
+message MyRecord {
+  map<int32, string> my_field = 1;
+}
+```
+
+...with this option off, it will be translated into:
+
+```json
+{
+  "type": "record",
+  "name": "MyRecord",
+  "fields": [
+    {
+      "name": "my_field",
+      "type": {
+        "type": "map",
+        "values": "string"
+      }
+    }
+  ]
+}
+```
+
+...but with it on, it will be translated into:
+
+```json
+{
+  "type": "record",
+  "name": "MyRecord",
+  "fields": [
+    {
+      "name": "my_field",
+      "type": {
+        "type": "record",
+        "fields": [
+          {
+            "name": "key",
+            "type": "int"
+          },
+          {
+            "name": "value",
+            "type": "string"
+          }
+        ]
+      }
+    }
+  ]
+}
+```
+
 ---
 
 To Do List:

diff --git a/avro/record.go b/avro/record.go
@@ -42,17 +42,19 @@ func (t Record) ToJSON(types *TypeRepo) (any, error) {
 	return jsonMap, nil
 }
 
-func RecordFromProto(proto *descriptorpb.DescriptorProto, namespace string) []NamedType {
+func RecordFromProto(proto *descriptorpb.DescriptorProto, namespace string, typeRepo *TypeRepo) []NamedType {
 	// Protobuf `map` is actually a record with a `map_entry` boolean set, and two fields ("key" and "value").
 	// Avro only supports maps with string keys, but Protobuf supports maps with any key type.
 	// If we detect a string key, we can turn this into an Avro map, but otherwise we have to leave it as a record.
-	if proto.Options.GetMapEntry() && proto.Field[0].GetType() == descriptorpb.FieldDescriptorProto_TYPE_STRING {
-		return []NamedType{
-			Map{
-				Namespace: namespace,
-			  Name: proto.GetName(),
-			  Values: FieldTypeFromProto(proto.Field[1]),
-  		},
+	if proto.Options.GetMapEntry() {
+		if !typeRepo.PreserveNonStringMaps || proto.Field[0].GetType() == descriptorpb.FieldDescriptorProto_TYPE_STRING {
+			return []NamedType{
+				Map{
+					Namespace: namespace,
+					Name:      proto.GetName(),
+					Values:    FieldTypeFromProto(proto.Field[1]),
+				},
+			}
 		}
 	}
 
@@ -61,7 +63,7 @@ func RecordFromProto(proto *descriptorpb.DescriptorProto, namespace string) []Na
 	nested := make([]NamedType, len(proto.NestedType))
 //	enums := make([]Enum, len(proto.EnumType))
 	for i, field := range proto.NestedType {
-		nested[i] = RecordFromProto(field, fmt.Sprintf("%s.%s", namespace, proto.GetName()))[0]
+		nested[i] = RecordFromProto(field, fmt.Sprintf("%s.%s", namespace, proto.GetName()), typeRepo)[0]
 	}
 	//for i, field := range proto.EnumType {
 	//	enums[i] = EnumFromProto(field)

diff --git a/avro/typeRepo.go b/avro/typeRepo.go
@@ -13,6 +13,7 @@ type TypeRepo struct {
   NamespaceMap map[string]string
   CollapseFields []string
   RemoveEnumPrefixes bool
+  PreserveNonStringMaps bool
 }
 
 func NewTypeRepo(params input.Params) *TypeRepo {
@@ -21,6 +22,7 @@ func NewTypeRepo(params input.Params) *TypeRepo {
     NamespaceMap: params.NamespaceMap,
     CollapseFields: params.CollapseFields,
     RemoveEnumPrefixes: params.RemoveEnumPrefixes,
+    PreserveNonStringMaps: params.PreserveNonStringMaps,
   }
 }
 

diff --git a/input/params.go b/input/params.go
@@ -13,6 +13,7 @@ type Params struct {
 	NamespaceMap map[string]string
 	CollapseFields []string
 	RemoveEnumPrefixes bool
+	PreserveNonStringMaps bool
 }
 
 func ReadRequest() (*pluginpb.CodeGeneratorRequest, error) {
@@ -60,6 +61,8 @@ func ParseParams(req *pluginpb.CodeGeneratorRequest) Params {
 			params.CollapseFields = strings.Split(v, ";")
 		} else if k == "remove_enum_prefixes" {
 			params.RemoveEnumPrefixes = v == "true"
+		} else if k == "preserve_non_string_maps" {
+			params.PreserveNonStringMaps = true
 		}
 	}
 	return params

diff --git a/main.go b/main.go
@@ -17,7 +17,7 @@ var params input.Params
 var typeRepo *avro.TypeRepo
 
 func processMessage(proto *descriptorpb.DescriptorProto, protoPackage string) {
-	records := avro.RecordFromProto(proto, protoPackage)
+	records := avro.RecordFromProto(proto, protoPackage, typeRepo)
 	for _, record := range records {
 		typeRepo.AddType(record)
 	}

diff --git a/main_test.go b/main_test.go
@@ -65,8 +65,13 @@ func runTest(t *testing.T, directory string, options map[string]string) {
     }
     protoc(t, args)
 
-    assertEqualFiles(t, workdir + "/testdata/" + directory, tmpdir)
-
+    testDir := workdir + "/testdata/" + directory
+    if os.Getenv("UPDATE_SNAPSHOTS") != "" {
+        cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cp %v/* %v", tmpdir, testDir))
+        cmd.Run()
+    } else {
+        assertEqualFiles(t, testDir, tmpdir)
+    }
 }
 
 func Test_Base(t *testing.T) {
@@ -85,6 +90,10 @@ func Test_NamespaceMap(t *testing.T) {
     runTest(t, "namespace_map", map[string]string{"namespace_map": "testdata:mynamespace"})
 }
 
+func Test_PreserveNonStringMaps(t *testing.T) {
+    runTest(t, "preserve_non_string_maps", map[string]string{"preserve_non_string_maps": "true"})
+}
+
 func assertEqualFiles(t *testing.T, original, generated string) {
     names, err := fileNames(original, false)
     if err != nil {

diff --git a/testdata/base/AOneOf.avsc b/testdata/base/AOneOf.avsc
@@ -4,7 +4,7 @@
   "namespace": "testdata",
   "fields": [
     {
-      "name": "a_oneof",
+      "name": "oneof_types",
       "type": [
         {
           "type": "record",

diff --git a/testdata/base/Foobar.avsc b/testdata/base/Foobar.avsc
@@ -110,26 +110,10 @@
     {
       "name": "a_yowza_map",
       "type": {
-        "type": "array",
-        "items": {
-          "type": "record",
-          "name": "AYowzaMapEntry",
-          "namespace": "testdata.Foobar",
-          "fields": [
-            {
-              "name": "key",
-              "type": "boolean",
-              "default": false
-            },
-            {
-              "name": "value",
-              "type": "testdata.Yowza",
-              "default": ""
-            }
-          ]
-        }
+        "type": "map",
+        "values": "testdata.Yowza"
       },
-      "default": []
+      "default": {}
     }
   ]
 }
diff --git a/testdata/base/Widget.avsc b/testdata/base/Widget.avsc
@@ -130,26 +130,10 @@
           {
             "name": "a_yowza_map",
             "type": {
-              "type": "array",
-              "items": {
-                "type": "record",
-                "name": "AYowzaMapEntry",
-                "namespace": "testdata.Foobar",
-                "fields": [
-                  {
-                    "name": "key",
-                    "type": "boolean",
-                    "default": false
-                  },
-                  {
-                    "name": "value",
-                    "type": "testdata.Yowza",
-                    "default": ""
-                  }
-                ]
-              }
+              "type": "map",
+              "values": "testdata.Yowza"
             },
-            "default": []
+            "default": {}
           }
         ]
       },
@@ -192,7 +176,7 @@
           "namespace": "testdata",
           "fields": [
             {
-              "name": "a_oneof",
+              "name": "oneof_types",
               "type": [
                 {
                   "type": "record",

diff --git a/testdata/collapse_fields/AOneOf.avsc b/testdata/collapse_fields/AOneOf.avsc
@@ -4,7 +4,7 @@
   "namespace": "testdata",
   "fields": [
     {
-      "name": "a_oneof",
+      "name": "oneof_types",
       "type": [
         {
           "type": "record",

diff --git a/testdata/collapse_fields/AYowzaMapEntry.avsc b/testdata/collapse_fields/AYowzaMapEntry.avsc
diff --git a/testdata/collapse_fields/Foobar.avsc b/testdata/collapse_fields/Foobar.avsc
@@ -110,26 +110,10 @@
     {
       "name": "a_yowza_map",
       "type": {
-        "type": "array",
-        "items": {
-          "type": "record",
-          "name": "AYowzaMapEntry",
-          "namespace": "testdata.Foobar",
-          "fields": [
-            {
-              "name": "key",
-              "type": "boolean",
-              "default": false
-            },
-            {
-              "name": "value",
-              "type": "testdata.Yowza",
-              "default": ""
-            }
-          ]
-        }
+        "type": "map",
+        "values": "testdata.Yowza"
       },
-      "default": []
+      "default": {}
     }
   ]
 }
diff --git a/testdata/collapse_fields/Widget.avsc b/testdata/collapse_fields/Widget.avsc
@@ -130,26 +130,10 @@
           {
             "name": "a_yowza_map",
             "type": {
-              "type": "array",
-              "items": {
-                "type": "record",
-                "name": "AYowzaMapEntry",
-                "namespace": "testdata.Foobar",
-                "fields": [
-                  {
-                    "name": "key",
-                    "type": "boolean",
-                    "default": false
-                  },
-                  {
-                    "name": "value",
-                    "type": "testdata.Yowza",
-                    "default": ""
-                  }
-                ]
-              }
+              "type": "map",
+              "values": "testdata.Yowza"
             },
-            "default": []
+            "default": {}
           }
         ]
       },
@@ -188,7 +172,7 @@
           "namespace": "testdata",
           "fields": [
             {
-              "name": "a_oneof",
+              "name": "oneof_types",
               "type": [
                 {
                   "type": "record",