From 65f2d6f676dd41387b04c3bd96927b98158b0554 Mon Sep 17 00:00:00 2001
From: Vladimir Vivien <vivienv@vmware.com>
Date: Thu, 19 Aug 2021 12:49:42 -0400
Subject: [PATCH] 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 <vivienv@vmware.com>
---
 exec/executor.go          | 33 ++++++++++++++++++++++++++++++
 exec/executor_test.go     | 42 +++++++++++++++++++++++----------------
 starlark/starlark_exec.go | 21 ++++++++++++++++++++
 3 files changed, 79 insertions(+), 17 deletions(-)

diff --git a/exec/executor.go b/exec/executor.go
index 0b105833..6cc1a270 100644
--- a/exec/executor.go
+++ b/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
+}
diff --git a/exec/executor_test.go b/exec/executor_test.go
index e537255d..d9a51154 100644
--- a/exec/executor_test.go
+++ b/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,7 +109,7 @@ 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 {
@@ -126,6 +117,23 @@ func TestExecute(t *testing.T) {
 				}
 			},
 		},
+		{
+			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 {
diff --git a/starlark/starlark_exec.go b/starlark/starlark_exec.go
index 81f06ae8..bcb259a0 100644
--- a/starlark/starlark_exec.go
+++ b/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)