Skip to content

Commit

Permalink
feat: Add support for Dataset property storageBillingModel (#2913)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@brendangreenley
brendangreenley committed Oct 20, 2023
1 parent 1b2bc18 commit f452cf4
Show file tree
Hide file tree
Showing 6 changed files with 70 additions and 0 deletions.
5 changes: 5 additions & 0 deletions google-cloud-bigquery/clirr-ignored-differences.xml
View file
Expand Up @@ -109,6 +109,11 @@
<className>com/google/cloud/bigquery/DatasetInfo*</className>
<method>*setExternalDatasetReference(*)</method>
</difference>
<difference>
<differenceType>7013</differenceType>
<className>com/google/cloud/bigquery/DatasetInfo*</className>
<method>*setStorageBillingModel(*)</method>
</difference>
<difference>
<differenceType>7013</differenceType>
<className>com/google/cloud/bigquery/StandardTableDefinition*</className>
Expand Down
6 changes: 6 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Dataset.java
View file
Expand Up @@ -158,6 +158,12 @@ public Builder setExternalDatasetReference(ExternalDatasetReference externalData
return this;
}

@Override
public Builder setStorageBillingModel(String storageBillingModel) {
infoBuilder.setStorageBillingModel(storageBillingModel);
return this;
}

@Override
public Dataset build() {
return new Dataset(bigquery, infoBuilder);
Expand Down
25 changes: 25 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java
View file
Expand Up @@ -74,6 +74,7 @@ public Dataset apply(DatasetInfo datasetInfo) {
private final Long defaultPartitionExpirationMs;
private final String defaultCollation;
private final ExternalDatasetReference externalDatasetReference;
private final String storageBillingModel;

/** A builder for {@code DatasetInfo} objects. */
public abstract static class Builder {
Expand Down Expand Up @@ -135,6 +136,12 @@ public abstract static class Builder {
public abstract Builder setExternalDatasetReference(
ExternalDatasetReference externalDatasetReference);

/**
* Optional. Storage billing model to be used for all tables in the dataset. Can be set to
* PHYSICAL. Default is LOGICAL.
*/
public abstract Builder setStorageBillingModel(String storageBillingModel);

/**
* The default encryption key for all tables in the dataset. Once this property is set, all
* newly-created partitioned tables in the dataset will have encryption key set to this value,
Expand Down Expand Up @@ -192,6 +199,7 @@ static final class BuilderImpl extends Builder {
private Long defaultPartitionExpirationMs;
private String defaultCollation;
private ExternalDatasetReference externalDatasetReference;
private String storageBillingModel;

BuilderImpl() {}

Expand All @@ -212,6 +220,7 @@ static final class BuilderImpl extends Builder {
this.defaultPartitionExpirationMs = datasetInfo.defaultPartitionExpirationMs;
this.defaultCollation = datasetInfo.defaultCollation;
this.externalDatasetReference = datasetInfo.externalDatasetReference;
this.storageBillingModel = datasetInfo.storageBillingModel;
}

BuilderImpl(com.google.api.services.bigquery.model.Dataset datasetPb) {
Expand Down Expand Up @@ -250,6 +259,7 @@ public Acl apply(Dataset.Access accessPb) {
this.externalDatasetReference =
ExternalDatasetReference.fromPb(datasetPb.getExternalDatasetReference());
}
this.storageBillingModel = datasetPb.getStorageBillingModel();
}

@Override
Expand Down Expand Up @@ -356,6 +366,12 @@ public Builder setExternalDatasetReference(ExternalDatasetReference externalData
return this;
}

@Override
public Builder setStorageBillingModel(String storageBillingModel) {
this.storageBillingModel = storageBillingModel;
return this;
}

@Override
public DatasetInfo build() {
return new DatasetInfo(this);
Expand All @@ -379,6 +395,7 @@ public DatasetInfo build() {
defaultPartitionExpirationMs = builder.defaultPartitionExpirationMs;
defaultCollation = builder.defaultCollation;
externalDatasetReference = builder.externalDatasetReference;
storageBillingModel = builder.storageBillingModel;
}

/** Returns the dataset identity. */
Expand Down Expand Up @@ -508,6 +525,10 @@ public String getDefaultCollation() {
return defaultCollation;
}

public String getStorageBillingModel() {
return storageBillingModel;
}

/**
* Returns information about the external metadata storage where the dataset is defined. Filled
* out when the dataset type is EXTERNAL.
Expand Down Expand Up @@ -540,6 +561,7 @@ public String toString() {
.add("defaultPartitionExpirationMs", defaultPartitionExpirationMs)
.add("defaultCollation", defaultCollation)
.add("externalDatasetReference", externalDatasetReference)
.add("storageBillingModel", storageBillingModel)
.toString();
}

Expand Down Expand Up @@ -621,6 +643,9 @@ public Dataset.Access apply(Acl acl) {
if (externalDatasetReference != null) {
datasetPb.setExternalDatasetReference(externalDatasetReference.toPb());
}
if (storageBillingModel != null) {
datasetPb.setStorageBillingModel(storageBillingModel);
}
return datasetPb;
}

Expand Down
6 changes: 6 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetInfoTest.java
View file
Expand Up @@ -58,6 +58,7 @@ public class DatasetInfoTest {
private static final DatasetId DATASET_ID_COMPLETE = DatasetId.of("project", "dataset");
private static final EncryptionConfiguration DATASET_ENCRYPTION_CONFIGURATION =
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
private static final String STORAGE_BILLING_MODEL = "LOGICAL";

private static final ExternalDatasetReference EXTERNAL_DATASET_REFERENCE =
ExternalDatasetReference.newBuilder()
Expand All @@ -79,6 +80,7 @@ public class DatasetInfoTest {
.setLabels(LABELS)
.setDefaultEncryptionConfiguration(DATASET_ENCRYPTION_CONFIGURATION)
.setDefaultPartitionExpirationMs(DEFAULT_PARTITION__EXPIRATION)
.setStorageBillingModel(STORAGE_BILLING_MODEL)
.build();
private static final DatasetInfo DATASET_INFO_COMPLETE =
DATASET_INFO
Expand Down Expand Up @@ -170,6 +172,7 @@ public void testBuilder() {
assertEquals(
EXTERNAL_DATASET_REFERENCE,
DATASET_INFO_COMPLETE_WITH_EXTERNAL_DATASET_REFERENCE.getExternalDatasetReference());
assertEquals(STORAGE_BILLING_MODEL, DATASET_INFO_COMPLETE.getStorageBillingModel());
}

@Test
Expand All @@ -190,6 +193,7 @@ public void testOf() {
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
assertTrue(datasetInfo.getLabels().isEmpty());
assertNull(datasetInfo.getExternalDatasetReference());
assertNull(datasetInfo.getStorageBillingModel());

datasetInfo = DatasetInfo.of(DATASET_ID);
assertEquals(DATASET_ID, datasetInfo.getDatasetId());
Expand All @@ -207,6 +211,7 @@ public void testOf() {
assertNull(datasetInfo.getDefaultPartitionExpirationMs());
assertTrue(datasetInfo.getLabels().isEmpty());
assertNull(datasetInfo.getExternalDatasetReference());
assertNull(datasetInfo.getStorageBillingModel());
}

@Test
Expand Down Expand Up @@ -243,5 +248,6 @@ private void compareDatasets(DatasetInfo expected, DatasetInfo value) {
assertEquals(
expected.getDefaultPartitionExpirationMs(), value.getDefaultPartitionExpirationMs());
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
assertEquals(expected.getStorageBillingModel(), value.getStorageBillingModel());
}
}
5 changes: 5 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/DatasetTest.java
View file
Expand Up @@ -66,6 +66,7 @@ public class DatasetTest {
private static final String SELF_LINK = "http://bigquery/p/d";
private static final DatasetInfo DATASET_INFO = DatasetInfo.newBuilder(DATASET_ID).build();
private static final Field FIELD = Field.of("FieldName", LegacySQLTypeName.INTEGER);
private static final String STORAGE_BILLING_MODEL = "LOGICAL";
private static final StandardTableDefinition TABLE_DEFINITION =
StandardTableDefinition.of(Schema.of(FIELD));
private static final ViewDefinition VIEW_DEFINITION = ViewDefinition.of("QUERY");
Expand Down Expand Up @@ -120,6 +121,7 @@ public void testBuilder() {
.setLocation(LOCATION)
.setSelfLink(SELF_LINK)
.setLabels(LABELS)
.setStorageBillingModel(STORAGE_BILLING_MODEL)
.build();
assertEquals(DATASET_ID, builtDataset.getDatasetId());
assertEquals(ACCESS_RULES, builtDataset.getAcl());
Expand All @@ -133,6 +135,7 @@ public void testBuilder() {
assertEquals(LOCATION, builtDataset.getLocation());
assertEquals(SELF_LINK, builtDataset.getSelfLink());
assertEquals(LABELS, builtDataset.getLabels());
assertEquals(STORAGE_BILLING_MODEL, builtDataset.getStorageBillingModel());
}

@Test
Expand Down Expand Up @@ -340,6 +343,7 @@ public void testExternalDatasetReference() {
.setSelfLink(SELF_LINK)
.setLabels(LABELS)
.setExternalDatasetReference(EXTERNAL_DATASET_REFERENCE)
.setStorageBillingModel(STORAGE_BILLING_MODEL)
.build();
assertEquals(
EXTERNAL_DATASET_REFERENCE,
Expand Down Expand Up @@ -369,5 +373,6 @@ private void compareDatasetInfo(DatasetInfo expected, DatasetInfo value) {
assertEquals(expected.getDefaultTableLifetime(), value.getDefaultTableLifetime());
assertEquals(expected.getLastModified(), value.getLastModified());
assertEquals(expected.getExternalDatasetReference(), value.getExternalDatasetReference());
assertEquals(expected.getStorageBillingModel(), value.getStorageBillingModel());
}
}
23 changes: 23 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Expand Up @@ -202,6 +202,7 @@ public class ITBigQueryTest {
private static final String ROUTINE_DATASET = RemoteBigQueryHelper.generateDatasetName();
private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
private static final String RANDOM_ID = UUID.randomUUID().toString().substring(0, 8);
private static final String STORAGE_BILLING_MODEL = "LOGICAL";
private static final String CLOUD_SAMPLES_DATA =
Optional.fromNullable(System.getenv("CLOUD_SAMPLES_DATA_BUCKET")).or("cloud-samples-data");
private static final Map<String, String> LABELS =
Expand Down Expand Up @@ -965,6 +966,7 @@ public void testGetDatasetWithSelectedFields() {
assertNull(dataset.getLastModified());
assertNull(dataset.getLocation());
assertNull(dataset.getSelfLink());
assertNull(dataset.getStorageBillingModel());
}

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

Map<String, String> updateLabels = new HashMap<>();
updateLabels.put("x", "y");
Expand All @@ -990,9 +993,11 @@ public void testUpdateDataset() {
.toBuilder()
.setDescription("Updated Description")
.setLabels(updateLabels)
.setStorageBillingModel("LOGICAL")
.build());
assertThat(updatedDataset.getDescription()).isEqualTo("Updated Description");
assertThat(updatedDataset.getLabels()).containsExactly("x", "y");
assertThat(updatedDataset.getStorageBillingModel()).isEqualTo("LOGICAL");

updatedDataset = bigquery.update(updatedDataset.toBuilder().setLabels(null).build());
assertThat(updatedDataset.getLabels()).isEmpty();
Expand Down Expand Up @@ -1022,6 +1027,7 @@ public void testUpdateDatasetWithSelectedFields() {
assertNull(updatedDataset.getLastModified());
assertNull(updatedDataset.getLocation());
assertNull(updatedDataset.getSelfLink());
assertNull(updatedDataset.getStorageBillingModel());
assertTrue(dataset.delete());
}

Expand Down Expand Up @@ -1283,6 +1289,23 @@ public void testCreateTableWithConstraints() {
bigquery.delete(tableId);
}

@Test
public void testCreateDatasetWithSpecifiedStorageBillingModel() {
String billingModelDataset = RemoteBigQueryHelper.generateDatasetName();
DatasetInfo info =
DatasetInfo.newBuilder(billingModelDataset)
.setDescription(DESCRIPTION)
.setStorageBillingModel(STORAGE_BILLING_MODEL)
.setLabels(LABELS)
.build();
bigquery.create(info);

Dataset dataset = bigquery.getDataset(DatasetId.of(billingModelDataset));
assertEquals(STORAGE_BILLING_MODEL, dataset.getStorageBillingModel());

RemoteBigQueryHelper.forceDelete(bigquery, billingModelDataset);
}

@Test
public void testCreateDatasetWithDefaultCollation() {
String collationDataset = RemoteBigQueryHelper.generateDatasetName();
Expand Down

0 comments on commit f452cf4

Please sign in to comment.