fix failing ProtoSchemaTranslator on proto3 optional fields (#32216)

* fix failing ProtoSchemaTranslator for proto3 optional fields * fix spotless checks
apache · Aug 27, 2024 · d1157d8 · d1157d8
1 parent 43963b6
commit d1157d8
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 2 deletions.
diff --git a/...protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoSchemaTranslator.java b/...protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoSchemaTranslator.java
@@ -207,7 +207,8 @@ static synchronized Schema getSchema(Descriptors.Descriptor descriptor) {
 
     for (Descriptors.FieldDescriptor fieldDescriptor : descriptor.getFields()) {
       int fieldDescriptorNumber = fieldDescriptor.getNumber();
-      if (!oneOfComponentFields.contains(fieldDescriptorNumber)) {
+      if (!(oneOfComponentFields.contains(fieldDescriptorNumber)
+          && fieldDescriptor.getRealContainingOneof() != null)) {
         // Store proto field number in metadata.
         FieldType fieldType = beamFieldTypeFromProtoField(fieldDescriptor);
         fields.add(

diff --git a/...obuf/src/test/java/org/apache/beam/sdk/extensions/protobuf/ProtoSchemaTranslatorTest.java b/...obuf/src/test/java/org/apache/beam/sdk/extensions/protobuf/ProtoSchemaTranslatorTest.java
@@ -52,6 +52,13 @@ public void testOptionalPrimitiveSchema() {
         ProtoSchemaTranslator.getSchema(Proto2SchemaMessages.OptionalPrimitive.class));
   }
 
+  @Test
+  public void testProto3OptionalPrimitiveSchema() {
+    assertEquals(
+        TestProtoSchemas.PROTO3_OPTIONAL_PRIMITIVE_SCHEMA,
+        ProtoSchemaTranslator.getSchema(Proto3SchemaMessages.OptionalPrimitive.class));
+  }
+
   @Test
   public void testRequiredPrimitiveSchema() {
     assertEquals(
@@ -169,6 +176,13 @@ public void testOptionalNestedSchema() {
         ProtoSchemaTranslator.getSchema(Proto2SchemaMessages.OptionalNested.class));
   }
 
+  @Test
+  public void testProto3OptionalNestedSchema() {
+    assertEquals(
+        TestProtoSchemas.PROTO3_OPTIONAL_NESTED_SCHEMA,
+        ProtoSchemaTranslator.getSchema(Proto3SchemaMessages.OptionalNested.class));
+  }
+
   @Test
   public void testRequiredNestedSchema() {
     assertEquals(
@@ -266,6 +280,7 @@ public void testOptionsMessageOnField() {
     assertEquals("foobar in field", optionMessage.getString("single_string"));
     assertEquals(Integer.valueOf(56), optionMessage.getInt32("single_int32"));
     assertEquals(Long.valueOf(78L), optionMessage.getInt64("single_int64"));
+    assertEquals("this is optional", optionMessage.getString("optional_string"));
   }
 
   @Test

diff --git a/...ions/protobuf/src/test/java/org/apache/beam/sdk/extensions/protobuf/TestProtoSchemas.java b/...ions/protobuf/src/test/java/org/apache/beam/sdk/extensions/protobuf/TestProtoSchemas.java
@@ -125,6 +125,20 @@ static Schema.Options withTypeName(String typeName) {
                       "proto2_schema_messages.OptionalPrimitive"))
           .build();
 
+  static final Schema PROTO3_OPTIONAL_PRIMITIVE_SCHEMA =
+      Schema.builder()
+          .addField(withFieldNumber("primitive_int32", FieldType.INT32, 1))
+          .addField(withFieldNumber("primitive_bool", FieldType.BOOLEAN, 2))
+          .addField(withFieldNumber("primitive_string", FieldType.STRING, 3))
+          .addField(withFieldNumber("primitive_bytes", FieldType.BYTES, 4))
+          .setOptions(
+              Schema.Options.builder()
+                  .setOption(
+                      SCHEMA_OPTION_META_TYPE_NAME,
+                      FieldType.STRING,
+                      "proto3_schema_messages.OptionalPrimitive"))
+          .build();
+
   static final Schema REQUIRED_PRIMITIVE_SCHEMA =
       Schema.builder()
           .addField(withFieldNumber("primitive_int32", FieldType.INT32, 1))
@@ -652,6 +666,14 @@ static Schema.Options withTypeName(String typeName) {
           .setOptions(withTypeName("proto2_schema_messages.OptionalNested"))
           .build();
 
+  static final Schema PROTO3_OPTIONAL_NESTED_SCHEMA =
+      Schema.builder()
+          .addField(
+              withFieldNumber("nested", FieldType.row(PROTO3_OPTIONAL_PRIMITIVE_SCHEMA), 1)
+                  .withNullable(true))
+          .setOptions(withTypeName("proto3_schema_messages.OptionalNested"))
+          .build();
+
   // A sample instance of the proto.
   static final OptionalNested OPTIONAL_NESTED =
       OptionalNested.newBuilder().setNested(OPTIONAL_PRIMITIVE_PROTO).build();

diff --git a/sdks/java/extensions/protobuf/src/test/proto/proto3_schema_messages.proto b/sdks/java/extensions/protobuf/src/test/proto/proto3_schema_messages.proto
@@ -180,6 +180,7 @@ message OptionMessage {
             single_string: "foobar in field"
             single_int32: 56
             single_int64: 78
+            optional_string: "this is optional"
         },
         (proto3_schema_options.field_option_repeated) = "field_string_1",
         (proto3_schema_options.field_option_repeated) = "field_string_2",
@@ -209,4 +210,15 @@ message SecondCircularNested {
 }
 
 message Empty {
-}
+}
+
+message OptionalPrimitive {
+    optional int32 primitive_int32 = 1;
+    optional bool primitive_bool = 2;
+    optional string primitive_string = 3;
+    optional bytes primitive_bytes = 4;
+}
+
+message OptionalNested {
+    optional OptionalPrimitive nested = 1;
+}
diff --git a/sdks/java/extensions/protobuf/src/test/proto/proto3_schema_options.proto b/sdks/java/extensions/protobuf/src/test/proto/proto3_schema_options.proto
@@ -90,4 +90,6 @@ message OptionTestMessage {
     }
     OptionEnum single_enum = 8;
     OptionTestSubMessage single_message = 9;
+
+    optional string optional_string = 10;
 }