bigquery: add table label

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java

@@ -25,7 +25,6 @@
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import java.io.Serializable;
 import java.util.List;
@@ -69,7 +68,7 @@ public Dataset apply(DatasetInfo datasetInfo) {
   private final Long lastModified;
   private final String location;
   private final String selfLink;
-  private final Map<String, String> labels;
+  private final Labels labels;
 
   /**
    * A builder for {@code DatasetInfo} objects.
@@ -157,7 +156,7 @@ static final class BuilderImpl extends Builder {
     private Long lastModified;
     private String location;
     private String selfLink;
-    private Map<String, String> labels;
+    private Labels labels = Labels.ZERO;
 
     BuilderImpl() {}
 
@@ -173,9 +172,7 @@ static final class BuilderImpl extends Builder {
       this.lastModified = datasetInfo.lastModified;
       this.location = datasetInfo.location;
       this.selfLink = datasetInfo.selfLink;
-      this.labels = datasetInfo.labels != null
-                    ? ImmutableMap.copyOf(datasetInfo.labels)
-                    : null;
+      this.labels = datasetInfo.labels;
     }
 
     BuilderImpl(com.google.api.services.bigquery.model.Dataset datasetPb) {
@@ -199,9 +196,7 @@ public Acl apply(Dataset.Access accessPb) {
       this.lastModified = datasetPb.getLastModifiedTime();
       this.location = datasetPb.getLocation();
       this.selfLink = datasetPb.getSelfLink();
-      this.labels = datasetPb.getLabels() != null
-                    ? ImmutableMap.copyOf(datasetPb.getLabels())
-                    : null;
+      this.labels = Labels.fromPb(datasetPb.getLabels());
     }
 
 
@@ -277,9 +272,16 @@ Builder setSelfLink(String selfLink) {
       return this;
     }
 
+    /**
+     * Sets the labels applied to this dataset.
+     *
+     * <p>When used with {@link BigQuery#update(DatasetInfo, DatasetOption...)}, setting {@code
+     * labels} to {@code null} removes all labels; otherwise all keys that are mapped to {@code
+     * null} values are removed and other keys are updated to their respective values.
+     */
     @Override
     public Builder setLabels(Map<String, String> labels) {
-      this.labels = ImmutableMap.copyOf(labels);
+      this.labels = Labels.fromUser(labels);
       return this;
     }
 
@@ -411,7 +413,7 @@ public String getSelfLink() {
    * @see <a href="https://cloud.google.com/bigquery/docs/labeling-datasets">Labeling Datasets</a>
    */
   public Map<String, String> getLabels() {
-    return labels;
+    return labels.userMap();
   }
 
   /**
@@ -494,9 +496,7 @@ public Dataset.Access apply(Acl acl) {
         }
       }));
     }
-    if (labels != null) {
-      datasetPb.setLabels(labels);
-    }
+    datasetPb.setLabels(labels.toPb());
     return datasetPb;
   }
 

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Labels.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import com.google.api.client.util.Data;
+import com.google.auto.value.AutoValue;
+import com.google.auto.value.extension.memoized.Memoized;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import java.io.Serializable;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+
+@AutoValue
+abstract class Labels implements Serializable {
+  private static final long serialVersionUID = 1L;
+  static final Labels ZERO = of(Collections.<String, String>emptyMap());
+
+  @Nullable
+  abstract Map<String, String> userMap();
+
+  @Memoized
+  @Nullable
+  Map<String, String> toPb() {
+    Map<String, String> userMap = userMap();
+    if (userMap == null) {
+      // converted to JSON null
+      return Data.nullOf(HashMap.class);
+    }
+    if (userMap.isEmpty()) {
+      // dropped from JSON
+      return null;
+    }
+    HashMap<String, String> pbMap = new HashMap<>();
+    for (Map.Entry<String, String> entry : userMap.entrySet()) {
+      String key = entry.getKey();
+      String val = entry.getValue();
+      if (val == null) {
+        val = Data.NULL_STRING;
+      }
+      pbMap.put(key, val);
+    }
+    return Collections.unmodifiableMap(pbMap);
+  }
+
+  private static Labels of(Map<String, String> userMap) {
+    Preconditions.checkArgument(
+        userMap == null || !userMap.containsKey(null), "null keys are not supported");
+    return new AutoValue_Labels(userMap);
+  }
+
+  static Labels fromUser(Map<String, String> map) {
+    if (map == null || map instanceof ImmutableMap) {
+      return of(map);
+    }
+    // Not ImmutableMap; we support null values!
+    return of(Collections.unmodifiableMap(new HashMap<>(map)));
+  }
+
+  static Labels fromPb(Map<String, String> pb) {
+    if (Data.isNull(pb)) {
+      return of(null);
+    }
+    if (pb == null || pb.isEmpty()) {
+      return of(Collections.<String, String>emptyMap());
+    }
+
+    HashMap<String, String> map = new HashMap<>();
+    for (Map.Entry<String, String> entry : pb.entrySet()) {
+      String key = entry.getKey();
+      String val = entry.getValue();
+      if (Data.isNull(val)) {
+        val = null;
+      }
+      map.put(key, val);
+    }
+    return of(Collections.unmodifiableMap(map));
+  }
+}

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Table.java

@@ -18,7 +18,6 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
-import com.google.api.gax.paging.Page;
 import com.google.cloud.bigquery.BigQuery.JobOption;
 import com.google.cloud.bigquery.BigQuery.TableDataListOption;
 import com.google.cloud.bigquery.BigQuery.TableOption;
@@ -27,6 +26,7 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -135,6 +135,12 @@ public TableInfo.Builder setEncryptionConfiguration(EncryptionConfiguration conf
       return this;
     }
 
+    @Override
+    public Builder setLabels(Map<String, String> labels) {
+      infoBuilder.setLabels(labels);
+      return this;
+    }
+
     @Override
     public Table build() {
       return new Table(bigquery, infoBuilder);

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java

@@ -20,12 +20,13 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.api.client.util.Data;
+import com.google.api.core.BetaApi;
 import com.google.api.services.bigquery.model.Table;
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
-
 import java.io.Serializable;
 import java.math.BigInteger;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -65,6 +66,7 @@ public Table apply(TableInfo tableInfo) {
   private final Long lastModifiedTime;
   private final TableDefinition definition;
   private final EncryptionConfiguration encryptionConfiguration;
+  private final Labels labels;
 
   /**
    * A builder for {@code TableInfo} objects.
@@ -100,7 +102,6 @@ public abstract static class Builder {
 
     abstract Builder setSelfLink(String selfLink);
 
-
     /**
      * Sets the table identity.
      */
@@ -114,6 +115,19 @@ public abstract static class Builder {
      */
     public abstract Builder setDefinition(TableDefinition definition);
 
+    /**
+     * Sets the labels applied to this table.
+     *
+     * <p>Unstable, because labels are <a
+     * href="https://cloud.google.com/bigquery/docs/reference/rest/v2/tables">experimental</a>.
+     *
+     * <p>When used with {@link BigQuery#update(TableInfo, TableOption...)}, setting {@code labels}
+     * to {@code null} removes all labels; otherwise all keys that are mapped to {@code null} values
+     * are removed and other keys are updated to their respective values.
+     */
+    @BetaApi
+    public abstract Builder setLabels(Map<String, String> labels);
+
     /**
      * Creates a {@code TableInfo} object.
      */
@@ -135,6 +149,7 @@ static class BuilderImpl extends Builder {
     private Long lastModifiedTime;
     private TableDefinition definition;
     private EncryptionConfiguration encryptionConfiguration;
+    private Labels labels = Labels.ZERO;
 
     BuilderImpl() {}
 
@@ -150,6 +165,7 @@ static class BuilderImpl extends Builder {
       this.lastModifiedTime = tableInfo.lastModifiedTime;
       this.definition = tableInfo.definition;
       this.encryptionConfiguration = tableInfo.encryptionConfiguration;
+      this.labels = tableInfo.labels;
     }
 
     BuilderImpl(Table tablePb) {
@@ -169,6 +185,7 @@ static class BuilderImpl extends Builder {
         this.encryptionConfiguration =
             new EncryptionConfiguration.Builder(tablePb.getEncryptionConfiguration()).build();
       }
+      this.labels = Labels.fromPb(tablePb.getLabels());
     }
 
     @Override
@@ -244,6 +261,12 @@ public Builder setEncryptionConfiguration(EncryptionConfiguration configuration)
     }
 
 
+    @Override
+    public Builder setLabels(Map<String, String> labels) {
+      this.labels = Labels.fromUser(labels);
+      return this;
+    }
+
     @Override
     public TableInfo build() {
       return new TableInfo(this);
@@ -262,6 +285,7 @@ public TableInfo build() {
     this.lastModifiedTime = builder.lastModifiedTime;
     this.definition = builder.definition;
     this.encryptionConfiguration = builder.encryptionConfiguration;
+    labels = builder.labels;
   }
 
 
@@ -351,6 +375,17 @@ public <T extends TableDefinition> T getDefinition() {
     return (T) definition;
   }
 
+  /**
+   * Return a map for labels applied to the table.
+   *
+   * <p>Unstable, because labels are <a
+   * href="https://cloud.google.com/bigquery/docs/reference/rest/v2/tables">experimental</a>.
+   */
+  @BetaApi
+  public Map<String, String> getLabels() {
+    return labels.userMap();
+  }
+
   /**
    * Returns a builder for the table object.
    */
@@ -372,6 +407,7 @@ public String toString() {
         .add("lastModifiedTime", lastModifiedTime)
         .add("definition", definition)
         .add("encryptionConfiguration", encryptionConfiguration)
+        .add("labels", labels)
         .toString();
   }
 
@@ -429,6 +465,7 @@ Table toPb() {
     if (encryptionConfiguration != null) {
       tablePb.setEncryptionConfiguration(encryptionConfiguration.toPb());
     }
+    tablePb.setLabels(labels.toPb());
     return tablePb;
   }
 

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetInfoTest.java

@@ -18,6 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -128,7 +129,8 @@ public void testOf() {
     assertNull(datasetInfo.getLastModified());
     assertNull(datasetInfo.getLocation());
     assertNull(datasetInfo.getSelfLink());
-    assertNull(datasetInfo.getLabels());
+    assertTrue(datasetInfo.getLabels().isEmpty());
+
     datasetInfo = DatasetInfo.of(DATASET_ID);
     assertEquals(DATASET_ID, datasetInfo.getDatasetId());
     assertNull(datasetInfo.getAcl());
@@ -141,7 +143,7 @@ public void testOf() {
     assertNull(datasetInfo.getLastModified());
     assertNull(datasetInfo.getLocation());
     assertNull(datasetInfo.getSelfLink());
-    assertNull(datasetInfo.getLabels());
+    assertTrue(datasetInfo.getLabels().isEmpty());
   }
 
   @Test