diff --git a/pubsub/integration_test.go b/pubsub/integration_test.go index 121cc2518ec..ae3356873e2 100644 --- 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 index 8a223a99f9a..78797f36475 100644 --- 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 }