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 CopyJobConfiguration. #6367

Merged
merged 2 commits 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.collect.ImmutableList;
import com.google.common.collect.Lists;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
Expand All @@ -40,6 +41,7 @@ public final class CopyJobConfiguration extends JobConfiguration {
private final JobInfo.CreateDisposition createDisposition;
private final JobInfo.WriteDisposition writeDisposition;
private final EncryptionConfiguration destinationEncryptionConfiguration;
private final Map<String, String> labels;

public static final class Builder
extends JobConfiguration.Builder<CopyJobConfiguration, Builder> {
Expand All @@ -49,6 +51,7 @@ public static final class Builder
private JobInfo.CreateDisposition createDisposition;
private JobInfo.WriteDisposition writeDisposition;
private EncryptionConfiguration destinationEncryptionConfiguration;
private Map<String, String> labels;

private Builder() {
super(Type.COPY);
Expand All @@ -61,6 +64,7 @@ private Builder(CopyJobConfiguration jobConfiguration) {
this.createDisposition = jobConfiguration.createDisposition;
this.writeDisposition = jobConfiguration.writeDisposition;
this.destinationEncryptionConfiguration = jobConfiguration.destinationEncryptionConfiguration;
this.labels = jobConfiguration.labels;
}

private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
Expand All @@ -87,6 +91,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
copyConfigurationPb.getDestinationEncryptionConfiguration())
.build();
}
if (configurationPb.getLabels() != null) {
this.labels = configurationPb.getLabels();
}
}

/** Sets the source tables to copy. */
Expand Down Expand Up @@ -131,6 +138,20 @@ public Builder setWriteDisposition(JobInfo.WriteDisposition writeDisposition) {
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 CopyJobConfiguration build() {
return new CopyJobConfiguration(this);
}
Expand All @@ -143,6 +164,7 @@ private CopyJobConfiguration(Builder builder) {
this.createDisposition = builder.createDisposition;
this.writeDisposition = builder.writeDisposition;
this.destinationEncryptionConfiguration = builder.destinationEncryptionConfiguration;
this.labels = builder.labels;
}

/** Returns the source tables to copy. */
Expand Down Expand Up @@ -181,6 +203,11 @@ public JobInfo.WriteDisposition getWriteDisposition() {
return writeDisposition;
}

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

@Override
public Builder toBuilder() {
return new Builder(this);
Expand All @@ -193,7 +220,8 @@ ToStringHelper toStringHelper() {
.add("destinationTable", destinationTable)
.add("destinationEncryptionConfiguration", destinationEncryptionConfiguration)
.add("createDisposition", createDisposition)
.add("writeDisposition", writeDisposition);
.add("writeDisposition", writeDisposition)
.add("labels", labels);
}

@Override
Expand All @@ -205,7 +233,12 @@ public boolean equals(Object obj) {
@Override
public int hashCode() {
return Objects.hash(
baseHashCode(), sourceTables, destinationTable, createDisposition, writeDisposition);
baseHashCode(),
sourceTables,
destinationTable,
createDisposition,
writeDisposition,
labels);
}

@Override
Expand All @@ -232,6 +265,8 @@ public TableId apply(TableId tableId) {
@Override
com.google.api.services.bigquery.model.JobConfiguration toPb() {
JobConfigurationTableCopy configurationPb = new JobConfigurationTableCopy();
com.google.api.services.bigquery.model.JobConfiguration jobConfiguration =
new com.google.api.services.bigquery.model.JobConfiguration();
configurationPb.setDestinationTable(destinationTable.toPb());
if (sourceTables.size() == 1) {
configurationPb.setSourceTable(sourceTables.get(0).toPb());
Expand All @@ -248,7 +283,11 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
configurationPb.setDestinationEncryptionConfiguration(
destinationEncryptionConfiguration.toPb());
}
return new com.google.api.services.bigquery.model.JobConfiguration().setCopy(configurationPb);
if (labels != null) {
jobConfiguration.setLabels(labels);
}
jobConfiguration.setCopy(configurationPb);
return jobConfiguration;
}

/** Creates a builder for a BigQuery Copy Job configuration given destination and source table. */
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@
import com.google.cloud.bigquery.JobInfo.WriteDisposition;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import java.util.List;
import java.util.Map;
import org.junit.Test;

public class CopyJobConfigurationTest {
Expand All @@ -40,16 +42,19 @@ public class CopyJobConfigurationTest {
private static final WriteDisposition WRITE_DISPOSITION = WriteDisposition.WRITE_APPEND;
private static final EncryptionConfiguration COPY_JOB_ENCRYPTION_CONFIGURATION =
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
private static final Map<String, String> LABELS = ImmutableMap.of("job-name", "copy");
private static final CopyJobConfiguration COPY_JOB_CONFIGURATION =
CopyJobConfiguration.newBuilder(DESTINATION_TABLE, SOURCE_TABLE)
.setCreateDisposition(CREATE_DISPOSITION)
.setWriteDisposition(WRITE_DISPOSITION)
.setDestinationEncryptionConfiguration(COPY_JOB_ENCRYPTION_CONFIGURATION)
.setLabels(LABELS)
.build();
private static final CopyJobConfiguration COPY_JOB_CONFIGURATION_MULTIPLE_TABLES =
CopyJobConfiguration.newBuilder(DESTINATION_TABLE, SOURCE_TABLES)
.setCreateDisposition(CREATE_DISPOSITION)
.setWriteDisposition(WRITE_DISPOSITION)
.setLabels(LABELS)
.build();

@Test
Expand Down Expand Up @@ -95,6 +100,7 @@ public void testBuilder() {
assertEquals(ImmutableList.of(SOURCE_TABLE), COPY_JOB_CONFIGURATION.getSourceTables());
assertEquals(CREATE_DISPOSITION, COPY_JOB_CONFIGURATION.getCreateDisposition());
assertEquals(WRITE_DISPOSITION, COPY_JOB_CONFIGURATION.getWriteDisposition());
assertEquals(LABELS, COPY_JOB_CONFIGURATION.getLabels());
}

@Test
Expand All @@ -105,6 +111,8 @@ public void testToPbAndFromPb() {
assertNull(COPY_JOB_CONFIGURATION.toPb().getQuery());
assertNull(COPY_JOB_CONFIGURATION.toPb().getCopy().getSourceTables());
assertNull(COPY_JOB_CONFIGURATION_MULTIPLE_TABLES.toPb().getCopy().getSourceTable());
assertNotNull(COPY_JOB_CONFIGURATION.getLabels());
assertNotNull(COPY_JOB_CONFIGURATION_MULTIPLE_TABLES.getLabels());
compareCopyJobConfiguration(
COPY_JOB_CONFIGURATION, CopyJobConfiguration.fromPb(COPY_JOB_CONFIGURATION.toPb()));
compareCopyJobConfiguration(
Expand Down Expand Up @@ -167,5 +175,6 @@ private void compareCopyJobConfiguration(
assertEquals(
expected.getDestinationEncryptionConfiguration(),
value.getDestinationEncryptionConfiguration());
assertEquals(expected.getLabels(), value.getLabels());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -1476,6 +1476,30 @@ public void testCopyJob() throws InterruptedException, TimeoutException {
assertTrue(remoteTable.delete());
}

@Test
public void testCopyJobWithLabels() throws InterruptedException {
String sourceTableName = "test_copy_job_source_table_label";
String destinationTableName = "test_copy_job_destination_table_label";
Map<String, String> labels = ImmutableMap.of("test_job_name", "test_copy_job");
TableId sourceTable = TableId.of(DATASET, sourceTableName);
StandardTableDefinition tableDefinition = StandardTableDefinition.of(TABLE_SCHEMA);
TableInfo tableInfo = TableInfo.of(sourceTable, tableDefinition);
Table createdTable = bigquery.create(tableInfo);
assertNotNull(createdTable);
TableId destinationTable = TableId.of(DATASET, destinationTableName);
CopyJobConfiguration configuration =
CopyJobConfiguration.newBuilder(destinationTable, sourceTable).setLabels(labels).build();
Job remoteJob = bigquery.create(JobInfo.of(configuration));
remoteJob = remoteJob.waitFor();
assertNull(remoteJob.getStatus().getError());
CopyJobConfiguration copyJobConfiguration = remoteJob.getConfiguration();
assertEquals(labels, copyJobConfiguration.getLabels());
Table remoteTable = bigquery.getTable(DATASET, destinationTableName);
assertNotNull(remoteTable);
assertTrue(createdTable.delete());
assertTrue(remoteTable.delete());
}

@Test
public void testQueryJob() throws InterruptedException, TimeoutException {
String tableName = "test_query_job_table";
Expand Down