From 184d8f9f629f5ef1ddf4abe64f5a7751b487a808 Mon Sep 17 00:00:00 2001
From: Tom Pantelis <tompantelis@gmail.com>
Date: Fri, 4 Oct 2024 15:01:38 -0400
Subject: [PATCH] Add IdentfyingLabels to CreateOrUpdateWithOptions

If the GenerateName field is set in the target resource, the
existing resource is searched for by the target's labels. However
this assumes the labels remain static for the lifetime of the
resource. If the labels are updated then the existing resource is
not found and a new one is created. To avoid this, allow the user
to specify which labels uniquely identify the resource by adding
an IdentfyingLabels field to CreateOrUpdateOptions.

Signed-off-by: Tom Pantelis <tompantelis@gmail.com>
---
 pkg/util/create_or_update.go      | 21 +++++++++++++++------
 pkg/util/create_or_update_test.go | 25 +++++++++++++++++--------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/pkg/util/create_or_update.go b/pkg/util/create_or_update.go
index 66015ee8..fea6caa1 100644
--- a/pkg/util/create_or_update.go
+++ b/pkg/util/create_or_update.go
@@ -70,6 +70,8 @@ type CreateOrUpdateOptions[T runtime.Object] struct {
 	Obj            T
 	MutateOnUpdate MutateFn[T]
 	MutateOnCreate MutateFn[T]
+	// IdentifyingLabels is used to find an existing resource if GenerateName is set in the target resource.
+	IdentifyingLabels map[string]string
 }
 
 func CreateOrUpdateWithOptions[T runtime.Object](ctx context.Context, options CreateOrUpdateOptions[T]) (OperationResult, T, error) {
@@ -121,7 +123,7 @@ func maybeCreateOrUpdate[T runtime.Object](ctx context.Context, options CreateOr
 	objMeta := resource.MustToMeta(options.Obj)
 
 	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
-		existing, err := getResource(ctx, options.Client, options.Obj)
+		existing, err := getResource(ctx, &options)
 		if apierrors.IsNotFound(err) {
 			if op != opCreate {
 				logger.V(log.LIBTRACE).Infof("Resource %q does not exist - not updating", objMeta.GetName())
@@ -204,16 +206,23 @@ func maybeCreateOrUpdate[T runtime.Object](ctx context.Context, options CreateOr
 }
 
 //nolint:wrapcheck // No need to wrap errors
-func getResource[T runtime.Object](ctx context.Context, client resource.Interface[T], obj T) (T, error) {
-	objMeta := resource.MustToMeta(obj)
+func getResource[T runtime.Object](ctx context.Context, options *CreateOrUpdateOptions[T]) (T, error) {
+	objMeta := resource.MustToMeta(options.Obj)
 
 	if objMeta.GetName() != "" || objMeta.GetGenerateName() == "" {
-		obj, err := client.Get(ctx, objMeta.GetName(), metav1.GetOptions{})
+		obj, err := options.Client.Get(ctx, objMeta.GetName(), metav1.GetOptions{})
 		return obj, err
 	}
 
-	list, err := client.List(ctx, metav1.ListOptions{
-		LabelSelector: labels.SelectorFromSet(objMeta.GetLabels()).String(),
+	var selectorLabels map[string]string
+	if len(options.IdentifyingLabels) > 0 {
+		selectorLabels = options.IdentifyingLabels
+	} else {
+		selectorLabels = objMeta.GetLabels()
+	}
+
+	list, err := options.Client.List(ctx, metav1.ListOptions{
+		LabelSelector: labels.SelectorFromSet(selectorLabels).String(),
 	})
 	if err != nil {
 		return *new(T), err
diff --git a/pkg/util/create_or_update_test.go b/pkg/util/create_or_update_test.go
index a18da61a..c77c9fbd 100644
--- a/pkg/util/create_or_update_test.go
+++ b/pkg/util/create_or_update_test.go
@@ -148,8 +148,23 @@ var _ = Describe("CreateOrUpdate function", func() {
 	t := newCreateOrUpdateTestDiver()
 
 	createOrUpdate := func(expResult util.OperationResult) error {
-		result, err := util.CreateOrUpdate[*unstructured.Unstructured](context.TODO(), resource.ForDynamic(t.client),
-			resource.MustToUnstructured(t.pod), t.mutateFn)
+		options := util.CreateOrUpdateOptions[*unstructured.Unstructured]{
+			Client:         resource.ForDynamic(t.client),
+			MutateOnUpdate: t.mutateFn,
+		}
+
+		if t.pod.GenerateName != "" {
+			options.IdentifyingLabels = map[string]string{}
+			for k, v := range t.pod.Labels {
+				options.IdentifyingLabels[k] = v
+			}
+
+			t.pod.Labels["new-label"] = "new-value"
+		}
+
+		options.Obj = resource.MustToUnstructured(t.pod)
+
+		result, _, err := util.CreateOrUpdateWithOptions[*unstructured.Unstructured](context.TODO(), options)
 		if err != nil && expResult != util.OperationResultNone {
 			return err
 		}
@@ -167,12 +182,6 @@ var _ = Describe("CreateOrUpdate function", func() {
 		})
 
 		Context("and a mutation function specified", func() {
-			BeforeEach(func() {
-				t.pod.Name = ""
-				t.pod.GenerateName = "name-prefix-"
-				t.pod.Labels = map[string]string{"label1": "value1", "label2": "value2"}
-			})
-
 			It("should invoke the function on create", func() {
 				result, created, err := util.CreateOrUpdateWithOptions[*unstructured.Unstructured](context.TODO(),
 					util.CreateOrUpdateOptions[*unstructured.Unstructured]{