[Java][BigQueryIO] Honor allowMissingRequiredFields when building Storage Write API messages

CHANGES.md

@@ -129,6 +129,7 @@
 * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)).
+* (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java

@@ -848,6 +848,13 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
     return Iterables.getOnlyElement(fileDescriptor.getMessageTypes());
   }
 
+  /**
+   * Given a BigQuery row map, returns a protocol-buffer message for the BigQuery Storage API.
+   *
+   * <p>When {@code allowMissingRequiredFields} is true, this method returns a partial
+   * protocol-buffer message if required fields are absent. When it is false, required-field
+   * validation and its conversion errors remain unchanged.
+   */
   public static @Nullable DynamicMessage messageFromMap(
       SchemaInformation schemaInformation,
       @Nullable Descriptor descriptor,
@@ -1042,7 +1049,8 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
       return null;
     }
     try {
-      return Preconditions.checkArgumentNotNull(builder).build();
+      DynamicMessage.Builder nonNullBuilder = Preconditions.checkArgumentNotNull(builder);
+      return allowMissingRequiredFields ? nonNullBuilder.buildPartial() : nonNullBuilder.build();
     } catch (Exception e) {
       throw new SchemaDoesntMatchException(
           "Couldn't convert schema for " + schemaInformation.getFullName(), e);
@@ -1053,6 +1061,10 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
    * Forwards {@code changeSequenceNum} to {@link #messageFromTableRow(SchemaInformation,
    * Descriptor, TableRow, boolean, boolean, TableRow, String, String, ErrorCollector)} via {@link
    * Long#toHexString}.
+   *
+   * <p>When {@code allowMissingRequiredFields} is true, this method returns a partial
+   * protocol-buffer message if required fields are absent. When it is false, required-field
+   * validation and its conversion errors remain unchanged.
    */
   public static @Nullable DynamicMessage messageFromTableRow(
       SchemaInformation schemaInformation,
@@ -1080,6 +1092,10 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
   /**
    * Given a BigQuery TableRow, returns a protocol-buffer message that can be used to write data
    * using the BigQuery Storage API.
+   *
+   * <p>When {@code allowMissingRequiredFields} is true, this method returns a partial
+   * protocol-buffer message if required fields are absent. When it is false, required-field
+   * validation and its conversion errors remain unchanged.
    */
   public static @Nullable DynamicMessage messageFromTableRow(
       SchemaInformation schemaInformation,
@@ -1202,7 +1218,7 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
       }
 
       try {
-        return builder.build();
+        return allowMissingRequiredFields ? builder.buildPartial() : builder.build();
       } catch (Exception e) {
         throw new SchemaDoesntMatchException(
             "Could convert schema for " + schemaInformation.getFullName(), e);

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java

@@ -19,6 +19,7 @@
 
 import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.TIMESTAMP_FORMATTER;
 import static org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.TYPE_MAP_PROTO_CONVERTERS;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -1903,6 +1904,184 @@ public void testNullRepeatedDescriptorFromTableSchema() throws Exception {
     assertTrue(repeatednof2.isEmpty());
   }
 
+  @Test
+  public void testAllowMissingRequiredFieldsMatchesStrictBuildForInitializedMessage()
+      throws Exception {
+    TableSchema tableSchema = requiredStringTableSchema();
+    SchemaInformation schemaInformation = SchemaInformation.fromTableSchema(tableSchema);
+    Descriptor optionalDescriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, false, false);
+    TableRow tableRow =
+        new TableRow().set("requiredvalue", "present").set("unknownvalue", "preserved");
+    TableRow strictUnknownFields = new TableRow();
+    TableRow allowMissingUnknownFields = new TableRow();
+
+    DynamicMessage strictMessage =
+        TableRowToStorageApiProto.messageFromTableRow(
+            schemaInformation,
+            optionalDescriptor,
+            tableRow,
+            true,
+            false,
+            strictUnknownFields,
+            null,
+            -1,
+            TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+    DynamicMessage allowMissingMessage =
+        TableRowToStorageApiProto.messageFromTableRow(
+            schemaInformation,
+            optionalDescriptor,
+            tableRow,
+            true,
+            true,
+            allowMissingUnknownFields,
+            null,
+            -1,
+            TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+
+    assertArrayEquals(strictMessage.toByteArray(), allowMissingMessage.toByteArray());
+    assertEquals(strictUnknownFields, allowMissingUnknownFields);
+    assertEquals(new TableRow().set("unknownvalue", "preserved"), strictUnknownFields);
+    assertTrue(strictMessage.isInitialized());
+    assertTrue(allowMissingMessage.isInitialized());
+  }
+
+  @Test
+  public void testStrictDescriptorRejectsMissingRequiredFieldWhenNotAllowed() throws Exception {
+    TableSchema tableSchema = requiredStringTableSchema();
+    SchemaInformation schemaInformation = SchemaInformation.fromTableSchema(tableSchema);
+    Descriptor strictDescriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, true, false);
+
+    try {
+      TableRowToStorageApiProto.messageFromTableRow(
+          schemaInformation,
+          strictDescriptor,
+          new TableRow().setF(Collections.emptyList()),
+          false,
+          false,
+          null,
+          null,
+          -1,
+          TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+      fail("Expected strict descriptor validation to reject the missing field");
+    } catch (TableRowToStorageApiProto.SchemaDoesntMatchException e) {
+      assertEquals(
+          "Could convert schema for . Exception: "
+              + "com.google.protobuf.UninitializedMessageException: Message missing required "
+              + "fields: requiredvalue",
+          e.getMessage());
+    }
+  }
+
+  @Test
+  public void testNestedStrictDescriptorPreservesExactFailureChain() throws Exception {
+    TableSchema tableSchema =
+        new TableSchema()
+            .setFields(
+                ImmutableList.of(
+                    new TableFieldSchema()
+                        .setName("nested")
+                        .setType("STRUCT")
+                        .setFields(
+                            ImmutableList.of(
+                                new TableFieldSchema()
+                                    .setName("requiredvalue")
+                                    .setType("STRING")
+                                    .setMode("REQUIRED")))));
+    SchemaInformation schemaInformation = SchemaInformation.fromTableSchema(tableSchema);
+    Descriptor strictDescriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, true, false);
+
+    try {
+      TableRowToStorageApiProto.messageFromTableRow(
+          schemaInformation,
+          strictDescriptor,
+          new TableRow().set("nested", new TableRow()),
+          false,
+          false,
+          null,
+          null,
+          -1,
+          TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+      fail("Expected nested strict descriptor validation to reject the missing field");
+    } catch (TableRowToStorageApiProto.SchemaDoesntMatchException e) {
+      assertEquals(
+          "Problem converting field nested expected type: STRUCT. Exception: "
+              + "org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto$"
+              + "SchemaMissingRequiredFieldException: Missing required fields: "
+              + "[nested.requiredvalue]",
+          e.getMessage());
+      Throwable cause = e;
+      while (cause != null
+          && !(cause instanceof TableRowToStorageApiProto.SchemaMissingRequiredFieldException)) {
+        cause = cause.getCause();
+      }
+      assertNotNull(cause);
+      assertEquals(
+          Collections.singleton("nested.requiredvalue"),
+          ((TableRowToStorageApiProto.SchemaMissingRequiredFieldException) cause)
+              .getMissingFields());
+    }
+  }
+
+  @Test
+  public void testAllowMissingRequiredFieldsReturnsPartialMessageFromMap() throws Exception {
+    TableSchema tableSchema = requiredStringTableSchema();
+    SchemaInformation schemaInformation = SchemaInformation.fromTableSchema(tableSchema);
+    Descriptor strictDescriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, true, false);
+
+    DynamicMessage message =
+        TableRowToStorageApiProto.messageFromTableRow(
+            schemaInformation,
+            strictDescriptor,
+            new TableRow(),
+            false,
+            true,
+            null,
+            null,
+            -1,
+            TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+
+    assertFalse(message.isInitialized());
+    assertTrue(message.findInitializationErrors().contains("requiredvalue"));
+  }
+
+  @Test
+  public void testAllowMissingRequiredFieldsReturnsPartialMessageFromPositionalRow()
+      throws Exception {
+    TableSchema tableSchema = requiredStringTableSchema();
+    SchemaInformation schemaInformation = SchemaInformation.fromTableSchema(tableSchema);
+    Descriptor strictDescriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, true, false);
+
+    DynamicMessage message =
+        TableRowToStorageApiProto.messageFromTableRow(
+            schemaInformation,
+            strictDescriptor,
+            new TableRow().setF(Collections.emptyList()),
+            false,
+            true,
+            null,
+            null,
+            -1,
+            TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT);
+
+    assertFalse(message.isInitialized());
+    assertTrue(message.findInitializationErrors().contains("requiredvalue"));
+  }
+
+  private static TableSchema requiredStringTableSchema() {
+    return new TableSchema()
+        .setFields(
+            ImmutableList.of(
+                new TableFieldSchema()
+                    .setName("requiredvalue")
+                    .setType("STRING")
+                    .setMode("REQUIRED")));
+  }
+
   @Test
   public void testIntegerTypeConversion() throws DescriptorValidationException {
     String intFieldName = "int_field";