diff --git a/bigquery/dataset.go b/bigquery/dataset.go index a10158f2d119..b2be93d9ba30 100644 --- a/bigquery/dataset.go +++ b/bigquery/dataset.go @@ -58,6 +58,12 @@ type DatasetMetadata struct { // More information: https://cloud.google.com/bigquery/docs/reference/standard-sql/collation-concepts DefaultCollation string + // Storage billing model to be used for all tables in the dataset. + // Can be set to PHYSICAL. Default is LOGICAL. + // Once you create a dataset with storage billing model set to physical bytes, you can't change it back to using logical bytes again. + // More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models + StorageBillingModel string + // These fields are read-only. CreationTime time.Time LastModifiedTime time.Time // When the dataset or any of its tables were modified. @@ -84,6 +90,14 @@ type DatasetTag struct { TagValue string } +const ( + // LogicalStorageBillingModel indicates billing for logical bytes. + LogicalStorageBillingModel = "" + + // PhysicalStorageBillingModel indicates billing for physical bytes. + PhysicalStorageBillingModel = "PHYSICAL" +) + func bqToDatasetTag(in *bq.DatasetTags) *DatasetTag { if in == nil { return nil @@ -117,6 +131,12 @@ type DatasetMetadataToUpdate struct { // created in the dataset. DefaultCollation optional.String + // Storage billing model to be used for all tables in the dataset. + // Can be set to PHYSICAL. Default is LOGICAL. + // Once you change a dataset's storage billing model to use physical bytes, you can't change it back to using logical bytes again. + // More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models + StorageBillingModel optional.String + // The entire access list. It is not possible to replace individual entries. Access []*AccessEntry @@ -187,7 +207,8 @@ func (dm *DatasetMetadata) toBQ() (*bq.Dataset, error) { ds.Location = dm.Location ds.DefaultTableExpirationMs = int64(dm.DefaultTableExpiration / time.Millisecond) ds.DefaultPartitionExpirationMs = int64(dm.DefaultPartitionExpiration / time.Millisecond) - ds.DefaultCollation = string(dm.DefaultCollation) + ds.DefaultCollation = dm.DefaultCollation + ds.StorageBillingModel = string(dm.StorageBillingModel) ds.Labels = dm.Labels var err error ds.Access, err = accessListToBQ(dm.Access) @@ -274,6 +295,7 @@ func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) { DefaultTableExpiration: time.Duration(d.DefaultTableExpirationMs) * time.Millisecond, DefaultPartitionExpiration: time.Duration(d.DefaultPartitionExpirationMs) * time.Millisecond, DefaultCollation: d.DefaultCollation, + StorageBillingModel: d.StorageBillingModel, DefaultEncryptionConfig: bqToEncryptionConfig(d.DefaultEncryptionConfiguration), Description: d.Description, Name: d.FriendlyName, @@ -363,6 +385,10 @@ func (dm *DatasetMetadataToUpdate) toBQ() (*bq.Dataset, error) { ds.DefaultCollation = optional.ToString(dm.DefaultCollation) forceSend("DefaultCollation") } + if dm.StorageBillingModel != nil { + ds.StorageBillingModel = optional.ToString(dm.StorageBillingModel) + forceSend("StorageBillingModel") + } if dm.DefaultEncryptionConfig != nil { ds.DefaultEncryptionConfiguration = dm.DefaultEncryptionConfig.toBQ() ds.DefaultEncryptionConfiguration.ForceSendFields = []string{"KmsKeyName"} diff --git a/bigquery/dataset_integration_test.go b/bigquery/dataset_integration_test.go index eb8b83ed82d9..91ec36c4bd62 100644 --- a/bigquery/dataset_integration_test.go +++ b/bigquery/dataset_integration_test.go @@ -277,6 +277,83 @@ func TestIntegration_DatasetUpdateDefaultCollation(t *testing.T) { } } +func TestIntegration_DatasetStorageBillingModel(t *testing.T) { + if client == nil { + t.Skip("Integration tests skipped") + } + + ctx := context.Background() + md, err := dataset.Metadata(ctx) + if err != nil { + t.Fatal(err) + } + if md.StorageBillingModel != LogicalStorageBillingModel { + t.Fatalf("got %q, want %q", md.StorageBillingModel, LogicalStorageBillingModel) + } + + ds := client.Dataset(datasetIDs.New()) + err = ds.Create(ctx, &DatasetMetadata{ + StorageBillingModel: PhysicalStorageBillingModel, + }) + if err != nil { + t.Fatal(err) + } + md, err = ds.Metadata(ctx) + if err != nil { + t.Fatal(err) + } + if md.StorageBillingModel != PhysicalStorageBillingModel { + t.Fatalf("got %q, want %q", md.StorageBillingModel, PhysicalStorageBillingModel) + } + if err := ds.Delete(ctx); err != nil { + t.Fatalf("deleting dataset %v: %v", ds, err) + } +} + +func TestIntegration_DatasetUpdateStorageBillingModel(t *testing.T) { + if client == nil { + t.Skip("Integration tests skipped") + } + + ctx := context.Background() + ds := client.Dataset(datasetIDs.New()) + err := ds.Create(ctx, &DatasetMetadata{ + StorageBillingModel: LogicalStorageBillingModel, + }) + if err != nil { + t.Fatal(err) + } + + md, err := ds.Metadata(ctx) + if md.StorageBillingModel != LogicalStorageBillingModel { + t.Fatalf("got %q, want %q", md.StorageBillingModel, LogicalStorageBillingModel) + } + + // Update the Storage billing model + md, err = ds.Update(ctx, DatasetMetadataToUpdate{ + StorageBillingModel: PhysicalStorageBillingModel, + }, "") + if err != nil { + t.Fatal(err) + } + if md.StorageBillingModel != PhysicalStorageBillingModel { + t.Fatalf("got %q, want %q", md.StorageBillingModel, PhysicalStorageBillingModel) + } + + // Omitting StorageBillingModel doesn't change it. + md, err = ds.Update(ctx, DatasetMetadataToUpdate{Name: "xyz"}, "") + if err != nil { + t.Fatal(err) + } + if md.StorageBillingModel != PhysicalStorageBillingModel { + t.Fatalf("got %q, want %q", md.StorageBillingModel, PhysicalStorageBillingModel) + } + + if err := ds.Delete(ctx); err != nil { + t.Fatalf("deleting dataset %v: %v", ds, err) + } +} + func TestIntegration_DatasetUpdateAccess(t *testing.T) { if client == nil { t.Skip("Integration tests skipped") diff --git a/bigquery/dataset_test.go b/bigquery/dataset_test.go index 8ed11fd9cc50..992aedece101 100644 --- a/bigquery/dataset_test.go +++ b/bigquery/dataset_test.go @@ -432,8 +432,9 @@ func TestBQToDatasetMetadata(t *testing.T) { DefaultEncryptionConfig: &EncryptionConfig{ KMSKeyName: "some_key", }, - Location: "EU", - Labels: map[string]string{"x": "y"}, + StorageBillingModel: LogicalStorageBillingModel, + Location: "EU", + Labels: map[string]string{"x": "y"}, Access: []*AccessEntry{ {Role: ReaderRole, Entity: "joe@example.com", EntityType: UserEmailEntity}, {Role: WriterRole, Entity: "users@example.com", EntityType: GroupEmailEntity}, @@ -466,6 +467,7 @@ func TestDatasetMetadataToUpdateToBQ(t *testing.T) { Name: "name", DefaultTableExpiration: time.Hour, DefaultPartitionExpiration: 24 * time.Hour, + StorageBillingModel: PhysicalStorageBillingModel, DefaultEncryptionConfig: &EncryptionConfig{ KMSKeyName: "some_key", }, @@ -482,12 +484,13 @@ func TestDatasetMetadataToUpdateToBQ(t *testing.T) { FriendlyName: "name", DefaultTableExpirationMs: 60 * 60 * 1000, DefaultPartitionExpirationMs: 24 * 60 * 60 * 1000, + StorageBillingModel: string(PhysicalStorageBillingModel), DefaultEncryptionConfiguration: &bq.EncryptionConfiguration{ KmsKeyName: "some_key", ForceSendFields: []string{"KmsKeyName"}, }, Labels: map[string]string{"label": "value"}, - ForceSendFields: []string{"Description", "FriendlyName"}, + ForceSendFields: []string{"Description", "FriendlyName", "StorageBillingModel"}, NullFields: []string{"Labels.del"}, } if diff := testutil.Diff(got, want); diff != "" {