Refactor starlark runtime ontop of runtimeutil

kyaml/fn/runtime/starlark/starlark.go

@@ -0,0 +1,224 @@
+// Copyright 2019 The Kubernetes Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package starlark
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+
+	"github.com/qri-io/starlib/util"
+	"go.starlark.net/starlark"
+	"sigs.k8s.io/kustomize/kyaml/comments"
+	"sigs.k8s.io/kustomize/kyaml/errors"
+	"sigs.k8s.io/kustomize/kyaml/fn/runtime/runtimeutil"
+	"sigs.k8s.io/kustomize/kyaml/kio/filters"
+	"sigs.k8s.io/kustomize/kyaml/yaml"
+)
+
+// Filter transforms a set of resources through the provided program
+type Filter struct {
+	Name string
+
+	// Program is a starlark script which will be run against the resources
+	Program string
+
+	// URL is the url of a starlark program to fetch and run
+	URL string
+
+	// Path is the path to a starlark program to read and run
+	Path string
+
+	runtimeutil.FunctionFilter
+
+	ids map[string]*yaml.RNode
+}
+
+func (sf *Filter) String() string {
+	return fmt.Sprintf(
+		"name: %v path: %v url: %v program: %v", sf.Name, sf.Path, sf.URL, sf.Program)
+}
+
+func (sf *Filter) Filter(nodes []*yaml.RNode) ([]*yaml.RNode, error) {
+	err := sf.setup()
+	if err != nil {
+		return nil, err
+	}
+	sf.FunctionFilter.Run = sf.Run
+
+	return sf.FunctionFilter.Filter(nodes)
+}
+
+func (sf *Filter) setup() error {
+	if sf.URL != "" && sf.Path != "" ||
+		sf.URL != "" && sf.Program != "" ||
+		sf.Path != "" && sf.Program != "" {
+		return errors.Errorf("Filter Path, Program and URL are mutually exclusive")
+	}
+
+	// read the program from a file
+	if sf.Path != "" {
+		b, err := ioutil.ReadFile(sf.Path)
+		if err != nil {
+			return err
+		}
+		sf.Program = string(b)
+	}
+
+	// read the program from a URL
+	if sf.URL != "" {
+		err := func() error {
+			resp, err := http.Get(sf.URL)
+			if err != nil {
+				return err
+			}
+			defer resp.Body.Close()
+			b, err := ioutil.ReadAll(resp.Body)
+			if err != nil {
+				return err
+			}
+			sf.Program = string(b)
+			return nil
+		}()
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (sf *Filter) Run(reader io.Reader, writer io.Writer) error {
+	// retain map of inputs to outputs by id so if the name is changed by the
+	// starlark program, we are able to match the same resources
+	value, err := sf.readResourceList(reader)
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	// run the starlark as program as transformation function
+	thread := &starlark.Thread{Name: sf.Name}
+
+	ctx := &Context{resourceList: value}
+	pd, err := ctx.predeclared()
+	if err != nil {
+		return errors.Wrap(err)
+	}
+	_, err = starlark.ExecFile(thread, sf.Name, sf.Program, pd)
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	return sf.writeResourceList(value, writer)
+}
+
+// inputToResourceList transforms input into a starlark.Value
+func (sf *Filter) readResourceList(reader io.Reader) (starlark.Value, error) {
+	// read and parse the inputs
+	rl := bytes.Buffer{}
+	_, err := rl.ReadFrom(reader)
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	rn, err := yaml.Parse(rl.String())
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+
+	// set the id on each node to map inputs to outputs
+	var id int
+	sf.ids = map[string]*yaml.RNode{}
+	items, err := rn.Pipe(yaml.Lookup("items"))
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	err = items.VisitElements(func(node *yaml.RNode) error {
+		id++
+		idStr := fmt.Sprintf("%v", id)
+		sf.ids[idStr] = node
+		return node.PipeE(yaml.SetAnnotation("config.k8s.io/id", idStr))
+	})
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+
+	// convert to a starlark value
+	b, err := yaml.Marshal(rn.Document()) // convert to bytes
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	var in map[string]interface{}
+	err = yaml.Unmarshal(b, &in) // convert to map[string]interface{}
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+	return util.Marshal(in) // convert to starlark value
+}
+
+// resourceListToOutput converts the output of the starlark program to the filter output
+func (sf *Filter) writeResourceList(value starlark.Value, writer io.Writer) error {
+	// convert the modified resourceList back into a slice of RNodes
+	// by first converting to a map[string]interface{}
+	out, err := util.Unmarshal(value)
+	if err != nil {
+		return errors.Wrap(err)
+	}
+	b, err := yaml.Marshal(out)
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	rl, err := yaml.Parse(string(b))
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	// preserve the comments from the input
+	items, err := rl.Pipe(yaml.Lookup("items"))
+	if err != nil {
+		return errors.Wrap(err)
+	}
+	err = items.VisitElements(func(node *yaml.RNode) error {
+		anID, err := node.Pipe(yaml.GetAnnotation("config.k8s.io/id"))
+		if err != nil {
+			return errors.Wrap(err)
+		}
+		if anID == nil {
+			return nil
+		}
+
+		var in *yaml.RNode
+		var found bool
+		if in, found = sf.ids[anID.YNode().Value]; !found {
+			return nil
+		}
+		if err := node.PipeE(yaml.ClearAnnotation("config.k8s.io/id")); err != nil {
+			return errors.Wrap(err)
+		}
+		if err := comments.CopyComments(in, node); err != nil {
+			return errors.Wrap(err)
+		}
+
+		// starlark will serialize the resources sorting the fields alphabetically,
+		// format them to have a better ordering
+		fmtFltr := filters.FormatFilter{}
+		if _, err := fmtFltr.Filter([]*yaml.RNode{node}); err != nil {
+			return errors.Wrap(err)
+		}
+		return nil
+	})
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	s, err := rl.String()
+	if err != nil {
+		return errors.Wrap(err)
+	}
+
+	_, err = writer.Write([]byte(s))
+	return err
+}