From cabb9c73326b8416fa548b61d16efb571dad057a Mon Sep 17 00:00:00 2001
From: Jeroen Op 't Eynde <jeroen@grafana.com>
Date: Tue, 5 Jan 2021 17:07:22 +0100
Subject: [PATCH] fix(cli): ensure TLACode works with EvalScript (#464)

---
 cmd/tk/jsonnet.go       |  3 ++-
 pkg/tanka/evaluators.go | 24 ++++++++++++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/cmd/tk/jsonnet.go b/cmd/tk/jsonnet.go
index 95c7382e0..5079416d7 100644
--- a/cmd/tk/jsonnet.go
+++ b/cmd/tk/jsonnet.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"github.com/go-clix/cli"
 
@@ -24,7 +25,7 @@ func evalCmd() *cli.Command {
 			JsonnetOpts: getJsonnetOpts(),
 		}
 		if *evalPattern != "" {
-			jsonnetOpts.EvalScript = "(import '%s')." + *evalPattern
+			jsonnetOpts.EvalScript = fmt.Sprintf(tanka.PatternEvalScript, *evalPattern)
 		}
 		raw, err := tanka.Eval(args[0], jsonnetOpts)
 
diff --git a/pkg/tanka/evaluators.go b/pkg/tanka/evaluators.go
index e016a5f18..b8f145090 100644
--- a/pkg/tanka/evaluators.go
+++ b/pkg/tanka/evaluators.go
@@ -2,6 +2,7 @@ package tanka
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/pkg/errors"
 
@@ -18,7 +19,24 @@ func EvalJsonnet(path string, opts jsonnet.Opts) (raw string, err error) {
 
 	// evaluate Jsonnet
 	if opts.EvalScript != "" {
-		evalScript := fmt.Sprintf(opts.EvalScript, entrypoint)
+		var tla []string
+		for k := range opts.TLACode {
+			tla = append(tla, k)
+		}
+		evalScript := fmt.Sprintf(`
+  local main = (import '%s');
+  %s
+`, entrypoint, opts.EvalScript)
+
+		if len(tla) != 0 {
+			tlaJoin := strings.Join(tla, ", ")
+			evalScript = fmt.Sprintf(`
+function(%s)
+  local main = (import '%s')(%s);
+  %s
+`, tlaJoin, entrypoint, tlaJoin, opts.EvalScript)
+		}
+
 		raw, err = jsonnet.Evaluate(path, evalScript, opts)
 		if err != nil {
 			return "", errors.Wrap(err, "evaluating jsonnet")
@@ -33,6 +51,8 @@ func EvalJsonnet(path string, opts jsonnet.Opts) (raw string, err error) {
 	return raw, nil
 }
 
+const PatternEvalScript = "main.%s"
+
 // EnvsOnlyEvalScript finds the Environment object (without its .data object)
 const EnvsOnlyEvalScript = `
 local noDataEnv(object) =
@@ -61,5 +81,5 @@ local noDataEnv(object) =
     else {}
   );
 
-noDataEnv(import '%s')
+noDataEnv(main)
 `