Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bigquery: Expose labels for ExtractJobConfiguration. #6376

Merged
merged 1 commit into from
Oct 2, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand All @@ -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<String, String> labels;

public static final class Builder
extends JobConfiguration.Builder<ExtractJobConfiguration, Builder> {
Expand All @@ -54,6 +56,7 @@ public static final class Builder
private String format;
private String compression;
private Boolean useAvroLogicalTypes;
private Map<String, String> labels;

private Builder() {
super(Type.EXTRACT);
Expand All @@ -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) {
Expand All @@ -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. */
Expand Down Expand Up @@ -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<String, String> labels) {
this.labels = labels;
return this;
}

public ExtractJobConfiguration build() {
return new ExtractJobConfiguration(this);
}
Expand All @@ -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. */
Expand Down Expand Up @@ -204,6 +226,11 @@ public Boolean getUseAvroLogicalTypes() {
return useAvroLogicalTypes;
}

/** Returns the labels associated with this job */
public Map<String, String> getLabels() {
return labels;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand All @@ -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
Expand All @@ -237,7 +265,8 @@ public int hashCode() {
fieldDelimiter,
format,
compression,
useAvroLogicalTypes);
useAvroLogicalTypes,
labels);
}

@Override
Expand All @@ -251,15 +280,20 @@ 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);
extractConfigurationPb.setFieldDelimiter(fieldDelimiter);
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;
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -37,19 +39,23 @@ 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<String, String> 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)
.setPrintHeader(PRINT_HEADER)
.setFieldDelimiter(FIELD_DELIMITER)
.setCompression(COMPRESSION)
.setFormat(FORMAT)
.setLabels(LABELS)
.build();
private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_AVRO =
ExtractJobConfiguration.newBuilder(TABLE_ID, DESTINATION_URI)
Expand All @@ -58,6 +64,7 @@ public class ExtractJobConfigurationTest {
.setCompression(COMPRESSION)
.setFormat(AVRO_FORMAT)
.setUseAvroLogicalTypes(USEAVROLOGICALTYPES)
.setLabels(LABELS)
.build();

@Test
Expand Down Expand Up @@ -113,20 +120,23 @@ 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());
assertEquals(FIELD_DELIMITER, EXTRACT_CONFIGURATION_ONE_URI.getFieldDelimiter());
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());
assertEquals(COMPRESSION, EXTRACT_CONFIGURATION_AVRO.getCompression());
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
Expand All @@ -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(
Expand Down Expand Up @@ -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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> 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";
Expand Down