From 4f05a71d1faa001867e357067dd39b01c6c6bb20 Mon Sep 17 00:00:00 2001
From: Jahvon Dockery <jahvondockery@gmail.com>
Date: Mon, 11 Oct 2021 11:24:39 -0400
Subject: [PATCH] Add merge function for cluster snapshot (#301)

* Add merge function for cluster snapshot
* changelog and test
* comment and reword
* comment:
* Merge remote-tracking branch 'origin/master' into cluster-snapshot-merge
* update changelog
---
 changelog/v0.21.1/add-merge-function.yaml |  4 ++
 pkg/resource/snapshot.go                  | 35 ++++++++++++++++
 pkg/resource/snapshot_test.go             | 49 +++++++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 changelog/v0.21.1/add-merge-function.yaml

diff --git a/changelog/v0.21.1/add-merge-function.yaml b/changelog/v0.21.1/add-merge-function.yaml
new file mode 100644
index 000000000..ec26eea35
--- /dev/null
+++ b/changelog/v0.21.1/add-merge-function.yaml
@@ -0,0 +1,4 @@
+changelog:
+  - type: NEW_FEATURE
+    description: Adds a `Merge` function to Snapshot and ClusterSnapshot so that two snapshots can be merged.
+    issueLink: https://github.com/solo-io/skv2/issues/302
diff --git a/pkg/resource/snapshot.go b/pkg/resource/snapshot.go
index 4f6220dc5..f7566c641 100644
--- a/pkg/resource/snapshot.go
+++ b/pkg/resource/snapshot.go
@@ -71,6 +71,25 @@ func (s Snapshot) Clone(selectors ...GVKSelectorFunc) Snapshot {
 	return clone
 }
 
+// Merges the Snapshot with a Snapshot passed in as an argument. The values
+// in the passed in Snapshot will take precedence when there is an object mapped
+// to the same gvk and name in both Snapshots.
+func (s Snapshot) Merge(toMerge Snapshot) Snapshot {
+	merged := s.Clone()
+	for gvk, objectsMap := range toMerge {
+		if _, ok := merged[gvk]; ok {
+			for name, object := range objectsMap {
+				// If there is already an object specified here, the object from toMerge
+				// will replace it
+				merged[gvk][name] = object
+			}
+		} else {
+			merged[gvk] = objectsMap
+		}
+	}
+	return merged
+}
+
 // ClusterSnapshot represents a set of snapshots partitioned by cluster
 type ClusterSnapshot map[string]Snapshot
 
@@ -128,3 +147,19 @@ func (cs ClusterSnapshot) Clone(selectors ...GVKSelectorFunc) ClusterSnapshot {
 	}
 	return clone
 }
+
+// Merges the ClusterSnapshot with a ClusterSnapshot passed in as an argument.
+// If a cluster exists in both ClusterSnapshots, then both Snapshots for the
+// cluster is merged; with the passed in ClusterSnapshot's corresponding Snapshot
+// taking precedence in case of conflicts.
+func (cs ClusterSnapshot) Merge(toMerge ClusterSnapshot) ClusterSnapshot {
+	merged := cs.Clone()
+	for cluster, snapshot := range toMerge {
+		if baseSnap, ok := merged[cluster]; ok {
+			merged[cluster] = baseSnap.Merge(snapshot)
+		} else {
+			merged[cluster] = snapshot
+		}
+	}
+	return merged
+}
diff --git a/pkg/resource/snapshot_test.go b/pkg/resource/snapshot_test.go
index 0fa5276e4..767f0b936 100644
--- a/pkg/resource/snapshot_test.go
+++ b/pkg/resource/snapshot_test.go
@@ -225,6 +225,55 @@ var _ = Describe("Snapshot", func() {
 			Expect(controllerutils.ObjectsEqual(copiedPaint, paint)).To(BeFalse())
 		})
 
+		It("will merge two snapshots properly", func() {
+			paint := &testv1.Paint{
+				ObjectMeta: metav1.ObjectMeta{Name: "paint", Namespace: "x"},
+			}
+			paintOverride := &testv1.Paint{
+				ObjectMeta: metav1.ObjectMeta{Name: "override", Namespace: "x"},
+			}
+			paint2 := &testv1.Paint{
+				ObjectMeta: metav1.ObjectMeta{Name: "paint2", Namespace: "y"},
+			}
+			name := types.NamespacedName{
+				Namespace: "x",
+				Name:      "paint",
+			}
+			name2 := types.NamespacedName{
+				Namespace: "y",
+				Name:      "paint2",
+			}
+			cluster1Name := "cluster1"
+			cluster2Name := "cluster2"
+			leftSnapshot := resource.ClusterSnapshot{
+				cluster1Name: resource.Snapshot{
+					testv1.PaintGVK: map[types.NamespacedName]resource.TypedObject{
+						name:  paint,
+						name2: paint2,
+					},
+				},
+			}
+			rightSnapshot := resource.ClusterSnapshot{
+				cluster1Name: resource.Snapshot{
+					testv1.PaintGVK: map[types.NamespacedName]resource.TypedObject{
+						name: paintOverride,
+					},
+				},
+				cluster2Name: resource.Snapshot{
+					testv1.PaintGVK: map[types.NamespacedName]resource.TypedObject{
+						name: paint,
+					},
+				},
+			}
+
+			snap := leftSnapshot.Merge(rightSnapshot)
+			Expect(snap[cluster1Name][testv1.PaintGVK][name].GetName()).
+				To(Equal(paintOverride.Name))
+			Expect(snap[cluster1Name][testv1.PaintGVK][name2].GetName()).
+				To(Equal(paint2.Name))
+			Expect(snap[cluster2Name][testv1.PaintGVK][name].GetName()).
+				To(Equal(paint.Name))
+		})
 	})
 
 })