From 215c4e07528afc29680195967cbbc87db50a493d Mon Sep 17 00:00:00 2001
From: Hongchao Deng <hongchaodeng1@gmail.com>
Date: Mon, 13 Sep 2021 08:19:16 -0400
Subject: [PATCH] Fix: patch any key using `retainKeys` strategy (#2280)

---
 pkg/cue/model/sets/operation.go      | 13 ++++++-------
 pkg/cue/model/sets/operation_test.go | 28 ++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/pkg/cue/model/sets/operation.go b/pkg/cue/model/sets/operation.go
index e7b3093f4..da5b88176 100644
--- a/pkg/cue/model/sets/operation.go
+++ b/pkg/cue/model/sets/operation.go
@@ -103,8 +103,8 @@ func listMergeProcess(field *ast.Field, key string, baseList, patchList *ast.Lis
 
 func strategyPatchHandle(baseNode ast.Node) interceptor {
 	return func(lnode ast.Node) (ast.Node, error) {
-		walker := newWalker(func(node ast.Node, ctx walkCtx) {
-			field, ok := node.(*ast.Field)
+		walker := newWalker(func(patchNode ast.Node, ctx walkCtx) {
+			field, ok := patchNode.(*ast.Field)
 			if !ok {
 				return
 			}
@@ -136,7 +136,7 @@ func strategyPatchHandle(baseNode ast.Node) interceptor {
 				}
 				listMergeProcess(field, key, baselist, val)
 
-			case *ast.StructLit:
+			default:
 				if !isStrategyRetainKeys(field) {
 					return
 				}
@@ -148,20 +148,19 @@ func strategyPatchHandle(baseNode ast.Node) interceptor {
 						for _, elt := range v.Elts {
 							if fe, ok := elt.(*ast.Field); ok &&
 								labelStr(fe.Label) == labelStr(field.Label) {
-								fe.Value = ast.NewStruct()
+								fe.Value = field.Value
 							}
 						}
-					case *ast.File:
+					case *ast.File: // For the top level element
 						for _, decl := range v.Decls {
 							if fe, ok := decl.(*ast.Field); ok &&
 								labelStr(fe.Label) == labelStr(field.Label) {
-								fe.Value = ast.NewStruct()
+								fe.Value = field.Value
 							}
 						}
 					}
 				}
 			}
-
 		})
 		walker.walk(lnode)
 		return lnode, nil
diff --git a/pkg/cue/model/sets/operation_test.go b/pkg/cue/model/sets/operation_test.go
index ac6694544..8af00345e 100644
--- a/pkg/cue/model/sets/operation_test.go
+++ b/pkg/cue/model/sets/operation_test.go
@@ -428,8 +428,32 @@ spec: {
 		envs: [{
 			name:  "e1"
 			value: "v2"
-		}]
-	}]
+		}, ...]
+	}, ...]
+}
+`}, {
+			base: `
+kind: "Old"
+metadata: {
+	name: "Old"
+	labels: keep: "true"
+}
+`,
+			patch: `// +patchStrategy=retainKeys
+kind: "New"
+metadata: {
+	// +patchStrategy=retainKeys
+	name: "New"
+}
+`,
+			result: `	// +patchStrategy=retainKeys
+kind: "New"
+metadata: {
+	// +patchStrategy=retainKeys
+	name: "New"
+	labels: {
+		keep: "true"
+	}
 }
 `},
 	}