This is an automated email from the ASF dual-hosted git repository.
ahmedabu98 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new e4c2f8c71a3 Honor allowMissingRequiredFields when building Storage
Write API messages (#40204)
e4c2f8c71a3 is described below
commit e4c2f8c71a344ba2e177904607c588499e3ac673
Author: Khaled Hammouda <[email protected]>
AuthorDate: Thu Oct 22 19:49:32 2026 -0400
Honor allowMissingRequiredFields when building Storage Write API messages
(#40204)
---
CHANGES.md | 1 +
.../io/gcp/bigquery/TableRowToStorageApiProto.java | 20 ++-
.../bigquery/TableRowToStorageApiProtoTest.java | 179 +++++++++++++++++++++
3 files changed, 198 insertions(+), 2 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 39d25b568a0..9d3f6fbac24 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -135,6 +135,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
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java
index 687796ed5b4..d06c3d764c4 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java
@@ -848,6 +848,13 @@ public class TableRowToStorageApiProto {
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 class TableRowToStorageApiProto {
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 class TableRowToStorageApiProto {
* 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 class TableRowToStorageApiProto {
/**
* 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 class TableRowToStorageApiProto {
}
try {
- return builder.build();
+ return allowMissingRequiredFields ? builder.buildPartial() :
builder.build();
} catch (Exception e) {
throw new SchemaDoesntMatchException(
"Could convert schema for " + schemaInformation.getFullName(), e);
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java
index 90735bae126..51c185fc7b9 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoTest.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.io.gcp.bigquery;
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 class TableRowToStorageApiProtoTest {
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";