Merge pull request #3270 from pwittrock/main

Improved fn framework support for patching

kyaml/fn/framework/containers.go

@@ -13,12 +13,18 @@ import (
 	"sigs.k8s.io/kustomize/kyaml/yaml/merge2"
 )
 
-// PatchTemplateContainers executes t as a template and patches each container in each resource
+// PatchContainersWithString executes t as a template and patches each container in each resource
 // with the result.
-func PatchTemplateContainers(resources []*yaml.RNode, t string, input interface{}, containers ...string) error {
+func PatchContainersWithString(resources []*yaml.RNode, t string, input interface{}, containers ...string) error {
 	resourcePatch := template.Must(template.New("containers").Parse(t))
+	return PatchContainersWithTemplate(resources, resourcePatch, input, containers...)
+}
+
+// PatchContainersWithTemplate executes t and patches each container in each resource
+// with the result.
+func PatchContainersWithTemplate(resources []*yaml.RNode, t *template.Template, input interface{}, containers ...string) error {
 	var b bytes.Buffer
-	if err := resourcePatch.Execute(&b, input); err != nil {
+	if err := t.Execute(&b, input); err != nil {
 		return errors.Wrap(err)
 	}
 	patch, err := yaml.Parse(b.String())

kyaml/fn/framework/example_test.go

@@ -893,8 +893,8 @@ metadata:
 	//     config.kubernetes.io/index: '1'
 }
 
-// ExamplePatchTemplateContainers_names patches all containers.
-func ExamplePatchTemplateContainers() {
+// ExamplePatchContainersWithString patches all containers.
+func ExamplePatchContainersWithString() {
 	resources, err := kio.ParseAll(`
 apiVersion: apps/v1
 kind: Deployment
@@ -943,7 +943,7 @@ spec:
 	}
 
 	input := struct{ Value string }{Value: "new-value"}
-	err = framework.PatchTemplateContainers(resources, `
+	err = framework.PatchContainersWithString(resources, `
 env:
   KEY: {{ .Value }}
 `, input)
@@ -1006,9 +1006,9 @@ env:
 	//  <nil>
 }
 
-// ExamplePatchTemplateContainers_names patches containers matching
+// PatchTemplateContainersWithString patches containers matching
 // a specific name.
-func ExamplePatchTemplateContainers_names() {
+func ExamplePatchContainersWithString_names() {
 	resources, err := kio.ParseAll(`
 apiVersion: apps/v1
 kind: Deployment
@@ -1057,7 +1057,7 @@ spec:
 	}
 
 	input := struct{ Value string }{Value: "new-value"}
-	err = framework.PatchTemplateContainers(resources, `
+	err = framework.PatchContainersWithString(resources, `
 env:
   KEY: {{ .Value }}
 `, input, "foo")

kyaml/fn/framework/framework.go

@@ -286,6 +286,16 @@ type TemplateCommand struct {
 	// PatchTemplates is a list of templates to render into Patches and apply.
 	PatchTemplates []PatchTemplate
 
+	// PatchTemplateFn returns a list of templates to render into Patches and apply.
+	// PatchTemplateFn is called after the ResourceList has been parsed.
+	PatchTemplatesFn func(*ResourceList) ([]PatchTemplate, error)
+
+	// PatchContainerTemplates applies patches to matching container fields
+	PatchContainerTemplates []ContainerPatchTemplate
+
+	// PatchContainerTemplates returns a list of PatchContainerTemplates
+	PatchContainerTemplatesFn func(*ResourceList) ([]ContainerPatchTemplate, error)
+
 	// TemplateFiles list of templates to read from disk which are appended
 	// to Templates.
 	TemplatesFiles []string
@@ -301,6 +311,13 @@ type TemplateCommand struct {
 	PostProcess func(*ResourceList) error
 }
 
+// ContainerPatchTemplate defines a patch to be applied to containers
+type ContainerPatchTemplate struct {
+	PatchTemplate
+
+	ContainerNames []string
+}
+
 func (tc TemplateCommand) doTemplate(t *template.Template, rl *ResourceList) error {
 	// invoke the template
 	var b bytes.Buffer
@@ -372,12 +389,39 @@ func (tc TemplateCommand) GetCommand() *cobra.Command {
 			}
 		}
 
+		if tc.PatchTemplatesFn != nil {
+			pt, err := tc.PatchTemplatesFn(&rl)
+			if err != nil {
+				return err
+			}
+			tc.PatchTemplates = append(tc.PatchTemplates, pt...)
+		}
+
 		for i := range tc.PatchTemplates {
 			if err := tc.PatchTemplates[i].Apply(&rl); err != nil {
 				return err
 			}
 		}
 
+		if tc.PatchContainerTemplatesFn != nil {
+			ct, err := tc.PatchContainerTemplatesFn(&rl)
+			if err != nil {
+				return err
+			}
+			tc.PatchContainerTemplates = append(tc.PatchContainerTemplates, ct...)
+		}
+		for i := range tc.PatchContainerTemplates {
+			ct := tc.PatchContainerTemplates[i]
+			matches, err := ct.Selector.GetMatches(&rl)
+			if err != nil {
+				return err
+			}
+			err = PatchContainersWithTemplate(matches, ct.Template, rl.FunctionConfig, ct.ContainerNames...)
+			if err != nil {
+				return err
+			}
+		}
+
 		var err error
 		if tc.MergeResources {
 			rl.Items, err = filters.MergeFilter{}.Filter(rl.Items)

kyaml/fn/framework/framework_test.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"text/template"
 
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
@@ -162,3 +163,204 @@ metadata:
     a: 'b'
 `), strings.TrimSpace(out.String()))
 }
+
+func TestCommand_PatchTemplateFn(t *testing.T) {
+	// TODO: make this test pass on windows -- currently failure seems spurious
+	testutil.SkipWindows(t)
+
+	type api = struct {
+		Spec struct {
+			A string `json:"a" yaml:"a"`
+		} `json:"spec" yaml:"spec"`
+	}
+	var config api
+
+	cmd := framework.TemplateCommand{
+		API: &config,
+		PatchTemplatesFn: func(_ *framework.ResourceList) ([]framework.PatchTemplate, error) {
+			return []framework.PatchTemplate{{
+				Selector: &framework.Selector{Names: []string{config.Spec.A}},
+				Template: template.Must(template.New("test").Parse(`
+metadata:
+  annotations:
+    baz: buz
+`)),
+			}}, nil
+		},
+	}.GetCommand()
+
+	cmd.SetIn(bytes.NewBufferString(`
+apiVersion: config.kubernetes.io/v1alpha1
+kind: ResourceList
+items:
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar1
+    namespace: default
+    annotations:
+      foo: bar1
+  spec:
+    replicas: 1
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar2
+    namespace: default
+    annotations:
+      foo: bar2
+  spec:
+    replicas: 1
+functionConfig:
+  apiVersion: example.com/v1alpha1
+  kind: Example
+  spec:
+    a: "bar1"
+`))
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+
+	require.NoError(t, cmd.Execute())
+
+	require.Equal(t, strings.TrimSpace(`
+apiVersion: config.kubernetes.io/v1alpha1
+kind: ResourceList
+items:
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar1
+    namespace: default
+    annotations:
+      foo: bar1
+      baz: buz
+      config.kubernetes.io/index: '0'
+  spec:
+    replicas: 1
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar2
+    namespace: default
+    annotations:
+      foo: bar2
+  spec:
+    replicas: 1
+functionConfig:
+  apiVersion: example.com/v1alpha1
+  kind: Example
+  spec:
+    a: "bar1"
+`), strings.TrimSpace(out.String()))
+}
+
+func TestCommand_PatchContainerTemplatesFn(t *testing.T) {
+	// TODO: make this test pass on windows -- currently failure seems spurious
+	testutil.SkipWindows(t)
+
+	type api = struct {
+		Spec struct {
+			A string `json:"a" yaml:"a"`
+		} `json:"spec" yaml:"spec"`
+	}
+	var config api
+
+	cmd := framework.TemplateCommand{
+		API: &config,
+		PatchContainerTemplatesFn: func(_ *framework.ResourceList) ([]framework.ContainerPatchTemplate, error) {
+			return []framework.ContainerPatchTemplate{{
+				PatchTemplate: framework.PatchTemplate{
+					Selector: &framework.Selector{Names: []string{config.Spec.A}},
+					Template: template.Must(template.New("test").Parse(`
+env:
+  key: Foo
+  value: Bar
+`))},
+			}}, nil
+		},
+	}.GetCommand()
+
+	cmd.SetIn(bytes.NewBufferString(`
+apiVersion: config.kubernetes.io/v1alpha1
+kind: ResourceList
+items:
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar1
+    namespace: default
+    annotations:
+      foo: bar1
+  spec:
+    template:
+      spec:
+        containers:
+        - name: foo
+        - name: bar
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar2
+    namespace: default
+    annotations:
+      foo: bar2
+  spec:
+    template:
+      spec:
+        containers:
+        - name: foo
+        - name: bar
+functionConfig:
+  apiVersion: example.com/v1alpha1
+  kind: Example
+  spec:
+    a: "bar1"
+`))
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+
+	require.NoError(t, cmd.Execute())
+
+	require.Equal(t, strings.TrimSpace(`
+apiVersion: config.kubernetes.io/v1alpha1
+kind: ResourceList
+items:
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar1
+    namespace: default
+    annotations:
+      foo: bar1
+  spec:
+    template:
+      spec:
+        containers:
+        - name: foo
+          env:
+            key: Foo
+            value: Bar
+        - name: bar
+          env:
+            key: Foo
+            value: Bar
+- apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: bar2
+    namespace: default
+    annotations:
+      foo: bar2
+  spec:
+    template:
+      spec:
+        containers:
+        - name: foo
+        - name: bar
+functionConfig:
+  apiVersion: example.com/v1alpha1
+  kind: Example
+  spec:
+    a: "bar1"
+`), strings.TrimSpace(out.String()))
+}