feat(pubsub): add support for snapshot labels (#6835)

hongalex · web-flow · commit c17851b5c3d8 · 2022-10-11T12:56:22.000-07:00
* feat(pubsub): add support for snapshot labels

* update comments and field construction

* fix typo in comment
diff --git a/pubsub/integration_test.go b/pubsub/integration_test.go
@@ -164,6 +164,15 @@ func TestIntegration_All(t *testing.T) {
 		t.Fatalf("CreateSnapshot error: %v", err)
 	}
 
+	labels := map[string]string{"foo": "bar"}
+	sc, err := snap.SetLabels(ctx, labels)
+	if err != nil {
+		t.Fatalf("Snapshot.SetLabels error: %v", err)
+	}
+	if diff := testutil.Diff(sc.Labels, labels); diff != "" {
+		t.Fatalf("\ngot: - want: +\n%s", diff)
+	}
+
 	timeoutCtx, cancel := context.WithTimeout(ctx, time.Minute)
 	defer cancel()
 	err = internal.Retry(timeoutCtx, gax.Backoff{}, func() (bool, error) {
diff --git a/pubsub/snapshot.go b/pubsub/snapshot.go
@@ -21,6 +21,7 @@ import (
 	"time"
 
 	pb "google.golang.org/genproto/googleapis/pubsub/v1"
+	fmpb "google.golang.org/genproto/protobuf/field_mask"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -42,11 +43,30 @@ func (s *Snapshot) ID() string {
 	return s.name[slash+1:]
 }
 
+// SetLabels sets or replaces the labels on a given snapshot.
+func (s *Snapshot) SetLabels(ctx context.Context, label map[string]string) (*SnapshotConfig, error) {
+	sc, err := s.c.subc.UpdateSnapshot(ctx, &pb.UpdateSnapshotRequest{
+		Snapshot: &pb.Snapshot{
+			Name:   s.name,
+			Labels: label,
+		},
+		UpdateMask: &fmpb.FieldMask{
+			Paths: []string{"labels"},
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	return toSnapshotConfig(sc, s.c)
+}
+
 // SnapshotConfig contains the details of a Snapshot.
 type SnapshotConfig struct {
 	*Snapshot
 	Topic      *Topic
 	Expiration time.Time
+	// The set of labels for the snapshot.
+	Labels map[string]string
 }
 
 // Snapshot creates a reference to a snapshot.
@@ -151,5 +171,6 @@ func toSnapshotConfig(snap *pb.Snapshot, c *Client) (*SnapshotConfig, error) {
 		Snapshot:   &Snapshot{c: c, name: snap.Name},
 		Topic:      newTopic(c, snap.Topic),
 		Expiration: exp,
+		Labels:     snap.Labels,
 	}, nil
 }
