Skip to content

Commit f452cf4

Browse files
brendangreenleybrendangreenley
authored
feat: Add support for Dataset property storageBillingModel (#2913)
Adds support for setting and getting the BigQuery v2 API Dataset property `storageBillingModel`, to allow end users to set the billing model depending on their use case. When no billing model is provided, the existing behavior (use the default billing model) continues.
1 parent 1b2bc18 commit f452cf4

File tree

6 files changed

+70
-0
lines changed

6 files changed

+70
-0
lines changed

‎google-cloud-bigquery/clirr-ignored-differences.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,11 @@
109109
<className>com/google/cloud/bigquery/DatasetInfo*</className>
110110
<method>*setExternalDatasetReference(*)</method>
111111
</difference>
112+
<difference>
113+
<differenceType>7013</differenceType>
114+
<className>com/google/cloud/bigquery/DatasetInfo*</className>
115+
<method>*setStorageBillingModel(*)</method>
116+
</difference>
112117
<difference>
113118
<differenceType>7013</differenceType>
114119
<className>com/google/cloud/bigquery/StandardTableDefinition*</className>

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,12 @@ public Builder setExternalDatasetReference(ExternalDatasetReference externalData
158158
return this;
159159
}
160160

161+
@Override
162+
public Builder setStorageBillingModel(String storageBillingModel) {
163+
infoBuilder.setStorageBillingModel(storageBillingModel);
164+
return this;
165+
}
166+
161167
@Override
162168
public Dataset build() {
163169
return new Dataset(bigquery, infoBuilder);

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ public Dataset apply(DatasetInfo datasetInfo) {
7474
private final Long defaultPartitionExpirationMs;
7575
private final String defaultCollation;
7676
private final ExternalDatasetReference externalDatasetReference;
77+
private final String storageBillingModel;
7778

7879
/** A builder for {@code DatasetInfo} objects. */
7980
public abstract static class Builder {
@@ -135,6 +136,12 @@ public abstract static class Builder {
135136
public abstract Builder setExternalDatasetReference(
136137
ExternalDatasetReference externalDatasetReference);
137138

139+
/**
140+
* Optional. Storage billing model to be used for all tables in the dataset. Can be set to
141+
* PHYSICAL. Default is LOGICAL.
142+
*/
143+
public abstract Builder setStorageBillingModel(String storageBillingModel);
144+
138145
/**
139146
* The default encryption key for all tables in the dataset. Once this property is set, all
140147
* newly-created partitioned tables in the dataset will have encryption key set to this value,
@@ -192,6 +199,7 @@ static final class BuilderImpl extends Builder {
192199
private Long defaultPartitionExpirationMs;
193200
private String defaultCollation;
194201
private ExternalDatasetReference externalDatasetReference;
202+
private String storageBillingModel;
195203

196204
BuilderImpl() {}
197205

@@ -212,6 +220,7 @@ static final class BuilderImpl extends Builder {
212220
this.defaultPartitionExpirationMs = datasetInfo.defaultPartitionExpirationMs;
213221
this.defaultCollation = datasetInfo.defaultCollation;
214222
this.externalDatasetReference = datasetInfo.externalDatasetReference;
223+
this.storageBillingModel = datasetInfo.storageBillingModel;
215224
}
216225

217226
BuilderImpl(com.google.api.services.bigquery.model.Dataset datasetPb) {
@@ -250,6 +259,7 @@ public Acl apply(Dataset.Access accessPb) {
250259
this.externalDatasetReference =
251260
ExternalDatasetReference.fromPb(datasetPb.getExternalDatasetReference());
252261
}
262+
this.storageBillingModel = datasetPb.getStorageBillingModel();
253263
}
254264

255265
@Override
@@ -356,6 +366,12 @@ public Builder setExternalDatasetReference(ExternalDatasetReference externalData
356366
return this;
357367
}
358368

369+
@Override
370+
public Builder setStorageBillingModel(String storageBillingModel) {
371+
this.storageBillingModel = storageBillingModel;
372+
return this;
373+
}
374+
359375
@Override
360376
public DatasetInfo build() {
361377
return new DatasetInfo(this);
@@ -379,6 +395,7 @@ public DatasetInfo build() {
379395
defaultPartitionExpirationMs = builder.defaultPartitionExpirationMs;
380396
defaultCollation = builder.defaultCollation;
381397
externalDatasetReference = builder.externalDatasetReference;
398+
storageBillingModel = builder.storageBillingModel;
382399
}
383400

384401
/** Returns the dataset identity. */
@@ -508,6 +525,10 @@ public String getDefaultCollation() {
508525
return defaultCollation;
509526
}
510527

528+
public String getStorageBillingModel() {
529+
return storageBillingModel;
530+
}
531+
511532
/**
512533
* Returns information about the external metadata storage where the dataset is defined. Filled
513534
* out when the dataset type is EXTERNAL.
@@ -540,6 +561,7 @@ public String toString() {
540561
.add("defaultPartitionExpirationMs", defaultPartitionExpirationMs)
541562
.add("defaultCollation", defaultCollation)
542563
.add("externalDatasetReference", externalDatasetReference)
564+
.add("storageBillingModel", storageBillingModel)
543565
.toString();
544566
}
545567

@@ -621,6 +643,9 @@ public Dataset.Access apply(Acl acl) {
621643
if (externalDatasetReference != null) {
622644
datasetPb.setExternalDatasetReference(externalDatasetReference.toPb());
623645
}
646+
if (storageBillingModel != null) {
647+
datasetPb.setStorageBillingModel(storageBillingModel);
648+
}
624649
return datasetPb;
625650
}
626651

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ public class DatasetInfoTest {
5858
private static final DatasetId DATASET_ID_COMPLETE = DatasetId.of("project", "dataset");
5959
private static final EncryptionConfiguration DATASET_ENCRYPTION_CONFIGURATION =
6060
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
61+
private static final String STORAGE_BILLING_MODEL = "LOGICAL";
6162

6263
private static final ExternalDatasetReference EXTERNAL_DATASET_REFERENCE =
6364
ExternalDatasetReference.newBuilder()
@@ -79,6 +80,7 @@ public class DatasetInfoTest {
7980
.setLabels(LABELS)
8081
.setDefaultEncryptionConfiguration(DATASET_ENCRYPTION_CONFIGURATION)
8182
.setDefaultPartitionExpirationMs(DEFAULT_PARTITION__EXPIRATION)
83+
.setStorageBillingModel(STORAGE_BILLING_MODEL)
8284
.build();
8385
private static final DatasetInfo DATASET_INFO_COMPLETE =
8486
DATASET_INFO
@@ -170,6 +172,7 @@ public void testBuilder() {
170172
assertEquals(
171173
EXTERNAL_DATASET_REFERENCE,
172174
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE.getExternalDatasetReference());
175+
assertEquals(STORAGE_BILLING_MODEL, DATASET_INFO_COMPLETE.getStorageBillingModel());
173176
}
174177

175178
@Test
@@ -190,6 +193,7 @@ public void testOf() {
190193
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
191194
assertTrue(datasetInfo.getLabels().isEmpty());
192195
assertNull(datasetInfo.getExternalDatasetReference());
196+
assertNull(datasetInfo.getStorageBillingModel());
193197

194198
datasetInfo = DatasetInfo.of(DATASET_ID);
195199
assertEquals(DATASET_ID, datasetInfo.getDatasetId());
@@ -207,6 +211,7 @@ public void testOf() {
207211
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
208212
assertTrue(datasetInfo.getLabels().isEmpty());
209213
assertNull(datasetInfo.getExternalDatasetReference());
214+
assertNull(datasetInfo.getStorageBillingModel());
210215
}
211216

212217
@Test
@@ -243,5 +248,6 @@ private void compareDatasets(DatasetInfo expected, DatasetInfo value) {
243248
assertEquals(
244249
expected.getDefaultPartitionExpirationMs(), value.getDefaultPartitionExpirationMs());
245250
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
251+
assertEquals(expected.getStorageBillingModel(), value.getStorageBillingModel());
246252
}
247253
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ public class DatasetTest {
6666
private static final String SELF_LINK = "http://bigquery/p/d";
6767
private static final DatasetInfo DATASET_INFO = DatasetInfo.newBuilder(DATASET_ID).build();
6868
private static final Field FIELD = Field.of("FieldName", LegacySQLTypeName.INTEGER);
69+
private static final String STORAGE_BILLING_MODEL = "LOGICAL";
6970
private static final StandardTableDefinition TABLE_DEFINITION =
7071
StandardTableDefinition.of(Schema.of(FIELD));
7172
private static final ViewDefinition VIEW_DEFINITION = ViewDefinition.of("QUERY");
@@ -120,6 +121,7 @@ public void testBuilder() {
120121
.setLocation(LOCATION)
121122
.setSelfLink(SELF_LINK)
122123
.setLabels(LABELS)
124+
.setStorageBillingModel(STORAGE_BILLING_MODEL)
123125
.build();
124126
assertEquals(DATASET_ID, builtDataset.getDatasetId());
125127
assertEquals(ACCESS_RULES, builtDataset.getAcl());
@@ -133,6 +135,7 @@ public void testBuilder() {
133135
assertEquals(LOCATION, builtDataset.getLocation());
134136
assertEquals(SELF_LINK, builtDataset.getSelfLink());
135137
assertEquals(LABELS, builtDataset.getLabels());
138+
assertEquals(STORAGE_BILLING_MODEL, builtDataset.getStorageBillingModel());
136139
}
137140

138141
@Test
@@ -340,6 +343,7 @@ public void testExternalDatasetReference() {
340343
.setSelfLink(SELF_LINK)
341344
.setLabels(LABELS)
342345
.setExternalDatasetReference(EXTERNAL_DATASET_REFERENCE)
346+
.setStorageBillingModel(STORAGE_BILLING_MODEL)
343347
.build();
344348
assertEquals(
345349
EXTERNAL_DATASET_REFERENCE,
@@ -369,5 +373,6 @@ private void compareDatasetInfo(DatasetInfo expected, DatasetInfo value) {
369373
assertEquals(expected.getDefaultTableLifetime(), value.getDefaultTableLifetime());
370374
assertEquals(expected.getLastModified(), value.getLastModified());
371375
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
376+
assertEquals(expected.getStorageBillingModel(), value.getStorageBillingModel());
372377
}
373378
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,7 @@ public class ITBigQueryTest {
202202
private static final String ROUTINE_DATASET = RemoteBigQueryHelper.generateDatasetName();
203203
private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
204204
private static final String RANDOM_ID = UUID.randomUUID().toString().substring(0, 8);
205+
private static final String STORAGE_BILLING_MODEL = "LOGICAL";
205206
private static final String CLOUD_SAMPLES_DATA =
206207
Optional.fromNullable(System.getenv("CLOUD_SAMPLES_DATA_BUCKET")).or("cloud-samples-data");
207208
private static final Map<String, String> LABELS =
@@ -965,6 +966,7 @@ public void testGetDatasetWithSelectedFields() {
965966
assertNull(dataset.getLastModified());
966967
assertNull(dataset.getLocation());
967968
assertNull(dataset.getSelfLink());
969+
assertNull(dataset.getStorageBillingModel());
968970
}
969971

970972
@Test
@@ -980,6 +982,7 @@ public void testUpdateDataset() {
980982
assertThat(dataset.getDatasetId().getDataset()).isEqualTo(OTHER_DATASET);
981983
assertThat(dataset.getDescription()).isEqualTo("Some Description");
982984
assertThat(dataset.getLabels()).containsExactly("a", "b");
985+
assertThat(dataset.getStorageBillingModel()).isNull();
983986

984987
Map<String, String> updateLabels = new HashMap<>();
985988
updateLabels.put("x", "y");
@@ -990,9 +993,11 @@ public void testUpdateDataset() {
990993
.toBuilder()
991994
.setDescription("Updated Description")
992995
.setLabels(updateLabels)
996+
.setStorageBillingModel("LOGICAL")
993997
.build());
994998
assertThat(updatedDataset.getDescription()).isEqualTo("Updated Description");
995999
assertThat(updatedDataset.getLabels()).containsExactly("x", "y");
1000+
assertThat(updatedDataset.getStorageBillingModel()).isEqualTo("LOGICAL");
9961001

9971002
updatedDataset = bigquery.update(updatedDataset.toBuilder().setLabels(null).build());
9981003
assertThat(updatedDataset.getLabels()).isEmpty();
@@ -1022,6 +1027,7 @@ public void testUpdateDatasetWithSelectedFields() {
10221027
assertNull(updatedDataset.getLastModified());
10231028
assertNull(updatedDataset.getLocation());
10241029
assertNull(updatedDataset.getSelfLink());
1030+
assertNull(updatedDataset.getStorageBillingModel());
10251031
assertTrue(dataset.delete());
10261032
}
10271033

@@ -1283,6 +1289,23 @@ public void testCreateTableWithConstraints() {
12831289
bigquery.delete(tableId);
12841290
}
12851291

1292+
@Test
1293+
public void testCreateDatasetWithSpecifiedStorageBillingModel() {
1294+
String billingModelDataset = RemoteBigQueryHelper.generateDatasetName();
1295+
DatasetInfo info =
1296+
DatasetInfo.newBuilder(billingModelDataset)
1297+
.setDescription(DESCRIPTION)
1298+
.setStorageBillingModel(STORAGE_BILLING_MODEL)
1299+
.setLabels(LABELS)
1300+
.build();
1301+
bigquery.create(info);
1302+
1303+
Dataset dataset = bigquery.getDataset(DatasetId.of(billingModelDataset));
1304+
assertEquals(STORAGE_BILLING_MODEL, dataset.getStorageBillingModel());
1305+
1306+
RemoteBigQueryHelper.forceDelete(bigquery, billingModelDataset);
1307+
}
1308+
12861309
@Test
12871310
public void testCreateDatasetWithDefaultCollation() {
12881311
String collationDataset = RemoteBigQueryHelper.generateDatasetName();

0 commit comments

Comments
 (0)