Bigquery: Expose labels for LoadJobConfiguration. (#6387)

* expose labels * modified test case
googleapis · Oct 2, 2019 · 620af66 · 620af66
1 parent 04b5c59
commit 620af66
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 3 deletions.
diff --git a/...s/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java b/...s/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java
@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Ints;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -50,6 +51,7 @@ public final class LoadJobConfiguration extends JobConfiguration implements Load
   private final TimePartitioning timePartitioning;
   private final Clustering clustering;
   private final Boolean useAvroLogicalTypes;
+  private final Map<String, String> labels;
 
   public static final class Builder extends JobConfiguration.Builder<LoadJobConfiguration, Builder>
       implements LoadConfiguration.Builder {
@@ -70,6 +72,7 @@ public static final class Builder extends JobConfiguration.Builder<LoadJobConfig
     private TimePartitioning timePartitioning;
     private Clustering clustering;
     private Boolean useAvroLogicalTypes;
+    private Map<String, String> labels;
 
     private Builder() {
       super(Type.LOAD);
@@ -93,6 +96,7 @@ private Builder(LoadJobConfiguration loadConfiguration) {
       this.timePartitioning = loadConfiguration.timePartitioning;
       this.clustering = loadConfiguration.clustering;
       this.useAvroLogicalTypes = loadConfiguration.useAvroLogicalTypes;
+      this.labels = loadConfiguration.labels;
     }
 
     private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -166,6 +170,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
                     loadConfigurationPb.getDestinationEncryptionConfiguration())
                 .build();
       }
+      if (configurationPb.getLabels() != null) {
+        this.labels = configurationPb.getLabels();
+      }
     }
 
     @Override
@@ -263,6 +270,20 @@ public Builder setSchemaUpdateOptions(List<JobInfo.SchemaUpdateOption> schemaUpd
       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;
+    }
+
     @Override
     public LoadJobConfiguration build() {
       return new LoadJobConfiguration(this);
@@ -286,6 +307,7 @@ private LoadJobConfiguration(Builder builder) {
     this.timePartitioning = builder.timePartitioning;
     this.clustering = builder.clustering;
     this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
+    this.labels = builder.labels;
   }
 
   @Override
@@ -378,6 +400,11 @@ public List<JobInfo.SchemaUpdateOption> getSchemaUpdateOptions() {
     return schemaUpdateOptions;
   }
 
+  /** Returns the labels associated with this job */
+  public Map<String, String> getLabels() {
+    return labels;
+  }
+
   @Override
   public Builder toBuilder() {
     return new Builder(this);
@@ -400,7 +427,8 @@ ToStringHelper toStringHelper() {
         .add("autodetect", autodetect)
         .add("timePartitioning", timePartitioning)
         .add("clustering", clustering)
-        .add("useAvroLogicalTypes", useAvroLogicalTypes);
+        .add("useAvroLogicalTypes", useAvroLogicalTypes)
+        .add("labels", labels);
   }
 
   @Override
@@ -425,6 +453,8 @@ LoadJobConfiguration setProjectId(String projectId) {
   @Override
   com.google.api.services.bigquery.model.JobConfiguration toPb() {
     JobConfigurationLoad loadConfigurationPb = new JobConfigurationLoad();
+    com.google.api.services.bigquery.model.JobConfiguration jobConfiguration =
+        new com.google.api.services.bigquery.model.JobConfiguration();
     loadConfigurationPb.setDestinationTable(destinationTable.toPb());
     if (createDisposition != null) {
       loadConfigurationPb.setCreateDisposition(createDisposition.toString());
@@ -482,8 +512,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;
   }
 
   /**

diff --git a/...ogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java b/...ogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.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 LoadJobConfigurationTest {
@@ -58,6 +60,8 @@ public class LoadJobConfigurationTest {
   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<String, String> LABELS =
+      ImmutableMap.of("test-job-name", "test-load-job");
   private static final LoadJobConfiguration LOAD_CONFIGURATION_CSV =
       LoadJobConfiguration.newBuilder(TABLE_ID, SOURCE_URIS)
           .setCreateDisposition(CREATE_DISPOSITION)
@@ -71,6 +75,7 @@ public class LoadJobConfigurationTest {
           .setDestinationEncryptionConfiguration(JOB_ENCRYPTION_CONFIGURATION)
           .setTimePartitioning(TIME_PARTITIONING)
           .setClustering(CLUSTERING)
+          .setLabels(LABELS)
           .build();
   private static final DatastoreBackupOptions BACKUP_OPTIONS =
       DatastoreBackupOptions.newBuilder()
@@ -86,6 +91,7 @@ public class LoadJobConfigurationTest {
           .setSchema(TABLE_SCHEMA)
           .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
           .setAutodetect(AUTODETECT)
+          .setLabels(LABELS)
           .build();
   private static final LoadJobConfiguration LOAD_CONFIGURATION_AVRO =
       LoadJobConfiguration.newBuilder(TABLE_ID, SOURCE_URIS)
@@ -101,6 +107,7 @@ public class LoadJobConfigurationTest {
           .setTimePartitioning(TIME_PARTITIONING)
           .setClustering(CLUSTERING)
           .setUseAvroLogicalTypes(USERAVROLOGICALTYPES)
+          .setLabels(LABELS)
           .build();
 
   @Test
@@ -216,5 +223,6 @@ private void compareLoadJobConfiguration(
     assertEquals(expected.getTimePartitioning(), value.getTimePartitioning());
     assertEquals(expected.getClustering(), value.getClustering());
     assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes());
+    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
@@ -279,6 +279,7 @@ public class ITBigQueryTest {
   public static void beforeClass() throws InterruptedException, TimeoutException {
     RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
     RemoteStorageHelper storageHelper = RemoteStorageHelper.create();
+    Map<String, String> labels = ImmutableMap.of("test-job-name", "test-load-job");
     bigquery = bigqueryHelper.getOptions().getService();
     storage = storageHelper.getOptions().getService();
     storage.create(BucketInfo.of(BUCKET));
@@ -302,10 +303,13 @@ public static void beforeClass() throws InterruptedException, TimeoutException {
                 TABLE_ID, "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json())
             .setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
             .setSchema(TABLE_SCHEMA)
+            .setLabels(labels)
             .build();
     Job job = bigquery.create(JobInfo.of(configuration));
     job = job.waitFor();
     assertNull(job.getStatus().getError());
+    LoadJobConfiguration loadJobConfiguration = job.getConfiguration();
+    assertEquals(labels, loadJobConfiguration.getLabels());
   }
 
   @AfterClass
@@ -1557,13 +1561,17 @@ public void testQueryJobWithDryRun() throws InterruptedException, TimeoutExcepti
   public void testExtractJob() throws InterruptedException, TimeoutException {
     String tableName = "test_export_job_table";
     TableId destinationTable = TableId.of(DATASET, tableName);
+    Map<String, String> labels = ImmutableMap.of("test-job-name", "test-load-extract-job");
     LoadJobConfiguration configuration =
         LoadJobConfiguration.newBuilder(destinationTable, "gs://" + BUCKET + "/" + LOAD_FILE)
             .setSchema(SIMPLE_SCHEMA)
+            .setLabels(labels)
             .build();
     Job remoteLoadJob = bigquery.create(JobInfo.of(configuration));
     remoteLoadJob = remoteLoadJob.waitFor();
     assertNull(remoteLoadJob.getStatus().getError());
+    LoadJobConfiguration loadJobConfiguration = remoteLoadJob.getConfiguration();
+    assertEquals(labels, loadJobConfiguration.getLabels());
 
     ExtractJobConfiguration extractConfiguration =
         ExtractJobConfiguration.newBuilder(destinationTable, "gs://" + BUCKET + "/" + EXTRACT_FILE)