fix: set-image with more matching rules

functions/go/set-image/custom/custom.go

@@ -31,7 +31,7 @@ func SetAdditionalFieldSpec(img *fn.SubObject, objects fn.KubeObjects, addImgFie
 		if err != nil {
 			ctx.ResultErr(err.Error(), obj)
 		}
-		objects[i] = newObj
+		*objects[i] = *newObj
 	}
 }
 

functions/go/set-image/transformer/images_transformer.go

@@ -2,19 +2,24 @@ package transformer
 
 import (
 	"fmt"
-
 	"github.com/GoogleContainerTools/kpt-functions-catalog/functions/go/set-image/custom"
 	"github.com/GoogleContainerTools/kpt-functions-catalog/functions/go/set-image/third_party/sigs.k8s.io/kustomize/api/image"
 	"github.com/GoogleContainerTools/kpt-functions-catalog/functions/go/set-image/third_party/sigs.k8s.io/kustomize/api/types"
 	"github.com/GoogleContainerTools/kpt-functions-sdk/go/fn"
 )
 
-// Image contains an image name, a new name, a new tag or digest, which will replace the original name and tag.
 type Image struct {
-	// Name is a tag-less image name.
+	// DEPRECATED
+	//Name is a tag-less image name.
 	Name string `json:"name,omitempty" yaml:"name,omitempty"`
 
-	// NewName is the value used to replace the original name.
+	// ContainerName is the Pod's container name. It is used to choose the matching containers to update the images.
+	ContainerName string `json:"containerName,omitempty" yaml:"containerName,omitempty"`
+
+	// ImageName is the image name. It is used to choose images for update.
+	ImageName string `json:"imageName,omitempty" yaml:"imageName,omitempty"`
+
+	// NewName is the new image name
 	NewName string `json:"newName,omitempty" yaml:"newName,omitempty"`
 
 	// NewTag is the value used to replace the original tag.
@@ -43,21 +48,28 @@ func (t SetImage) Run(ctx *fn.Context, functionConfig *fn.KubeObject, items fn.K
 	if err != nil {
 		ctx.ResultErrAndDie(err.Error(), nil)
 	}
-	err = t.validateInput()
-	if err != nil {
-		ctx.ResultErrAndDie(err.Error(), nil)
+
+	items = items.WhereNot(fn.IsLocalConfig)
+
+	res := t.validateInput(items)
+	if res != nil {
+		if res.Severity == fn.Error {
+			ctx.ResultErrAndDie(res.Message, nil)
+		}
+		ctx.Result(res.Message, res.Severity)
+		return
 	}
 
 	for _, o := range items {
-		switch o.GetKind() {
-		case "Pod":
-			if err = t.setPodContainers(o); err != nil {
-				ctx.ResultErr(err.Error(), o)
-			}
-		case "Deployment", "StatefulSet", "ReplicaSet", "DaemonSet", "PodTemplate":
-			if err = t.setPodSpecContainers(o); err != nil {
-				ctx.ResultErr(err.Error(), o)
-			}
+		containers := getContainers(o)
+		warnings, err := t.updateContainerImages(containers)
+
+		for _, w := range warnings {
+			ctx.ResultWarn(w.Message, o)
+		}
+
+		if err != nil {
+			ctx.ResultErr(err.Error(), o)
 		}
 	}
 
@@ -74,7 +86,11 @@ func (t *SetImage) configDefaultData() error {
 	for key, val := range t.DataFromDefaultConfig {
 		switch key {
 		case "name":
-			t.Image.Name = val
+			t.Image.ImageName = val
+		case "containerName":
+			t.Image.ContainerName = val
+		case "imageName":
+			t.Image.ImageName = val
 		case "newName":
 			t.Image.NewName = val
 		case "newTag":
@@ -89,68 +105,144 @@ func (t *SetImage) configDefaultData() error {
 }
 
 // validateInput validates the inputs passed into via the functionConfig
-func (t *SetImage) validateInput() error {
-	// TODO: support container name and only one argument input in the next PR
-	if t.Image.Name == "" {
-		return fmt.Errorf("must specify `name`")
+func (t *SetImage) validateInput(items fn.KubeObjects) *fn.Result {
+	// if user does not input image name or container name, there should be only one image name to select from
+	if t.Image.Name == "" && t.Image.ImageName == "" && t.Image.ContainerName == "" {
+		if !t.isImageNameUnique(items) {
+			return fn.GeneralResult("must specify `imageName`, resources contain non-unique image names", fn.Error)
+		}
+	}
+	if t.Image.Name != "" && t.Image.ImageName != "" && t.Image.Name != t.Image.ImageName {
+		return fn.GeneralResult("must not fill `imageName` and `name` at same time, their values should be equal", fn.Error)
 	}
 	if t.Image.NewName == "" && t.Image.NewTag == "" && t.Image.Digest == "" {
-		return fmt.Errorf("must specify one of `newName`, `newTag`, or `digest`")
+		return fn.GeneralResult("must specify one of `newName`, `newTag`, or `digest`", fn.Error)
+	}
+	if len(items) == 0 {
+		return fn.GeneralResult("no input resources", fn.Info)
 	}
 	return nil
 }
 
-// updateContainerImages updates the images inside containers, return potential error
-func (t *SetImage) updateContainerImages(pod *fn.SubObject) error {
-	var containers fn.SliceSubObjects
-	containers = append(containers, pod.GetSlice("iniContainers")...)
-	containers = append(containers, pod.GetSlice("containers")...)
+// matchImage takes the resources image name and container name, return if there is a match and potential warning
+func matchImage(oldImageName string, oldContainer string, newImage *Image) (bool, *fn.Result) {
+	// name would be deprecated, means image name
+	if newImage.Name != "" {
+		newImage.ImageName = newImage.Name
+	}
+
+	// match image name, no container name
+	if newImage.ImageName != "" && newImage.ContainerName == "" {
+		if !image.IsImageMatched(oldImageName, newImage.ImageName) {
+			return false, nil
+		}
+	}
+	// match container name, no image name
+	if newImage.ImageName == "" && newImage.ContainerName != "" {
+		if oldContainer != newImage.ContainerName {
+			return false, nil
+		}
+	}
+	// match both
+	if newImage.ImageName != "" && newImage.ContainerName != "" {
+		if !image.IsImageMatched(oldImageName, newImage.ImageName) {
+			return false, nil
+		}
+		//if only image name match, container name does not match, provide a warning
+		if oldContainer != newImage.ContainerName {
+			msg := fmt.Sprintf("container name `%v` does not match `%v`, only image name matches", newImage.ContainerName, oldContainer)
+			warning := &fn.Result{
+				Message:  msg,
+				Severity: fn.Warning,
+			}
+			return true, warning
+		}
+	}
+	return true, nil
+}
+
+// isImageNameUnique checks if there is only one image name in all resources, return true if name is unique
+func (t *SetImage) isImageNameUnique(items fn.KubeObjects) bool {
+	curImageName := ""
+	for _, o := range items {
+		containers := getContainers(o)
+		for _, c := range containers {
+			oldValue := c.NestedStringOrDie("image")
+			imageName, _, _ := image.Split(oldValue)
+			if curImageName == "" {
+				curImageName = imageName
+			} else if curImageName != imageName {
+				return false
+			}
+		}
+	}
+	return true
+}
 
+// updateContainerImages updates the images inside containers, return warnings and error
+func (t *SetImage) updateContainerImages(containers fn.SliceSubObjects) (fn.Results, error) {
+	var warningList fn.Results
 	for _, o := range containers {
 		oldValue := o.NestedStringOrDie("image")
-		if !image.IsImageMatched(oldValue, t.Image.Name) {
+		oldContainer := o.NestedStringOrDie("name")
+
+		matched, warning := matchImage(oldValue, oldContainer, &t.Image)
+		if !matched {
 			continue
 		}
+
 		newName := getNewImageName(oldValue, t.Image)
 		if oldValue == newName {
 			continue
 		}
 
+		if warning != nil {
+			warningList = append(warningList, warning)
+		}
+
 		if err := o.SetNestedString(newName, "image"); err != nil {
-			return err
+			return warningList, err
 		}
 		t.resultCount += 1
 	}
+	return warningList, nil
+}
+
+// getContainers gets the containers inside kubeObject
+func getContainers(o *fn.KubeObject) fn.SliceSubObjects {
+	switch o.GetKind() {
+	case "Pod":
+		return getPodContainers(o)
+	case "Deployment", "StatefulSet", "ReplicaSet", "DaemonSet", "PodTemplate":
+		return getPodSpecContainers(o)
+	}
 	return nil
 }
 
-func (t *SetImage) setPodSpecContainers(o *fn.KubeObject) error {
+// getPodContainers gets the containers from pod
+func getPodContainers(o *fn.KubeObject) fn.SliceSubObjects {
 	spec := o.GetMap("spec")
 	if spec == nil {
 		return nil
 	}
-	template := spec.GetMap("template")
-	if template == nil {
-		return nil
-	}
-	podSpec := template.GetMap("spec")
-	err := t.updateContainerImages(podSpec)
-	if err != nil {
-		return err
-	}
-	return nil
+	return append(spec.GetSlice("iniContainers"), spec.GetSlice("containers")...)
 }
 
-func (t *SetImage) setPodContainers(o *fn.KubeObject) error {
+// getPodSpecContainers gets the containers from podSpec
+func getPodSpecContainers(o *fn.KubeObject) fn.SliceSubObjects {
 	spec := o.GetMap("spec")
 	if spec == nil {
 		return nil
 	}
-	err := t.updateContainerImages(spec)
-	if err != nil {
-		return err
+	template := spec.GetMap("template")
+	if template == nil {
+		return nil
 	}
-	return nil
+	podSpec := template.GetMap("spec")
+	if podSpec == nil {
+		return nil
+	}
+	return append(podSpec.GetSlice("iniContainers"), podSpec.GetSlice("containers")...)
 }
 
 // getNewImageName return the new name for image field

tests/set-image/containerName-mismatch/.expected/diff.patch

@@ -0,0 +1,14 @@
+diff --git a/resources.yaml b/resources.yaml
+index f674e0b..ae7d829 100644
+--- a/resources.yaml
++++ b/resources.yaml
+@@ -5,6 +5,6 @@ metadata:
+ spec:
+   containers:
+     - name: server
+-      image: nginx:1.20.2
++      image: bar:v1.0
+     - name: store
+-      image: nginx:14.1
+\ No newline at end of file
++      image: bar:v1.0

tests/set-image/containerName-mismatch/.expected/results.yaml

@@ -0,0 +1,19 @@
+apiVersion: kpt.dev/v1
+kind: FunctionResultList
+metadata:
+  name: fnresults
+exitCode: 0
+items:
+  - image: gcr.io/kpt-fn/set-image:unstable
+    exitCode: 0
+    results:
+      - message: container name `server` does not match `store`, only image name matches
+        severity: warning
+        resourceRef:
+          apiVersion: v1
+          kind: Pod
+          name: app1
+        file:
+          path: resources.yaml
+      - message: 'summary: updated a total of 2 image(s)'
+        severity: info

tests/set-image/containerName-mismatch/Kptfile

@@ -0,0 +1,12 @@
+apiVersion: kpt.dev/v1
+kind: Kptfile
+metadata:
+  name: example
+pipeline:
+  mutators:
+    - image: gcr.io/kpt-fn/set-image:unstable
+      configMap:
+        imageName: nginx
+        containerName: server
+        newName: bar
+        newTag: v1.0

tests/set-image/containerName-mismatch/resources.yaml

@@ -0,0 +1,10 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: app1
+spec:
+  containers:
+    - name: server
+      image: nginx:1.20.2
+    - name: store
+      image: nginx:14.1

tests/set-image/empty-resource/.expected/results.yaml

@@ -0,0 +1,11 @@
+apiVersion: kpt.dev/v1
+kind: FunctionResultList
+metadata:
+  name: fnresults
+exitCode: 0
+items:
+  - image: gcr.io/kpt-fn/set-image:unstable
+    exitCode: 0
+    results:
+      - message: no input resources
+        severity: info

tests/set-image/empty-resource/.krmignore

@@ -0,0 +1 @@
+.expected

tests/set-image/empty-resource/Kptfile

@@ -0,0 +1,14 @@
+apiVersion: kpt.dev/v1
+kind: Kptfile
+metadata:
+  name: example
+  annotations:
+    config.kubernetes.io/local-config: "true"
+pipeline:
+  mutators:
+    - image: gcr.io/kpt-fn/set-image:unstable
+      configMap:
+        imageName: nginx
+        containerName: server
+        newName: bar
+        newTag: v1.0

tests/set-image/no-image-name/.expected/results.yaml → tests/set-image/nonunique-image-with-no-image-name-input/.expected/results.yaml

@@ -5,10 +5,10 @@ metadata:
 exitCode: 1
 items:
   - image: gcr.io/kpt-fn/set-image:unstable
-    stderr: 'failed to evaluate function: function is terminated: must specify `name`'
+    stderr: 'failed to evaluate function: function is terminated: must specify `imageName`, resources contain non-unique image names'
     exitCode: 1
     results:
-      - message: must specify `name`
+      - message: must specify `imageName`, resources contain non-unique image names
         severity: error
-      - message: 'function is terminated: must specify `name`'
+      - message: 'function is terminated: must specify `imageName`, resources contain non-unique image names'
         severity: error

tests/set-image/nonunique-image-with-no-image-name-input/.krmignore

@@ -0,0 +1 @@
+.expected

tests/set-image/nonunique-image-with-no-image-name-input/resources.yaml

@@ -0,0 +1,10 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: app1
+spec:
+  containers:
+    - name: server
+      image: nginx:1.20.2
+    - name: store
+      image: postgres:14.1