Hide the Jsonnet Go implementation behind an interface (#914)

* Hide the Jsonnet Go implementation behind an interface

Tanka currently evals jsonnet using the Go native code.
However, some other implementations have come up in the past years that could be worth using (ex: https://github.com/CertainLach/jrsonnet, which is much faster)

In this PR is the first step: I create an interface where all the jsonnet eval code happens. The Go Jsonnet implementation is now hidden behind this interface.

The setting can either be passed as a global flag or as an env spec attribute to be used when exporting (`spec.exportJsonnetImplementation`)

* PR review changes
`VM` -> `Evaluator`
Add some docstrings
Add some comments

* Select the Jsonnet implementation when constructing the loader

Add the `path` as a struct member of the Go-jsonnet implementation. We
have this available when we construct the loader. This lets us then drop
the same argument from `EvaluateAnonymousSnippet()` in the interface -
as not all implementations need it.

---------

Co-authored-by: Iain Lane <iain.lane@grafana.com>

cmd/tk/flags.go

@@ -58,12 +58,14 @@ func labelSelectorFlag(fs *pflag.FlagSet) func() labels.Selector {
 func jsonnetFlags(fs *pflag.FlagSet) func() tanka.JsonnetOpts {
 	getExtCode, getTLACode := cliCodeParser(fs)
 	maxStack := fs.Int("max-stack", 0, "Jsonnet VM max stack. The default value is the value set in the go-jsonnet library. Increase this if you get: max stack frames exceeded")
+	jsonnetImplementation := fs.String("jsonnet-implementation", "go", "Only go is supported for now.")
 
 	return func() tanka.JsonnetOpts {
 		return tanka.JsonnetOpts{
-			MaxStack: *maxStack,
-			ExtCode:  getExtCode(),
-			TLACode:  getTLACode(),
+			MaxStack:              *maxStack,
+			ExtCode:               getExtCode(),
+			TLACode:               getTLACode(),
+			JsonnetImplementation: *jsonnetImplementation,
 		}
 	}
 }

pkg/jsonnet/eval.go

@@ -9,8 +9,9 @@ import (
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/goimpl"
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/types"
 	"github.com/grafana/tanka/pkg/jsonnet/jpath"
-	"github.com/grafana/tanka/pkg/jsonnet/native"
 )
 
 // Modifier allows to set optional parameters on the Jsonnet VM.
@@ -31,12 +32,13 @@ func (i *InjectedCode) Set(key, value string) {
 
 // Opts are additional properties for the Jsonnet VM
 type Opts struct {
-	MaxStack    int
-	ExtCode     InjectedCode
-	TLACode     InjectedCode
-	ImportPaths []string
-	EvalScript  string
-	CachePath   string
+	JsonnetImplementation string
+	MaxStack              int
+	ExtCode               InjectedCode
+	TLACode               InjectedCode
+	ImportPaths           []string
+	EvalScript            string
+	CachePath             string
 
 	CachePathRegexes []*regexp.Regexp
 }
@@ -75,76 +77,51 @@ func (o Opts) Clone() Opts {
 	}
 }
 
-// MakeVM returns a Jsonnet VM with some extensions of Tanka, including:
-// - extended importer
-// - extCode and tlaCode applied
-// - native functions registered
-func MakeVM(opts Opts) *jsonnet.VM {
-	vm := jsonnet.MakeVM()
-	vm.Importer(NewExtendedImporter(opts.ImportPaths))
-
-	for k, v := range opts.ExtCode {
-		vm.ExtCode(k, v)
-	}
-	for k, v := range opts.TLACode {
-		vm.TLACode(k, v)
-	}
-
-	for _, nf := range native.Funcs() {
-		vm.NativeFunction(nf)
-	}
-
-	if opts.MaxStack > 0 {
-		vm.MaxStack = opts.MaxStack
-	}
-
-	return vm
-}
-
 // EvaluateFile evaluates the Jsonnet code in the given file and returns the
 // result in JSON form. It disregards opts.ImportPaths in favor of automatically
 // resolving these according to the specified file.
-func EvaluateFile(jsonnetFile string, opts Opts) (string, error) {
-	evalFunc := func(vm *jsonnet.VM) (string, error) {
-		return vm.EvaluateFile(jsonnetFile)
+func EvaluateFile(impl types.JsonnetImplementation, jsonnetFile string, opts Opts) (string, error) {
+	evalFunc := func(evaluator types.JsonnetEvaluator) (string, error) {
+		return evaluator.EvaluateFile(jsonnetFile)
 	}
 	data, err := os.ReadFile(jsonnetFile)
 	if err != nil {
 		return "", err
 	}
-	return evaluateSnippet(evalFunc, jsonnetFile, string(data), opts)
+	return evaluateSnippet(impl, evalFunc, jsonnetFile, string(data), opts)
 }
 
 // Evaluate renders the given jsonnet into a string
 // If cache options are given, a hash from the data will be computed and
 // the resulting string will be cached for future retrieval
-func Evaluate(path, data string, opts Opts) (string, error) {
-	evalFunc := func(vm *jsonnet.VM) (string, error) {
-		return vm.EvaluateAnonymousSnippet(path, data)
+func Evaluate(path string, impl types.JsonnetImplementation, data string, opts Opts) (string, error) {
+	evalFunc := func(evaluator types.JsonnetEvaluator) (string, error) {
+		return evaluator.EvaluateAnonymousSnippet(data)
 	}
-	return evaluateSnippet(evalFunc, path, data, opts)
+	return evaluateSnippet(impl, evalFunc, path, data, opts)
 }
 
-type evalFunc func(vm *jsonnet.VM) (string, error)
+type evalFunc func(evaluator types.JsonnetEvaluator) (string, error)
 
-func evaluateSnippet(evalFunc evalFunc, path, data string, opts Opts) (string, error) {
+func evaluateSnippet(jsonnetImpl types.JsonnetImplementation, evalFunc evalFunc, path, data string, opts Opts) (string, error) {
 	var cache *FileEvalCache
 	if opts.CachePath != "" && opts.PathIsCached(path) {
 		cache = NewFileEvalCache(opts.CachePath)
 	}
 
-	// Create VM
 	jpath, _, _, err := jpath.Resolve(path, false)
 	if err != nil {
 		return "", errors.Wrap(err, "resolving import paths")
 	}
 	opts.ImportPaths = jpath
-	vm := MakeVM(opts)
+	evaluator := jsonnetImpl.MakeEvaluator(opts.ImportPaths, opts.ExtCode, opts.TLACode, opts.MaxStack)
+	// We're using the go implementation to deal with imports because we're not evaluating, we're reading the AST
+	importVM := goimpl.MakeRawVM(opts.ImportPaths, opts.ExtCode, opts.TLACode, opts.MaxStack)
 
 	var hash string
 	if cache != nil {
 		startTime := time.Now()
-		if hash, err = getSnippetHash(vm, path, data); err != nil {
+		if hash, err = getSnippetHash(importVM, path, data); err != nil {
 			return "", err
 		}
 		cacheLog := log.Debug().Str("path", path).Str("hash", hash).Dur("duration_ms", time.Since(startTime))
@@ -157,7 +134,7 @@ func evaluateSnippet(evalFunc evalFunc, path, data string, opts Opts) (string, e
 		cacheLog.Bool("cache_hit", false).Msg("computed snippet hash")
 	}
 
-	content, err := evalFunc(vm)
+	content, err := evalFunc(evaluator)
 	if err != nil {
 		return "", err
 	}

pkg/jsonnet/eval_test.go

@@ -5,10 +5,13 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/goimpl"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+var jsonnetImpl = &goimpl.JsonnetGoImplementation{}
+
 const importTreeResult = `[
    {
       "breed": "apple",
@@ -51,13 +54,13 @@ const thisFileResult = `{
 // To be consistent with the jsonnet executable,
 // when evaluating a file, `std.thisFile` should point to the given path
 func TestEvaluateFile(t *testing.T) {
-	result, err := EvaluateFile("testdata/thisFile/main.jsonnet", Opts{})
+	result, err := EvaluateFile(jsonnetImpl, "testdata/thisFile/main.jsonnet", Opts{})
 	assert.NoError(t, err)
 	assert.Equal(t, thisFileResult, result)
 }
 
 func TestEvaluateFileDoesntExist(t *testing.T) {
-	result, err := EvaluateFile("testdata/doesnt-exist/main.jsonnet", Opts{})
+	result, err := EvaluateFile(jsonnetImpl, "testdata/doesnt-exist/main.jsonnet", Opts{})
 	assert.EqualError(t, err, "open testdata/doesnt-exist/main.jsonnet: no such file or directory")
 	assert.Equal(t, "", result)
 }
@@ -69,10 +72,10 @@ func TestEvaluateFileWithCaching(t *testing.T) {
 	cachePath := filepath.Join(tmp, "cache") // Should be created during caching
 
 	// Evaluate two files
-	result, err := EvaluateFile("testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
+	result, err := EvaluateFile(jsonnetImpl, "testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, thisFileResult, result)
-	result, err = EvaluateFile("testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
+	result, err = EvaluateFile(jsonnetImpl, "testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, importTreeResult, result)
 
@@ -82,10 +85,10 @@ func TestEvaluateFileWithCaching(t *testing.T) {
 	assert.Len(t, readCache, 2)
 
 	// Evaluate two files again, same result
-	result, err = EvaluateFile("testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
+	result, err = EvaluateFile(jsonnetImpl, "testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, thisFileResult, result)
-	result, err = EvaluateFile("testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
+	result, err = EvaluateFile(jsonnetImpl, "testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, importTreeResult, result)
 
@@ -95,10 +98,10 @@ func TestEvaluateFileWithCaching(t *testing.T) {
 	}
 
 	// Evaluate two files again, modified cache is returned instead of the actual result
-	result, err = EvaluateFile("testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
+	result, err = EvaluateFile(jsonnetImpl, "testdata/thisFile/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, "BYfdlr1ZOVwiOfbd89JYTcK-eRQh05bi8ky3k1vVW5o=.json", result)
-	result, err = EvaluateFile("testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
+	result, err = EvaluateFile(jsonnetImpl, "testdata/importTree/main.jsonnet", Opts{CachePath: cachePath})
 	assert.NoError(t, err)
 	assert.Equal(t, "R_3hy-dRfOwXN-fezQ50ZF4dnrFcBcbQ9LztR_XWzJA=.json", result)
 }

pkg/jsonnet/find_importers_test.go

@@ -228,7 +228,7 @@ func TestFindImportersForFiles(t *testing.T) {
 		if filepath.Base(file) != jpath.DefaultEntrypoint {
 			continue
 		}
-		_, err := EvaluateFile(file, Opts{})
+		_, err := EvaluateFile(jsonnetImpl, file, Opts{})
 		require.NoError(t, err, "failed to eval %s", file)
 	}
 

pkg/jsonnet/implementations/goimpl/impl.go

@@ -0,0 +1,32 @@
+package goimpl
+
+import (
+	"github.com/google/go-jsonnet"
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/types"
+)
+
+type JsonnetGoVM struct {
+	vm *jsonnet.VM
+
+	path string
+}
+
+func (vm *JsonnetGoVM) EvaluateAnonymousSnippet(snippet string) (string, error) {
+	return vm.vm.EvaluateAnonymousSnippet(vm.path, snippet)
+}
+
+func (vm *JsonnetGoVM) EvaluateFile(filename string) (string, error) {
+	return vm.vm.EvaluateFile(filename)
+}
+
+type JsonnetGoImplementation struct {
+	Path string
+}
+
+func (i *JsonnetGoImplementation) MakeEvaluator(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) types.JsonnetEvaluator {
+	return &JsonnetGoVM{
+		vm: MakeRawVM(importPaths, extCode, tlaCode, maxStack),
+
+		path: i.Path,
+	}
+}

pkg/jsonnet/importer.go → pkg/jsonnet/implementations/goimpl/importer.go

@@ -1,4 +1,4 @@
-package jsonnet
+package goimpl
 
 import (
 	"path/filepath"
@@ -8,10 +8,10 @@ import (
 
 const locationInternal = "<internal>"
 
-// ExtendedImporter wraps jsonnet.FileImporter to add additional functionality:
+// extendedImporter wraps jsonnet.FileImporter to add additional functionality:
 // - `import "file.yaml"`
 // - `import "tk"`
-type ExtendedImporter struct {
+type extendedImporter struct {
 	loaders    []importLoader    // for loading jsonnet from somewhere. First one that returns non-nil is used
 	processors []importProcessor // for post-processing (e.g. yaml -> json)
 }
@@ -24,10 +24,10 @@ type importLoader func(importedFrom, importedPath string) (c *jsonnet.Contents,
 // further
 type importProcessor func(contents, foundAt string) (c *jsonnet.Contents, err error)
 
-// NewExtendedImporter returns a new instance of ExtendedImporter with the
+// newExtendedImporter returns a new instance of ExtendedImporter with the
 // correct jpaths set up
-func NewExtendedImporter(jpath []string) *ExtendedImporter {
-	return &ExtendedImporter{
+func newExtendedImporter(jpath []string) *extendedImporter {
+	return &extendedImporter{
 		loaders: []importLoader{
 			tkLoader,
 			newFileLoader(&jsonnet.FileImporter{
@@ -38,7 +38,7 @@ func NewExtendedImporter(jpath []string) *ExtendedImporter {
 }
 
 // Import implements the functionality offered by the ExtendedImporter
-func (i *ExtendedImporter) Import(importedFrom, importedPath string) (contents jsonnet.Contents, foundAt string, err error) {
+func (i *extendedImporter) Import(importedFrom, importedPath string) (contents jsonnet.Contents, foundAt string, err error) {
 	// load using loader
 	for _, loader := range i.loaders {
 		c, f, err := loader(importedFrom, importedPath)

pkg/jsonnet/tk.libsonnet.go → pkg/jsonnet/implementations/goimpl/tk.libsonnet.go

@@ -1,4 +1,4 @@
-package jsonnet
+package goimpl
 
 import jsonnet "github.com/google/go-jsonnet"
 

pkg/jsonnet/implementations/goimpl/vm.go

@@ -0,0 +1,33 @@
+package goimpl
+
+import (
+	"github.com/google/go-jsonnet"
+	"github.com/grafana/tanka/pkg/jsonnet/native"
+)
+
+// MakeRawVM returns a Jsonnet VM with some extensions of Tanka, including:
+// - extended importer
+// - extCode and tlaCode applied
+// - native functions registered
+// This is exposed because Go is used for advanced use cases, like finding transitive imports or linting.
+func MakeRawVM(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) *jsonnet.VM {
+	vm := jsonnet.MakeVM()
+	vm.Importer(newExtendedImporter(importPaths))
+
+	for k, v := range extCode {
+		vm.ExtCode(k, v)
+	}
+	for k, v := range tlaCode {
+		vm.TLACode(k, v)
+	}
+
+	for _, nf := range native.Funcs() {
+		vm.NativeFunction(nf)
+	}
+
+	if maxStack > 0 {
+		vm.MaxStack = maxStack
+	}
+
+	return vm
+}

pkg/jsonnet/implementations/types/types.go

@@ -0,0 +1,13 @@
+package types
+
+// JsonnetEvaluator represents a struct that can evaluate Jsonnet code
+// It is configured with import paths, external code and top-level arguments
+type JsonnetEvaluator interface {
+	EvaluateAnonymousSnippet(snippet string) (string, error)
+	EvaluateFile(filename string) (string, error)
+}
+
+// JsonnetImplementation is a factory for JsonnetEvaluator
+type JsonnetImplementation interface {
+	MakeEvaluator(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) JsonnetEvaluator
+}

pkg/jsonnet/imports.go

@@ -15,8 +15,8 @@ import (
 	"github.com/google/go-jsonnet/toolutils"
 	"github.com/pkg/errors"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/goimpl"
 	"github.com/grafana/tanka/pkg/jsonnet/jpath"
-	"github.com/grafana/tanka/pkg/jsonnet/native"
 )
 
 var importsRegexp = regexp.MustCompile(`import(str)?\s+['"]([^'"%()]+)['"]`)
@@ -48,12 +48,7 @@ func TransitiveImports(dir string) ([]string, error) {
 		return nil, errors.Wrap(err, "resolving JPATH")
 	}
 
-	vm := jsonnet.MakeVM()
-	vm.Importer(NewExtendedImporter(jpath))
-	for _, nf := range native.Funcs() {
-		vm.NativeFunction(nf)
-	}
-
+	vm := goimpl.MakeRawVM(jpath, nil, nil, 0)
 	node, err := jsonnet.SnippetToAST(filepath.Base(entrypoint), string(sonnet))
 	if err != nil {
 		return nil, errors.Wrap(err, "creating Jsonnet AST")

pkg/jsonnet/imports_test.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementations/goimpl"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -53,7 +54,7 @@ func BenchmarkGetSnippetHash(b *testing.B) {
 			// Create a VM. It's important to reuse the same VM
 			// While there is a caching mechanism that normally shouldn't be shared in a benchmark iteration,
 			// it's useful to evaluate its impact here, because the caching will also improve the evaluation performance afterwards.
-			vm := MakeVM(Opts{ImportPaths: []string{tempDir}})
+			vm := goimpl.MakeRawVM([]string{tempDir}, nil, nil, 0)
 			content, err := os.ReadFile(filepath.Join(tempDir, "main.jsonnet"))
 			require.NoError(b, err)