Support for script module preloading

This patch introduces the support for Starlark code module preload.
This feature allows scripts to reuse code in other scripts loaded
ahead of time allowing for libraries to be used by the main script.

Signed-off-by: Vladimir Vivien <[email protected]>

exec/executor.go

@@ -4,6 +4,7 @@
 package exec
 
 import (
+	"fmt"
 	"io"
 	"os"
 
@@ -36,3 +37,35 @@ func Execute(name string, source io.Reader, args ArgMap) error {
 func ExecuteFile(file *os.File, args ArgMap) error {
 	return Execute(file.Name(), file, args)
 }
+
+type StarlarkModule struct {
+	Name   string
+	Source io.Reader
+}
+
+func ExecuteWithModules(name string, source io.Reader, args ArgMap, modules ...StarlarkModule) error {
+	star := starlark.New()
+
+	if args != nil {
+		starStruct, err := starlark.NewGoValue(args).ToStarlarkStruct("args")
+		if err != nil {
+			return err
+		}
+
+		star.AddPredeclared("args", starStruct)
+	}
+
+	// load modules
+	for _, mod := range modules {
+		if err := star.Preload(mod.Name, mod.Source); err != nil {
+			return fmt.Errorf("module load: %w", err)
+		}
+	}
+
+	err := star.Exec(name, source)
+	if err != nil {
+		return fmt.Errorf("exec failed: %w", err)
+	}
+
+	return nil
+}

exec/executor_test.go

@@ -17,7 +17,7 @@ var (
 	support *testcrashd.TestSupport
 )
 
-func TestMain(m *testing.M) {
+func setupTestSupport() {
 	test, err := testcrashd.Init()
 	if err != nil {
 		logrus.Fatal(err)
@@ -36,17 +36,17 @@ func TestMain(m *testing.M) {
 	if err != nil {
 		logrus.Fatal(err)
 	}
+}
 
-	result := m.Run()
-
+func teardownTestSupport() {
 	if err := support.TearDown(); err != nil {
 		logrus.Fatal(err)
 	}
-
-	os.Exit(result)
 }
 
-func TestKindScript(t *testing.T) {
+func TestExampleScripts(t *testing.T) {
+	setupTestSupport()
+
 	tests := []struct {
 		name       string
 		scriptPath string
@@ -80,16 +80,6 @@ func TestKindScript(t *testing.T) {
 				"ssh_port":    support.PortValue(),
 			},
 		},
-		//{
-		//	name:       "kube-nodes provider",
-		//	scriptPath: "../examples/kube-nodes-provider.crsh",
-		//	args: ArgMap{
-		//		"kubecfg":  getTestKubeConf(),
-		//		"ssh_port": testSSHPort,
-		//		"username": getUsername(),
-		//		"key_path": getPrivateKey(),
-		//	},
-		//},
 		{
 			name:       "kind-capi-bootstrap",
 			scriptPath: "../examples/kind-capi-bootstrap.crsh",
@@ -109,6 +99,7 @@ func TestKindScript(t *testing.T) {
 			}
 		})
 	}
+	teardownTestSupport()
 }
 
 func TestExecute(t *testing.T) {
@@ -118,14 +109,31 @@ func TestExecute(t *testing.T) {
 		exec   func(t *testing.T, script string)
 	}{
 		{
-			name:   "run_local",
+			name:   "execute single script",
 			script: `result = run_local("echo 'Hello World!'")`,
 			exec: func(t *testing.T, script string) {
 				if err := Execute("run_local", strings.NewReader(script), ArgMap{}); err != nil {
 					t.Fatal(err)
 				}
 			},
 		},
+		{
+			name:   "execute with modules",
+			script: `result = multiply(2, 3)`,
+			exec: func(t *testing.T, script string) {
+				mod := `
+def multiply(x, y):
+    log (msg="{} * {} = {}".format(x,y,x*y))
+`
+				if err := ExecuteWithModules(
+					"multiply",
+					strings.NewReader(script),
+					ArgMap{},
+					StarlarkModule{Name: "lib", Source: strings.NewReader(mod)}); err != nil {
+					t.Fatal(err)
+				}
+			},
+		},
 	}
 
 	for _, test := range tests {

starlark/starlark_exec.go

@@ -34,6 +34,27 @@ func (e *Executor) AddPredeclared(name string, value starlark.Value) {
 	}
 }
 
+// Preload loads parse and load code modules to be used in other scripts.
+// This call should take place prior to calling the Exec call for main script execution.
+// The result of the loaded module are added to the global predeclared
+// values used in Exec call.
+func (e *Executor) Preload(name string, source io.Reader) error {
+	result, err := starlark.ExecFile(e.thread, name, source, e.predecs)
+	if err != nil {
+		if evalErr, ok := err.(*starlark.EvalError); ok {
+			return fmt.Errorf(evalErr.Backtrace())
+		}
+		return err
+	}
+
+	// add result to predeclared
+	for k, v := range result {
+		e.predecs[k] = v
+	}
+
+	return nil
+}
+
 func (e *Executor) Exec(name string, source io.Reader) error {
 	if err := setupLocalDefaults(e.thread); err != nil {
 		return fmt.Errorf("failed to setup defaults: %s", err)