From bc220e9455e349bdba543c0466e72255ef7c1357 Mon Sep 17 00:00:00 2001 From: Khaled Hammouda Date: Wed, 21 Oct 2026 13:08:23 -0400 Subject: [PATCH] Honor allowMissingRequiredFields when building Storage Write API messages --- CHANGES.md | 1 + .../bigquery/TableRowToStorageApiProto.java | 20 +- .../TableRowToStorageApiProtoTest.java | 179 ++++++++++++++++++ 3 files changed, 198 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 74dcb41af94b..dc166eb4d64e 100644 --- a/CHANGES.md +++ b/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 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 687796ed5b47..d06c3d764c4f 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 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. + * + *

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}. + * + *

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. + * + *

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); 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 90735bae1261..51c185fc7b9e 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 @@ 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";