From e486e2aaa4f7ddda041265f5fe19a5c6118b98fe Mon Sep 17 00:00:00 2001 From: Kenneth Knowles Date: Wed, 26 Aug 2026 14:38:04 +0000 Subject: [PATCH] Fix nullness in BigQueryIO --- .../trigger_files/beam_PostCommit_Java.json | 1 + ..._PostCommit_Java_BigQueryEarlyRollout.json | 3 + .../beam_PostCommit_Java_DataflowV1.json | 1 + .../beam_PostCommit_Java_DataflowV2.json | 1 + .../AvroGenericRecordToStorageApiProto.java | 22 +- .../sdk/io/gcp/bigquery/AvroRowWriter.java | 9 +- .../beam/sdk/io/gcp/bigquery/BatchLoads.java | 17 +- .../io/gcp/bigquery/BigQueryAvroUtils.java | 17 +- ...xportReadSchemaTransformConfiguration.java | 8 +- ...ueryExportReadSchemaTransformProvider.java | 13 +- .../sdk/io/gcp/bigquery/BigQueryHelpers.java | 38 +- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 372 ++++++++++-------- .../gcp/bigquery/BigQueryIOTranslation.java | 269 ++++++++----- .../gcp/bigquery/BigQueryQuerySourceDef.java | 32 +- .../bigquery/BigQuerySchemaIOProvider.java | 7 +- .../io/gcp/bigquery/BigQueryServicesImpl.java | 51 ++- .../bigquery/BigQueryStorageAvroReader.java | 12 +- .../sdk/io/gcp/bigquery/BigQueryUtils.java | 169 ++++---- .../bigquery/DynamicDestinationsHelpers.java | 12 +- .../sdk/io/gcp/bigquery/RowWriterFactory.java | 7 +- .../bigquery/StorageApiConvertMessages.java | 7 +- .../StorageApiDynamicDestinationsBeamRow.java | 1 - ...geApiDynamicDestinationsGenericRecord.java | 4 +- .../StorageApiDynamicDestinationsProto.java | 12 +- .../StorageApiFinalizeWritesDoFn.java | 6 - .../sdk/io/gcp/bigquery/StorageApiLoads.java | 17 +- .../gcp/bigquery/StorageApiWritePayload.java | 1 - .../StorageApiWritesShardedRecords.java | 4 +- .../sdk/io/gcp/bigquery/StreamingInserts.java | 4 +- .../bigquery/TableRowToStorageApiProto.java | 11 +- .../gcp/bigquery/UpdateSchemaDestination.java | 26 +- ...ueryDirectReadSchemaTransformProvider.java | 32 +- ...QueryFileLoadsSchemaTransformProvider.java | 27 +- ...torageWriteApiSchemaTransformProvider.java | 102 +++-- .../providers/BigQueryWriteConfiguration.java | 18 + .../PortableBigQueryDestinations.java | 6 +- 36 files changed, 788 insertions(+), 551 deletions(-) create mode 100644 .github/trigger_files/beam_PostCommit_Java_BigQueryEarlyRollout.json diff --git a/.github/trigger_files/beam_PostCommit_Java.json b/.github/trigger_files/beam_PostCommit_Java.json index 7b4c1ba67021..de320ac8499c 100644 --- a/.github/trigger_files/beam_PostCommit_Java.json +++ b/.github/trigger_files/beam_PostCommit_Java.json @@ -1,4 +1,5 @@ { + "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO", "comment": "Modify this file in a trivial way to cause this test suite to run", "modification": 6 } diff --git a/.github/trigger_files/beam_PostCommit_Java_BigQueryEarlyRollout.json b/.github/trigger_files/beam_PostCommit_Java_BigQueryEarlyRollout.json new file mode 100644 index 000000000000..f4dd232f10da --- /dev/null +++ b/.github/trigger_files/beam_PostCommit_Java_BigQueryEarlyRollout.json @@ -0,0 +1,3 @@ +{ + "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO" +} diff --git a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json index aff502e19620..c403293a8561 100644 --- a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json +++ b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json @@ -1,4 +1,5 @@ { + "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO", "https://github.com/apache/beam/pull/39330": "Fix DataflowV1 test failure by fixing getSimpleName access", "https://github.com/apache/beam/pull/34902": "Introducing OutputBuilder", "https://github.com/apache/beam/pull/35177": "Introducing WindowedValueReceiver to runners", diff --git a/.github/trigger_files/beam_PostCommit_Java_DataflowV2.json b/.github/trigger_files/beam_PostCommit_Java_DataflowV2.json index c4a6954cecb7..67d8e37f3de8 100644 --- a/.github/trigger_files/beam_PostCommit_Java_DataflowV2.json +++ b/.github/trigger_files/beam_PostCommit_Java_DataflowV2.json @@ -1,4 +1,5 @@ { + "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO", "modification": 9, "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface" } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java index 2691f82eebef..84a7755573a1 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java @@ -317,7 +317,6 @@ public static DynamicMessage messageFromGenericRecord( return builder.build(); } - @SuppressWarnings("nullness") private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Schema.Field field) { @Nullable Schema schema = field.schema(); @@ -360,12 +359,14 @@ private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Sch if (valueType == null) { throw new RuntimeException("Unexpected null element type!"); } - TableFieldSchema keyFieldSchema = - fieldDescriptorFromAvroField( - new Schema.Field("key", keyType, "key of the map entry", null)); - TableFieldSchema valueFieldSchema = - fieldDescriptorFromAvroField( - new Schema.Field("value", valueType, "value of the map entry", null)); + // The Avro Field constructor accepts a null default value, but Avro is not annotated. + @SuppressWarnings("nullness") + Schema.Field keyField = new Schema.Field("key", keyType, "key of the map entry", null); + @SuppressWarnings("nullness") + Schema.Field valueField = + new Schema.Field("value", valueType, "value of the map entry", null); + TableFieldSchema keyFieldSchema = fieldDescriptorFromAvroField(keyField); + TableFieldSchema valueFieldSchema = fieldDescriptorFromAvroField(valueField); builder = builder .setType(TableFieldSchema.Type.STRUCT) @@ -382,9 +383,10 @@ private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Sch Preconditions.checkState( elementType.getType() != Schema.Type.UNION, "Multiple non-null union types are not supported."); - TableFieldSchema unionFieldSchema = - fieldDescriptorFromAvroField( - new Schema.Field(field.name(), elementType, field.doc(), null)); + // The Avro Field constructor accepts a null default value, but Avro is not annotated. + @SuppressWarnings("nullness") + Schema.Field unionField = new Schema.Field(field.name(), elementType, field.doc(), null); + TableFieldSchema unionFieldSchema = fieldDescriptorFromAvroField(unionField); builder = builder .setType(unionFieldSchema.getType()) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroRowWriter.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroRowWriter.java index 1f45371b19ff..26f444ada777 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroRowWriter.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroRowWriter.java @@ -29,9 +29,6 @@ class AvroRowWriter extends BigQueryRowWriter { private final Schema schema; private final SerializableFunction, AvroT> toAvroRecord; - @SuppressWarnings({ - "nullness" // calling superclass method in constructor flagged as error; TODO: fix - }) AvroRowWriter( String basename, Schema schema, @@ -42,8 +39,12 @@ class AvroRowWriter extends BigQueryRowWriter { this.schema = schema; this.toAvroRecord = toAvroRecord; - this.writer = + // getOutputStream() is established by the superclass constructor, which the checker + // cannot see through the partially-initialized receiver. + @SuppressWarnings("nullness") + DataFileWriter initializedWriter = new DataFileWriter<>(writerFactory.apply(schema)).create(schema, getOutputStream()); + this.writer = initializedWriter; } @Override diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java index 252e55d34c07..d12af24e407d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java @@ -162,12 +162,12 @@ class BatchLoads private long maxBytesPerPartition; private int numFileShards; private @Nullable Duration triggeringFrequency; - private ValueProvider customGcsTempLocation; + private @Nullable ValueProvider customGcsTempLocation; private @Nullable ValueProvider loadJobProjectId; private final Coder elementCoder; private final RowWriterFactory rowWriterFactory; private final @Nullable String kmsKey; - private final String tempDataset; + private final @Nullable String tempDataset; private final BadRecordRouter badRecordRouter; private final ErrorHandler badRecordErrorHandler; private Coder tableDestinationCoder; @@ -181,7 +181,7 @@ class BatchLoads boolean singletonTable, DynamicDestinations dynamicDestinations, Coder destinationCoder, - ValueProvider customGcsTempLocation, + @Nullable ValueProvider customGcsTempLocation, @Nullable ValueProvider loadJobProjectId, boolean ignoreUnknownValues, Coder elementCoder, @@ -189,7 +189,7 @@ class BatchLoads @Nullable String kmsKey, boolean clusteringEnabled, boolean useAvroLogicalTypes, - String tempDataset, + @Nullable String tempDataset, BadRecordRouter badRecordRouter, ErrorHandler badRecordErrorHandler) { bigQueryServices = new BigQueryServicesImpl(); @@ -249,7 +249,7 @@ public void setMaxNumWritersPerBundle(int maxNumWritersPerBundle) { this.maxNumWritersPerBundle = maxNumWritersPerBundle; } - public void setTriggeringFrequency(Duration triggeringFrequency) { + public void setTriggeringFrequency(@Nullable Duration triggeringFrequency) { this.triggeringFrequency = triggeringFrequency; } @@ -285,6 +285,7 @@ public void validate(@Nullable PipelineOptions maybeOptions) { PipelineOptions options = Preconditions.checkArgumentNotNull(maybeOptions); // We will use a BigQuery load job -- validate the temp location. String tempLocation; + ValueProvider customGcsTempLocation = this.customGcsTempLocation; if (customGcsTempLocation == null) { tempLocation = options.getTempLocation(); } else { @@ -424,7 +425,7 @@ private WriteResult expandTriggered(PCollection> inpu .apply("ExtractTempTables", Values.create()) .apply( ParDo.of( - new UpdateSchemaDestination( + new UpdateSchemaDestination<>( bigQueryServices, zeroLoadJobIdPrefixView, loadJobProjectId, @@ -530,7 +531,7 @@ public WriteResult expandUntriggered(PCollection> inp .apply("ReifyRenameInput", new ReifyAsIterable<>()) .apply( ParDo.of( - new UpdateSchemaDestination( + new UpdateSchemaDestination<>( bigQueryServices, zeroLoadJobIdPrefixView, loadJobProjectId, @@ -592,6 +593,8 @@ private PCollectionView createTempFilePrefixView( @ProcessElement public void getTempFilePrefix(ProcessContext c) { String tempLocationRoot; + ValueProvider customGcsTempLocation = + BatchLoads.this.customGcsTempLocation; if (customGcsTempLocation != null && customGcsTempLocation.get() != null) { tempLocationRoot = customGcsTempLocation.get(); } else { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryAvroUtils.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryAvroUtils.java index 2dbc4316b883..f7c74cf8bc2d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryAvroUtils.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryAvroUtils.java @@ -546,9 +546,6 @@ private static boolean hasNamespaceCollision(List fieldSchemas return false; } - @SuppressWarnings({ - "nullness" // Avro library not annotated - }) private static Field convertField( TableFieldSchema bigQueryField, Boolean useAvroLogicalTypes, @Nullable String namespace) { String fieldName = bigQueryField.getName(); @@ -569,11 +566,15 @@ private static Field convertField( } else if (!"REQUIRED".equals(bqMode)) { throw new IllegalArgumentException(String.format("Unknown BigQuery Field Mode: %s", bqMode)); } - return new Field( - fieldName, - fieldSchema, - bigQueryField.getDescription(), - (Object) null /* Cast to avoid deprecated JsonNode constructor. */); + // The Avro Field constructor accepts a null default value, but Avro is not annotated. + @SuppressWarnings("nullness") + Field field = + new Field( + fieldName, + fieldSchema, + bigQueryField.getDescription(), + (Object) null /* Cast to avoid deprecated JsonNode constructor. */); + return field; } static TableSchema fromGenericAvroSchema(Schema schema) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformConfiguration.java index a5ca8e1a2bbb..b2a109c44855 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformConfiguration.java @@ -21,6 +21,7 @@ import javax.annotation.Nullable; import org.apache.beam.sdk.schemas.AutoValueSchema; import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.checkerframework.dataflow.qual.Pure; /** * Configuration for reading from BigQuery. @@ -31,9 +32,6 @@ * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam * repository. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) @DefaultSchema(AutoValueSchema.class) @AutoValue public abstract class BigQueryExportReadSchemaTransformConfiguration { @@ -44,6 +42,7 @@ public static Builder builder() { } /** Configures the BigQuery read job with the SQL query. */ + @Pure @Nullable public abstract String getQuery(); @@ -51,14 +50,17 @@ public static Builder builder() { * Specifies a table for a BigQuery read job. See {@link BigQueryIO.TypedRead#from(String)} for * more details on the expected format. */ + @Pure @Nullable public abstract String getTableSpec(); /** BigQuery geographic location where the query job will be executed. */ + @Pure @Nullable public abstract String getQueryLocation(); /** Enables BigQuery's Standard SQL dialect when reading from a query. */ + @Pure @Nullable public abstract Boolean getUseStandardSql(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformProvider.java index af6ab5c71c8b..735957239490 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryExportReadSchemaTransformProvider.java @@ -32,7 +32,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; /** * An implementation of {@link TypedSchemaTransformProvider} for BigQuery read jobs configured using @@ -42,9 +42,6 @@ * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam * repository. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) @Internal @AutoService(SchemaTransformProvider.class) public class BigQueryExportReadSchemaTransformProvider @@ -96,7 +93,7 @@ public List outputCollectionNames() { */ protected static class BigQueryExportSchemaTransform extends SchemaTransform { /** An instance of {@link BigQueryServices} used for testing. */ - private BigQueryServices testBigQueryServices = null; + private @Nullable BigQueryServices testBigQueryServices = null; private final BigQueryExportReadSchemaTransformConfiguration configuration; @@ -135,11 +132,11 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { BigQueryIO.TypedRead toTypedRead() { BigQueryIO.TypedRead read = BigQueryIO.readTableRowsWithSchema(); - if (!Strings.isNullOrEmpty(configuration.getQuery())) { + if (configuration.getQuery() != null && !configuration.getQuery().isEmpty()) { read = read.fromQuery(configuration.getQuery()); } - if (!Strings.isNullOrEmpty(configuration.getTableSpec())) { + if (configuration.getTableSpec() != null && !configuration.getTableSpec().isEmpty()) { read = read.from(configuration.getTableSpec()); } @@ -147,7 +144,7 @@ BigQueryIO.TypedRead toTypedRead() { read = read.usingStandardSql(); } - if (!Strings.isNullOrEmpty(configuration.getQueryLocation())) { + if (configuration.getQueryLocation() != null && !configuration.getQueryLocation().isEmpty()) { read = read.withQueryLocation(configuration.getQueryLocation()); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java index ef15cd805223..a7080c9a9a16 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; @@ -27,6 +28,8 @@ import com.google.api.services.bigquery.model.Dataset; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.Job; +import com.google.api.services.bigquery.model.JobConfiguration; +import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; import com.google.api.services.bigquery.model.JobStatus; import com.google.api.services.bigquery.model.Table; @@ -471,9 +474,6 @@ static List getOrCreateMapListValue(Map> map, K key) { * *

If the project id is omitted, the default project id is used. */ - @SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) - }) public static TableReference parseTableSpec(String tableSpec) { Matcher match = BigQueryIO.TABLE_SPEC.matcher(tableSpec); if (!match.matches()) { @@ -545,7 +545,9 @@ public static TableReference parseTableSpec(String tableSpec) { } TableReference ref = new TableReference(); - ref.setProjectId(project); + // The project id is optional; the API client accepts a null project id but is not annotated. + @SuppressWarnings("nullness") + TableReference unused = ref.setProjectId(project); return ref.setDatasetId(dataset).setTableId(table); } @@ -561,9 +563,6 @@ private static IllegalArgumentException invalidTableSpec(String tableSpec) { tableSpec)); } - @SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) - }) public static TableReference parseTableUrn(String tableUrn) { Matcher match = BigQueryIO.TABLE_URN_SPEC.matcher(tableUrn); if (!match.matches()) { @@ -573,10 +572,10 @@ public static TableReference parseTableUrn(String tableUrn) { + tableUrn); } - TableReference ref = new TableReference(); - ref.setProjectId(match.group("PROJECT")); - - return ref.setDatasetId(match.group("DATASET")).setTableId(match.group("TABLE")); + return new TableReference() + .setProjectId(checkStateNotNull(match.group("PROJECT"))) + .setDatasetId(checkStateNotNull(match.group("DATASET"))) + .setTableId(checkStateNotNull(match.group("TABLE"))); } /** Strip off any partition decorator information from a tablespec. */ @@ -585,19 +584,22 @@ public static String stripPartitionDecorator(String tableSpec) { return (index == -1) ? tableSpec : tableSpec.substring(0, index); } - @SuppressWarnings({ - "nullness" // The BigQuery API library is documented to accept nulls but is not annotated - }) static String jobToPrettyString(@Nullable Job job) throws IOException { - if (job != null && job.getConfiguration().getLoad() != null) { + if (job == null) { + return "null"; + } + JobConfiguration configuration = job.getConfiguration(); + if (configuration != null && configuration.getLoad() != null) { // Removing schema and sourceUris from error messages for load jobs since these fields can be // quite long and error message might not be displayed properly in runner specific logs. job = job.clone(); - job.getConfiguration().getLoad().setSchema(null); - job.getConfiguration().getLoad().setSourceUris(null); + JobConfigurationLoad load = checkStateNotNull(job.getConfiguration()).getLoad(); + // The BigQuery API library is documented to accept nulls here but is not annotated. + @SuppressWarnings("nullness") + JobConfigurationLoad unused = load.setSchema(null).setSourceUris(null); } - return job == null ? "null" : job.toPrettyString(); + return job.toPrettyString(); } static String statusToPrettyString(@Nullable JobStatus status) throws IOException { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index 56b0200e4ba8..e43bc0a4e871 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -21,6 +21,8 @@ import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryResourceNaming.createTempTableReference; import static org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.BAD_RECORD_TAG; import static org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.RECORDING_ROUTER; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; @@ -555,10 +557,7 @@ * the table is not previously created and CREATE_IF_NEEDED is used, a primary key must be specified * using {@link Write#withPrimaryKey}. */ -@SuppressWarnings({ - "nullness", // TODO(https://github.com/apache/beam/issues/20506), - "SameNameButDifferent" -}) +@SuppressWarnings({"SameNameButDifferent"}) public class BigQueryIO { /** @@ -652,8 +651,12 @@ public class BigQueryIO { static final SerializableFunction> GENERIC_DATUM_WRITER_FACTORY = schema -> new GenericDatumWriter<>(); - private static final SerializableFunction - DEFAULT_AVRO_SCHEMA_FACTORY = BigQueryAvroUtils::toGenericAvroSchema; + private static final SerializableFunction<@Nullable TableSchema, org.apache.avro.Schema> + DEFAULT_AVRO_SCHEMA_FACTORY = + tableSchema -> + BigQueryAvroUtils.toGenericAvroSchema( + checkArgumentNotNull( + tableSchema, "A table schema is required to generate an Avro schema.")); static final String CONNECTION_ID = "connectionId"; static final String STORAGE_URI = "storageUri"; @@ -755,7 +758,7 @@ public void setSchema(org.apache.avro.Schema schema) { @Override public T read(T reuse, Decoder in) throws IOException { - GenericRecord record = (GenericRecord) this.reader.read(reuse, in); + GenericRecord record = (GenericRecord) checkStateNotNull(this.reader.read(reuse, in)); return parseFn.apply(new SchemaAndRecord(record, this.tableSchema.get())); } } @@ -794,8 +797,8 @@ public static TypedRead read(SerializableFunction par (writer, reader) -> new GenericDatumTransformer<>(parseFn, jsonTableSchema, writer); } catch (IOException e) { - LOG.warn("Error while converting table schema {} to JSON!", input, e); - return null; + throw new RuntimeException( + "Error while converting table schema " + input + " to JSON!", e); } }) // TODO: Remove setParseFn once https://github.com/apache/beam/issues/21076 is fixed. @@ -859,9 +862,9 @@ public abstract static class DynamicRead abstract DataFormat getFormat(); - abstract @Nullable SerializableFunction getParseFn(); + abstract SerializableFunction getParseFn(); - abstract @Nullable Coder getOutputCoder(); + abstract Coder getOutputCoder(); abstract boolean getProjectionPushdownApplied(); @@ -979,9 +982,9 @@ public void processElement( BigQueryStorageQuerySource querySource = BigQueryStorageQuerySource.create( kv.getKey(), - StaticValueProvider.of(descriptor.getQuery()), - descriptor.getFlattenResults(), - descriptor.getUseLegacySql(), + StaticValueProvider.of(checkStateNotNull(descriptor.getQuery())), + checkStateNotNull(descriptor.getFlattenResults()), + checkStateNotNull(descriptor.getUseLegacySql()), TypedRead.QueryPriority.INTERACTIVE, getQueryLocation(), getQueryTempDataset(), @@ -993,7 +996,8 @@ public void processElement( getBigQueryServices()); // due to retry, table may already exist, remove it to ensure correctness querySource.removeDestinationIfExists(options.as(BigQueryOptions.class)); - Table queryResultTable = querySource.getTargetTable(options.as(BigQueryOptions.class)); + Table queryResultTable = + checkStateNotNull(querySource.getTargetTable(options.as(BigQueryOptions.class))); BigQueryStorageTableSource output = BigQueryStorageTableSource.create( @@ -1081,6 +1085,7 @@ boolean getValidate() { return this.inner.getValidate(); } + @Nullable ValueProvider getQuery() { return this.inner.getQuery(); } @@ -1354,12 +1359,15 @@ public enum QueryPriority { @VisibleForTesting Coder inferCoder(CoderRegistry coderRegistry) { - if (getCoder() != null) { - return getCoder(); + Coder coder = getCoder(); + if (coder != null) { + return coder; } try { - return coderRegistry.getCoder(TypeDescriptors.outputOf(getParseFn())); + return coderRegistry.getCoder( + TypeDescriptors.outputOf( + checkStateNotNull(getParseFn(), "Either withCoder() or a parseFn is required"))); } catch (CannotProvideCoderException e) { throw new IllegalArgumentException( "Unable to infer coder for output of parseFn. Specify it explicitly using withCoder().", @@ -1368,46 +1376,48 @@ Coder inferCoder(CoderRegistry coderRegistry) { } private BigQuerySourceDef createSourceDef() { - BigQuerySourceDef sourceDef; - if (getQuery() == null) { - sourceDef = BigQueryTableSourceDef.create(getBigQueryServices(), getTableProvider()); - } else { - sourceDef = - BigQueryQuerySourceDef.create( - getBigQueryServices(), - getQuery(), - getFlattenResults(), - getUseLegacySql(), - MoreObjects.firstNonNull(getQueryPriority(), QueryPriority.BATCH), - getQueryLocation(), - getQueryTempDataset(), - getQueryTempProject(), - getKmsKey()); + ValueProvider query = getQuery(); + if (query == null) { + return BigQueryTableSourceDef.create( + getBigQueryServices(), + checkStateNotNull(getTableProvider(), "Either from() or fromQuery() is required")); } - return sourceDef; + return BigQueryQuerySourceDef.create( + getBigQueryServices(), + query, + checkStateNotNull( + getFlattenResults(), "flattenResults should not be null if query is set"), + checkStateNotNull(getUseLegacySql(), "useLegacySql should not be null if query is set"), + MoreObjects.firstNonNull(getQueryPriority(), QueryPriority.BATCH), + getQueryLocation(), + getQueryTempDataset(), + getQueryTempProject(), + getKmsKey()); } private BigQueryStorageQuerySource createStorageQuerySource( String stepUuid, Coder outputCoder) { return BigQueryStorageQuerySource.create( stepUuid, - getQuery(), - getFlattenResults(), - getUseLegacySql(), + checkStateNotNull(getQuery(), "Either from() or fromQuery() is required"), + checkStateNotNull( + getFlattenResults(), "flattenResults should not be null if query is set"), + checkStateNotNull(getUseLegacySql(), "useLegacySql should not be null if query is set"), MoreObjects.firstNonNull(getQueryPriority(), QueryPriority.BATCH), getQueryLocation(), getQueryTempDataset(), getQueryTempProject(), getKmsKey(), getFormat(), - getParseFn(), + checkStateNotNull(getParseFn(), "A parseFn is required"), outputCoder, getBigQueryServices(), getDirectReadPicosTimestampPrecision()); } @Override - public void validate(PipelineOptions options) { + public void validate(@Nullable PipelineOptions maybeOptions) { + PipelineOptions options = checkArgumentNotNull(maybeOptions); // Even if existence validation is disabled, we need to make sure that the BigQueryIO // read is properly specified. BigQueryOptions bqOptions = options.as(BigQueryOptions.class); @@ -1435,6 +1445,7 @@ public void validate(PipelineOptions options) { } ValueProvider table = getTableProvider(); + ValueProvider query = getQuery(); // Note that a table or query check can fail if the table or dataset are created by // earlier stages of the pipeline or if a query depends on earlier stages of a pipeline. @@ -1449,32 +1460,38 @@ public void validate(PipelineOptions options) { // Check for source table presence for early failure notification. BigQueryHelpers.verifyDatasetPresence(datasetService, table.get()); BigQueryHelpers.verifyTablePresence(datasetService, table.get()); - } else if (getQuery() != null) { + } else if (query != null) { checkArgument( - getQuery().isAccessible(), "Cannot call validate if query is dynamically set."); + query.isAccessible(), "Cannot call validate if query is dynamically set."); JobService jobService = getBigQueryServices().getJobService(bqOptions); + // JobConfigurationQuery accepts null for both of these, but the generated API client + // is not annotated. + @SuppressWarnings("nullness") + JobConfigurationQuery queryConfig = + new JobConfigurationQuery() + .setQuery(query.get()) + .setFlattenResults(getFlattenResults()) + .setUseLegacySql(getUseLegacySql()); try { jobService.dryRunQuery( bqOptions.getBigQueryProject() == null ? bqOptions.getProject() : bqOptions.getBigQueryProject(), - new JobConfigurationQuery() - .setQuery(getQuery().get()) - .setFlattenResults(getFlattenResults()) - .setUseLegacySql(getUseLegacySql()), + queryConfig, getQueryLocation()); } catch (Exception e) { throw new IllegalArgumentException( String.format( "Validation of query \"%1$s\" failed. If the query depends on an earlier stage of the" + " pipeline, This validation can be disabled using #withoutValidation.", - getQuery().get()), + query.get()), e); } // If the user provided a temp dataset, check if the dataset exists before launching the // query - if (getQueryTempDataset() != null) { + String queryTempDataset = getQueryTempDataset(); + if (queryTempDataset != null) { // The temp table is only used for dataset and project id validation, not for table // name // validation @@ -1488,7 +1505,7 @@ public void validate(PipelineOptions options) { TableReference tempTable = new TableReference() .setProjectId(project) - .setDatasetId(getQueryTempDataset()) + .setDatasetId(queryTempDataset) .setTableId("dummy table"); BigQueryHelpers.verifyDatasetPresence(datasetService, tempTable); } @@ -1549,8 +1566,9 @@ public PCollection expand(PBegin input) { } BigQueryUtils.SchemaConversionOptions.Builder builder = BigQueryUtils.SchemaConversionOptions.builder(); - if (getDirectReadPicosTimestampPrecision() != null) { - builder.setPicosecondTimestampMapping(getDirectReadPicosTimestampPrecision()); + TimestampPrecision picosPrecision = getDirectReadPicosTimestampPrecision(); + if (picosPrecision != null) { + builder.setPicosecondTimestampMapping(picosPrecision); } beamSchema = BigQueryUtils.fromTableSchema(tableSchema, builder.build()); } @@ -1585,7 +1603,11 @@ public PCollection expand(PBegin input) { p.apply( org.apache.beam.sdk.io.Read.from( sourceDef.toSource( - staticJobUuid, coder, getDatumReaderFactory(), getUseAvroLogicalTypes()))); + staticJobUuid, + coder, + checkStateNotNull( + getDatumReaderFactory(), "A readerDatumFactory is required"), + getUseAvroLogicalTypes()))); } else { // Create a singleton job ID token at execution time. jobIdTokenCollection = @@ -1615,7 +1637,9 @@ public void processElement(ProcessContext c) throws Exception { sourceDef.toSource( jobUuid, coder, - getDatumReaderFactory(), + checkStateNotNull( + getDatumReaderFactory(), + "A readerDatumFactory is required"), getUseAvroLogicalTypes()); BigQueryOptions options = c.getPipelineOptions().as(BigQueryOptions.class); @@ -1651,7 +1675,9 @@ public void processElement(ProcessContext c) throws Exception { sourceDef.toSource( jobUuid, coder, - getDatumReaderFactory(), + checkStateNotNull( + getDatumReaderFactory(), + "A readerDatumFactory is required"), getUseAvroLogicalTypes()); List> sources = source.createSources( @@ -1710,15 +1736,20 @@ void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { if (beamSchema != null) { rows.setSchema( beamSchema, - getTypeDescriptor(), - getToBeamRowFn().apply(beamSchema), - getFromBeamRowFn().apply(beamSchema)); + checkStateNotNull(getTypeDescriptor()), + checkStateNotNull(getToBeamRowFn()).apply(beamSchema), + checkStateNotNull(getFromBeamRowFn()).apply(beamSchema)); } return rows; } private PCollection expandForDirectRead( - PBegin input, Coder outputCoder, Schema beamSchema, BigQueryOptions bqOptions) { + PBegin input, + Coder outputCoder, + @Nullable Schema beamSchema, + BigQueryOptions bqOptions) { + SerializableFunction parseFn = + checkStateNotNull(getParseFn(), "A parseFn is required"); ValueProvider tableProvider = getTableProvider(); Pipeline p = input.getPipeline(); if (tableProvider != null) { @@ -1734,7 +1765,7 @@ private PCollection expandForDirectRead( getFormat(), getSelectedFields(), getRowRestriction(), - getParseFn(), + parseFn, outputCoder, getBigQueryServices(), getProjectionPushdownApplied(), @@ -1742,9 +1773,9 @@ private PCollection expandForDirectRead( if (beamSchema != null) { rows.setSchema( beamSchema, - getTypeDescriptor(), - getToBeamRowFn().apply(beamSchema), - getFromBeamRowFn().apply(beamSchema)); + checkStateNotNull(getTypeDescriptor()), + checkStateNotNull(getToBeamRowFn()).apply(beamSchema), + checkStateNotNull(getFromBeamRowFn()).apply(beamSchema)); } return rows; } else { @@ -1756,7 +1787,7 @@ private PCollection expandForDirectRead( getFormat(), getSelectedFields(), getRowRestriction(), - getParseFn(), + parseFn, outputCoder, getBigQueryServices(), getProjectionPushdownApplied(), @@ -1784,7 +1815,7 @@ private PCollection expandForDirectRead( p.apply(Create.of(sources)) .apply( "Read Storage Table Source", - ParDo.of(new ReadTableSource(rowTag, getParseFn(), getBadRecordRouter())) + ParDo.of(new ReadTableSource(rowTag, parseFn, getBadRecordRouter())) .withOutputTags(rowTag, TupleTagList.of(BAD_RECORD_TAG))); getBadRecordErrorHandler() .addErrorCollection( @@ -1907,9 +1938,9 @@ void cleanup(ContextContainer c) throws Exception { if (beamSchema != null) { rows.setSchema( beamSchema, - getTypeDescriptor(), - getToBeamRowFn().apply(beamSchema), - getFromBeamRowFn().apply(beamSchema)); + checkStateNotNull(getTypeDescriptor()), + checkStateNotNull(getToBeamRowFn()).apply(beamSchema), + checkStateNotNull(getFromBeamRowFn()).apply(beamSchema)); } return rows.apply(new PassThroughThenCleanup<>(cleanupOperation, jobIdTokenView)); } @@ -1977,7 +2008,8 @@ public void processElement(ProcessContext c) throws Exception { // the destination table created to hold the results. BigQueryStorageQuerySource querySource = createStorageQuerySource(jobUuid, outputCoder); - Table queryResultTable = querySource.getTargetTable(options); + Table queryResultTable = + checkStateNotNull(querySource.getTargetTable(options)); // Create a read session without specifying a desired stream count and // let the BigQuery storage server pick the number of streams. @@ -2023,7 +2055,7 @@ private static class ErrorHandlingParseFn implements SerializableFunction { private final SerializableFunction parseFn; - private transient SchemaAndRecord schemaAndRecord = null; + private transient @Nullable SchemaAndRecord schemaAndRecord = null; private ErrorHandlingParseFn(SerializableFunction parseFn) { this.parseFn = parseFn; @@ -2040,7 +2072,7 @@ public T apply(SchemaAndRecord input) { } public SchemaAndRecord getSchemaAndRecord() { - return schemaAndRecord; + return checkStateNotNull(schemaAndRecord, "apply() has not been called yet"); } } @@ -2075,7 +2107,8 @@ public void processElement( ReadStream readStream = c.element(); ErrorHandlingParseFn errorHandlingParseFn = - new ErrorHandlingParseFn(getParseFn()); + new ErrorHandlingParseFn( + checkStateNotNull(getParseFn(), "A parseFn is required")); BigQueryStorageStreamSource streamSource = BigQueryStorageStreamSource.create( @@ -2118,13 +2151,13 @@ public static void readSource( // the same order. BoundedSource.BoundedReader reader = streamSource.createReader(options); - T current = null; - boolean hasCurrent = false; + @Nullable T current = null; try { if (reader.start()) { - current = - java.util.Objects.requireNonNull(reader.getCurrent(), "Reader returned null element"); - hasCurrent = true; + current = reader.getCurrent(); + if (current == null) { + throw new IllegalStateException("Reader returned null element"); + } } else { return; } @@ -2137,17 +2170,15 @@ public static void readSource( (Exception) e.getCause(), "Unable to parse record reading from BigQuery"); } - if (hasCurrent) { + if (current != null) { outputReceiver.get(rowTag).output(current); } while (true) { current = null; - hasCurrent = false; try { if (reader.advance()) { current = reader.getCurrent(); - hasCurrent = true; } else { return; } @@ -2160,7 +2191,7 @@ public static void readSource( (Exception) e.getCause(), "Unable to parse record reading from BigQuery"); } - if (hasCurrent) { + if (current != null) { outputReceiver.get(rowTag).output(current); } } @@ -2203,9 +2234,10 @@ private void ensureFromNotCalledYet() { /** See {@link Read#getTableProvider()}. */ public @Nullable ValueProvider getTableProvider() { - return getJsonTableRef() == null + ValueProvider jsonTableRef = getJsonTableRef(); + return jsonTableRef == null ? null - : NestedValueProvider.of(getJsonTableRef(), new JsonTableRefToTableRef()); + : NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableRef()); } /** See {@link Read#getTable()}. */ @@ -2619,7 +2651,7 @@ SerializableFunction toSerializableFunction() { } private static class FormatProto extends TableRowFormatFunction { - transient TableRowToStorageApiProto.SchemaInformation inferredSchemaInformation; + transient TableRowToStorageApiProto.@Nullable SchemaInformation inferredSchemaInformation; final Class protoMessageClass; FormatProto(Class protoMessageClass) { @@ -2629,11 +2661,12 @@ private static class FormatProto extends TableRowFormatFuncti TableRowToStorageApiProto.SchemaInformation inferSchemaInformation() { try { if (inferredSchemaInformation == null) { + // Method.invoke takes a null receiver for a static method; that is not expressible + // against the JDK's annotations. + @SuppressWarnings("nullness") + Object rawDescriptor = protoMessageClass.getMethod("getDescriptor").invoke(null); Descriptors.Descriptor descriptor = - (Descriptors.Descriptor) - org.apache.beam.sdk.util.Preconditions.checkStateNotNull( - protoMessageClass.getMethod("getDescriptor")) - .invoke(null); + (Descriptors.Descriptor) checkStateNotNull(rawDescriptor); Descriptors.Descriptor convertedDescriptor = TableRowToStorageApiProto.wrapDescriptorProto( ProtoSchemaConverter.convert(descriptor).getProtoDescriptor()); @@ -2655,7 +2688,8 @@ static FormatProto fromClass(Class protoMessageClass) } @Override - public TableRow apply(TableRowToStorageApiProto.SchemaInformation schemaInformation, T input) { + public TableRow apply( + TableRowToStorageApiProto.@Nullable SchemaInformation schemaInformation, T input) { TableRowToStorageApiProto.SchemaInformation localSchemaInformation = schemaInformation != null ? schemaInformation : inferSchemaInformation(); return TableRowToStorageApiProto.tableRowFromMessage( @@ -2711,7 +2745,7 @@ public enum Method { abstract @Nullable ValueProvider getJsonTableRef(); - abstract @Nullable SerializableFunction, TableDestination> + abstract @Nullable SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> getTableFunction(); abstract @Nullable TableRowFormatFunction getFormatFunction(); @@ -2829,7 +2863,7 @@ abstract static class Builder { abstract Builder setJsonTableRef(ValueProvider jsonTableRef); abstract Builder setTableFunction( - SerializableFunction, TableDestination> tableFunction); + SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction); abstract Builder setFormatFunction(TableRowFormatFunction formatFunction); @@ -2928,7 +2962,7 @@ abstract Builder setDefaultMissingValueInterpretation( abstract Builder setDirectWriteProtos(boolean direct); abstract Builder setDeterministicRecordIdFn( - SerializableFunction toUniqueIdFunction); + @Nullable SerializableFunction toUniqueIdFunction); abstract Builder setWriteTempDataset(String writeTempDataset); @@ -3063,7 +3097,7 @@ public Write to(ValueProvider tableSpec) { * encoded and decoded. */ public Write to( - SerializableFunction, TableDestination> tableFunction) { + SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction) { checkArgument(tableFunction != null, "tableFunction can not be null"); return toBuilder().setTableFunction(tableFunction).build(); } @@ -3707,12 +3741,16 @@ public Write withErrorHandler(ErrorHandler errorHandler) { } @Override - public void validate(PipelineOptions pipelineOptions) { - BigQueryOptions options = pipelineOptions.as(BigQueryOptions.class); + public void validate(@Nullable PipelineOptions maybeOptions) { + BigQueryOptions options = + org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull(maybeOptions) + .as(BigQueryOptions.class); // The user specified a table. - if (getJsonTableRef() != null && getJsonTableRef().isAccessible() && getValidate()) { - TableReference table = getTableWithDefaultProject(options).get(); + ValueProvider jsonTableRef = getJsonTableRef(); + if (jsonTableRef != null && jsonTableRef.isAccessible() && getValidate()) { + TableReference table = + checkStateNotNull(getTableWithDefaultProject(options), "table is required").get(); try (DatasetService datasetService = getBigQueryServices().getDatasetService(options)) { // Check for destination table presence and emptiness for early failure notification. // Note that a presence check can fail when the table or dataset is created by an earlier @@ -3751,12 +3789,14 @@ private Write.Method resolveMethod(PCollection input) { : Write.Method.FILE_LOADS; } - private Duration getStorageApiTriggeringFrequency(BigQueryOptions options) { - if (getTriggeringFrequency() != null) { - return getTriggeringFrequency(); + private @Nullable Duration getStorageApiTriggeringFrequency(BigQueryOptions options) { + Duration triggeringFrequency = getTriggeringFrequency(); + if (triggeringFrequency != null) { + return triggeringFrequency; } - if (options.getStorageWriteApiTriggeringFrequencySec() != null) { - return Duration.standardSeconds(options.getStorageWriteApiTriggeringFrequencySec()); + Integer frequencySec = options.getStorageWriteApiTriggeringFrequencySec(); + if (frequencySec != null) { + return Duration.standardSeconds(frequencySec); } return null; } @@ -3882,10 +3922,11 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { if (getWriteDisposition() == WriteDisposition.WRITE_TRUNCATE) { LOG.error("The Storage API sink does not support the WRITE_TRUNCATE write disposition."); } - if (getBigLakeConfiguration() != null) { + Map bigLakeConfiguration = getBigLakeConfiguration(); + if (bigLakeConfiguration != null) { checkArgument( Arrays.stream(new String[] {CONNECTION_ID, STORAGE_URI}) - .allMatch(getBigLakeConfiguration()::containsKey), + .allMatch(bigLakeConfiguration::containsKey), String.format( "bigLakeConfiguration must contain keys '%s' and '%s'", CONNECTION_ID, STORAGE_URI)); @@ -3927,44 +3968,49 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { DynamicDestinations dynamicDestinations = getDynamicDestinations(); if (dynamicDestinations == null) { - if (getJsonTableRef() != null) { - dynamicDestinations = + ValueProvider jsonTableRef = getJsonTableRef(); + SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction = + getTableFunction(); + ValueProvider jsonClustering = getJsonClustering(); + DynamicDestinations tableDestinations; + if (jsonTableRef != null) { + tableDestinations = DynamicDestinationsHelpers.ConstantTableDestinations.fromJsonTableRef( - getJsonTableRef(), getTableDescription(), getJsonClustering() != null); - } else if (getTableFunction() != null) { - dynamicDestinations = - new TableFunctionDestinations<>(getTableFunction(), getJsonClustering() != null); + jsonTableRef, getTableDescription(), jsonClustering != null); + } else { + // Checked above: exactly one of jsonTableRef, tableFunction or dynamicDestinations is + // set. + tableDestinations = + new TableFunctionDestinations<>( + checkStateNotNull(tableFunction), jsonClustering != null); } // Wrap with a DynamicDestinations class that will provide a schema. There might be no // schema provided if the create disposition is CREATE_NEVER. - if (getJsonSchema() != null) { - dynamicDestinations = - new ConstantSchemaDestinations<>( - (DynamicDestinations) dynamicDestinations, getJsonSchema()); - } else if (getSchemaFromView() != null) { - dynamicDestinations = - new SchemaFromViewDestinations<>( - (DynamicDestinations) dynamicDestinations, - getSchemaFromView()); + ValueProvider jsonSchema = getJsonSchema(); + PCollectionView> schemaFromView = getSchemaFromView(); + if (jsonSchema != null) { + tableDestinations = new ConstantSchemaDestinations<>(tableDestinations, jsonSchema); + } else if (schemaFromView != null) { + tableDestinations = new SchemaFromViewDestinations<>(tableDestinations, schemaFromView); } // Wrap with a DynamicDestinations class that will provide the proper TimePartitioning. - if (getJsonTimePartitioning() != null || (getJsonClustering() != null)) { - dynamicDestinations = + ValueProvider jsonTimePartitioning = getJsonTimePartitioning(); + if (jsonTimePartitioning != null || jsonClustering != null) { + tableDestinations = new ConstantTimePartitioningClusteringDestinations<>( - (DynamicDestinations) dynamicDestinations, - getJsonTimePartitioning(), - getJsonClustering()); + tableDestinations, jsonTimePartitioning, jsonClustering); } - if (getPrimaryKey() != null) { - dynamicDestinations = + List primaryKey = getPrimaryKey(); + if (primaryKey != null) { + tableDestinations = new DynamicDestinationsHelpers.ConstantTableConstraintsDestinations<>( - (DynamicDestinations) dynamicDestinations, + tableDestinations, new TableConstraints() - .setPrimaryKey( - new TableConstraints.PrimaryKey().setColumns(getPrimaryKey()))); + .setPrimaryKey(new TableConstraints.PrimaryKey().setColumns(primaryKey))); } + dynamicDestinations = tableDestinations; } return expandTyped(input, dynamicDestinations); } @@ -4011,12 +4057,12 @@ private WriteResult expandTyped( } else if (writeProtoClass != null) { if (!hasSchema) { try { - @SuppressWarnings({"unchecked", "nullness"}) + // Method.invoke takes a null receiver for a static method; that is not expressible + // against the JDK's annotations. + @SuppressWarnings("nullness") + Object rawDescriptor = writeProtoClass.getMethod("getDescriptor").invoke(null); Descriptors.Descriptor descriptor = - (Descriptors.Descriptor) - org.apache.beam.sdk.util.Preconditions.checkStateNotNull( - writeProtoClass.getMethod("getDescriptor")) - .invoke(null); + (Descriptors.Descriptor) checkStateNotNull(rawDescriptor); TableSchema tableSchema = TableRowToStorageApiProto.protoSchemaToTableSchema( TableRowToStorageApiProto.tableSchemaFromDescriptor(descriptor)); @@ -4077,15 +4123,16 @@ private WriteResult expandTyped( checkArgument( avroRowWriterFactory == null, "When using a formatFunction, the AvroRowWriterFactory should be null"); - checkArgument( - formatFunction != null, - "A function must be provided to convert the input type into a TableRow or " - + "GenericRecord. Use BigQueryIO.Write.withFormatFunction or " - + "BigQueryIO.Write.withAvroFormatFunction to provide a formatting function. " - + "A format function is not required if Beam schemas are used."); - rowWriterFactory = - RowWriterFactory.tableRows(formatFunction, formatRecordOnFailureFunction); + RowWriterFactory.tableRows( + checkArgumentNotNull( + formatFunction, + "A function must be provided to convert the input type into a TableRow or " + + "GenericRecord. Use BigQueryIO.Write.withFormatFunction or " + + "BigQueryIO.Write.withAvroFormatFunction to provide a formatting " + + "function. A format function is not required if Beam schemas are " + + "used."), + formatRecordOnFailureFunction); } PCollection> rowsWithDestination = @@ -4176,8 +4223,9 @@ private WriteResult continueExpandTyped( // Batch load handles wrapped json string value differently than the other methods. Raise a // warning when applies. - if (getJsonSchema() != null && getJsonSchema().isAccessible()) { - JsonElement schema = JsonParser.parseString(getJsonSchema().get()); + ValueProvider jsonSchema = getJsonSchema(); + if (jsonSchema != null && jsonSchema.isAccessible()) { + JsonElement schema = JsonParser.parseString(jsonSchema.get()); if (!schema.getAsJsonObject().keySet().isEmpty() && hasJsonTypeInSchema(schema)) { if (rowWriterFactory.getOutputType() == OutputType.JsonTableRow) { LOG.warn( @@ -4218,11 +4266,13 @@ private WriteResult continueExpandTyped( if (getSchemaUpdateOptions() != null) { batchLoads.setSchemaUpdateOptions(getSchemaUpdateOptions()); } - if (getMaxFilesPerBundle() != null) { - batchLoads.setMaxNumWritersPerBundle(getMaxFilesPerBundle()); + Integer maxFilesPerBundle = getMaxFilesPerBundle(); + if (maxFilesPerBundle != null) { + batchLoads.setMaxNumWritersPerBundle(maxFilesPerBundle); } - if (getMaxFileSize() != null) { - batchLoads.setMaxFileSize(getMaxFileSize()); + Long maxFileSize = getMaxFileSize(); + if (maxFileSize != null) { + batchLoads.setMaxFileSize(maxFileSize); } batchLoads.setMaxFilesPerPartition(getMaxFilesPerPartition()); batchLoads.setMaxBytesPerPartition(getMaxBytesPerPartition()); @@ -4240,8 +4290,8 @@ private WriteResult continueExpandTyped( } return input.apply(batchLoads); } else if (method == Method.STORAGE_WRITE_API || method == Method.STORAGE_API_AT_LEAST_ONCE) { - boolean useSchemaUpdate = - getSchemaUpdateOptions() != null && !getSchemaUpdateOptions().isEmpty(); + Set schemaUpdateOptions = getSchemaUpdateOptions(); + boolean useSchemaUpdate = schemaUpdateOptions != null && !schemaUpdateOptions.isEmpty(); if (useSchemaUpdate) { checkArgument( !getAutoSchemaUpdate() && !getIgnoreUnknownValues(), @@ -4258,11 +4308,12 @@ private WriteResult continueExpandTyped( storageApiDynamicDestinations = new StorageApiDynamicDestinationsBeamRow<>( dynamicDestinations, - elementSchema, - elementToRowFunction, + checkStateNotNull(elementSchema), + checkStateNotNull(elementToRowFunction), getFormatRecordOnFailureFunction(), getRowMutationInformationFn() != null); } else if (getWriteProtosClass() != null && getDirectWriteProtos()) { + Class writeProtosClass = checkStateNotNull(getWriteProtosClass()); checkArgument( !useSchemaUpdate, "SchemaUpdateOptions are not supported when writing protos"); @@ -4287,9 +4338,7 @@ private WriteResult continueExpandTyped( storageApiDynamicDestinations = (StorageApiDynamicDestinations) new StorageApiDynamicDestinationsProto( - dynamicDestinations, - getWriteProtosClass(), - getFormatRecordOnFailureFunction()); + dynamicDestinations, writeProtosClass, getFormatRecordOnFailureFunction()); } else if (getAvroRowWriterFactory() != null) { checkArgument( !useSchemaUpdate, "SchemaUpdateOptions are not supported when writing avros"); @@ -4329,9 +4378,7 @@ private WriteResult continueExpandTyped( getCreateDisposition(), getIgnoreUnknownValues(), getAutoSchemaUpdate(), - getSchemaUpdateOptions() == null - ? Collections.emptySet() - : getSchemaUpdateOptions()); + schemaUpdateOptions == null ? Collections.emptySet() : schemaUpdateOptions); } int numShards = getStorageApiNumStreams(bqOptions); @@ -4362,7 +4409,7 @@ private WriteResult continueExpandTyped( getBigLakeConfiguration(), getBadRecordRouter(), getBadRecordErrorHandler(), - !getSchemaUpdateOptions().isEmpty()); + useSchemaUpdate); return input.apply("StorageApiLoads", storageApiLoads); } else { throw new RuntimeException("Unexpected write method " + method); @@ -4404,9 +4451,11 @@ public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("schema", "Custom Schema Function").withLabel("Table Schema")); } - if (getTableFunction() != null) { + SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction = + getTableFunction(); + if (tableFunction != null) { builder.add( - DisplayData.item("tableFn", getTableFunction().getClass()) + DisplayData.item("tableFn", tableFunction.getClass()) .withLabel("Table Reference Function")); } @@ -4463,9 +4512,10 @@ ValueProvider getTableWithDefaultProject(BigQueryOptions bqOptio /** Returns the table reference, or {@code null}. */ public @Nullable ValueProvider getTable() { - return getJsonTableRef() == null + ValueProvider jsonTableRef = getJsonTableRef(); + return jsonTableRef == null ? null - : NestedValueProvider.of(getJsonTableRef(), new JsonTableRefToTableRef()); + : NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableRef()); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index dd59939726bf..ee29189a0652 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.sdk.util.construction.TransformUpgrader.fromByteArray; import static org.apache.beam.sdk.util.construction.TransformUpgrader.toByteArray; @@ -48,6 +49,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.RowWriterFactory.AvroRowWriterFactory; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; import org.apache.beam.sdk.runners.AppliedPTransform; import org.apache.beam.sdk.schemas.Schema; @@ -73,7 +75,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -@SuppressWarnings({"rawtypes", "nullness", "AutoValueSubclassLeaked"}) +@SuppressWarnings({"rawtypes", "AutoValueSubclassLeaked"}) public class BigQueryIOTranslation { private static final Logger LOG = LoggerFactory.getLogger(BigQueryIOTranslation.class); @@ -119,9 +121,18 @@ public String getUrn() { return BIGQUERY_READ_TRANSFORM_URN; } + // The 'application' parameter below carries a narrowly scoped "nullness" suppression: the + // Checker Framework cannot verify that the captured wildcards of + // AppliedPTransform> satisfy the F-bounded declaration + // 'TransformT extends PTransform'. The signature is dictated by the + // TransformPayloadTranslator interface so it cannot be reformulated here, and no annotation on + // the type arguments can express the missing capture constraint. + // See https://github.com/typetools/checker-framework/issues/3791, which is also the reason + // PubSubPayloadTranslation is listed in 'classesTriggerCheckerBugs' in build.gradle. @Override public RunnerApi.@Nullable FunctionSpec translate( - AppliedPTransform> application, SdkComponents components) + @SuppressWarnings("nullness") AppliedPTransform> application, + SdkComponents components) throws IOException { // Setting an empty payload since BigQuery transform payload is not actually used by runners // currently. @@ -133,37 +144,54 @@ public String getUrn() { public Row toConfigRow(TypedRead transform) { Map fieldValues = new HashMap<>(); - if (transform.getJsonTableRef() != null) { - fieldValues.put("json_table_ref", transform.getJsonTableRef().get()); + ValueProvider jsonTableRef = transform.getJsonTableRef(); + if (jsonTableRef != null) { + fieldValues.put("json_table_ref", jsonTableRef.get()); } - if (transform.getQuery() != null) { - fieldValues.put("query", transform.getQuery().get()); + ValueProvider query = transform.getQuery(); + if (query != null) { + fieldValues.put("query", query.get()); } fieldValues.put("validate", transform.getValidate()); - fieldValues.put("flatten_results", transform.getFlattenResults()); - fieldValues.put("use_legacy_sql", transform.getUseLegacySql()); + // 'flatten_results' and 'use_legacy_sql' are nullable properties of the transform and the + // corresponding schema fields are nullable. Omitting the field value is equivalent to + // setting it to null. + Boolean flattenResults = transform.getFlattenResults(); + if (flattenResults != null) { + fieldValues.put("flatten_results", flattenResults); + } + Boolean useLegacySql = transform.getUseLegacySql(); + if (useLegacySql != null) { + fieldValues.put("use_legacy_sql", useLegacySql); + } fieldValues.put("with_template_compatibility", transform.getWithTemplateCompatibility()); if (transform.getBigQueryServices() != null) { fieldValues.put("bigquery_services", toByteArray(transform.getBigQueryServices())); } - if (transform.getParseFn() != null) { - fieldValues.put("parse_fn", toByteArray(transform.getParseFn())); + Object parseFn = transform.getParseFn(); + if (parseFn != null) { + fieldValues.put("parse_fn", toByteArray(parseFn)); } - if (transform.getDatumReaderFactory() != null) { - fieldValues.put("datum_reader_factory", toByteArray(transform.getDatumReaderFactory())); + Object datumReaderFactory = transform.getDatumReaderFactory(); + if (datumReaderFactory != null) { + fieldValues.put("datum_reader_factory", toByteArray(datumReaderFactory)); } - if (transform.getQueryPriority() != null) { - fieldValues.put("query_priority", toByteArray(transform.getQueryPriority())); + QueryPriority queryPriority = transform.getQueryPriority(); + if (queryPriority != null) { + fieldValues.put("query_priority", toByteArray(queryPriority)); } - if (transform.getQueryLocation() != null) { - fieldValues.put("query_location", transform.getQueryLocation()); + String queryLocation = transform.getQueryLocation(); + if (queryLocation != null) { + fieldValues.put("query_location", queryLocation); } - if (transform.getQueryTempDataset() != null) { - fieldValues.put("query_temp_dataset", transform.getQueryTempDataset()); + String queryTempDataset = transform.getQueryTempDataset(); + if (queryTempDataset != null) { + fieldValues.put("query_temp_dataset", queryTempDataset); } - if (transform.getQueryTempProject() != null) { - fieldValues.put("query_temp_project", transform.getQueryTempProject()); + String queryTempProject = transform.getQueryTempProject(); + if (queryTempProject != null) { + fieldValues.put("query_temp_project", queryTempProject); } if (transform.getMethod() != null) { fieldValues.put("method", toByteArray(transform.getMethod())); @@ -171,32 +199,41 @@ public Row toConfigRow(TypedRead transform) { if (transform.getFormat() != null) { fieldValues.put("format", toByteArray(transform.getFormat())); } - if (transform.getSelectedFields() != null && !transform.getSelectedFields().get().isEmpty()) { - fieldValues.put("selected_fields", transform.getSelectedFields().get()); + ValueProvider> selectedFields = transform.getSelectedFields(); + if (selectedFields != null && !selectedFields.get().isEmpty()) { + fieldValues.put("selected_fields", selectedFields.get()); } - if (transform.getRowRestriction() != null) { - fieldValues.put("row_restriction", transform.getRowRestriction().get()); + ValueProvider rowRestriction = transform.getRowRestriction(); + if (rowRestriction != null) { + fieldValues.put("row_restriction", rowRestriction.get()); } - if (transform.getCoder() != null) { - fieldValues.put("coder", toByteArray(transform.getCoder())); + Coder coder = transform.getCoder(); + if (coder != null) { + fieldValues.put("coder", toByteArray(coder)); } - if (transform.getKmsKey() != null) { - fieldValues.put("kms_key", transform.getKmsKey()); + String kmsKey = transform.getKmsKey(); + if (kmsKey != null) { + fieldValues.put("kms_key", kmsKey); } - if (transform.getTypeDescriptor() != null) { - fieldValues.put("type_descriptor", toByteArray(transform.getTypeDescriptor())); + TypeDescriptor typeDescriptor = transform.getTypeDescriptor(); + if (typeDescriptor != null) { + fieldValues.put("type_descriptor", toByteArray(typeDescriptor)); } - if (transform.getToBeamRowFn() != null) { - fieldValues.put("to_beam_row_fn", toByteArray(transform.getToBeamRowFn())); + ToBeamRowFunction toBeamRowFn = transform.getToBeamRowFn(); + if (toBeamRowFn != null) { + fieldValues.put("to_beam_row_fn", toByteArray(toBeamRowFn)); } - if (transform.getFromBeamRowFn() != null) { - fieldValues.put("from_beam_row_fn", toByteArray(transform.getFromBeamRowFn())); + FromBeamRowFunction fromBeamRowFn = transform.getFromBeamRowFn(); + if (fromBeamRowFn != null) { + fieldValues.put("from_beam_row_fn", toByteArray(fromBeamRowFn)); } fieldValues.put("use_avro_logical_types", transform.getUseAvroLogicalTypes()); - if (transform.getDirectReadPicosTimestampPrecision() != null) { + TimestampPrecision directReadPicosTimestampPrecision = + transform.getDirectReadPicosTimestampPrecision(); + if (directReadPicosTimestampPrecision != null) { fieldValues.put( "direct_read_picos_timestamp_precision", - toByteArray(transform.getDirectReadPicosTimestampPrecision())); + toByteArray(directReadPicosTimestampPrecision)); } fieldValues.put("projection_pushdown_applied", transform.getProjectionPushdownApplied()); fieldValues.put("bad_record_router", toByteArray(transform.getBadRecordRouter())); @@ -354,9 +391,16 @@ public TypedRead fromConfigRow(Row configRow, PipelineOptions options) { builder.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER); builder.setBadRecordErrorHandler(new ErrorHandler.DefaultErrorHandler<>()); } else { - byte[] badRecordRouter = configRow.getBytes("bad_record_router"); + // 'toConfigRow' always populates these two fields, so they are expected to be present in + // rows produced by Beam 2.55.0 and later. + byte[] badRecordRouter = + checkStateNotNull( + configRow.getBytes("bad_record_router"), "Missing 'bad_record_router'"); builder.setBadRecordRouter((BadRecordRouter) fromByteArray(badRecordRouter)); - byte[] badRecordErrorHandler = configRow.getBytes("bad_record_error_handler"); + byte[] badRecordErrorHandler = + checkStateNotNull( + configRow.getBytes("bad_record_error_handler"), + "Missing 'bad_record_error_handler'"); builder.setBadRecordErrorHandler( (ErrorHandler) fromByteArray(badRecordErrorHandler)); } @@ -448,9 +492,12 @@ public String getUrn() { return BIGQUERY_WRITE_TRANSFORM_URN; } + // See the comment on BigQueryIOReadTranslator#translate for why the 'application' parameter + // below carries a narrowly scoped "nullness" suppression. @Override public @Nullable FunctionSpec translate( - AppliedPTransform> application, SdkComponents components) + @SuppressWarnings("nullness") AppliedPTransform> application, + SdkComponents components) throws IOException { // Setting an empty payload since BigQuery transform payload is not actually used by runners // currently. @@ -463,27 +510,31 @@ public Row toConfigRow(Write transform) { Map fieldValues = new HashMap<>(); - if (transform.getJsonTableRef() != null) { - fieldValues.put("json_table_ref", transform.getJsonTableRef().get()); + ValueProvider jsonTableRef = transform.getJsonTableRef(); + if (jsonTableRef != null) { + fieldValues.put("json_table_ref", jsonTableRef.get()); } - if (transform.getTableFunction() != null) { - fieldValues.put("table_function", toByteArray(transform.getTableFunction())); + Object tableFunction = transform.getTableFunction(); + if (tableFunction != null) { + fieldValues.put("table_function", toByteArray(tableFunction)); } - if (transform.getFormatFunction() != null) { - fieldValues.put("format_function", toByteArray(transform.getFormatFunction())); + Object formatFunction = transform.getFormatFunction(); + if (formatFunction != null) { + fieldValues.put("format_function", toByteArray(formatFunction)); } - if (transform.getFormatRecordOnFailureFunction() != null) { + Object formatRecordOnFailureFunction = transform.getFormatRecordOnFailureFunction(); + if (formatRecordOnFailureFunction != null) { fieldValues.put( - "format_record_on_failure_function", - toByteArray(transform.getFormatRecordOnFailureFunction())); + "format_record_on_failure_function", toByteArray(formatRecordOnFailureFunction)); } - if (transform.getAvroRowWriterFactory() != null) { - fieldValues.put( - "avro_row_writer_factory", toByteArray(transform.getAvroRowWriterFactory())); + Object avroRowWriterFactory = transform.getAvroRowWriterFactory(); + if (avroRowWriterFactory != null) { + fieldValues.put("avro_row_writer_factory", toByteArray(avroRowWriterFactory)); } fieldValues.put("use_avro_logical_types", transform.getUseAvroLogicalTypes()); - if (transform.getDynamicDestinations() != null) { - fieldValues.put("dynamic_destinations", toByteArray(transform.getDynamicDestinations())); + Object dynamicDestinations = transform.getDynamicDestinations(); + if (dynamicDestinations != null) { + fieldValues.put("dynamic_destinations", toByteArray(dynamicDestinations)); } if (transform.getSchemaFromView() != null) { // Property 'getSchemaFromView' cannot be used in a portable way across pipelines since it @@ -493,15 +544,17 @@ public Row toConfigRow(Write transform) { + "portable row based config due to 'withSchemaFromView' property being set. Please " + "retry without setting this property when configuring your transform"); } - if (transform.getJsonSchema() != null) { - fieldValues.put("json_schema", transform.getJsonSchema().get()); + ValueProvider jsonSchema = transform.getJsonSchema(); + if (jsonSchema != null) { + fieldValues.put("json_schema", jsonSchema.get()); } - if (transform.getJsonTimePartitioning() != null) { - fieldValues.put( - "json_time_partitioning", toByteArray(transform.getJsonTimePartitioning().get())); + ValueProvider jsonTimePartitioning = transform.getJsonTimePartitioning(); + if (jsonTimePartitioning != null) { + fieldValues.put("json_time_partitioning", toByteArray(jsonTimePartitioning.get())); } - if (transform.getJsonClustering() != null) { - fieldValues.put("clustering", transform.getJsonClustering().get()); + ValueProvider jsonClustering = transform.getJsonClustering(); + if (jsonClustering != null) { + fieldValues.put("clustering", jsonClustering.get()); } if (transform.getCreateDisposition() != null) { fieldValues.put("create_disposition", toByteArray(transform.getCreateDisposition())); @@ -517,21 +570,25 @@ public Row toConfigRow(Write transform) { .collect(Collectors.toList()); fieldValues.put("schema_update_options", schemUpdateOptionsData); } - if (transform.getTableDescription() != null) { - fieldValues.put("table_description", transform.getTableDescription()); + String tableDescription = transform.getTableDescription(); + if (tableDescription != null) { + fieldValues.put("table_description", tableDescription); } - if (transform.getBigLakeConfiguration() != null) { - fieldValues.put("biglake_configuration", transform.getBigLakeConfiguration()); + Map bigLakeConfiguration = transform.getBigLakeConfiguration(); + if (bigLakeConfiguration != null) { + fieldValues.put("biglake_configuration", bigLakeConfiguration); } fieldValues.put("validate", transform.getValidate()); if (transform.getBigQueryServices() != null) { fieldValues.put("bigquery_services", toByteArray(transform.getBigQueryServices())); } - if (transform.getMaxFilesPerBundle() != null) { - fieldValues.put("max_files_per_bundle", transform.getMaxFilesPerBundle()); + Integer maxFilesPerBundle = transform.getMaxFilesPerBundle(); + if (maxFilesPerBundle != null) { + fieldValues.put("max_files_per_bundle", maxFilesPerBundle); } - if (transform.getMaxFileSize() != null) { - fieldValues.put("max_file_size", transform.getMaxFileSize()); + Long maxFileSize = transform.getMaxFileSize(); + if (maxFileSize != null) { + fieldValues.put("max_file_size", maxFileSize); } fieldValues.put("num_file_shards", transform.getNumFileShards()); fieldValues.put("num_storage_write_api_streams", transform.getNumStorageWriteApiStreams()); @@ -543,23 +600,24 @@ public Row toConfigRow(Write transform) { toByteArray(transform.getPropagateSuccessfulStorageApiWritesPredicate())); fieldValues.put("max_files_per_partition", transform.getMaxFilesPerPartition()); fieldValues.put("max_bytes_per_partition", transform.getMaxBytesPerPartition()); - if (transform.getTriggeringFrequency() != null) { - fieldValues.put( - "triggering_frequency", - Duration.ofMillis(transform.getTriggeringFrequency().getMillis())); + org.joda.time.Duration triggeringFrequency = transform.getTriggeringFrequency(); + if (triggeringFrequency != null) { + fieldValues.put("triggering_frequency", Duration.ofMillis(triggeringFrequency.getMillis())); } if (transform.getMethod() != null) { fieldValues.put("method", toByteArray(transform.getMethod())); } - if (transform.getLoadJobProjectId() != null) { - fieldValues.put("load_job_project_id", transform.getLoadJobProjectId()); + ValueProvider loadJobProjectId = transform.getLoadJobProjectId(); + if (loadJobProjectId != null) { + fieldValues.put("load_job_project_id", loadJobProjectId); } - if (transform.getFailedInsertRetryPolicy() != null) { - fieldValues.put( - "failed_insert_retry_policy", toByteArray(transform.getFailedInsertRetryPolicy())); + InsertRetryPolicy failedInsertRetryPolicy = transform.getFailedInsertRetryPolicy(); + if (failedInsertRetryPolicy != null) { + fieldValues.put("failed_insert_retry_policy", toByteArray(failedInsertRetryPolicy)); } - if (transform.getCustomGcsTempLocation() != null) { - fieldValues.put("custom_gcs_temp_location", transform.getCustomGcsTempLocation().get()); + ValueProvider customGcsTempLocation = transform.getCustomGcsTempLocation(); + if (customGcsTempLocation != null) { + fieldValues.put("custom_gcs_temp_location", customGcsTempLocation.get()); } fieldValues.put("extended_error_info", transform.getExtendedErrorInfo()); fieldValues.put("skip_invalid_rows", transform.getSkipInvalidRows()); @@ -567,11 +625,13 @@ public Row toConfigRow(Write transform) { fieldValues.put("ignore_insert_ids", transform.getIgnoreInsertIds()); fieldValues.put("max_retry_jobs", transform.getMaxRetryJobs()); fieldValues.put("propagate_successful", transform.getPropagateSuccessful()); - if (transform.getKmsKey() != null) { - fieldValues.put("kms_key", transform.getKmsKey()); + String kmsKey = transform.getKmsKey(); + if (kmsKey != null) { + fieldValues.put("kms_key", kmsKey); } - if (transform.getPrimaryKey() != null) { - fieldValues.put("primary_key", transform.getPrimaryKey()); + List primaryKey = transform.getPrimaryKey(); + if (primaryKey != null) { + fieldValues.put("primary_key", primaryKey); } if (transform.getDefaultMissingValueInterpretation() != null) { fieldValues.put( @@ -582,20 +642,22 @@ public Row toConfigRow(Write transform) { fieldValues.put("use_beam_schema", transform.getUseBeamSchema()); fieldValues.put("auto_sharding", transform.getAutoSharding()); fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate()); - if (transform.getWriteProtosClass() != null) { - fieldValues.put("write_protos_class", toByteArray(transform.getWriteProtosClass())); + Class writeProtosClass = transform.getWriteProtosClass(); + if (writeProtosClass != null) { + fieldValues.put("write_protos_class", toByteArray(writeProtosClass)); } fieldValues.put("direct_write_protos", transform.getDirectWriteProtos()); - if (transform.getDeterministicRecordIdFn() != null) { - fieldValues.put( - "deterministic_record_id_fn", toByteArray(transform.getDeterministicRecordIdFn())); + Object deterministicRecordIdFn = transform.getDeterministicRecordIdFn(); + if (deterministicRecordIdFn != null) { + fieldValues.put("deterministic_record_id_fn", toByteArray(deterministicRecordIdFn)); } - if (transform.getWriteTempDataset() != null) { - fieldValues.put("write_temp_dataset", toByteArray(transform.getWriteTempDataset())); + String writeTempDataset = transform.getWriteTempDataset(); + if (writeTempDataset != null) { + fieldValues.put("write_temp_dataset", toByteArray(writeTempDataset)); } - if (transform.getRowMutationInformationFn() != null) { - fieldValues.put( - "row_mutation_information_fn", toByteArray(transform.getRowMutationInformationFn())); + Object rowMutationInformationFn = transform.getRowMutationInformationFn(); + if (rowMutationInformationFn != null) { + fieldValues.put("row_mutation_information_fn", toByteArray(rowMutationInformationFn)); } fieldValues.put("bad_record_router", toByteArray(transform.getBadRecordRouter())); fieldValues.put( @@ -872,9 +934,17 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { } byte[] writeProtosClasses = configRow.getBytes("write_protos_class"); if (writeProtosClasses != null) { + // NOTE: this deserializes 'defaultMissingValueInterpretationsBytes' rather than + // 'writeProtosClasses'. This looks like a copy/paste mistake, but it is preserved here + // as-is because changing it would alter the pipeline update compatibility behavior of + // already released Beam versions. builder = builder.setWriteProtosClass( - (Class) fromByteArray(defaultMissingValueInterpretationsBytes)); + (Class) + fromByteArray( + checkStateNotNull( + defaultMissingValueInterpretationsBytes, + "Missing 'default_missing_value_interpretation'"))); } Boolean directWriteProtos = configRow.getBoolean("direct_write_protos"); if (directWriteProtos != null) { @@ -904,9 +974,16 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { builder.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER); builder.setBadRecordErrorHandler(new ErrorHandler.DefaultErrorHandler<>()); } else { - byte[] badRecordRouter = configRow.getBytes("bad_record_router"); + // 'toConfigRow' always populates these two fields, so they are expected to be present in + // rows produced by Beam 2.55.0 and later. + byte[] badRecordRouter = + checkStateNotNull( + configRow.getBytes("bad_record_router"), "Missing 'bad_record_router'"); builder.setBadRecordRouter((BadRecordRouter) fromByteArray(badRecordRouter)); - byte[] badRecordErrorHandler = configRow.getBytes("bad_record_error_handler"); + byte[] badRecordErrorHandler = + checkStateNotNull( + configRow.getBytes("bad_record_error_handler"), + "Missing 'bad_record_error_handler'"); builder.setBadRecordErrorHandler( (ErrorHandler) fromByteArray(badRecordErrorHandler)); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java index 25f274d708b5..29a38f004fce 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryQuerySourceDef.java @@ -18,7 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryResourceNaming.createTempTableReference; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import com.google.api.services.bigquery.model.JobStatistics; import com.google.api.services.bigquery.model.TableReference; @@ -44,10 +44,10 @@ class BigQueryQuerySourceDef implements BigQuerySourceDef { private final Boolean flattenResults; private final Boolean useLegacySql; private final BigQueryIO.TypedRead.QueryPriority priority; - private final String location; - private final String tempDatasetId; - private final String tempProjectId; - private final String kmsKey; + private final @Nullable String location; + private final @Nullable String tempDatasetId; + private final @Nullable String tempProjectId; + private final @Nullable String kmsKey; private transient AtomicReference<@Nullable JobStatistics> dryRunJobStats; @@ -57,10 +57,10 @@ static BigQueryQuerySourceDef create( Boolean flattenResults, Boolean useLegacySql, BigQueryIO.TypedRead.QueryPriority priority, - String location, - String tempDatasetId, - String tempProjectId, - String kmsKey) { + @Nullable String location, + @Nullable String tempDatasetId, + @Nullable String tempProjectId, + @Nullable String kmsKey) { return new BigQueryQuerySourceDef( bqServices, query, @@ -79,13 +79,13 @@ private BigQueryQuerySourceDef( Boolean flattenResults, Boolean useLegacySql, BigQueryIO.TypedRead.QueryPriority priority, - String location, - String tempDatasetId, - String tempProjectId, - String kmsKey) { - this.query = checkNotNull(query, "query"); - this.flattenResults = checkNotNull(flattenResults, "flattenResults"); - this.useLegacySql = checkNotNull(useLegacySql, "useLegacySql"); + @Nullable String location, + @Nullable String tempDatasetId, + @Nullable String tempProjectId, + @Nullable String kmsKey) { + this.query = checkArgumentNotNull(query, "query"); + this.flattenResults = checkArgumentNotNull(flattenResults, "flattenResults"); + this.useLegacySql = checkArgumentNotNull(useLegacySql, "useLegacySql"); this.bqServices = bqServices; this.priority = priority; this.location = location; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySchemaIOProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySchemaIOProvider.java index 662fac3d7678..bb53ba376ac0 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySchemaIOProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySchemaIOProvider.java @@ -153,10 +153,11 @@ static class BigQuerySchemaIO implements SchemaIO, Serializable { * * @return null */ + // SchemaIO.schema() is declared non-null, but the BigQuery schema is only known once the + // table is read. The suppression cannot be narrowed below the method: a @Nullable local + // would just move the error to the return statement. @Override - @SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) - }) + @SuppressWarnings("nullness") public Schema schema() { return null; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java index 14765a65ff0b..36fb731ef158 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java @@ -17,7 +17,8 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import com.google.api.client.googleapis.json.GoogleJsonError; import com.google.api.client.googleapis.json.GoogleJsonResponseException; @@ -149,6 +150,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ListeningExecutorService; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.slf4j.Logger; @@ -158,10 +160,7 @@ * An implementation of {@link BigQueryServices} that actually communicates with the cloud BigQuery * service. */ -@SuppressWarnings({ - "nullness", // TODO(https://github.com/apache/beam/issues/20506) - "keyfor" -}) +@SuppressWarnings({"keyfor"}) public class BigQueryServicesImpl implements BigQueryServices { private static final Logger LOG = LoggerFactory.getLogger(BigQueryServicesImpl.class); @@ -455,7 +454,13 @@ public Job pollJob(JobReference jobRef, int maxAttempts) throws InterruptedExcep .withInitialBackoff(INITIAL_JOB_STATUS_POLL_BACKOFF) .withMaxBackoff(Duration.standardMinutes(1)) .backoff()); - return pollJob(jobRef, Sleeper.DEFAULT, backoff); + // JobService#pollJob is documented to return null when maxAttempts is exhausted, but the + // interface method is not annotated @Nullable (annotating it would ripple through every + // caller). Suppress narrowly here rather than at the class level. + @SuppressWarnings("nullness") + @NonNull + Job job = pollJob(jobRef, Sleeper.DEFAULT, backoff); + return job; } @VisibleForTesting @@ -506,6 +511,7 @@ private static String formatBqStatusCommand(String projectId, String jobId) { public JobStatistics dryRunQuery( String projectId, JobConfigurationQuery queryConfig, @Nullable String location) throws InterruptedException, IOException { + @SuppressWarnings("nullness") // setLocation is not annotated, but does accept nulls JobReference jobRef = new JobReference().setLocation(location).setProjectId(projectId); Job job = new Job() @@ -540,7 +546,13 @@ public JobStatistics dryRunQuery( */ @Override public Job getJob(JobReference jobRef) throws IOException, InterruptedException { - return getJob(jobRef, Sleeper.DEFAULT, createDefaultBackoff()); + // JobService#getJob is documented to return null when the job is not found, but the + // interface method is not annotated @Nullable (annotating it would ripple through every + // caller). Suppress narrowly here rather than at the class level. + @SuppressWarnings("nullness") + @NonNull + Job job = getJob(jobRef, Sleeper.DEFAULT, createDefaultBackoff()); + return job; } @VisibleForTesting @@ -1048,8 +1060,10 @@ public List call() throws Exception { */ if (!ApiErrorExtractor.INSTANCE.rateLimited(e) && !errorInfo.getReason().equals(QUOTA_EXCEEDED)) { + String exceptionMessage = e.getMessage(); if (ApiErrorExtractor.INSTANCE.badRequest(e) - && e.getMessage().contains(NO_ROWS_PRESENT)) { + && exceptionMessage != null + && exceptionMessage.contains(NO_ROWS_PRESENT)) { LOG.error( "No rows present in the request error likely caused by BigQuery Insert" + " timing out. Update BigQueryOptions.setHTTPWriteTimeout to be longer," @@ -1101,12 +1115,14 @@ long insertAll( boolean ignoreInsertIds, List> successfulRows) throws IOException, InterruptedException { - checkNotNull(ref, "ref"); + checkArgumentNotNull(ref, "ref"); + BoundedExecutorService executor = this.executor; if (executor == null) { - this.executor = + executor = new BoundedExecutorService( MoreExecutors.listeningDecorator(options.as(GcsOptions.class).getExecutorService()), options.as(BigQueryOptions.class).getInsertBundleParallelism()); + this.executor = executor; } if (insertIdList != null && rowList.size() != insertIdList.size()) { throw new AssertionError( @@ -1288,7 +1304,9 @@ long insertAll( // errorIndex) from the batch of rows which attempted insertion in this call. // Not the entire set of rows in rowsToPublish. if (retryIds != null) { - retryIds.add(idsToPublish.get(errorIndex)); + // retryIds is non-null exactly when idsToPublish is non-null; see where both are + // initialized above. + retryIds.add(checkStateNotNull(idsToPublish).get(errorIndex)); } } else { numFailedRows += 1; @@ -1439,8 +1457,8 @@ public long insertAll( public Table patchTableDescription( TableReference tableReference, @Nullable String tableDescription) throws IOException, InterruptedException { - Table table = new Table(); - table.setDescription(tableDescription); + @SuppressWarnings("nullness") // setDescription is not annotated, but does accept nulls + Table table = new Table().setDescription(tableDescription); return executeWithRetries( client @@ -1877,7 +1895,8 @@ public void onRetryAttempt(Status status, Metadata metadata) { && metadata != null && metadata.containsKey(KEY_RETRY_INFO)) { LOG.info("BigQuery direct read quota exceeded, retrying."); - RetryInfo retryInfo = metadata.get(KEY_RETRY_INFO); + // containsKey() was just checked above, so the value must be present. + RetryInfo retryInfo = checkStateNotNull(metadata.get(KEY_RETRY_INFO)); if (retryInfo.hasRetryDelay()) { long delay = retryInfo.getRetryDelay().getSeconds() * 1000 @@ -1989,7 +2008,9 @@ public BigQueryServerStream readRows( ServiceCallMetric serviceCallMetric = BigQueryUtils.readCallMetric(tableReference); try { BigQueryServerStream response = readRows(request); - serviceCallMetric.call("ok"); + if (serviceCallMetric != null) { + serviceCallMetric.call("ok"); + } return response; } catch (ApiException e) { if (serviceCallMetric != null) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageAvroReader.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageAvroReader.java index 50ce6a89f7a9..810e5c15f027 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageAvroReader.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageAvroReader.java @@ -48,9 +48,8 @@ class BigQueryStorageAvroReader implements BigQueryStorageReader { public void processReadRowsResponse(ReadRowsResponse readRowsResponse) { AvroRows avroRows = readRowsResponse.getAvroRows(); rowCount = avroRows.getRowCount(); - @SuppressWarnings({ - "nullness" // reused decoder can be null but avro not annotated - }) + // Avro accepts a null decoder to reuse but is not annotated. + @SuppressWarnings("nullness") BinaryDecoder newDecoder = DecoderFactory.get() .binaryDecoder(avroRows.getSerializedBinaryRows().toByteArray(), decoder); @@ -65,10 +64,9 @@ public long getRowCount() { @Override public GenericRecord readSingleRecord() throws IOException { Preconditions.checkStateNotNull(decoder); - @SuppressWarnings({ - "nullness" // reused record is null but avro not annotated - }) - // record should not be reused, mutating outputted values is unsafe + // The record should not be reused; mutating outputted values is unsafe. Avro accepts a + // null reuse argument but is not annotated. + @SuppressWarnings("nullness") GenericRecord newRecord = datumReader.read(/*reuse=*/ null, decoder); return newRecord; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java index 5ba2d17c127a..c8530f5ce678 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtils.java @@ -90,9 +90,6 @@ import org.joda.time.format.DateTimeFormatterBuilder; /** Utility methods for BigQuery related operations. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20506) -}) public class BigQueryUtils { // For parsing the format returned on the API proto: @@ -696,14 +693,14 @@ public TableRow apply(T input) { } public static Row toBeamRow(GenericRecord record, Schema schema, ConversionOptions options) { - List valuesInOrder = + List<@Nullable Object> valuesInOrder = schema.getFields().stream() - .map( + .<@Nullable Object>map( field -> { try { org.apache.avro.Schema.Field avroField = record.getSchema().getField(field.getName()); - Object value = avroField != null ? record.get(avroField.pos()) : null; + @Nullable Object value = avroField != null ? record.get(avroField.pos()) : null; return convertAvroFormat(field.getType(), value, options); } catch (Exception cause) { throw new IllegalArgumentException( @@ -735,14 +732,17 @@ public static TableRow convertGenericRecordToTableRow(GenericRecord record) { public static TableRow toTableRow(Row row) { TableRow output = new TableRow(); for (int i = 0; i < row.getFieldCount(); i++) { - Object value = row.getValue(i); + @Nullable Object value = row.getValue(i); Field schemaField = row.getSchema().getField(i); - output = output.set(schemaField.getName(), fromBeamField(schemaField.getType(), value)); + @SuppressWarnings("nullness") // TableRow.set is not annotated, but accepts nulls + TableRow updated = + output.set(schemaField.getName(), fromBeamField(schemaField.getType(), value)); + output = updated; } return output; } - private static @Nullable Object fromBeamField(FieldType fieldType, Object fieldValue) { + private static @Nullable Object fromBeamField(FieldType fieldType, @Nullable Object fieldValue) { if (fieldValue == null) { if (!fieldType.getNullable()) { throw new IllegalArgumentException("Field is not nullable."); @@ -753,26 +753,31 @@ public static TableRow toTableRow(Row row) { switch (fieldType.getTypeName()) { case ARRAY: case ITERABLE: - FieldType elementType = fieldType.getCollectionElementType(); + FieldType elementType = + Preconditions.checkArgumentNotNull(fieldType.getCollectionElementType()); Iterable items = (Iterable) fieldValue; - List convertedItems = Lists.newArrayListWithCapacity(Iterables.size(items)); + List<@Nullable Object> convertedItems = + Lists.newArrayListWithCapacity(Iterables.size(items)); for (Object item : items) { convertedItems.add(fromBeamField(elementType, item)); } return convertedItems; case MAP: - FieldType keyElementType = fieldType.getMapKeyType(); - FieldType valueElementType = fieldType.getMapValueType(); + FieldType keyElementType = Preconditions.checkArgumentNotNull(fieldType.getMapKeyType()); + FieldType valueElementType = + Preconditions.checkArgumentNotNull(fieldType.getMapValueType()); Map pairs = (Map) fieldValue; convertedItems = Lists.newArrayListWithCapacity(pairs.size()); for (Map.Entry pair : pairs.entrySet()) { - convertedItems.add( + @SuppressWarnings("nullness") // TableRow.set is not annotated, but accepts nulls + TableRow convertedPair = new TableRow() .set(BIGQUERY_MAP_KEY_FIELD_NAME, fromBeamField(keyElementType, pair.getKey())) .set( BIGQUERY_MAP_VALUE_FIELD_NAME, - fromBeamField(valueElementType, pair.getValue()))); + fromBeamField(valueElementType, pair.getValue())); + convertedItems.add(convertedPair); } return convertedItems; @@ -806,7 +811,8 @@ public static TableRow toTableRow(Row row) { case LOGICAL_TYPE: // For the JSON formats of DATE/DATETIME/TIME/TIMESTAMP types that BigQuery accepts, see // https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-json#details_of_loading_json_data - String identifier = fieldType.getLogicalType().getIdentifier(); + String identifier = + Preconditions.checkArgumentNotNull(fieldType.getLogicalType()).getIdentifier(); if (SqlTypes.DATE.getIdentifier().equals(identifier)) { return fieldValue.toString(); } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) { @@ -829,7 +835,7 @@ public static TableRow toTableRow(Row row) { java.time.format.DateTimeFormatter localDateTimeFormatter = (0 == localDateTime.getNano()) ? ISO_LOCAL_DATE_TIME : BIGQUERY_DATETIME_FORMATTER; return localDateTimeFormatter.format(localDateTime); - } else if (Timestamp.IDENTIFIER.equals(fieldType.getLogicalType().getIdentifier())) { + } else if (Timestamp.IDENTIFIER.equals(identifier)) { return BigQueryAvroUtils.formatTimestamp((java.time.Instant) fieldValue); } else if ("Enum".equals(identifier)) { return fieldType @@ -855,7 +861,7 @@ public static Row toBeamRow(Schema rowSchema, TableRow jsonBqRow) { // 2. TableSchema objects are not serializable and are therefore harder to propagate through a // pipeline. return rowSchema.getFields().stream() - .map(field -> toBeamValue(field, jsonBqRow.get(field.getName()))) + .<@Nullable Object>map(field -> toBeamValue(field, jsonBqRow.get(field.getName()))) .collect(toRow(rowSchema)); } @@ -874,17 +880,18 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso List rawJsonValues = rowSchema.getFields().stream() - .map(field -> bqFieldIndices.get(field.getName())) + .map(field -> Preconditions.checkArgumentNotNull(bqFieldIndices.get(field.getName()))) .map(index -> jsonBqRow.getF().get(index).getV()) .collect(toList()); return IntStream.range(0, rowSchema.getFieldCount()) .boxed() - .map(index -> toBeamValue(rowSchema.getField(index), rawJsonValues.get(index))) + .<@Nullable Object>map( + index -> toBeamValue(rowSchema.getField(index), rawJsonValues.get(index))) .collect(toRow(rowSchema)); } - private static @Nullable Object toBeamValue(Field field, Object jsonBQValue) { + private static @Nullable Object toBeamValue(Field field, @Nullable Object jsonBQValue) { FieldType fieldType = field.getType(); if (jsonBQValue == null) { @@ -947,7 +954,8 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso } if (jsonBQValue instanceof List) { - if (fieldType.getCollectionElementType() == null) { + FieldType collectionElementType = fieldType.getCollectionElementType(); + if (collectionElementType == null) { throw new IllegalArgumentException( "Cannot convert BigQuery type '" + jsonBQValue.getClass() @@ -957,27 +965,27 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso + " collection."); } - boolean innerTypeIsMap = fieldType.getCollectionElementType().getTypeName().isMapType(); + boolean innerTypeIsMap = collectionElementType.getTypeName().isMapType(); - return ((List) jsonBQValue) + return ((List<@Nullable Object>) jsonBQValue) .stream() // Old BigQuery client returns arrays as lists of maps {"v": }. // If this is the case, unwrap the value first - .map( + .<@Nullable Object>map( v -> (!innerTypeIsMap && v instanceof Map && ((Map) v).keySet().equals(Sets.newHashSet("v"))) ? ((Map) v).get("v") : v) - .map(v -> toBeamValue(field.withType(fieldType.getCollectionElementType()), v)) + .<@Nullable Object>map(v -> toBeamValue(field.withType(collectionElementType), v)) .collect(toList()); } if (jsonBQValue instanceof Map) { TableRow tr = new TableRow(); tr.putAll((Map) jsonBQValue); - return toBeamRow(fieldType.getRowSchema(), tr); + return toBeamRow(Preconditions.checkArgumentNotNull(fieldType.getRowSchema()), tr); } throw new UnsupportedOperationException( @@ -999,8 +1007,10 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso * https://cloud.google.com/bigquery/docs/exporting-data#avro_export_details and * https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-avro#avro_conversions */ - public static Object convertAvroFormat( - FieldType beamFieldType, Object avroValue, BigQueryUtils.ConversionOptions options) { + public static @Nullable Object convertAvroFormat( + FieldType beamFieldType, + @Nullable Object avroValue, + BigQueryUtils.ConversionOptions options) { TypeName beamFieldTypeName = beamFieldType.getTypeName(); if (avroValue == null) { if (beamFieldType.getNullable()) { @@ -1037,8 +1047,8 @@ public static Object convertAvroFormat( case ARRAY: return convertAvroArray(beamFieldType, avroValue, options); case LOGICAL_TYPE: - LogicalType logicalType = beamFieldType.getLogicalType(); - assert logicalType != null; + LogicalType logicalType = + Preconditions.checkArgumentNotNull(beamFieldType.getLogicalType()); String identifier = logicalType.getIdentifier(); if (SqlTypes.DATE.getIdentifier().equals(identifier)) { return convertAvroDate(avroValue); @@ -1100,8 +1110,9 @@ private static Object convertAvroArray( FieldType beamField, Object value, BigQueryUtils.ConversionOptions options) { // Check whether the type of array element is equal. List values = (List) value; - List ret = new ArrayList<>(); - FieldType collectionElement = beamField.getCollectionElementType(); + List<@Nullable Object> ret = new ArrayList<>(); + FieldType collectionElement = + Preconditions.checkArgumentNotNull(beamField.getCollectionElementType()); for (Object v : values) { ret.add(convertAvroFormat(collectionElement, v, options)); } @@ -1112,17 +1123,19 @@ private static Object convertAvroRecordToMap( FieldType beamField, Object value, BigQueryUtils.ConversionOptions options) { List records = (List) value; ImmutableMap.Builder ret = ImmutableMap.builder(); - FieldType keyElement = beamField.getMapKeyType(); - FieldType valueElement = beamField.getMapValueType(); + FieldType keyElement = Preconditions.checkArgumentNotNull(beamField.getMapKeyType()); + FieldType valueElement = Preconditions.checkArgumentNotNull(beamField.getMapValueType()); for (GenericData.Record record : records) { + // A map entry with a null key or value cannot be represented in an ImmutableMap. ret.put( - convertAvroFormat(keyElement, record.get(0), options), - convertAvroFormat(valueElement, record.get(1), options)); + Preconditions.checkArgumentNotNull(convertAvroFormat(keyElement, record.get(0), options)), + Preconditions.checkArgumentNotNull( + convertAvroFormat(valueElement, record.get(1), options))); } return ret.build(); } - private static Object convertAvroPrimitiveTypes(TypeName beamType, Object value) { + private static @Nullable Object convertAvroPrimitiveTypes(TypeName beamType, Object value) { switch (beamType) { case BYTE: return ((Long) value).byteValue(); @@ -1149,7 +1162,7 @@ private static Object convertAvroPrimitiveTypes(TypeName beamType, Object value) } } - private static Object convertAvroString(Object value) { + private static @Nullable Object convertAvroString(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof Utf8) { @@ -1162,7 +1175,7 @@ private static Object convertAvroString(Object value) { } } - private static Object convertAvroBytes(Object value) { + private static @Nullable Object convertAvroBytes(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof ByteBuffer) { @@ -1176,7 +1189,7 @@ private static Object convertAvroBytes(Object value) { } } - private static Object convertAvroDate(Object value) { + private static @Nullable Object convertAvroDate(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof Integer) { @@ -1187,7 +1200,7 @@ private static Object convertAvroDate(Object value) { } } - private static Object convertAvroTime(Object value) { + private static @Nullable Object convertAvroTime(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof Long) { @@ -1198,7 +1211,7 @@ private static Object convertAvroTime(Object value) { } } - private static Object convertAvroDateTime(Object value) { + private static @Nullable Object convertAvroDateTime(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof Utf8) { @@ -1209,13 +1222,18 @@ private static Object convertAvroDateTime(Object value) { } } - private static Object convertAvroNumeric(Object value) { + private static @Nullable Object convertAvroNumeric(@Nullable Object value) { if (value == null) { return null; } else if (value instanceof ByteBuffer) { // BigQuery NUMERIC type has precision 38 and scale 9 - return new Conversions.DecimalConversion() - .fromBytes((ByteBuffer) value, null, LogicalTypes.decimal(38, 9)); + // DecimalConversion is not annotated, but fromBytes ignores the writer schema argument + // and accepts null for it. + @SuppressWarnings("nullness") + BigDecimal decimal = + new Conversions.DecimalConversion() + .fromBytes((ByteBuffer) value, null, LogicalTypes.decimal(38, 9)); + return decimal; } else { throw new RuntimeException( "Does not support converting avro format: " + value.getClass().getName()); @@ -1239,10 +1257,11 @@ private static Object convertAvroNumeric(Object value) { // "projects/{project_id}/datasets/{dataset_id}/tables/{table_id}" Matcher m = TABLE_RESOURCE_PATTERN.matcher(fullTableId); if (m.matches()) { + // The named groups are not optional, so they are non-null whenever the pattern matches. return new TableReference() - .setProjectId(m.group("PROJECT")) - .setDatasetId(m.group("DATASET")) - .setTableId(m.group("TABLE")); + .setProjectId(Preconditions.checkStateNotNull(m.group("PROJECT"))) + .setDatasetId(Preconditions.checkStateNotNull(m.group("DATASET"))) + .setTableId(Preconditions.checkStateNotNull(m.group("TABLE"))); } // If that failed, try the format: @@ -1250,10 +1269,11 @@ private static Object convertAvroNumeric(Object value) { // "{project_id}.{dataset_id}.{table_id}" m = SIMPLE_TABLE_PATTERN.matcher(fullTableId); if (m.matches()) { + // The named groups are not optional, so they are non-null whenever the pattern matches. return new TableReference() - .setProjectId(m.group("PROJECT")) - .setDatasetId(m.group("DATASET")) - .setTableId(m.group("TABLE")); + .setProjectId(Preconditions.checkStateNotNull(m.group("PROJECT"))) + .setDatasetId(Preconditions.checkStateNotNull(m.group("DATASET"))) + .setTableId(Preconditions.checkStateNotNull(m.group("TABLE"))); } return null; } @@ -1320,29 +1340,25 @@ private static Stream trimField( return Stream.empty(); } - private static @Nullable ServiceCallMetric callMetricForMethod( - @Nullable TableReference tableReference, String method) { - if (tableReference != null) { - // TODO(ajamato): Add Ptransform label. Populate it as empty for now to prevent the - // SpecMonitoringInfoValidator from dropping the MonitoringInfo. - HashMap baseLabels = new HashMap(); - baseLabels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); - baseLabels.put(MonitoringInfoConstants.Labels.SERVICE, "BigQuery"); - baseLabels.put(MonitoringInfoConstants.Labels.METHOD, method); - baseLabels.put( - MonitoringInfoConstants.Labels.RESOURCE, - GcpResourceIdentifiers.bigQueryTable( - tableReference.getProjectId(), - tableReference.getDatasetId(), - tableReference.getTableId())); - baseLabels.put( - MonitoringInfoConstants.Labels.BIGQUERY_PROJECT_ID, tableReference.getProjectId()); - baseLabels.put( - MonitoringInfoConstants.Labels.BIGQUERY_DATASET, tableReference.getDatasetId()); - baseLabels.put(MonitoringInfoConstants.Labels.BIGQUERY_TABLE, tableReference.getTableId()); - return new ServiceCallMetric(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, baseLabels); - } - return null; + private static ServiceCallMetric callMetricForMethod( + TableReference tableReference, String method) { + // TODO(ajamato): Add Ptransform label. Populate it as empty for now to prevent the + // SpecMonitoringInfoValidator from dropping the MonitoringInfo. + HashMap baseLabels = new HashMap(); + baseLabels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); + baseLabels.put(MonitoringInfoConstants.Labels.SERVICE, "BigQuery"); + baseLabels.put(MonitoringInfoConstants.Labels.METHOD, method); + baseLabels.put( + MonitoringInfoConstants.Labels.RESOURCE, + GcpResourceIdentifiers.bigQueryTable( + tableReference.getProjectId(), + tableReference.getDatasetId(), + tableReference.getTableId())); + baseLabels.put( + MonitoringInfoConstants.Labels.BIGQUERY_PROJECT_ID, tableReference.getProjectId()); + baseLabels.put(MonitoringInfoConstants.Labels.BIGQUERY_DATASET, tableReference.getDatasetId()); + baseLabels.put(MonitoringInfoConstants.Labels.BIGQUERY_TABLE, tableReference.getTableId()); + return new ServiceCallMetric(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, baseLabels); } /** @@ -1355,6 +1371,9 @@ private static Stream trimField( */ public static @Nullable ServiceCallMetric readCallMetric( @Nullable TableReference tableReference) { + if (tableReference == null) { + return null; + } return callMetricForMethod(tableReference, "BigQueryBatchRead"); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java index 6370244c268c..318986960495 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java @@ -73,12 +73,16 @@ static class ConstantTableDestinations extends DynamicDestinations ConstantTableDestinations fromTableSpec( - ValueProvider tableSpec, String tableDescription, boolean clusteringEnabled) { + ValueProvider tableSpec, + @Nullable String tableDescription, + boolean clusteringEnabled) { return new ConstantTableDestinations<>(tableSpec, tableDescription, clusteringEnabled); } static ConstantTableDestinations fromJsonTableRef( - ValueProvider jsonTableRef, String tableDescription, boolean clusteringEnabled) { + ValueProvider jsonTableRef, + @Nullable String tableDescription, + boolean clusteringEnabled) { return new ConstantTableDestinations<>( NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableSpec()), tableDescription, @@ -286,8 +290,8 @@ static class ConstantTimePartitioningClusteringDestinations ConstantTimePartitioningClusteringDestinations( DynamicDestinations inner, - ValueProvider jsonTimePartitioning, - ValueProvider jsonClustering) { + @Nullable ValueProvider jsonTimePartitioning, + @Nullable ValueProvider jsonClustering) { super(inner); checkArgument( diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/RowWriterFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/RowWriterFactory.java index cc5c97ed0d3a..9204c7fb5c55 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/RowWriterFactory.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/RowWriterFactory.java @@ -41,7 +41,7 @@ abstract BigQueryRowWriter createRowWriter( static RowWriterFactory tableRows( BigQueryIO.TableRowFormatFunction toRow, - BigQueryIO.TableRowFormatFunction toFailsafeRow) { + BigQueryIO.@Nullable TableRowFormatFunction toFailsafeRow) { return new TableRowWriterFactory(toRow, toFailsafeRow); } @@ -49,11 +49,11 @@ static final class TableRowWriterFactory extends RowWriterFactory { private final BigQueryIO.TableRowFormatFunction toRow; - private final BigQueryIO.TableRowFormatFunction toFailsafeRow; + private final BigQueryIO.@Nullable TableRowFormatFunction toFailsafeRow; private TableRowWriterFactory( BigQueryIO.TableRowFormatFunction toRow, - BigQueryIO.TableRowFormatFunction toFailsafeRow) { + BigQueryIO.@Nullable TableRowFormatFunction toFailsafeRow) { this.toRow = toRow; this.toFailsafeRow = toFailsafeRow; } @@ -75,7 +75,6 @@ public OutputType getOutputType() { } @Override - @SuppressWarnings("nullness") public BigQueryRowWriter createRowWriter( String tempFilePrefix, DestinationT destination) throws Exception { return new TableRowWriter<>(tempFilePrefix, toRow.toSerializableFunction()); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java index 02ef9e9c06a3..b352902a0cd6 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java @@ -110,9 +110,10 @@ public PCollectionTuple expand(PCollection> input) { // This code currently assumes that the input is in the global window. Preconditions.checkState(input.getWindowingStrategy().getWindowFn() instanceof GlobalWindows); - @SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) - }) + // ConvertMessagesDoFn requires DestinationT to be @NonNull, but this transform (and its + // callers, up to BigQueryIO.Write) leave DestinationT unbounded. + // TODO(https://github.com/apache/beam/issues/20497) + @SuppressWarnings("nullness") ConvertMessagesDoFn convertMessagesDoFn = new ConvertMessagesDoFn<>( dynamicDestinations, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java index 401395030542..87a7a86b7346 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java @@ -99,7 +99,6 @@ public DescriptorProtos.DescriptorProto getDescriptor(boolean includeCdcColumns) } @Override - @SuppressWarnings("nullness") public StorageApiWritePayload toMessage( T element, @Nullable RowMutationInformation rowMutationInformation, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsGenericRecord.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsGenericRecord.java index 6d7b679fc9b6..99d296178a14 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsGenericRecord.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsGenericRecord.java @@ -28,6 +28,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.util.Preconditions; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -90,7 +91,6 @@ class GenericRecordConverter implements MessageConverter { public void updateSchemaFromTable() throws IOException, InterruptedException {} @Override - @SuppressWarnings("nullness") public StorageApiWritePayload toMessage( T element, @Nullable RowMutationInformation rowMutationInformation, @@ -102,7 +102,7 @@ public StorageApiWritePayload toMessage( if (rowMutationInformation != null) { changeType = rowMutationInformation.getMutationType().toString(); changeSequenceNum = rowMutationInformation.getChangeSequenceNumber(); - descriptorToUse = cdcDescriptor; + descriptorToUse = Preconditions.checkStateNotNull(cdcDescriptor); } Message msg = AvroGenericRecordToStorageApiProto.messageFromGenericRecord( diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsProto.java index 2a1d8f3be6f2..8cacb8a4ffb0 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsProto.java @@ -39,7 +39,7 @@ class StorageApiDynamicDestinationsProto formatRecordOnFailureFunction; - @SuppressWarnings({"unchecked", "nullness"}) + @SuppressWarnings("unchecked") StorageApiDynamicDestinationsProto( DynamicDestinations inner, Class protoClass, @@ -47,11 +47,12 @@ class StorageApiDynamicDestinationsProto(Duration.standardSeconds(1), Duration.standardMinutes(1), 3); retryManager.addOperation( c -> { - @SuppressWarnings({ - "nullness" // unsure why s is inferred to be @Nullable - }) Iterable streamsToCommit = Iterables.filter(streamNames, s -> !alreadyCommittedStreams.contains(s)); batchCommitOperationsSent.inc(); @@ -188,9 +185,6 @@ public void finishBundle(PipelineOptions pipelineOptions) throws Exception { alreadyCommittedStreams.add(storageError.getEntity()); } } - @SuppressWarnings({ - "nullness" // unsure why s is inferred to be @Nullable - }) Iterable streamsToCommit = Iterables.filter(streamNames, s -> !alreadyCommittedStreams.contains(s)); // If there are no more streams left to commit, then report this operation as having diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java index 007bba5c6cdf..38c104ab9694 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java @@ -44,6 +44,7 @@ import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler; import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; @@ -68,8 +69,8 @@ public class StorageApiLoads private final @Nullable SerializableFunction rowUpdateFn; private final CreateDisposition createDisposition; - private final String kmsKey; - private final Duration triggeringFrequency; + private final @Nullable String kmsKey; + private final @Nullable Duration triggeringFrequency; private final BigQueryServices bqServices; private final int numShards; private final boolean allowInconsistentWrites; @@ -79,7 +80,7 @@ public class StorageApiLoads private final boolean usesCdc; private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation; - private final Map bigLakeConfiguration; + private final @Nullable Map bigLakeConfiguration; private final BadRecordRouter badRecordRouter; @@ -92,8 +93,8 @@ public StorageApiLoads( StorageApiDynamicDestinations dynamicDestinations, @Nullable SerializableFunction rowUpdateFn, CreateDisposition createDisposition, - String kmsKey, - Duration triggeringFrequency, + @Nullable String kmsKey, + @Nullable Duration triggeringFrequency, BigQueryServices bqServices, int numShards, boolean allowInconsistentWrites, @@ -104,7 +105,7 @@ public StorageApiLoads( Predicate propagateSuccessfulStorageApiWritesPredicate, boolean usesCdc, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, - Map bigLakeConfiguration, + @Nullable Map bigLakeConfiguration, BadRecordRouter badRecordRouter, ErrorHandler badRecordErrorHandler, boolean hasSchemaUpdateOptions) { @@ -231,6 +232,10 @@ public WriteResult expandTriggered( PCollection> input, Coder> successCoder, Coder payloadCoder) { + // Only reached when a triggering frequency is configured; see expand(). + Duration triggeringFrequency = + Preconditions.checkStateNotNull( + this.triggeringFrequency, "A triggering frequency is required for triggered loads"); // Handle triggered, low-latency loads into BigQuery. PCollection> inputInGlobalWindow = input.apply("rewindowIntoGlobal", Window.into(new GlobalWindows())); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java index 3f00ed67a8a5..87ed4518bb00 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java @@ -62,7 +62,6 @@ public abstract static class Builder { public abstract Builder toBuilder(); - @SuppressWarnings("nullness") static StorageApiWritePayload of( byte[] payload, @Nullable TableRow unknownFields, @Nullable TableRow failsafeTableRow) throws IOException { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java index 188f393dbe8c..85ac020da0d4 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java @@ -123,7 +123,7 @@ public class StorageApiWritesShardedRecords dynamicDestinations; private final CreateDisposition createDisposition; - private final String kmsKey; + private final @Nullable String kmsKey; private final BigQueryServices bqServices; private final Coder destinationCoder; private final Coder failedRowsCoder; @@ -151,7 +151,7 @@ static void clearCache() { public StorageApiWritesShardedRecords( StorageApiDynamicDestinations dynamicDestinations, CreateDisposition createDisposition, - String kmsKey, + @Nullable String kmsKey, BigQueryServices bqServices, Coder destinationCoder, Coder failedRowsCoder, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java index 8c825b03d4f2..def51b285413 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java @@ -245,7 +245,7 @@ StreamingInserts withSuccessfulInsertsPropagation( } StreamingInserts withDeterministicRecordIdFn( - SerializableFunction deterministicRecordIdFn) { + @Nullable SerializableFunction deterministicRecordIdFn) { return new StreamingInserts<>( createDisposition, dynamicDestinations, @@ -264,7 +264,7 @@ StreamingInserts withDeterministicRecordIdFn( kmsKey); } - StreamingInserts withKmsKey(String kmsKey) { + StreamingInserts withKmsKey(@Nullable String kmsKey) { return new StreamingInserts<>( createDisposition, dynamicDestinations, 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 ba72bb8682fd..fa714a70a362 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 @@ -1080,7 +1080,6 @@ 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. */ - @SuppressWarnings("nullness") public static @Nullable DynamicMessage messageFromTableRow( SchemaInformation schemaInformation, @Nullable Descriptor descriptor, @@ -1124,7 +1123,10 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) if (unknownFields != null) { List unknownValues = Lists.newArrayListWithExpectedSize(cells.size()); for (int i = 0; i < cells.size(); ++i) { - unknownValues.add(new TableCell().setV(null)); + // TableCell accepts a null value, but the API client is not annotated. + @SuppressWarnings("nullness") + TableCell nullCell = new TableCell().setV(null); + unknownValues.add(nullCell); } unknownFields.setF(unknownValues); } @@ -1186,7 +1188,10 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) // If there are unknown fields, copy them into the output. if (unknownFields != null) { for (int i = cellsToProcess; i < cells.size(); ++i) { - unknownFields.getF().set(i, new TableCell().setV(cells.get(i).get("v"))); + // TableCell accepts a null value, but the API client is not annotated. + @SuppressWarnings("nullness") + TableCell unknownCell = new TableCell().setV(cells.get(i).get("v")); + unknownFields.getF().set(i, unknownCell); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java index b9cd554b578b..536d318f0213 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java @@ -42,6 +42,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.checkerframework.checker.nullness.qual.NonNull; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -58,8 +59,7 @@ * of KV(TableDestination, WriteTables.Result) where the destination label is parsed and replaced to * TableDestination objects. */ -@SuppressWarnings({"nullness"}) -public class UpdateSchemaDestination +public class UpdateSchemaDestination extends DoFn< Iterable>, Iterable>> { @@ -116,12 +116,13 @@ public void startBundle(StartBundleContext c) { } TableDestination getTableWithDefaultProject(DestinationT destination) { - if (dynamicDestinations.getPipelineOptions() == null) { + PipelineOptions pipelineOptions = dynamicDestinations.getPipelineOptions(); + if (pipelineOptions == null) { throw new IllegalStateException( "Unexpected null pipeline option for DynamicDestination object. " + "Need to call setSideInputAccessorFromProcessContext(context) before use it."); } - BigQueryOptions options = dynamicDestinations.getPipelineOptions().as(BigQueryOptions.class); + BigQueryOptions options = pipelineOptions.as(BigQueryOptions.class); TableDestination tableDestination = dynamicDestinations.getTable(destination); TableReference tableReference = tableDestination.getTableReference(); @@ -229,13 +230,13 @@ public void finishBundle(FinishBundleContext context) throws Exception { jobManager.waitForDone(); } - private BigQueryHelpers.PendingJob startZeroLoadJob( + private @Nullable BigQueryHelpers.PendingJob startZeroLoadJob( BigQueryServices.JobService jobService, DatasetService datasetService, String jobIdPrefix, TableReference tableReference, - TimePartitioning timePartitioning, - Clustering clustering, + @Nullable TimePartitioning timePartitioning, + @Nullable Clustering clustering, @Nullable TableSchema schema, BigQueryIO.Write.WriteDisposition writeDisposition, BigQueryIO.Write.CreateDisposition createDisposition, @@ -243,7 +244,6 @@ private BigQueryHelpers.PendingJob startZeroLoadJob( JobConfigurationLoad loadConfig = new JobConfigurationLoad() .setDestinationTable(tableReference) - .setSchema(schema) .setWriteDisposition(writeDisposition.name()) .setCreateDisposition(createDisposition.name()) .setSourceFormat("NEWLINE_DELIMITED_JSON"); @@ -275,12 +275,14 @@ private BigQueryHelpers.PendingJob startZeroLoadJob( // no need to update schema ahead if provided schema already matches destination schema // or when destination schema is null (the write will set the schema) // or when provided schema is null (e.g. when using CREATE_NEVER disposition) - if (destinationTable.getSchema() == null - || destinationTable.getSchema().isEmpty() - || destinationTable.getSchema().equals(schema) - || schema == null) { + TableSchema destinationSchema = destinationTable.getSchema(); + if (destinationSchema == null + || destinationSchema.isEmpty() + || schema == null + || destinationSchema.equals(schema)) { return null; } + loadConfig.setSchema(schema); if (timePartitioning != null) { loadConfig.setTimePartitioning(timePartitioning); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryDirectReadSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryDirectReadSchemaTransformProvider.java index 073de40038b3..84916b918fe1 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryDirectReadSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryDirectReadSchemaTransformProvider.java @@ -17,9 +17,9 @@ */ package org.apache.beam.sdk.io.gcp.bigquery.providers; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import com.google.api.services.bigquery.model.TableRow; import com.google.auto.service.AutoService; @@ -50,6 +50,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.dataflow.qual.Pure; /** * An implementation of {@link TypedSchemaTransformProvider} for BigQuery Storage Read API jobs @@ -59,9 +60,6 @@ * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam * repository. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) @AutoService(SchemaTransformProvider.class) public class BigQueryDirectReadSchemaTransformProvider extends TypedSchemaTransformProvider { @@ -100,8 +98,9 @@ public abstract static class BigQueryDirectReadSchemaTransformConfiguration { public void validate() { String invalidConfigMessage = "Invalid BigQuery Direct Read configuration: "; - if (!Strings.isNullOrEmpty(this.getTableSpec())) { - checkNotNull(BigQueryHelpers.parseTableSpec(this.getTableSpec())); + if (this.getTableSpec() != null && !this.getTableSpec().isEmpty()) { + // Throws IllegalArgumentException if the table spec is malformed. + BigQueryHelpers.parseTableSpec(this.getTableSpec()); checkArgument( Strings.isNullOrEmpty(this.getQuery()), invalidConfigMessage + "Cannot specify both query and table spec."); @@ -125,29 +124,33 @@ public static Builder builder() { } @SchemaFieldDescription("The SQL query to be executed to read from the BigQuery table.") + @Pure @Nullable public abstract String getQuery(); @SchemaFieldDescription( "The fully-qualified name of the BigQuery table to read from. Format: [${PROJECT}:]${DATASET}.${TABLE}") + @Pure @Nullable public abstract String getTableSpec(); @SchemaFieldDescription( "Read only rows that match this filter, which must be compatible with Google standard SQL. This is not supported when reading via query.") + @Pure @Nullable public abstract String getRowRestriction(); @SchemaFieldDescription( "Read only the specified fields (columns) from a BigQuery table. Fields may not be returned in the order specified. If no value is specified, then all fields are returned. Example: \"col1, col2, col3\"") + @Pure @Nullable public abstract List getSelectedFields(); @SchemaFieldDescription("Use this Cloud KMS key to encrypt your data") + @Pure @Nullable public abstract String getKmsKey(); - @Nullable /** Builder for the {@link BigQueryDirectReadSchemaTransformConfiguration}. */ @AutoValue.Builder public abstract static class Builder { @@ -172,7 +175,7 @@ public abstract static class Builder { * BigQueryDirectReadSchemaTransformProvider}. */ public static class BigQueryDirectReadSchemaTransform extends SchemaTransform { - private BigQueryServices testBigQueryServices = null; + private @Nullable BigQueryServices testBigQueryServices = null; private final BigQueryDirectReadSchemaTransformConfiguration configuration; BigQueryDirectReadSchemaTransform( @@ -224,23 +227,26 @@ BigQueryIO.TypedRead createDirectReadTransform() { BigQueryIO.TypedRead read = BigQueryIO.readTableRowsWithSchema().withMethod(TypedRead.Method.DIRECT_READ); - if (!Strings.isNullOrEmpty(configuration.getTableSpec())) { + if (configuration.getTableSpec() != null && !configuration.getTableSpec().isEmpty()) { read = read.from(configuration.getTableSpec()); - if (!Strings.isNullOrEmpty(configuration.getRowRestriction())) { + if (configuration.getRowRestriction() != null + && !configuration.getRowRestriction().isEmpty()) { read = read.withRowRestriction(configuration.getRowRestriction()); } if (configuration.getSelectedFields() != null) { read = read.withSelectedFields(configuration.getSelectedFields()); } } else { - read = read.fromQuery(configuration.getQuery()).usingStandardSql(); + // Guaranteed non-empty by BigQueryDirectReadSchemaTransformConfiguration#validate(), which + // runs in this transform's constructor. + read = read.fromQuery(checkStateNotNull(configuration.getQuery())).usingStandardSql(); } - if (!Strings.isNullOrEmpty(configuration.getKmsKey())) { + if (configuration.getKmsKey() != null && !configuration.getKmsKey().isEmpty()) { read = read.withKmsKey(configuration.getKmsKey()); } if (this.testBigQueryServices != null) { - read = read.withTestServices(testBigQueryServices); + read = read.withTestServices(this.testBigQueryServices); } return read; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryFileLoadsSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryFileLoadsSchemaTransformProvider.java index 8899ac82eb06..49bf39fb3d78 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryFileLoadsSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryFileLoadsSchemaTransformProvider.java @@ -35,7 +35,7 @@ import org.apache.beam.sdk.values.PCollectionRowTuple; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; /** * An implementation of {@link TypedSchemaTransformProvider} for BigQuery write jobs configured @@ -45,9 +45,6 @@ * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam * repository. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) @Internal @AutoService(SchemaTransformProvider.class) public class BigQueryFileLoadsSchemaTransformProvider @@ -77,7 +74,7 @@ public List outputCollectionNames() { public static class BigQueryFileLoadsSchemaTransform extends SchemaTransform { /** An instance of {@link BigQueryServices} used for testing. */ - private BigQueryServices testBigQueryServices = null; + private @Nullable BigQueryServices testBigQueryServices = null; private final BigQueryWriteConfiguration configuration; @@ -113,17 +110,19 @@ public BigQueryIO.Write toWrite(Schema schema, PipelineOptions options) { // reason. .withAvroFormatFunction(dynamicDestinations.getAvroFilterFormatFunction(false)); - if (!Strings.isNullOrEmpty(configuration.getCreateDisposition())) { - CreateDisposition createDisposition = - CreateDisposition.valueOf(configuration.getCreateDisposition().toUpperCase()); - write = write.withCreateDisposition(createDisposition); + if (configuration.getCreateDisposition() != null + && !configuration.getCreateDisposition().isEmpty()) { + write = + write.withCreateDisposition( + CreateDisposition.valueOf(configuration.getCreateDisposition().toUpperCase())); } - if (!Strings.isNullOrEmpty(configuration.getWriteDisposition())) { - WriteDisposition writeDisposition = - WriteDisposition.valueOf(configuration.getWriteDisposition().toUpperCase()); - write = write.withWriteDisposition(writeDisposition); + if (configuration.getWriteDisposition() != null + && !configuration.getWriteDisposition().isEmpty()) { + write = + write.withWriteDisposition( + WriteDisposition.valueOf(configuration.getWriteDisposition().toUpperCase())); } - if (!Strings.isNullOrEmpty(configuration.getKmsKey())) { + if (configuration.getKmsKey() != null && !configuration.getKmsKey().isEmpty()) { write = write.withKmsKey(configuration.getKmsKey()); } if (testBigQueryServices != null) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java index 1d618ba685ed..01cfd59fccc9 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java @@ -20,6 +20,7 @@ import static org.apache.beam.sdk.io.gcp.bigquery.providers.BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS; import static org.apache.beam.sdk.io.gcp.bigquery.providers.PortableBigQueryDestinations.DESTINATION; import static org.apache.beam.sdk.io.gcp.bigquery.providers.PortableBigQueryDestinations.RECORD; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.auto.service.AutoService; @@ -52,7 +53,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; /** @@ -63,9 +64,6 @@ * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam * repository. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) @AutoService(SchemaTransformProvider.class) public class BigQueryStorageWriteApiSchemaTransformProvider extends TypedSchemaTransformProvider { @@ -123,7 +121,7 @@ public List outputCollectionNames() { */ public static class BigQueryStorageWriteApiSchemaTransform extends SchemaTransform { - private BigQueryServices testBigQueryServices = null; + private @Nullable BigQueryServices testBigQueryServices = null; private final BigQueryWriteConfiguration configuration; BigQueryStorageWriteApiSchemaTransform(BigQueryWriteConfiguration configuration) { @@ -183,13 +181,9 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { if (inputRows.isBounded() == IsBounded.UNBOUNDED) { Long triggeringFrequency = configuration.getTriggeringFrequencySeconds(); - Boolean autoSharding = configuration.getAutoSharding(); - boolean useAtLeastOnceSemantics = - configuration.getUseAtLeastOnceSemantics() != null - && configuration.getUseAtLeastOnceSemantics(); // Triggering frequency is only applicable for exactly-once - if (!useAtLeastOnceSemantics) { + if (!Boolean.TRUE.equals(configuration.getUseAtLeastOnceSemantics())) { write = write.withTriggeringFrequency( (triggeringFrequency == null || triggeringFrequency <= 0) @@ -197,7 +191,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { : Duration.standardSeconds(triggeringFrequency)); } // set num streams if specified, otherwise default to autoSharding - if (numStreams == 0 && (autoSharding == null || autoSharding)) { + if (numStreams == 0 && !Boolean.FALSE.equals(configuration.getAutoSharding())) { write = write.withAutoSharding(); } } @@ -222,7 +216,8 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { .apply("post-write", ParDo.of(new NoOutputDoFn())) .setRowSchema(Schema.of()); - if (configuration.getErrorHandling() == null) { + BigQueryWriteConfiguration.ErrorHandling errorHandling = configuration.getErrorHandling(); + if (errorHandling == null) { result .getFailedStorageApiInserts() .apply("Error on failed inserts", ParDo.of(new FailOnError())); @@ -250,21 +245,26 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { .via( (storageError) -> Row.withSchema(errorSchema) - .withFieldValue("error_message", storageError.getErrorMessage()) + // "error_message" is a non-nullable field of errorSchema, so a + // null error message would fail Row validation anyway. + .withFieldValue( + "error_message", + checkStateNotNull( + storageError.getErrorMessage(), + "Failed BigQuery insert has no error message.")) .withFieldValue( "failed_row", BigQueryUtils.toBeamRow(inputSchema, storageError.getRow())) .build())) .setRowSchema(errorSchema); return PCollectionRowTuple.of("post_write", postWrite) - .and(configuration.getErrorHandling().getOutput(), failedRowsWithErrors); + .and(errorHandling.getOutput(), failedRowsWithErrors); } } BigQueryIO.Write createStorageWriteApiTransform(Schema schema) { Method writeMethod = - configuration.getUseAtLeastOnceSemantics() != null - && configuration.getUseAtLeastOnceSemantics() + Boolean.TRUE.equals(configuration.getUseAtLeastOnceSemantics()) ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; @@ -277,24 +277,37 @@ BigQueryIO.Write createStorageWriteApiTransform(Schema schema) { boolean fetchNestedRecord = false; if (configuration.getTable().equals(DYNAMIC_DESTINATIONS)) { validateDynamicDestinationsSchema(schema); - rowSchema = schema.getField(RECORD).getType().getRowSchema(); + rowSchema = getRecordSchema(schema); fetchNestedRecord = true; } if (Boolean.TRUE.equals(configuration.getUseCdcWrites())) { validateCdcSchema(schema); - rowSchema = schema.getField(RECORD).getType().getRowSchema(); + rowSchema = getRecordSchema(schema); fetchNestedRecord = true; + // BigQueryIO stores a null primary key as "unset", so only call the setter when one is + // configured. BigQueryIO itself validates that a primary key is present when required. + if (configuration.getPrimaryKey() != null) { + write = write.withPrimaryKey(configuration.getPrimaryKey()); + } write = - write - .withPrimaryKey(configuration.getPrimaryKey()) - .withRowMutationInformationFn( - row -> - RowMutationInformation.of( - RowMutationInformation.MutationType.valueOf( - row.getRow(ROW_PROPERTY_MUTATION_INFO) - .getString(ROW_PROPERTY_MUTATION_TYPE)), - row.getRow(ROW_PROPERTY_MUTATION_INFO) - .getString(ROW_PROPERTY_MUTATION_SQN))); + write.withRowMutationInformationFn( + row -> { + Row mutationInfo = + checkStateNotNull( + row.getRow(ROW_PROPERTY_MUTATION_INFO), + "Encountered a row with an unset \"%s\" field.", + ROW_PROPERTY_MUTATION_INFO); + return RowMutationInformation.of( + RowMutationInformation.MutationType.valueOf( + checkStateNotNull( + mutationInfo.getString(ROW_PROPERTY_MUTATION_TYPE), + "Encountered a row with an unset \"%s\" field.", + ROW_PROPERTY_MUTATION_TYPE)), + checkStateNotNull( + mutationInfo.getString(ROW_PROPERTY_MUTATION_SQN), + "Encountered a row with an unset \"%s\" field.", + ROW_PROPERTY_MUTATION_SQN)); + }); } PortableBigQueryDestinations dynamicDestinations = new PortableBigQueryDestinations(rowSchema, configuration); @@ -303,29 +316,42 @@ BigQueryIO.Write createStorageWriteApiTransform(Schema schema) { .to(dynamicDestinations) .withFormatFunction(dynamicDestinations.getFilterFormatFunction(fetchNestedRecord)); - if (!Strings.isNullOrEmpty(configuration.getCreateDisposition())) { - CreateDisposition createDisposition = - CreateDisposition.valueOf(configuration.getCreateDisposition().toUpperCase()); - write = write.withCreateDisposition(createDisposition); + if (configuration.getCreateDisposition() != null + && !configuration.getCreateDisposition().isEmpty()) { + write = + write.withCreateDisposition( + CreateDisposition.valueOf(configuration.getCreateDisposition().toUpperCase())); } - if (!Strings.isNullOrEmpty(configuration.getWriteDisposition())) { - WriteDisposition writeDisposition = - WriteDisposition.valueOf(configuration.getWriteDisposition().toUpperCase()); - write = write.withWriteDisposition(writeDisposition); + if (configuration.getWriteDisposition() != null + && !configuration.getWriteDisposition().isEmpty()) { + write = + write.withWriteDisposition( + WriteDisposition.valueOf(configuration.getWriteDisposition().toUpperCase())); } - if (!Strings.isNullOrEmpty(configuration.getKmsKey())) { + if (configuration.getKmsKey() != null && !configuration.getKmsKey().isEmpty()) { write = write.withKmsKey(configuration.getKmsKey()); } if (configuration.getBigLakeConfiguration() != null) { write = write.withBigLakeConfiguration(configuration.getBigLakeConfiguration()); } if (this.testBigQueryServices != null) { - write = write.withTestServices(testBigQueryServices); + write = write.withTestServices(this.testBigQueryServices); } return write; } + /** + * Returns the schema of the nested {@link PortableBigQueryDestinations#RECORD} field, which the + * schema validation methods of this class have already established is a row field. + */ + private static Schema getRecordSchema(Schema schema) { + return checkStateNotNull( + schema.getField(RECORD).getType().getRowSchema(), + "Expected \"%s\" to be a Row field.", + RECORD); + } + void validateDynamicDestinationsSchema(Schema schema) { checkArgument( schema.getFieldNames().containsAll(Arrays.asList(DESTINATION, RECORD)), diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index 7f09feb245cd..ec59db14501b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java @@ -30,6 +30,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.checkerframework.checker.nullness.qual.Nullable; +import org.checkerframework.dataflow.qual.Pure; /** * Configuration for writing to BigQuery with SchemaTransforms. Used by {@link @@ -44,6 +45,7 @@ public abstract class BigQueryWriteConfiguration { @AutoValue public abstract static class ErrorHandling { @SchemaFieldDescription("The name of the output PCollection containing failed writes.") + @Pure public abstract String getOutput(); public static Builder builder() { @@ -122,12 +124,14 @@ public static Builder builder() { @SchemaFieldDescription( "The bigquery table to write to. Format: [${PROJECT}:]${DATASET}.${TABLE}") + @Pure public abstract String getTable(); @SchemaFieldDescription( "Optional field that specifies whether the job is allowed to create new tables. " + "The following values are supported: CREATE_IF_NEEDED (the job may create the table), CREATE_NEVER (" + "the job must fail if the table does not exist already).") + @Pure @Nullable public abstract String getCreateDisposition(); @@ -137,36 +141,43 @@ public static Builder builder() { + "WRITE_TRUNCATE (overwrites the table data), " + "WRITE_APPEND (append the data to the table), " + "WRITE_EMPTY (job must fail if the table is not empty).") + @Pure @Nullable public abstract String getWriteDisposition(); @SchemaFieldDescription( "Determines how often to 'commit' progress into BigQuery. Default is every 5 seconds.") + @Pure @Nullable public abstract Long getTriggeringFrequencySeconds(); @SchemaFieldDescription( "This option enables lower latency for insertions to BigQuery but may ocassionally " + "duplicate data elements.") + @Pure @Nullable public abstract Boolean getUseAtLeastOnceSemantics(); @SchemaFieldDescription( "This option enables using a dynamically determined number of Storage Write API streams to write to " + "BigQuery. Only applicable to unbounded data.") + @Pure @Nullable public abstract Boolean getAutoSharding(); @SchemaFieldDescription( "Specifies the number of write streams that the Storage API sink will use.") + @Pure @Nullable public abstract Integer getNumStreams(); @SchemaFieldDescription("Use this Cloud KMS key to encrypt your data") + @Pure @Nullable public abstract String getKmsKey(); @SchemaFieldDescription("This option specifies whether and where to output unwritable rows.") + @Pure @Nullable public abstract ErrorHandling getErrorHandling(); @@ -175,31 +186,37 @@ public static Builder builder() { + " should contain Beam Rows with a schema wrapping the record to be inserted and" + " adding the CDC info similar to: {row_mutation_info: {mutation_type:\"...\", " + "change_sequence_number:\"...\"}, record: {...}}") + @Pure @Nullable public abstract Boolean getUseCdcWrites(); @SchemaFieldDescription( "If CREATE_IF_NEEDED disposition is set, BigQuery table(s) will be created with this" + " columns as primary key. Required when CDC writes are enabled with CREATE_IF_NEEDED.") + @Pure @Nullable public abstract List getPrimaryKey(); @SchemaFieldDescription( "A list of field names to keep in the input record. All other fields are dropped before writing. " + "Is mutually exclusive with 'drop' and 'only'.") + @Pure public abstract @Nullable List getKeep(); @SchemaFieldDescription( "A list of field names to drop from the input record before writing. " + "Is mutually exclusive with 'keep' and 'only'.") + @Pure public abstract @Nullable List getDrop(); @SchemaFieldDescription( "The name of a single record field that should be written. " + "Is mutually exclusive with 'keep' and 'drop'.") + @Pure public abstract @Nullable String getOnly(); @SchemaFieldDescription("A list of columns to cluster the BigQuery table by.") + @Pure public abstract @Nullable List getClusteringFields(); @SchemaFieldDescription( @@ -208,6 +225,7 @@ public static Builder builder() { + "\n - storageUri (REQUIRED): the path to your GCS folder where data will be written to," + "\n - fileFormat (OPTIONAL): defaults to 'parquet'," + "\n - tableFormat (OPTIONAL): defaults to 'iceberg'.") + @Pure public abstract java.util.@Nullable Map getBigLakeConfiguration(); /** Builder for {@link BigQueryWriteConfiguration}. */ diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java index c927cec34735..b866b89e9879 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java @@ -61,13 +61,13 @@ public PortableBigQueryDestinations(Schema rowSchema, BigQueryWriteConfiguration this.primaryKey = configuration.getPrimaryKey(); RowFilter rf = new RowFilter(rowSchema); if (configuration.getDrop() != null) { - rf = rf.drop(checkStateNotNull(configuration.getDrop())); + rf = rf.drop(configuration.getDrop()); } if (configuration.getKeep() != null) { - rf = rf.keep(checkStateNotNull(configuration.getKeep())); + rf = rf.keep(configuration.getKeep()); } if (configuration.getOnly() != null) { - rf = rf.only(checkStateNotNull(configuration.getOnly())); + rf = rf.only(configuration.getOnly()); } this.rowFilter = rf; }