Bigquery: Expose labels for CopyJobConfiguration. (#6367)

* fix labels * modified code
googleapis · Oct 2, 2019 · 8d569df · 8d569df
1 parent 8d16214
commit 8d569df
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 3 deletions.
diff --git a/...s/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java b/...s/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java
@@ -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;
 
 /**
@@ -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> {
@@ -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);
@@ -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) {
@@ -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. */
@@ -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);
     }
@@ -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. */
@@ -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);
@@ -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
@@ -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
@@ -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());
@@ -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. */

diff --git a/...ogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java b/...ogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java
@@ -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 {
@@ -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
@@ -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
@@ -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(
@@ -167,5 +175,6 @@ private void compareCopyJobConfiguration(
     assertEquals(
         expected.getDestinationEncryptionConfiguration(),
         value.getDestinationEncryptionConfiguration());
+    assertEquals(expected.getLabels(), value.getLabels());
   }
 }
diff --git a/...ents/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/...ents/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -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";