From 1f35044894c7395326d7a5b787a2406d5cd92cb1 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Mon, 13 Jun 2022 12:29:51 -0700
Subject: [PATCH] feat(bigquery): add support for dataset tags (#6114)

---
 bigquery/dataset.go      | 32 ++++++++++++++++++++++++++++++++
 bigquery/dataset_test.go |  8 ++++++++
 2 files changed, 40 insertions(+)

diff --git a/bigquery/dataset.go b/bigquery/dataset.go
index 7c23b286e47..a9deeb3295e 100644
--- a/bigquery/dataset.go
+++ b/bigquery/dataset.go
@@ -50,11 +50,37 @@ type DatasetMetadata struct {
 	LastModifiedTime time.Time // When the dataset or any of its tables were modified.
 	FullID           string    // The full dataset ID in the form projectID:datasetID.
 
+	// The tags associated with this dataset. Tag keys are
+	// globally unique, and managed via the resource manager API.
+	// More information: https://cloud.google.com/resource-manager/docs/tags/tags-overview
+	Tags []*DatasetTag
+
 	// ETag is the ETag obtained when reading metadata. Pass it to Dataset.Update to
 	// ensure that the metadata hasn't changed since it was read.
 	ETag string
 }
 
+// DatasetTag is a representation of a single tag key/value.
+type DatasetTag struct {
+	// TagKey is the namespaced friendly name of the tag key, e.g.
+	// "12345/environment" where 12345 is org id.
+	TagKey string
+
+	// TagValue is the friendly short name of the tag value, e.g.
+	// "production".
+	TagValue string
+}
+
+func bqToDatasetTag(in *bq.DatasetTags) *DatasetTag {
+	if in == nil {
+		return nil
+	}
+	return &DatasetTag{
+		TagKey:   in.TagKey,
+		TagValue: in.TagValue,
+	}
+}
+
 // DatasetMetadataToUpdate is used when updating a dataset's metadata.
 // Only non-nil fields will be updated.
 type DatasetMetadataToUpdate struct {
@@ -230,6 +256,12 @@ func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) {
 		}
 		dm.Access = append(dm.Access, e)
 	}
+	for _, bqTag := range d.Tags {
+		tag := bqToDatasetTag(bqTag)
+		if tag != nil {
+			dm.Tags = append(dm.Tags, tag)
+		}
+	}
 	return dm, nil
 }
 
diff --git a/bigquery/dataset_test.go b/bigquery/dataset_test.go
index 1d49c681512..2ce72b9d92a 100644
--- a/bigquery/dataset_test.go
+++ b/bigquery/dataset_test.go
@@ -413,6 +413,10 @@ func TestBQToDatasetMetadata(t *testing.T) {
 				},
 			},
 		},
+		Tags: []*bq.DatasetTags{
+			{TagKey: "tag1", TagValue: "value1"},
+			{TagKey: "tag2", TagValue: "value2"},
+		},
 		Etag: "etag",
 	}
 	want := &DatasetMetadata{
@@ -437,6 +441,10 @@ func TestBQToDatasetMetadata(t *testing.T) {
 				},
 			},
 		},
+		Tags: []*DatasetTag{
+			{TagKey: "tag1", TagValue: "value1"},
+			{TagKey: "tag2", TagValue: "value2"},
+		},
 		ETag: "etag",
 	}
 	got, err := bqToDatasetMetadata(q, client)