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 <[email protected]>
Reviewed-by: Sai Cheemalapati <[email protected]>

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.

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")

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 {

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)
+	}
+}