From 7dbd0df6e46a0760c28925ea65eed9c306c22f21 Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Wed, 13 May 2020 20:51:32 +0530 Subject: [PATCH] feat: add labels for writechannelconfiguration (#347) --- .../bigquery/WriteChannelConfiguration.java | 32 ++++++++++++++--- .../WriteChannelConfigurationTest.java | 7 ++++ .../cloud/bigquery/it/ITBigQueryTest.java | 34 +++++++++++++++++++ 3 files changed, 69 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java index b99f71879..479a68099 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java @@ -28,6 +28,7 @@ import com.google.common.primitives.Ints; import java.io.Serializable; import java.util.List; +import java.util.Map; import java.util.Objects; /** @@ -53,6 +54,7 @@ public final class WriteChannelConfiguration implements LoadConfiguration, Seria private final TimePartitioning timePartitioning; private final Clustering clustering; private final Boolean useAvroLogicalTypes; + private final Map labels; public static final class Builder implements LoadConfiguration.Builder { @@ -70,6 +72,7 @@ public static final class Builder implements LoadConfiguration.Builder { private TimePartitioning timePartitioning; private Clustering clustering; private Boolean useAvroLogicalTypes; + private Map labels; private Builder() {} @@ -89,6 +92,7 @@ private Builder(WriteChannelConfiguration writeChannelConfiguration) { this.timePartitioning = writeChannelConfiguration.timePartitioning; this.clustering = writeChannelConfiguration.clustering; this.useAvroLogicalTypes = writeChannelConfiguration.useAvroLogicalTypes; + this.labels = writeChannelConfiguration.labels; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -162,6 +166,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur this.clustering = Clustering.fromPb(loadConfigurationPb.getClustering()); } this.useAvroLogicalTypes = loadConfigurationPb.getUseAvroLogicalTypes(); + if (configurationPb.getLabels() != null) { + this.labels = configurationPb.getLabels(); + } } @Override @@ -250,6 +257,11 @@ public Builder setUseAvroLogicalTypes(Boolean useAvroLogicalTypes) { return this; } + public Builder setLabels(Map labels) { + this.labels = labels; + return this; + } + @Override public WriteChannelConfiguration build() { return new WriteChannelConfiguration(this); @@ -271,6 +283,7 @@ protected WriteChannelConfiguration(Builder builder) { this.timePartitioning = builder.timePartitioning; this.clustering = builder.clustering; this.useAvroLogicalTypes = builder.useAvroLogicalTypes; + this.labels = builder.labels; } @Override @@ -355,6 +368,10 @@ public Boolean getUseAvroLogicalTypes() { return useAvroLogicalTypes; } + public Map getLabels() { + return labels; + } + @Override public Builder toBuilder() { return new Builder(this); @@ -375,7 +392,8 @@ MoreObjects.ToStringHelper toStringHelper() { .add("autodetect", autodetect) .add("timePartitioning", timePartitioning) .add("clustering", clustering) - .add("useAvroLogicalTypes", useAvroLogicalTypes); + .add("useAvroLogicalTypes", useAvroLogicalTypes) + .add("labels", labels); } @Override @@ -405,7 +423,8 @@ public int hashCode() { autodetect, timePartitioning, clustering, - useAvroLogicalTypes); + useAvroLogicalTypes, + labels); } WriteChannelConfiguration setProjectId(String projectId) { @@ -416,6 +435,8 @@ WriteChannelConfiguration setProjectId(String projectId) { } com.google.api.services.bigquery.model.JobConfiguration toPb() { + com.google.api.services.bigquery.model.JobConfiguration jobConfiguration = + new com.google.api.services.bigquery.model.JobConfiguration(); JobConfigurationLoad loadConfigurationPb = new JobConfigurationLoad(); loadConfigurationPb.setDestinationTable(destinationTable.toPb()); if (createDisposition != null) { @@ -471,8 +492,11 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { loadConfigurationPb.setClustering(clustering.toPb()); } loadConfigurationPb.setUseAvroLogicalTypes(useAvroLogicalTypes); - return new com.google.api.services.bigquery.model.JobConfiguration() - .setLoad(loadConfigurationPb); + if (labels != null) { + jobConfiguration.setLabels(labels); + } + jobConfiguration.setLoad(loadConfigurationPb); + return jobConfiguration; } static WriteChannelConfiguration fromPb( diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java index d78d3995a..dbee1094b 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java @@ -23,8 +23,10 @@ import com.google.cloud.bigquery.JobInfo.WriteDisposition; import com.google.cloud.bigquery.TimePartitioning.Type; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import java.nio.charset.StandardCharsets; import java.util.List; +import java.util.Map; import org.junit.Test; public class WriteChannelConfigurationTest { @@ -55,6 +57,8 @@ public class WriteChannelConfigurationTest { private static final TimePartitioning TIME_PARTITIONING = TimePartitioning.of(Type.DAY); private static final Clustering CLUSTERING = Clustering.newBuilder().setFields(ImmutableList.of("Foo", "Bar")).build(); + private static final Map LABELS = + ImmutableMap.of("test-job-name", "test-write-channel"); private static final WriteChannelConfiguration LOAD_CONFIGURATION_CSV = WriteChannelConfiguration.newBuilder(TABLE_ID) .setCreateDisposition(CREATE_DISPOSITION) @@ -68,6 +72,7 @@ public class WriteChannelConfigurationTest { .setAutodetect(AUTODETECT) .setTimePartitioning(TIME_PARTITIONING) .setClustering(CLUSTERING) + .setLabels(LABELS) .build(); private static final DatastoreBackupOptions BACKUP_OPTIONS = @@ -151,6 +156,7 @@ public void testBuilder() { assertEquals(IGNORE_UNKNOWN_VALUES, LOAD_CONFIGURATION_CSV.ignoreUnknownValues()); assertEquals(MAX_BAD_RECORDS, LOAD_CONFIGURATION_CSV.getMaxBadRecords()); assertEquals(TABLE_SCHEMA, LOAD_CONFIGURATION_CSV.getSchema()); + assertEquals(LABELS, LOAD_CONFIGURATION_CSV.getLabels()); assertEquals(BACKUP_OPTIONS, LOAD_CONFIGURATION_BACKUP.getDatastoreBackupOptions()); assertEquals(SCHEMA_UPDATE_OPTIONS, LOAD_CONFIGURATION_CSV.getSchemaUpdateOptions()); assertEquals(SCHEMA_UPDATE_OPTIONS, LOAD_CONFIGURATION_BACKUP.getSchemaUpdateOptions()); @@ -218,5 +224,6 @@ private void compareLoadConfiguration( assertEquals(expected.getTimePartitioning(), value.getTimePartitioning()); assertEquals(expected.getClustering(), value.getClustering()); assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes()); + assertEquals(expected.getLabels(), value.getLabels()); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 9e334c8a8..cd69d8afd 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -2105,6 +2105,40 @@ public void testInsertFromFile() throws InterruptedException, IOException, Timeo assertTrue(bigquery.delete(tableId)); } + @Test + public void testInsertFromFileWithLabels() + throws InterruptedException, IOException, TimeoutException { + String destinationTableName = "test_insert_from_file_table_with_labels"; + TableId tableId = TableId.of(DATASET, destinationTableName); + WriteChannelConfiguration configuration = + WriteChannelConfiguration.newBuilder(tableId) + .setFormatOptions(FormatOptions.json()) + .setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED) + .setSchema(TABLE_SCHEMA) + .setLabels(LABELS) + .build(); + TableDataWriteChannel channel = bigquery.writer(configuration); + try { + // A zero byte write should not throw an exception. + assertEquals(0, channel.write(ByteBuffer.wrap("".getBytes(StandardCharsets.UTF_8)))); + } finally { + // Force the channel to flush by calling `close`. + channel.close(); + } + channel = bigquery.writer(configuration); + try { + channel.write(ByteBuffer.wrap(JSON_CONTENT.getBytes(StandardCharsets.UTF_8))); + } finally { + channel.close(); + } + Job job = channel.getJob().waitFor(); + LoadJobConfiguration jobConfiguration = job.getConfiguration(); + assertEquals(TABLE_SCHEMA, jobConfiguration.getSchema()); + assertEquals(LABELS, jobConfiguration.getLabels()); + assertNull(job.getStatus().getError()); + assertTrue(bigquery.delete(tableId)); + } + @Test public void testLocation() throws Exception { String location = "EU";