From bb60c29672f3736adf2b23c11957c3fa4b7e0995 Mon Sep 17 00:00:00 2001
From: Donny Xia <donnyxia@google.com>
Date: Mon, 28 Sep 2020 16:37:49 -0700
Subject: [PATCH] add GetElementByKey

---
 kyaml/yaml/fns.go      | 43 ++++++++++++++++++++++++++++++++++++++++--
 kyaml/yaml/fns_test.go | 41 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/kyaml/yaml/fns.go b/kyaml/yaml/fns.go
index ffc1e4d6f..50a1980a3 100644
--- a/kyaml/yaml/fns.go
+++ b/kyaml/yaml/fns.go
@@ -57,6 +57,16 @@ type ElementSetter struct {
 	Value string `yaml:"value,omitempty"`
 }
 
+// isMappingNode returns whether node is a mapping node
+func (e ElementSetter) isMappingNode(node *RNode) bool {
+	return ErrorIfInvalid(node, yaml.MappingNode) == nil
+}
+
+// isMappingSetter returns is this setter intended to set a mapping node
+func (e ElementSetter) isMappingSetter() bool {
+	return e.Key != "" && e.Value != ""
+}
+
 func (e ElementSetter) Filter(rn *RNode) (*RNode, error) {
 	if err := ErrorIfInvalid(rn, SequenceNode); err != nil {
 		return nil, err
@@ -73,6 +83,11 @@ func (e ElementSetter) Filter(rn *RNode) (*RNode, error) {
 		if IsMissingOrNull(newNode) || IsEmptyMap(newNode) {
 			continue
 		}
+		// keep non-mapping node in the Content when we want to match a mapping.
+		if !e.isMappingNode(newNode) && e.isMappingSetter() {
+			newContent = append(newContent, elem)
+			continue
+		}
 
 		// check if this is the element we are matching
 		val, err := newNode.Pipe(FieldMatcher{Name: e.Key, StringValue: e.Value})
@@ -165,8 +180,13 @@ func MatchElement(field, value string) ElementMatcher {
 	return ElementMatcher{FieldName: field, FieldValue: value}
 }
 
+func GetElementByKey(key string) ElementMatcher {
+	return ElementMatcher{FieldName: key, MatchAnyValue: true}
+}
+
 // ElementMatcher returns the first element from a Sequence matching the
-// specified field's value.
+// specified field's value. If there's no match, and no configuration error,
+// the matcher returns nil, nil.
 type ElementMatcher struct {
 	Kind string `yaml:"kind,omitempty"`
 
@@ -180,12 +200,20 @@ type ElementMatcher struct {
 
 	// Create will create the Element if it is not found
 	Create *RNode `yaml:"create,omitempty"`
+
+	// MatchAnyValue indicates that matcher should only consider the key and ignore
+	// the actual value in the list. FieldValue must be empty when NoValue is
+	// set to true.
+	MatchAnyValue bool `yaml:"noValue,omitempty"`
 }
 
 func (e ElementMatcher) Filter(rn *RNode) (*RNode, error) {
 	if err := ErrorIfInvalid(rn, yaml.SequenceNode); err != nil {
 		return nil, err
 	}
+	if e.MatchAnyValue && e.FieldValue != "" {
+		return nil, fmt.Errorf("FieldValue must be empty when NoValue is set to true")
+	}
 
 	// SequenceNode Content is a slice of ScalarNodes.  Each ScalarNode has a
 	// YNode containing the primitive data.
@@ -207,7 +235,18 @@ func (e ElementMatcher) Filter(rn *RNode) (*RNode, error) {
 		// cast the entry to a RNode so we can operate on it
 		elem := NewRNode(rn.Content()[i])
 
-		field, err := elem.Pipe(MatchField(e.FieldName, e.FieldValue))
+		// only check mapping node
+		if err := ErrorIfInvalid(elem, yaml.MappingNode); err != nil {
+			continue
+		}
+
+		var field *RNode
+		var err error
+		if e.MatchAnyValue {
+			field, err = elem.Pipe(Get(e.FieldName))
+		} else {
+			field, err = elem.Pipe(MatchField(e.FieldName, e.FieldValue))
+		}
 		if IsFoundOrError(field, err) {
 			return elem, err
 		}
diff --git a/kyaml/yaml/fns_test.go b/kyaml/yaml/fns_test.go
index ad5b9937b..1b4e327c5 100644
--- a/kyaml/yaml/fns_test.go
+++ b/kyaml/yaml/fns_test.go
@@ -57,6 +57,47 @@ func TestAppend(t *testing.T) {
 	assert.Nil(t, rn)
 }
 
+func TestGetElementByKey(t *testing.T) {
+	node, err := Parse(`
+- b: c
+- i
+- d: e
+- f: g
+- f: h
+`)
+	assert.NoError(t, err)
+
+	rn, err := node.Pipe(GetElementByKey("b"))
+	assert.NoError(t, err)
+	assert.Equal(t, "b: c\n", assertNoErrorString(t)(rn.String()))
+
+	rn, err = node.Pipe(GetElementByKey("f"))
+	assert.NoError(t, err)
+	assert.Equal(t, "f: g\n", assertNoErrorString(t)(rn.String()))
+}
+
+func TestElementMatcherWithNoValue(t *testing.T) {
+	node, err := Parse(`
+- a: c
+- b: ""
+`)
+	assert.NoError(t, err)
+
+	rn, err := node.Pipe(ElementMatcher{FieldName: "b"})
+	assert.NoError(t, err)
+	assert.Equal(t, "b: \"\"\n", assertNoErrorString(t)(rn.String()))
+
+	rn, err = node.Pipe(ElementMatcher{FieldName: "a"})
+	assert.NoError(t, err)
+	assert.Nil(t, rn)
+
+	rn, err = node.Pipe(ElementMatcher{FieldName: "a", MatchAnyValue: true})
+	assert.NoError(t, err)
+	assert.Equal(t, "a: c\n", assertNoErrorString(t)(rn.String()))
+
+	_, err = node.Pipe(ElementMatcher{FieldName: "a", FieldValue: "c", MatchAnyValue: true})
+	assert.Errorf(t, err, "FieldValue must be empty when NoValue is set to true")
+}
 func TestClearField_Fn(t *testing.T) {
 	node, err := Parse(NodeSampleData)
 	assert.NoError(t, err)