From 74d277ff90b20fed8d0700c343f3760082f17dc1 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Thu, 7 Sep 2017 16:32:41 -0400 Subject: [PATCH] bigquery: support update of dataset labels The design follows the one for storage buckets (https://godoc.org/cloud.google.com/go/storage#BucketAttrsToUpdate). Change-Id: I1f1baac05d5e5a9d804cd894b603800c404d0d60 Reviewed-on: https://code-review.googlesource.com/16650 Reviewed-by: kokoro Reviewed-by: Sai Cheemalapati --- bigquery/dataset.go | 21 +++++++++++++++++++++ bigquery/integration_test.go | 31 +++++++++++++++++++++++++++++++ bigquery/service.go | 35 ++++++++++++++++++++++++++--------- bigquery/service_test.go | 23 +++++++++++++++++++++++ 4 files changed, 101 insertions(+), 9 deletions(-) diff --git a/bigquery/dataset.go b/bigquery/dataset.go index 66e87771c6fd..04e4a643ef49 100644 --- a/bigquery/dataset.go +++ b/bigquery/dataset.go @@ -52,6 +52,27 @@ type DatasetMetadataToUpdate struct { // DefaultTableExpiration is the the default expiration time for new tables. // If set to time.Duration(0), new tables never expire. DefaultTableExpiration optional.Duration + + setLabels map[string]string + deleteLabels map[string]bool +} + +// SetLabel causes a label to be added or modified when dm is used +// in a call to Dataset.Update. +func (dm *DatasetMetadataToUpdate) SetLabel(name, value string) { + if dm.setLabels == nil { + dm.setLabels = map[string]string{} + } + dm.setLabels[name] = value +} + +// DeleteLabel causes a label to be deleted when dm is used in a +// call to Dataset.Update. +func (dm *DatasetMetadataToUpdate) DeleteLabel(name string) { + if dm.deleteLabels == nil { + dm.deleteLabels = map[string]bool{} + } + dm.deleteLabels[name] = true } // Dataset creates a handle to a BigQuery dataset in the client's project. diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go index 0f42b9cadc01..99604e631b0b 100644 --- a/bigquery/integration_test.go +++ b/bigquery/integration_test.go @@ -303,6 +303,37 @@ func TestIntegration_DatasetUpdateDefaultExpiration(t *testing.T) { } } +func TestIntegration_DatasetUpdateLabels(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) + } + // TODO(jba): use a separate dataset for each test run so + // tests don't interfere with each other. + var dm DatasetMetadataToUpdate + dm.SetLabel("label", "value") + md, err = dataset.Update(ctx, dm, "") + if err != nil { + t.Fatal(err) + } + if got, want := md.Labels["label"], "value"; got != want { + t.Errorf("got %q, want %q", got, want) + } + dm = DatasetMetadataToUpdate{} + dm.DeleteLabel("label") + md, err = dataset.Update(ctx, dm, "") + if err != nil { + t.Fatal(err) + } + if _, ok := md.Labels["label"]; ok { + t.Error("label still present after deletion") + } +} + func TestIntegration_Tables(t *testing.T) { if client == nil { t.Skip("Integration tests skipped") diff --git a/bigquery/service.go b/bigquery/service.go index 1726cef1cac6..dc3f3946cbe4 100644 --- a/bigquery/service.go +++ b/bigquery/service.go @@ -689,6 +689,20 @@ func (s *bigqueryService) insertDataset(ctx context.Context, datasetID, projectI } func (s *bigqueryService) patchDataset(ctx context.Context, projectID, datasetID string, dm *DatasetMetadataToUpdate, etag string) (*DatasetMetadata, error) { + ds := bqDatasetFromMetadata(dm) + call := s.s.Datasets.Patch(projectID, datasetID, ds).Context(ctx) + setClientHeader(call.Header()) + if etag != "" { + call.Header().Set("If-Match", etag) + } + ds2, err := call.Do() + if err != nil { + return nil, err + } + return bqDatasetToMetadata(ds2), nil +} + +func bqDatasetFromMetadata(dm *DatasetMetadataToUpdate) *bq.Dataset { ds := &bq.Dataset{} forceSend := func(field string) { ds.ForceSendFields = append(ds.ForceSendFields, field) @@ -711,16 +725,19 @@ func (s *bigqueryService) patchDataset(ctx context.Context, projectID, datasetID ds.DefaultTableExpirationMs = int64(dur.Seconds() * 1000) } } - call := s.s.Datasets.Patch(projectID, datasetID, ds).Context(ctx) - setClientHeader(call.Header()) - if etag != "" { - call.Header().Set("If-Match", etag) - } - ds2, err := call.Do() - if err != nil { - return nil, err + if dm.setLabels != nil || dm.deleteLabels != nil { + ds.Labels = map[string]string{} + for k, v := range dm.setLabels { + ds.Labels[k] = v + } + if len(ds.Labels) == 0 && len(dm.deleteLabels) > 0 { + forceSend("Labels") + } + for l := range dm.deleteLabels { + ds.NullFields = append(ds.NullFields, "Labels."+l) + } } - return bqDatasetToMetadata(ds2), nil + return ds } func (s *bigqueryService) deleteDataset(ctx context.Context, datasetID, projectID string) error { diff --git a/bigquery/service_test.go b/bigquery/service_test.go index 1944da4fc4dc..79f190a31f96 100644 --- a/bigquery/service_test.go +++ b/bigquery/service_test.go @@ -83,3 +83,26 @@ func TestBQTableToMetadata(t *testing.T) { } } } + +func TestBQDatasetFromMetadata(t *testing.T) { + dm := DatasetMetadataToUpdate{ + Description: "desc", + Name: "name", + DefaultTableExpiration: time.Hour, + } + dm.SetLabel("label", "value") + dm.DeleteLabel("del") + + got := bqDatasetFromMetadata(&dm) + want := &bq.Dataset{ + Description: "desc", + FriendlyName: "name", + DefaultTableExpirationMs: 60 * 60 * 1000, + Labels: map[string]string{"label": "value"}, + ForceSendFields: []string{"Description", "FriendlyName"}, + NullFields: []string{"Labels.del"}, + } + if diff := testutil.Diff(got, want); diff != "" { + t.Errorf("-got, +want:\n%s", diff) + } +}