Hide the Jsonnet Go implementation behind an interface

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)