From 925693ded7216770545e6ee334f00b0e86da83ca Mon Sep 17 00:00:00 2001
From: Tamal Saha <tamal@appscode.com>
Date: Tue, 10 Dec 2024 14:02:32 -0800
Subject: [PATCH] Add PatchE method

Signed-off-by: Tamal Saha <tamal@appscode.com>
---
 client/client.go | 66 ++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 61 insertions(+), 5 deletions(-)

diff --git a/client/client.go b/client/client.go
index b1607e34d..321ec3abc 100644
--- a/client/client.go
+++ b/client/client.go
@@ -66,9 +66,10 @@ func NewUncachedClient(cfg *rest.Config, funcs ...func(*runtime.Scheme) error) (
 }
 
 type (
-	TransformFunc       func(obj client.Object, createOp bool) client.Object
-	TransformFuncE      func(obj client.Object, createOp bool) (client.Object, error)
-	TransformStatusFunc func(obj client.Object) client.Object
+	TransformFunc  func(obj client.Object, createOp bool) client.Object
+	TransformFuncE func(obj client.Object, createOp bool) (client.Object, error)
+	PatchFunc      func(obj client.Object) client.Object
+	PatchFuncE     func(obj client.Object) (client.Object, error)
 )
 
 func CreateOrPatchE(ctx context.Context, c client.Client, obj client.Object, transform TransformFuncE, opts ...client.PatchOption) (kutil.VerbType, error) {
@@ -147,6 +148,51 @@ func CreateOrPatch(ctx context.Context, c client.Client, obj client.Object, tran
 	}, opts...)
 }
 
+func PatchE(ctx context.Context, c client.Client, obj client.Object, transform PatchFuncE, opts ...client.PatchOption) (kutil.VerbType, error) {
+	gvk, err := apiutil.GVKForObject(obj, c.Scheme())
+	if err != nil {
+		return kutil.VerbUnchanged, errors.Wrapf(err, "failed to get GVK for object %T", obj)
+	}
+
+	_, unstructuredObj := obj.(*unstructured.Unstructured)
+
+	var patch client.Patch
+	if isOfficialTypes(gvk.Group) && !unstructuredObj {
+		patch = client.StrategicMergeFrom(obj)
+	} else {
+		patch = client.MergeFrom(obj)
+	}
+	mod, err := transform(obj.DeepCopyObject().(client.Object))
+	if err != nil {
+		return kutil.VerbUnchanged, err
+	}
+	err = c.Patch(ctx, mod, patch, opts...)
+	if err != nil {
+		return kutil.VerbUnchanged, err
+	}
+
+	vt := kutil.VerbUnchanged
+	if mod.GetGeneration() > 0 {
+		if obj.GetGeneration() != mod.GetGeneration() {
+			vt = kutil.VerbPatched
+		}
+	} else {
+		// Secret, ServiceAccount etc resources do not use metadata.generation
+		if meta.ObjectHash(obj) != meta.ObjectHash(mod) {
+			vt = kutil.VerbPatched
+		}
+	}
+	assign(obj, mod)
+	return vt, nil
+}
+
+func Patch(ctx context.Context, c client.Client, obj client.Object, transform PatchFunc, opts ...client.PatchOption) (kutil.VerbType, error) {
+	return PatchE(ctx, c, obj, func(obj client.Object) (client.Object, error) {
+		transform(obj)
+		return obj, nil
+	}, opts...)
+}
+
 func assign(target, src any) {
 	srcValue := reflect.ValueOf(src)
 	if srcValue.Kind() == reflect.Pointer {
@@ -155,7 +201,7 @@ func assign(target, src any) {
 	reflect.ValueOf(target).Elem().Set(srcValue)
 }
 
-func PatchStatus(ctx context.Context, c client.Client, obj client.Object, transform TransformStatusFunc, opts ...client.SubResourcePatchOption) (kutil.VerbType, error) {
+func PatchStatusE(ctx context.Context, c client.Client, obj client.Object, transform PatchFuncE, opts ...client.SubResourcePatchOption) (kutil.VerbType, error) {
 	cur := obj.DeepCopyObject().(client.Object)
 	key := types.NamespacedName{
 		Namespace: cur.GetNamespace(),
@@ -172,7 +218,10 @@ func PatchStatus(ctx context.Context, c client.Client, obj client.Object, transf
 	//   - application/merge-patch+json,
 	//   - application/apply-patch+yaml
 	patch := client.MergeFrom(cur)
-	mod := transform(cur.DeepCopyObject().(client.Object))
+	mod, err := transform(cur.DeepCopyObject().(client.Object))
+	if err != nil {
+		return kutil.VerbUnchanged, err
+	}
 	err = c.Status().Patch(ctx, mod, patch, opts...)
 	if err != nil {
 		return kutil.VerbUnchanged, err
@@ -181,6 +230,13 @@ func PatchStatus(ctx context.Context, c client.Client, obj client.Object, transf
 	return kutil.VerbPatched, nil
 }
 
+func PatchStatus(ctx context.Context, c client.Client, obj client.Object, transform PatchFunc, opts ...client.SubResourcePatchOption) (kutil.VerbType, error) {
+	return PatchStatusE(ctx, c, obj, func(obj client.Object) (client.Object, error) {
+		transform(obj)
+		return obj, nil
+	}, opts...)
+}
+
 func isOfficialTypes(group string) bool {
 	return !strings.ContainsRune(group, '.')
 }