feat(jsonnet): add support for importing all files in a directory

pkg/jsonnet/implementations/goimpl/vm.go

@@ -25,6 +25,10 @@ func MakeRawVM(importPaths []string, extCode map[string]string, tlaCode map[stri
 		vm.NativeFunction(nf)
 	}
 
+	for _, nvf := range native.VMFuncs(vm) {
+		vm.NativeFunction(nvf)
+	}
+
 	if maxStack > 0 {
 		vm.MaxStack = maxStack
 	}

pkg/jsonnet/native/funcs.go

@@ -6,14 +6,18 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 	"regexp"
+	"slices"
 	"strings"
 
 	jsonnet "github.com/google/go-jsonnet"
 	"github.com/google/go-jsonnet/ast"
 	"github.com/grafana/tanka/pkg/helm"
 	"github.com/grafana/tanka/pkg/kustomize"
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
 	yaml "gopkg.in/yaml.v3"
 )
 
@@ -42,6 +46,14 @@ func Funcs() []*jsonnet.NativeFunction {
 	}
 }
 
+// VMFuncs returns a slice of functions similar to Funcs but are passed the jsonnet VM
+// for in-line evaluation
+func VMFuncs(vm *jsonnet.VM) []*jsonnet.NativeFunction {
+	return []*jsonnet.NativeFunction{
+		importFiles(vm),
+	}
+}
+
 // parseJSON wraps `json.Unmarshal` to convert a json string into a dict
 func parseJSON() *jsonnet.NativeFunction {
 	return &jsonnet.NativeFunction{
@@ -178,3 +190,87 @@ func regexSubst() *jsonnet.NativeFunction {
 		},
 	}
 }
+
+type importFilesOpts struct {
+	CalledFrom string   `json:"calledFrom"`
+	Exclude    []string `json:"exclude"`
+	Extension  string   `json:"extension"`
+}
+
+func parseImportOpts(data interface{}) (*importFilesOpts, error) {
+	c, err := json.Marshal(data)
+	if err != nil {
+		return nil, err
+	}
+
+	// default extension to `.libsonnet`
+	opts := importFilesOpts{
+		Extension: ".libsonnet",
+	}
+	if err := json.Unmarshal(c, &opts); err != nil {
+		return nil, err
+	}
+	if opts.CalledFrom == "" {
+		return nil, fmt.Errorf("importFiles: `opts.calledFrom` is unset or empty\nTanka needs this to find your directory.")
+	}
+	// Make sure calledFrom is inside the current working directory
+	cwd, err := os.Getwd()
+	if err != nil {
+		return nil, fmt.Errorf("importFiles: failed to get current working directory: %s", err)
+	}
+	calledFromAbs, err := filepath.Abs(opts.CalledFrom)
+	if err != nil {
+		return nil, fmt.Errorf("importFiles: failed to get absolute path to `opts.calledFrom`: %s", err)
+	}
+	if !strings.HasPrefix(calledFromAbs, cwd) {
+		return nil, fmt.Errorf("importFiles: `opts.calledFrom` must be a subdirectory of the current working directory: %s", cwd)
+	}
+	return &opts, nil
+}
+
+// importFiles imports and evaluates all matching jsonnet files in the given relative directory
+func importFiles(vm *jsonnet.VM) *jsonnet.NativeFunction {
+	return &jsonnet.NativeFunction{
+		Name:   "importFiles",
+		Params: ast.Identifiers{"directory", "opts"},
+		Func: func(data []interface{}) (interface{}, error) {
+			dir, ok := data[0].(string)
+			if !ok {
+				return nil, fmt.Errorf("first argument 'directory' must be of 'string' type, got '%T' instead", data[0])
+			}
+			opts, err := parseImportOpts(data[1])
+			if err != nil {
+				return nil, err
+			}
+			dirPath := filepath.Join(filepath.Dir(opts.CalledFrom), dir)
+			imports := make(map[string]interface{})
+			err = filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
+				if err != nil {
+					return err
+				}
+				if info.IsDir() || !strings.HasSuffix(info.Name(), opts.Extension) {
+					return nil
+				}
+				if slices.Contains(opts.Exclude, info.Name()) {
+					return nil
+				}
+				log.Debug().Msgf("importFiles: parsing file %s", info.Name())
+				resultStr, err := vm.EvaluateFile(path)
+				if err != nil {
+					return fmt.Errorf("importFiles: failed to evaluate %s: %s", path, err)
+				}
+				var result interface{}
+				err = json.Unmarshal([]byte(resultStr), &result)
+				if err != nil {
+					return err
+				}
+				imports[info.Name()] = result
+				return nil
+			})
+			if err != nil {
+				return nil, err
+			}
+			return imports, nil
+		},
+	}
+}

pkg/jsonnet/native/funcs_test.go

@@ -3,8 +3,13 @@ package native
 import (
 	"encoding/json"
 	"fmt"
+	"maps"
+	"os"
+	"path/filepath"
+	"slices"
 	"testing"
 
+	jsonnet "github.com/google/go-jsonnet"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -21,6 +26,20 @@ func callNative(name string, data []interface{}) (res interface{}, err error, ca
 	return nil, nil, fmt.Errorf("could not find native function %s", name)
 }
 
+// callVMNative calls a native function used by jsonnet VM that requires access to the VM resource
+func callVMNative(name string, data []interface{}) (res interface{}, err error, callerr error) {
+	vm := jsonnet.MakeVM()
+	for _, fun := range VMFuncs(vm) {
+		if fun.Name == name {
+			// Call the function
+			ret, err := fun.Func(data)
+			return ret, err, nil
+		}
+	}
+
+	return nil, nil, fmt.Errorf("could not find VM native function %s", name)
+}
+
 func TestSha256(t *testing.T) {
 	ret, err, callerr := callNative("sha256", []interface{}{"foo"})
 
@@ -208,3 +227,46 @@ func TestRegexSubstInvalid(t *testing.T) {
 	assert.Empty(t, ret)
 	assert.NotEmpty(t, err)
 }
+
+func TestImportFiles(t *testing.T) {
+	cwd, err := os.Getwd()
+	assert.NoError(t, err)
+	tempDir := t.TempDir()
+	defer func() {
+		if err := os.Chdir(cwd); err != nil {
+			panic(err)
+		}
+	}()
+	err = os.Chdir(tempDir)
+	assert.NoError(t, err)
+	importDirName := "imports"
+	importDir := filepath.Join(tempDir, importDirName)
+	err = os.Mkdir(importDir, 0750)
+	assert.NoError(t, err)
+	importMap := map[string]string{
+		"test1.libsonnet": "{ test: 1 }",
+		"test2.libsonnet": "{ test: 2 }",
+	}
+	excludeMap := map[string]string{
+		"skip1.libsonnet": `{ test: error "should not be included" }`,
+		"skip2.libsonnet": `{ test: error "should not be included" }`,
+	}
+	expectedJson := `{"test1.libsonnet":{"test":1},"test2.libsonnet":{"test":2}}`
+	allMap := make(map[string]string)
+	maps.Copy(allMap, importMap)
+	maps.Copy(allMap, excludeMap)
+	for fName, content := range allMap {
+		fPath := filepath.Join(importDir, fName)
+		err = os.WriteFile(fPath, []byte(content), 0644)
+		assert.NoError(t, err)
+	}
+	opts := make(map[string]interface{})
+	opts["calledFrom"] = filepath.Join(tempDir, "main.jsonnet")
+	opts["exclude"] = slices.Collect(maps.Keys(excludeMap))
+	ret, err, callerr := callVMNative("importFiles", []interface{}{importDirName, opts})
+	assert.NoError(t, err)
+	assert.Nil(t, callerr)
+	retJson, err := json.Marshal(ret)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedJson, string(retJson))
+}