Merge pull request #2320 from pwittrock/starlark

Import environment and openAPI into starlark fn runtime

kyaml/starlark/context.go

@@ -0,0 +1,95 @@
+// Copyright 2019 The Kubernetes Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package starlark
+
+import (
+	"encoding/json"
+	"os"
+	"strings"
+
+	"github.com/qri-io/starlib/util"
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+	"sigs.k8s.io/kustomize/kyaml/errors"
+	"sigs.k8s.io/kustomize/kyaml/openapi"
+	"sigs.k8s.io/kustomize/kyaml/yaml"
+)
+
+type Context struct {
+	resourceList starlark.Value
+}
+
+func (c *Context) predeclared() (starlark.StringDict, error) {
+	e, err := env()
+	if err != nil {
+		return nil, err
+	}
+	oa, err := oa()
+	if err != nil {
+		return nil, err
+	}
+	dict := starlark.StringDict{
+		"resource_list": c.resourceList,
+		"open_api":      oa,
+		"environment":   e,
+	}
+
+	return starlark.StringDict{
+		"ctx": starlarkstruct.FromStringDict(starlarkstruct.Default, dict),
+	}, nil
+}
+
+func oa() (starlark.Value, error) {
+	return interfaceToValue(openapi.Schema())
+}
+
+func env() (starlark.Value, error) {
+	env := map[string]interface{}{}
+	for _, e := range os.Environ() {
+		pair := strings.SplitN(e, "=", 2)
+		if len(pair) < 2 {
+			continue
+		}
+		env[pair[0]] = pair[1]
+	}
+	value, err := util.Marshal(env)
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	return value, nil
+}
+
+func nodeToValue(node *yaml.RNode) (starlark.Value, error) {
+	s, err := node.String()
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	var in map[string]interface{}
+	if err := yaml.Unmarshal([]byte(s), &in); err != nil {
+		return nil, errors.Wrap(err)
+	}
+	value, err := util.Marshal(in)
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	return value, nil
+}
+
+func interfaceToValue(i interface{}) (starlark.Value, error) {
+	b, err := json.Marshal(i)
+	if err != nil {
+		return nil, err
+	}
+
+	var in map[string]interface{}
+	if err := yaml.Unmarshal(b, &in); err != nil {
+		return nil, errors.Wrap(err)
+	}
+
+	value, err := util.Marshal(in)
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	return value, nil
+}

kyaml/starlark/example_test.go

@@ -50,7 +50,7 @@ def run(items):
   for item in items:
     item["metadata"]["annotations"]["foo"] = "bar"
 
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 	}
 
@@ -139,7 +139,7 @@ def run(items, value):
   for item in items:
     item["metadata"]["annotations"]["foo"] = value
 
-run(resourceList["items"], resourceList["functionConfig"]["spec"]["value"])
+run(ctx.resource_list["items"], ctx.resource_list["functionConfig"]["spec"]["value"])
 `,
 		FunctionConfig: fc,
 	}
@@ -233,7 +233,7 @@ def run(items):
   for item in items:
     item["metadata"]["annotations"]["foo"] = "bar"
 
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `), 0600)
 	if err != nil {
 		log.Fatal(err)

kyaml/starlark/starlark.go

@@ -84,8 +84,15 @@ func (sf *Filter) Filter(input []*yaml.RNode) ([]*yaml.RNode, error) {
 
 	// run the starlark as program as transformation function
 	thread := &starlark.Thread{Name: sf.Name}
-	predeclared := starlark.StringDict{"resourceList": value}
+
-	_, err = starlark.ExecFile(thread, sf.Name, sf.Program, predeclared)
+	ctx := &Context{
+		resourceList: value,
+	}
+	pd, err := ctx.predeclared()
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	_, err = starlark.ExecFile(thread, sf.Name, sf.Program, pd)
 	if err != nil {
 		return nil, errors.Wrap(err)
 	}
@@ -151,18 +158,7 @@ func (sf *Filter) inputToResourceList(
 
 	// convert the ResourceList into a starlark dictionary by
 	// first converting it into a map[string]interface{}
-	s, err := resourceList.String()
+	value, err := nodeToValue(resourceList)
-	if err != nil {
-		return nil, nil, errors.Wrap(err)
-	}
-	var in map[string]interface{}
-	if err := yaml.Unmarshal([]byte(s), &in); err != nil {
-		return nil, nil, errors.Wrap(err)
-	}
-	value, err := util.Marshal(in)
-	if err != nil {
-		return nil, nil, errors.Wrap(err)
-	}
 	return value, ids, err
 }
 

kyaml/starlark/starlark_test.go

@@ -24,6 +24,7 @@ func TestFilter_Filter(t *testing.T) {
 		script                 string
 		expected               string
 		expectedFunctionConfig string
+		env                    map[string]string
 	}{
 		{
 			name: "add_annotation",
@@ -46,7 +47,7 @@ def run(r):
   for resource in r:
     resource["metadata"]["annotations"]["foo"] = "bar"
 
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 			expected: `
 apiVersion: apps/v1
@@ -62,6 +63,83 @@ spec:
       - name: nginx
         # head comment
         image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+		},
+		{
+			name: "add_annotation_from_env",
+			input: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+			script: `
+def run(r):
+  for resource in r:
+    resource["metadata"]["annotations"]["foo"] = ctx.environment["ANNOTATION"]
+
+run(ctx.resource_list["items"])
+`,
+			env: map[string]string{"ANNOTATION": "annotation-value"},
+			expected: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+  annotations:
+    foo: annotation-value
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+		},
+		{
+			name: "add_annotation_from_open_api",
+			input: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+			script: `
+def run(r):
+  for resource in r:
+    resource["metadata"]["annotations"]["foo"] = ctx.open_api["definitions"]["io.k8s.api.apps.v1.Deployment"]["description"]
+
+run(ctx.resource_list["items"])
+`,
+			expected: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+  annotations:
+    foo: Deployment enables declarative updates for Pods and ReplicaSets.
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
 `,
 		},
 		{
@@ -87,7 +165,7 @@ def run(r):
   for resource in r:
     resource["metadata"]["annotations"]["foo"] = "bar"
 
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 			expected: `
 apiVersion: apps/v1
@@ -103,6 +181,45 @@ spec:
       - name: nginx
         # head comment
         image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+		},
+		{
+			name: "delete_annotation",
+			input: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+  annotations:
+    foo: baz
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
+`,
+			script: `
+# set the foo annotation on each resource
+def run(r):
+  for resource in r:
+    resource["metadata"]["annotations"].pop("foo")
+
+run(ctx.resource_list["items"])
+`,
+			expected: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        # head comment
+        image: nginx:1.8.1 # {"$ref": "#/definitions/io.k8s.cli.substitutions.image"}
 `,
 		},
 		{
@@ -140,7 +257,7 @@ def run(r):
   for resource in r:
     resource["metadata"]["annotations"]["foo"] = "bar"
 
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 			expected: `
 apiVersion: apps/v1
@@ -196,7 +313,7 @@ def run(r):
   },
 }
   r.append(d)
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 			expected: `
 apiVersion: apps/v1
@@ -233,7 +350,7 @@ metadata:
 			script: `
 def run(r):
   r.pop()
-run(resourceList["items"])
+run(ctx.resource_list["items"])
 `,
 			expected: `
 apiVersion: apps/v1
@@ -268,8 +385,8 @@ def run(r, an):
   for resource in r:
     resource["metadata"]["annotations"]["foo"] = an
 
-an = resourceList["functionConfig"]["spec"]["value"]
+an = ctx.resource_list["functionConfig"]["spec"]["value"]
-run(resourceList["items"], an)
+run(ctx.resource_list["items"], an)
 `,
 			expected: `
 apiVersion: apps/v1
@@ -319,9 +436,9 @@ def run(r, an):
   for resource in r:
     resource["metadata"]["annotations"]["foo"] = an
 
-an = resourceList["functionConfig"]["spec"]["value"]
+an = ctx.resource_list["functionConfig"]["spec"]["value"]
-run(resourceList["items"], an)
+run(ctx.resource_list["items"], an)
-resourceList["functionConfig"]["spec"]["value"] = "updated"
+ctx.resource_list["functionConfig"]["spec"]["value"] = "updated"
 `,
 			expected: `
 apiVersion: apps/v1
@@ -348,6 +465,10 @@ spec:
 	for i := range tests {
 		test := tests[i]
 		t.Run(test.name, func(t *testing.T) {
+			os.Clearenv()
+			for k, v := range test.env {
+				os.Setenv(k, v)
+			}
 			f := &Filter{Name: test.name, Program: test.script}
 
 			if test.functionConfig != "" {