Skip to content

Commit 8c0e636

Browse files
rouzwawirouzwawi
committed
Expose BigQuery Dataset labels
1 parent a831272 commit 8c0e636

6 files changed

Lines changed: 69 additions & 7 deletions

File tree

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ enum DatasetField implements FieldSelector {
5151
ETAG("etag"),
5252
FRIENDLY_NAME("friendlyName"),
5353
ID("id"),
54+
LABELS("labels"),
5455
LAST_MODIFIED_TIME("lastModifiedTime"),
5556
LOCATION("location"),
5657
SELF_LINK("selfLink");

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Dataset.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,10 @@
2323
import com.google.cloud.bigquery.BigQuery.DatasetOption;
2424
import com.google.cloud.bigquery.BigQuery.TableListOption;
2525
import com.google.cloud.bigquery.BigQuery.TableOption;
26-
2726
import java.io.IOException;
2827
import java.io.ObjectInputStream;
2928
import java.util.List;
29+
import java.util.Map;
3030
import java.util.Objects;
3131

3232
/**
@@ -136,6 +136,12 @@ Builder setSelfLink(String selfLink) {
136136
return this;
137137
}
138138

139+
@Override
140+
public Builder setLabels(Map<String, String> labels) {
141+
infoBuilder.setLabels(labels);
142+
return this;
143+
}
144+
139145
@Override
140146
public Dataset build() {
141147
return new Dataset(bigquery, infoBuilder);

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,11 @@
2525
import com.google.common.base.Function;
2626
import com.google.common.base.MoreObjects;
2727
import com.google.common.collect.ImmutableList;
28+
import com.google.common.collect.ImmutableMap;
2829
import com.google.common.collect.Lists;
29-
3030
import java.io.Serializable;
3131
import java.util.List;
32+
import java.util.Map;
3233
import java.util.Objects;
3334

3435
/**
@@ -68,6 +69,7 @@ public Dataset apply(DatasetInfo datasetInfo) {
6869
private final Long lastModified;
6970
private final String location;
7071
private final String selfLink;
72+
private Map<String, String> labels;
7173

7274
/**
7375
* A builder for {@code DatasetInfo} objects.
@@ -134,6 +136,8 @@ public abstract static class Builder {
134136

135137
abstract Builder setSelfLink(String selfLink);
136138

139+
public abstract Builder setLabels(Map<String, String> labels);
140+
137141
/**
138142
* Creates a {@code DatasetInfo} object.
139143
*/
@@ -153,6 +157,7 @@ static final class BuilderImpl extends Builder {
153157
private Long lastModified;
154158
private String location;
155159
private String selfLink;
160+
private Map<String, String> labels;
156161

157162
BuilderImpl() {}
158163

@@ -168,6 +173,9 @@ static final class BuilderImpl extends Builder {
168173
this.lastModified = datasetInfo.lastModified;
169174
this.location = datasetInfo.location;
170175
this.selfLink = datasetInfo.selfLink;
176+
this.labels = datasetInfo.labels != null
177+
? ImmutableMap.copyOf(datasetInfo.labels)
178+
: null;
171179
}
172180

173181
BuilderImpl(com.google.api.services.bigquery.model.Dataset datasetPb) {
@@ -191,6 +199,9 @@ public Acl apply(Dataset.Access accessPb) {
191199
this.lastModified = datasetPb.getLastModifiedTime();
192200
this.location = datasetPb.getLocation();
193201
this.selfLink = datasetPb.getSelfLink();
202+
this.labels = datasetPb.getLabels() != null
203+
? ImmutableMap.copyOf(datasetPb.getLabels())
204+
: null;
194205
}
195206

196207

@@ -266,6 +277,12 @@ Builder setSelfLink(String selfLink) {
266277
return this;
267278
}
268279

280+
@Override
281+
public Builder setLabels(Map<String, String> labels) {
282+
this.labels = ImmutableMap.copyOf(labels);
283+
return this;
284+
}
285+
269286
@Override
270287
public DatasetInfo build() {
271288
return new DatasetInfo(this);
@@ -284,6 +301,7 @@ public DatasetInfo build() {
284301
lastModified = builder.lastModified;
285302
location = builder.location;
286303
selfLink = builder.selfLink;
304+
labels = builder.labels;
287305
}
288306

289307

@@ -387,6 +405,15 @@ public String getSelfLink() {
387405
return selfLink;
388406
}
389407

408+
/**
409+
* Return a map for labels applied to the dataset.
410+
*
411+
* @see <a href="https://cloud.google.com/bigquery/docs/labeling-datasets">Labeling Datasets</a>
412+
*/
413+
public Map<String, String> getLabels() {
414+
return labels;
415+
}
416+
390417
/**
391418
* Returns a builder for the dataset object.
392419
*/
@@ -408,6 +435,7 @@ public String toString() {
408435
.add("location", location)
409436
.add("selfLink", selfLink)
410437
.add("acl", acl)
438+
.add("labels", labels)
411439
.toString();
412440
}
413441

@@ -466,6 +494,9 @@ public Dataset.Access apply(Acl acl) {
466494
}
467495
}));
468496
}
497+
if (labels != null) {
498+
datasetPb.setLabels(labels);
499+
}
469500
return datasetPb;
470501
}
471502

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@
2020
import static org.junit.Assert.assertNull;
2121

2222
import com.google.common.collect.ImmutableList;
23-
24-
import org.junit.Test;
25-
23+
import com.google.common.collect.ImmutableMap;
2624
import java.util.List;
25+
import java.util.Map;
26+
import org.junit.Test;
2727

2828
public class DatasetInfoTest {
2929

@@ -33,6 +33,9 @@ public class DatasetInfoTest {
3333
private static final List<Acl> ACCESS_RULES_COMPLETE = ImmutableList.of(
3434
Acl.of(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER),
3535
Acl.of(new Acl.View(TableId.of("project", "dataset", "table"))));
36+
private static final Map<String, String> LABELS = ImmutableMap.of(
37+
"example-label1", "example-value1",
38+
"example-label2", "example-value2");
3639
private static final Long CREATION_TIME = System.currentTimeMillis();
3740
private static final Long DEFAULT_TABLE_EXPIRATION = CREATION_TIME + 100;
3841
private static final String DESCRIPTION = "description";
@@ -55,6 +58,7 @@ public class DatasetInfoTest {
5558
.setLastModified(LAST_MODIFIED)
5659
.setLocation(LOCATION)
5760
.setSelfLink(SELF_LINK)
61+
.setLabels(LABELS)
5862
.build();
5963
private static final DatasetInfo DATASET_INFO_COMPLETE = DATASET_INFO.toBuilder()
6064
.setDatasetId(DATASET_ID_COMPLETE)
@@ -106,6 +110,7 @@ public void testBuilder() {
106110
assertEquals(LAST_MODIFIED, DATASET_INFO_COMPLETE.getLastModified());
107111
assertEquals(LOCATION, DATASET_INFO_COMPLETE.getLocation());
108112
assertEquals(SELF_LINK, DATASET_INFO_COMPLETE.getSelfLink());
113+
assertEquals(LABELS, DATASET_INFO_COMPLETE.getLabels());
109114
}
110115

111116

@@ -123,6 +128,7 @@ public void testOf() {
123128
assertNull(datasetInfo.getLastModified());
124129
assertNull(datasetInfo.getLocation());
125130
assertNull(datasetInfo.getSelfLink());
131+
assertNull(datasetInfo.getLabels());
126132
datasetInfo = DatasetInfo.of(DATASET_ID);
127133
assertEquals(DATASET_ID, datasetInfo.getDatasetId());
128134
assertNull(datasetInfo.getAcl());
@@ -135,6 +141,7 @@ public void testOf() {
135141
assertNull(datasetInfo.getLastModified());
136142
assertNull(datasetInfo.getLocation());
137143
assertNull(datasetInfo.getSelfLink());
144+
assertNull(datasetInfo.getLabels());
138145
}
139146

140147
@Test
@@ -162,5 +169,6 @@ private void compareDatasets(DatasetInfo expected, DatasetInfo value) {
162169
assertEquals(expected.getCreationTime(), value.getCreationTime());
163170
assertEquals(expected.getDefaultTableLifetime(), value.getDefaultTableLifetime());
164171
assertEquals(expected.getLastModified(), value.getLastModified());
172+
assertEquals(expected.getLabels(), value.getLabels());
165173
}
166174
}

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,10 @@
3333
import com.google.api.gax.paging.Page;
3434
import com.google.cloud.PageImpl;
3535
import com.google.common.collect.ImmutableList;
36+
import com.google.common.collect.ImmutableMap;
3637
import com.google.common.collect.Iterables;
3738

39+
import java.util.Map;
3840
import org.junit.After;
3941
import org.junit.Test;
4042

@@ -46,6 +48,9 @@ public class DatasetTest {
4648
private static final List<Acl> ACCESS_RULES = ImmutableList.of(
4749
Acl.of(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER),
4850
Acl.of(new Acl.View(TableId.of("dataset", "table"))));
51+
private static final Map<String, String> LABELS = ImmutableMap.of(
52+
"example-label1", "example-value1",
53+
"example-label2", "example-value2");
4954
private static final Long CREATION_TIME = System.currentTimeMillis();
5055
private static final Long DEFAULT_TABLE_EXPIRATION = CREATION_TIME + 100;
5156
private static final String DESCRIPTION = "description";
@@ -106,6 +111,7 @@ public void testBuilder() {
106111
.setLastModified(LAST_MODIFIED)
107112
.setLocation(LOCATION)
108113
.setSelfLink(SELF_LINK)
114+
.setLabels(LABELS)
109115
.build();
110116
assertEquals(DATASET_ID, builtDataset.getDatasetId());
111117
assertEquals(ACCESS_RULES, builtDataset.getAcl());
@@ -118,6 +124,7 @@ public void testBuilder() {
118124
assertEquals(LAST_MODIFIED, builtDataset.getLastModified());
119125
assertEquals(LOCATION, builtDataset.getLocation());
120126
assertEquals(SELF_LINK, builtDataset.getSelfLink());
127+
assertEquals(LABELS, builtDataset.getLabels());
121128
}
122129

123130

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@
8383
import java.util.HashSet;
8484
import java.util.Iterator;
8585
import java.util.List;
86+
import java.util.Map;
8687
import java.util.Set;
8788
import java.util.concurrent.ExecutionException;
8889
import java.util.concurrent.TimeUnit;
@@ -103,6 +104,9 @@ public class ITBigQueryTest {
103104
private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
104105
private static final String DESCRIPTION = "Test dataset";
105106
private static final String OTHER_DATASET = RemoteBigQueryHelper.generateDatasetName();
107+
private static final Map<String, String> LABELS = ImmutableMap.of(
108+
"example-label1", "example-value1",
109+
"example-label2", "example-value2");
106110
private static final Field TIMESTAMP_FIELD_SCHEMA =
107111
Field.newBuilder("TimestampField", Field.Type.timestamp())
108112
.setMode(Field.Mode.NULLABLE)
@@ -221,7 +225,10 @@ public static void beforeClass() throws InterruptedException, TimeoutException {
221225
.setContentType("application/json")
222226
.build(),
223227
JSON_CONTENT.getBytes(StandardCharsets.UTF_8));
224-
DatasetInfo info = DatasetInfo.newBuilder(DATASET).setDescription(DESCRIPTION).build();
228+
DatasetInfo info = DatasetInfo.newBuilder(DATASET)
229+
.setDescription(DESCRIPTION)
230+
.setLabels(LABELS)
231+
.build();
225232
bigquery.create(info);
226233
LoadJobConfiguration configuration = LoadJobConfiguration.newBuilder(
227234
TABLE_ID, "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json())
@@ -265,6 +272,7 @@ public void testGetDataset() {
265272
assertEquals(bigquery.getOptions().getProjectId(), dataset.getDatasetId().getProject());
266273
assertEquals(DATASET, dataset.getDatasetId().getDataset());
267274
assertEquals(DESCRIPTION, dataset.getDescription());
275+
assertEquals(LABELS, dataset.getLabels());
268276
assertNotNull(dataset.getAcl());
269277
assertNotNull(dataset.getEtag());
270278
assertNotNull(dataset.getGeneratedId());
@@ -275,9 +283,10 @@ public void testGetDataset() {
275283
@Test
276284
public void testGetDatasetWithSelectedFields() {
277285
Dataset dataset = bigquery.getDataset(DATASET,
278-
DatasetOption.fields(DatasetField.CREATION_TIME));
286+
DatasetOption.fields(DatasetField.CREATION_TIME, DatasetField.LABELS));
279287
assertEquals(bigquery.getOptions().getProjectId(), dataset.getDatasetId().getProject());
280288
assertEquals(DATASET, dataset.getDatasetId().getDataset());
289+
assertEquals(LABELS, dataset.getLabels());
281290
assertNotNull(dataset.getCreationTime());
282291
assertNull(dataset.getDescription());
283292
assertNull(dataset.getDefaultTableLifetime());

0 commit comments

Comments
 (0)