diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java index 777e23c6ac17..6aed6ecf0766 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java @@ -25,6 +25,7 @@ import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import java.util.List; +import java.util.Map; import java.util.Objects; /** @@ -43,6 +44,7 @@ public final class ExtractJobConfiguration extends JobConfiguration { private final String format; private final String compression; private final Boolean useAvroLogicalTypes; + private final Map labels; public static final class Builder extends JobConfiguration.Builder { @@ -54,6 +56,7 @@ public static final class Builder private String format; private String compression; private Boolean useAvroLogicalTypes; + private Map labels; private Builder() { super(Type.EXTRACT); @@ -68,6 +71,7 @@ private Builder(ExtractJobConfiguration jobInfo) { this.format = jobInfo.format; this.compression = jobInfo.compression; this.useAvroLogicalTypes = jobInfo.useAvroLogicalTypes; + this.labels = jobInfo.labels; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -80,6 +84,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur this.format = extractConfigurationPb.getDestinationFormat(); this.compression = extractConfigurationPb.getCompression(); this.useAvroLogicalTypes = extractConfigurationPb.getUseAvroLogicalTypes(); + if (configurationPb.getLabels() != null) { + this.labels = configurationPb.getLabels(); + } } /** Sets the table to export. */ @@ -146,6 +153,20 @@ public Builder setUseAvroLogicalTypes(Boolean useAvroLogicalTypes) { return this; } + /** + * The labels associated with this job. You can use these to organize and group your jobs. Label + * keys and values can be no longer than 63 characters, can only contain lowercase letters, + * numeric characters, underscores and dashes. International characters are allowed. Label + * values are optional. Label keys must start with a letter and each label in the list must have + * a different key. + * + * @param labels labels or {@code null} for none + */ + public Builder setLabels(Map labels) { + this.labels = labels; + return this; + } + public ExtractJobConfiguration build() { return new ExtractJobConfiguration(this); } @@ -160,6 +181,7 @@ private ExtractJobConfiguration(Builder builder) { this.format = builder.format; this.compression = builder.compression; this.useAvroLogicalTypes = builder.useAvroLogicalTypes; + this.labels = builder.labels; } /** Returns the table to export. */ @@ -204,6 +226,11 @@ public Boolean getUseAvroLogicalTypes() { return useAvroLogicalTypes; } + /** Returns the labels associated with this job */ + public Map getLabels() { + return labels; + } + @Override public Builder toBuilder() { return new Builder(this); @@ -218,7 +245,8 @@ ToStringHelper toStringHelper() { .add("printHeader", printHeader) .add("fieldDelimiter", fieldDelimiter) .add("compression", compression) - .add("useAvroLogicalTypes", useAvroLogicalTypes); + .add("useAvroLogicalTypes", useAvroLogicalTypes) + .add("labels", labels); } @Override @@ -237,7 +265,8 @@ public int hashCode() { fieldDelimiter, format, compression, - useAvroLogicalTypes); + useAvroLogicalTypes, + labels); } @Override @@ -251,6 +280,8 @@ ExtractJobConfiguration setProjectId(String projectId) { @Override com.google.api.services.bigquery.model.JobConfiguration toPb() { JobConfigurationExtract extractConfigurationPb = new JobConfigurationExtract(); + com.google.api.services.bigquery.model.JobConfiguration jobConfiguration = + new com.google.api.services.bigquery.model.JobConfiguration(); extractConfigurationPb.setDestinationUris(destinationUris); extractConfigurationPb.setSourceTable(sourceTable.toPb()); extractConfigurationPb.setPrintHeader(printHeader); @@ -258,8 +289,11 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { extractConfigurationPb.setDestinationFormat(format); extractConfigurationPb.setCompression(compression); extractConfigurationPb.setUseAvroLogicalTypes(useAvroLogicalTypes); - return new com.google.api.services.bigquery.model.JobConfiguration() - .setExtract(extractConfigurationPb); + if (labels != null) { + jobConfiguration.setLabels(labels); + } + jobConfiguration.setExtract(extractConfigurationPb); + return jobConfiguration; } /** diff --git a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java index 1611145777c1..8a1a650372e0 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java +++ b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java @@ -21,7 +21,9 @@ import static org.junit.Assert.assertNull; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import java.util.List; +import java.util.Map; import org.junit.Test; public class ExtractJobConfigurationTest { @@ -37,12 +39,15 @@ public class ExtractJobConfigurationTest { private static final Boolean PRINT_HEADER = true; private static final String COMPRESSION = "GZIP"; private static final Boolean USEAVROLOGICALTYPES = true; + private static final Map LABELS = + ImmutableMap.of("test-job-name", "test-extract-job"); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URIS) .setPrintHeader(PRINT_HEADER) .setFieldDelimiter(FIELD_DELIMITER) .setCompression(COMPRESSION) .setFormat(FORMAT) + .setLabels(LABELS) .build(); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_ONE_URI = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URI) @@ -50,6 +55,7 @@ public class ExtractJobConfigurationTest { .setFieldDelimiter(FIELD_DELIMITER) .setCompression(COMPRESSION) .setFormat(FORMAT) + .setLabels(LABELS) .build(); private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_AVRO = ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URI) @@ -58,6 +64,7 @@ public class ExtractJobConfigurationTest { .setCompression(COMPRESSION) .setFormat(AVRO_FORMAT) .setUseAvroLogicalTypes(USEAVROLOGICALTYPES) + .setLabels(LABELS) .build(); @Test @@ -113,6 +120,7 @@ public void testBuilder() { assertEquals(COMPRESSION, EXTRACT_CONFIGURATION.getCompression()); assertEquals(PRINT_HEADER, EXTRACT_CONFIGURATION.printHeader()); assertEquals(FORMAT, EXTRACT_CONFIGURATION.getFormat()); + assertEquals(LABELS, EXTRACT_CONFIGURATION.getLabels()); assertEquals(TABLE_ID, EXTRACT_CONFIGURATION_ONE_URI.getSourceTable()); assertEquals( ImmutableList.of(DESTINATION_URI), EXTRACT_CONFIGURATION_ONE_URI.getDestinationUris()); @@ -120,6 +128,7 @@ public void testBuilder() { assertEquals(COMPRESSION, EXTRACT_CONFIGURATION_ONE_URI.getCompression()); assertEquals(PRINT_HEADER, EXTRACT_CONFIGURATION_ONE_URI.printHeader()); assertEquals(FORMAT, EXTRACT_CONFIGURATION_ONE_URI.getFormat()); + assertEquals(LABELS, EXTRACT_CONFIGURATION_ONE_URI.getLabels()); assertEquals( ImmutableList.of(DESTINATION_URI), EXTRACT_CONFIGURATION_AVRO.getDestinationUris()); assertEquals(FIELD_DELIMITER, EXTRACT_CONFIGURATION_AVRO.getFieldDelimiter()); @@ -127,6 +136,7 @@ public void testBuilder() { assertEquals(PRINT_HEADER, EXTRACT_CONFIGURATION_AVRO.printHeader()); assertEquals(AVRO_FORMAT, EXTRACT_CONFIGURATION_AVRO.getFormat()); assertEquals(USEAVROLOGICALTYPES, EXTRACT_CONFIGURATION_AVRO.getUseAvroLogicalTypes()); + assertEquals(LABELS, EXTRACT_CONFIGURATION_AVRO.getLabels()); } @Test @@ -135,6 +145,7 @@ public void testToPbAndFromPb() { assertNull(EXTRACT_CONFIGURATION.toPb().getCopy()); assertNull(EXTRACT_CONFIGURATION.toPb().getLoad()); assertNull(EXTRACT_CONFIGURATION.toPb().getQuery()); + assertNotNull(EXTRACT_CONFIGURATION.toPb().getLabels()); compareExtractJobConfiguration( EXTRACT_CONFIGURATION, ExtractJobConfiguration.fromPb(EXTRACT_CONFIGURATION.toPb())); compareExtractJobConfiguration( @@ -182,5 +193,6 @@ private void compareExtractJobConfiguration( assertEquals(expected.printHeader(), value.printHeader()); assertEquals(expected.getFieldDelimiter(), value.getFieldDelimiter()); assertEquals(expected.getFormat(), value.getFormat()); + assertEquals(expected.getLabels(), value.getLabels()); } } diff --git a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 58869a3848d4..5906730d5607 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -1556,6 +1556,32 @@ public void testExtractJob() throws InterruptedException, TimeoutException { assertTrue(bigquery.delete(destinationTable)); } + @Test + public void testExtractJobWithLabels() throws InterruptedException, TimeoutException { + String tableName = "test_export_job_table_label"; + Map labels = ImmutableMap.of("test_job_name", "test_export_job"); + TableId destinationTable = TableId.of(DATASET, tableName); + LoadJobConfiguration configuration = + LoadJobConfiguration.newBuilder(destinationTable, "gs://" + BUCKET + "/" + LOAD_FILE) + .setSchema(SIMPLE_SCHEMA) + .build(); + Job remoteLoadJob = bigquery.create(JobInfo.of(configuration)); + remoteLoadJob = remoteLoadJob.waitFor(); + assertNull(remoteLoadJob.getStatus().getError()); + + ExtractJobConfiguration extractConfiguration = + ExtractJobConfiguration.newBuilder(destinationTable, "gs://" + BUCKET + "/" + EXTRACT_FILE) + .setLabels(labels) + .setPrintHeader(false) + .build(); + Job remoteExtractJob = bigquery.create(JobInfo.of(extractConfiguration)); + remoteExtractJob = remoteExtractJob.waitFor(); + assertNull(remoteExtractJob.getStatus().getError()); + ExtractJobConfiguration extractJobConfiguration = remoteExtractJob.getConfiguration(); + assertEquals(labels, extractJobConfiguration.getLabels()); + assertTrue(bigquery.delete(destinationTable)); + } + @Test public void testCancelJob() throws InterruptedException, TimeoutException { String destinationTableName = "test_cancel_query_job_table";