fix(bigquery/storage/managedwriter): correct enum processing in Norma…

…lizeDescriptor (#5811)
googleapis · Mar 30, 2022 · 52cf48e · 52cf48e
1 parent 7076391
commit 52cf48e
Show file tree

Hide file tree

Showing 6 changed files with 448 additions and 23 deletions.
diff --git a/bigquery/storage/managedwriter/adapt/protoconversion.go b/bigquery/storage/managedwriter/adapt/protoconversion.go
@@ -352,7 +352,7 @@ func normalizeDescriptorInternal(in protoreflect.MessageDescriptor, visitedTypes
 			} else {
 				enumDP := protodesc.ToEnumDescriptorProto(inField.Enum())
 				enumDP.Name = proto.String(enumName)
-				resultDP.NestedType = append(resultDP.NestedType, &descriptorpb.DescriptorProto{
+				root.NestedType = append(resultDP.NestedType, &descriptorpb.DescriptorProto{
 					Name:     proto.String(enclosingTypeName),
 					EnumType: []*descriptorpb.EnumDescriptorProto{enumDP},
 				})

diff --git a/bigquery/storage/managedwriter/adapt/protoconversion_test.go b/bigquery/storage/managedwriter/adapt/protoconversion_test.go
@@ -621,6 +621,88 @@ func TestNormalizeDescriptor(t *testing.T) {
 				},
 			},
 		},
+		{
+			description: "WithExternalEnum",
+			in:          (&testdata.ExternalEnumMessage{}).ProtoReflect().Descriptor(),
+			want: &descriptorpb.DescriptorProto{
+				Name: proto.String("testdata_ExternalEnumMessage"),
+				Field: []*descriptorpb.FieldDescriptorProto{
+					{
+						Name:     proto.String("msg_a"),
+						JsonName: proto.String("msgA"),
+						Number:   proto.Int32(1),
+						TypeName: proto.String("testdata_EnumMsgA"),
+						Type:     descriptorpb.FieldDescriptorProto_TYPE_MESSAGE.Enum(),
+						Label:    descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
+					},
+					{
+						Name:     proto.String("msg_b"),
+						JsonName: proto.String("msgB"),
+						Number:   proto.Int32(2),
+						TypeName: proto.String("testdata_EnumMsgB"),
+						Type:     descriptorpb.FieldDescriptorProto_TYPE_MESSAGE.Enum(),
+						Label:    descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
+					},
+				},
+				NestedType: []*descriptorpb.DescriptorProto{
+					{
+						Name: proto.String("testdata_ExtEnum_E"),
+						EnumType: []*descriptorpb.EnumDescriptorProto{
+							{
+								Name: proto.String("ExtEnum"),
+								Value: []*descriptorpb.EnumValueDescriptorProto{
+									{
+										Name:   proto.String("UNDEFINED"),
+										Number: proto.Int32(0),
+									},
+									{
+										Name:   proto.String("THING"),
+										Number: proto.Int32(1),
+									},
+									{
+										Name:   proto.String("OTHER_THING"),
+										Number: proto.Int32(2),
+									},
+								},
+							},
+						},
+					},
+					{
+						Name: proto.String("testdata_EnumMsgA"),
+						Field: []*descriptorpb.FieldDescriptorProto{
+							{
+								Name:     proto.String("foo"),
+								JsonName: proto.String("foo"),
+								Number:   proto.Int32(1),
+								Type:     descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(),
+								Label:    descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
+							},
+							{
+								Name:     proto.String("bar"),
+								JsonName: proto.String("bar"),
+								Number:   proto.Int32(2),
+								Type:     descriptorpb.FieldDescriptorProto_TYPE_ENUM.Enum(),
+								TypeName: proto.String("testdata_ExtEnum_E.ExtEnum"),
+								Label:    descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
+							},
+						},
+					},
+					{
+						Name: proto.String("testdata_EnumMsgB"),
+						Field: []*descriptorpb.FieldDescriptorProto{
+							{
+								Name:     proto.String("baz"),
+								JsonName: proto.String("baz"),
+								Number:   proto.Int32(1),
+								Type:     descriptorpb.FieldDescriptorProto_TYPE_ENUM.Enum(),
+								TypeName: proto.String("testdata_ExtEnum_E.ExtEnum"),
+								Label:    descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {

diff --git a/bigquery/storage/managedwriter/integration_test.go b/bigquery/storage/managedwriter/integration_test.go
@@ -730,6 +730,25 @@ func TestIntegration_ProtoNormalization(t *testing.T) {
 			descriptor := (mesg).ProtoReflect().Descriptor()
 			testProtoNormalization(ctx, t, mwClient, bqClient, dataset, schema, descriptor, b)
 		})
+		t.Run("WithExternalEnum", func(t *testing.T) {
+			t.Parallel()
+			schema := testdata.ExternalEnumMessageSchema
+			mesg := &testdata.ExternalEnumMessage{
+				MsgA: &testdata.EnumMsgA{
+					Foo: proto.String("foo"),
+					Bar: testdata.ExtEnum_THING.Enum(),
+				},
+				MsgB: &testdata.EnumMsgB{
+					Baz: testdata.ExtEnum_OTHER_THING.Enum(),
+				},
+			}
+			b, err := proto.Marshal(mesg)
+			if err != nil {
+				t.Fatalf("proto.Marshal: %v", err)
+			}
+			descriptor := (mesg).ProtoReflect().Descriptor()
+			testProtoNormalization(ctx, t, mwClient, bqClient, dataset, schema, descriptor, b)
+		})
 	})
 }