Merge pull request GoogleCloudPlatform#2285 from gemmahou/gen-mapper

chore: Handle list of enum during mapper generation
gemmahou · Jul 15, 2024 · e8c8b0d · e8c8b0d
2 parents 9de4abf + 79b1c88
commit e8c8b0d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 7 deletions.
diff --git a/dev/tools/proto-to-mapper/Makefile b/dev/tools/proto-to-mapper/Makefile
@@ -17,7 +17,6 @@ generate-pb: install-protoc-linux
 		./third_party/googleapis/google/monitoring/v3/*.proto \
 		./third_party/googleapis/google/monitoring/dashboard/v1/*.proto \
 		./third_party/googleapis/google/devtools/cloudbuild/*/*.proto \
-		./third_party/googleapis/google/cloud/dataform/*/*.proto \
 		-o build/googleapis.pb
 
 install-protoc-linux:

diff --git a/dev/tools/proto-to-mapper/main.go b/dev/tools/proto-to-mapper/main.go
@@ -380,7 +380,7 @@ func (v *visitor) writeMapFunctionsForPair(out io.Writer, pair *typePair) {
 			}
 
 			if protoField.Cardinality() == protoreflect.Repeated {
-				useSliceFromProto := ""
+				useSliceFromProtoFunction := ""
 				useCustomMethod := false
 
 				switch protoField.Kind() {
@@ -390,12 +390,22 @@ func (v *visitor) writeMapFunctionsForPair(out io.Writer, pair *typePair) {
 					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
 
 					functionName := krmElemTypeName + "_FromProto"
-					useSliceFromProto = functionName
+					useSliceFromProtoFunction = functionName
 				case protoreflect.StringKind:
 					if krmField.Type != "[]string" {
 						useCustomMethod = true
-						// useSliceFromProto = fmt.Sprintf("%s_%s_FromProto", goTypeName, protoFieldName)
+						// useSliceFromProtoFunction = fmt.Sprintf("%s_%s_FromProto", goTypeName, protoFieldName)
 					}
+				case protoreflect.EnumKind:
+					krmElemTypeName := krmField.Type
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "*")
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
+
+					functionName := "Enum_FromProto"
+					useSliceFromProtoFunction = fmt.Sprintf("%s(mapCtx, in.%s)",
+						functionName,
+						krmFieldName,
+					)
 
 				case
 					protoreflect.FloatKind,
@@ -406,16 +416,16 @@ func (v *visitor) writeMapFunctionsForPair(out io.Writer, pair *typePair) {
 					protoreflect.Uint32Kind,
 					protoreflect.Uint64Kind,
 					protoreflect.BytesKind:
-					useSliceFromProto = ""
+					useSliceFromProtoFunction = ""
 				default:
 					klog.Fatalf("unhandled kind %q for repeated field %v", protoField.Kind(), protoField)
 				}
 
-				if useSliceFromProto != "" {
+				if useSliceFromProtoFunction != "" {
 					fmt.Fprintf(out, "\tout.%s = Slice_FromProto(mapCtx, in.%s, %s)\n",
 						krmFieldName,
 						krmFieldName,
-						useSliceFromProto,
+						useSliceFromProtoFunction,
 					)
 				} else if useCustomMethod {
 					methodName := fmt.Sprintf("%s_%s_FromProto", goTypeName, protoFieldName)
@@ -514,6 +524,19 @@ func (v *visitor) writeMapFunctionsForPair(out io.Writer, pair *typePair) {
 						//useSliceToProtoFunction = fmt.Sprintf("%s_%s_ToProto", goTypeName, protoFieldName)
 					}
 
+				case protoreflect.EnumKind:
+					krmElemTypeName := krmField.Type
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "*")
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
+
+					protoTypeName := "pb." + protoNameForEnum(protoField.Enum())
+					functionName := "Enum_ToProto"
+					useSliceToProtoFunction = fmt.Sprintf("%s[%s](mapCtx, in.%s)",
+						functionName,
+						protoTypeName,
+						krmFieldName,
+					)
+
 				case protoreflect.FloatKind,
 					protoreflect.DoubleKind,
 					protoreflect.BoolKind,