From fa5aaa40ac9d8f25a0b258b71d6a8450b4409b9d Mon Sep 17 00:00:00 2001
From: ashmrtn <ashmrtnz@alcion.ai>
Date: Tue, 22 Nov 2022 18:23:58 -0800
Subject: [PATCH] feat(snapshots): Add labels to snapshot checkpoint manifests
 (#2548)

* Apply labels to checkpoint snapshot manifests

This allows looking up checkpoint snapshot manifests by label if
desired.

* Fixup other calls to Upload function

* Fix missed Upload call

* Revert code to pass checkpoint labels as arg

This reverts commits
* bf844ff0683fd99bc8efe2d1ccd9741c354cce72
* 76438c9bc1c5ba0b81aca8669358accad92441f7
* d748b6361b7364a52734d612e5da84e8c4ff5a15

* Add checkpoint labels to checkpoint manifests

Allow finer-grained lookups on checkpoints by adding labels to them.
Labels are specified in the Uploader struct and apply to all checkpoints
for a snapshot.

* Test checkpoint labels.

Be strict about testing and have a separate copy of the labels.
---
 snapshot/snapshotfs/upload.go      |  4 ++++
 snapshot/snapshotfs/upload_test.go | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/snapshot/snapshotfs/upload.go b/snapshot/snapshotfs/upload.go
index 307e3683e..efbb7eacb 100644
--- a/snapshot/snapshotfs/upload.go
+++ b/snapshot/snapshotfs/upload.go
@@ -95,6 +95,9 @@ type Uploader struct {
 	// When set to true, do not ignore any files, regardless of policy settings.
 	DisableIgnoreRules bool
 
+	// Labels to apply to every checkpoint made for this snapshot.
+	CheckpointLabels map[string]string
+
 	repo repo.RepositoryWriter
 
 	// stats must be allocated on heap to enforce 64-bit alignment due to atomic access on ARM.
@@ -506,6 +509,7 @@ func (u *Uploader) checkpointRoot(ctx context.Context, cp *checkpointRegistry, p
 	man.EndTime = fs.UTCTimestampFromTime(u.repo.Time())
 	man.StartTime = man.EndTime
 	man.IncompleteReason = IncompleteReasonCheckpoint
+	man.Tags = u.CheckpointLabels
 
 	if _, err := snapshot.SaveSnapshot(ctx, u.repo, &man); err != nil {
 		return errors.Wrap(err, "error saving checkpoint snapshot")
diff --git a/snapshot/snapshotfs/upload_test.go b/snapshot/snapshotfs/upload_test.go
index db927b840..9ef7d5626 100644
--- a/snapshot/snapshotfs/upload_test.go
+++ b/snapshot/snapshotfs/upload_test.go
@@ -636,6 +636,18 @@ func TestUploadWithCheckpointing(t *testing.T) {
 		Path:     "path",
 	}
 
+	labels := map[string]string{
+		"shape": "square",
+		"color": "red",
+	}
+
+	// Be paranoid and make a copy of the labels in the uploader so we know stuff
+	// didn't change.
+	u.CheckpointLabels = make(map[string]string, len(labels))
+	for k, v := range labels {
+		u.CheckpointLabels[k] = v
+	}
+
 	// inject a action into mock filesystem to trigger and wait for checkpoints at few places.
 	// the places are not important, what's important that those are 3 separate points in time.
 	dirsToCheckpointAt := []*mockfs.Directory{
@@ -671,6 +683,8 @@ func TestUploadWithCheckpointing(t *testing.T) {
 		if got, want := sn.IncompleteReason, IncompleteReasonCheckpoint; got != want {
 			t.Errorf("unexpected incompleteReason %q, want %q", got, want)
 		}
+
+		assert.Equal(t, labels, sn.Tags)
 	}
 }