Lazy expansion of json payloads that we get from user executors

cmd/venom/root/root.go

@@ -18,7 +18,7 @@ func New() *cobra.Command {
 	return rootCmd
 }
 
-//AddCommands adds child commands to the root command rootCmd.
+// AddCommands adds child commands to the root command rootCmd.
 func addCommands(cmd *cobra.Command) {
 	cmd.AddCommand(run.Cmd)
 	cmd.AddCommand(version.Cmd)

process_testcase.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"os"
+	"reflect"
 	"regexp"
 	"strconv"
 	"strings"
@@ -574,13 +576,43 @@ func processVariableAssignments(ctx context.Context, tcName string, tcVars H, ra
 
 	for varname, assignment := range stepAssignment.Assignments {
 		Debug(ctx, "Processing %s assignment", varname)
+		lazilyLoadJSON := os.Getenv(LAZY_JSON_EXPANSION_FLAG) == FLAG_ENABLED
+		if lazilyLoadJSON {
+			if strings.Contains(assignment.From, "json") {
+				_, has := tcVars[assignment.From]
+				if !has {
+					key := getKeyForLookup(assignment.From)
+					varValue, has := tcVars[key]
+					if !has {
+						if assignment.Default == nil {
+							err := fmt.Errorf("%s reference not found and tried to create json for %s", assignment.From, key)
+							Error(ctx, "%v", err)
+							return nil, true, err
+						}
+						tcVars[key] = assignment.Default
+						result.Add(varname, assignment.Default)
+					} else {
+						Debug(ctx, "process json %v", key)
+						varjson, err := UnmarshalJSON(key, fmt.Sprintf("%v", varValue))
+						if err != nil {
+							err := fmt.Errorf("%s could not parse json for %s", assignment.From, key)
+							Error(ctx, "%v", err)
+							return nil, true, err
+						}
+						tcVars.AddAll(varjson)
+					}
+				}
+
+			}
+		}
 		varValue, has := tcVars[assignment.From]
+
 		if !has {
 			varValue, has = tcVars[tcName+"."+assignment.From]
 			if !has {
 				if assignment.Default == nil {
 					err := fmt.Errorf("%s reference not found in %s", assignment.From, strings.Join(tcVarsKeys, "\n"))
-					Info(ctx, "%v", err)
+					Error(ctx, "%v", err)
 					return nil, true, err
 				}
 				varValue = assignment.Default
@@ -613,3 +645,61 @@ func processVariableAssignments(ctx context.Context, tcName string, tcVars H, ra
 	}
 	return result, true, nil
 }
+
+// UnmarshalJSON tries to unmarshal the value as a json object and
+// creates a map with the key as the name of the `key` that we are
+// passing as argument and the value as the unmarshalled json object.
+func UnmarshalJSON(key string, value string) (map[string]interface{}, error) {
+	result := make(map[string]interface{})
+	var outJSON interface{}
+	if err := JSONUnmarshal([]byte(value), &outJSON); err == nil {
+		result[key+"json"] = outJSON
+		initialDump, err := DumpStringPreserveCase(outJSON)
+		if err != nil {
+			return nil, errors.Wrapf(err, "unable to compute result")
+		}
+		// Now we have to dump this object, but the key will change if this is a array or not
+		if reflect.ValueOf(outJSON).Kind() == reflect.Slice {
+			prefix := key + "json"
+			splitPrefix := strings.Split(prefix, ".")
+			prefix += "." + splitPrefix[len(splitPrefix)-1]
+			for ko, vo := range initialDump {
+				result[prefix+ko] = vo
+			}
+		} else {
+			outJSONDump := map[string]interface{}{}
+			for ko, vo := range initialDump {
+				if !strings.Contains(ko, "__Type__") && !strings.Contains(ko, "__Len__") {
+					outJSONDump[key+"json."+ko] = vo
+				} else {
+					outJSONDump[ko] = vo
+				}
+			}
+			for ko, vo := range outJSONDump {
+				result[ko] = vo
+			}
+		}
+	}
+	//make it compatible with the one before (ie all lowercase)
+	for k, v := range result {
+		result[strings.ToLower(k)] = v
+	}
+
+	return result, nil
+}
+
+// getKeyForLookup we need the first json key in order to process the json blob
+func getKeyForLookup(originalKey string) string {
+	parts := strings.Split(originalKey, ".")
+	keyparts := []string{}
+	for _, part := range parts {
+		if strings.HasSuffix(part, "json") {
+			keyparts = append(keyparts, strings.Replace(part, "json", "", 1))
+			break
+		} else {
+			keyparts = append(keyparts, part)
+		}
+	}
+	key := strings.Join(keyparts, ".")
+	return key
+}

process_testcase_test.go

@@ -3,6 +3,7 @@ package venom
 import (
 	"context"
 	"fmt"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -44,3 +45,56 @@ script: echo 'foo'
 	assert.Nil(t, result)
 	assert.Empty(t, result)
 }
+
+func TestProcessVariableAssignmentsWithJSON(t *testing.T) {
+	InitTestLogger(t)
+	assign := AssignStep{}
+	assign.Assignments = make(map[string]Assignment)
+	assign.Assignments["name"] = Assignment{
+		From: "payloadjson.users.users0.name",
+	}
+
+	b, _ := yaml.Marshal(assign)
+	t.Log("\n" + string(b))
+
+	tcVars := H{"payload": "{\"users\": [{\"name\":\"Tester\"}]}"}
+	os.Setenv(LAZY_JSON_EXPANSION_FLAG, FLAG_ENABLED)
+	result, is, err := processVariableAssignments(context.TODO(), "", tcVars, b)
+	assert.True(t, is)
+	assert.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.Equal(t, "Tester", result["name"])
+	t.Log(result)
+	os.Unsetenv(LAZY_JSON_EXPANSION_FLAG)
+}
+
+func TestProcessJsonBlobWithObject(t *testing.T) {
+	InitTestLogger(t)
+	items, err := UnmarshalJSON("test", "{\"key\":123,\"another\":\"one\"}")
+	assert.NoError(t, err)
+	assert.NotNil(t, items)
+	assert.Contains(t, items, "testjson.key")
+	assert.Contains(t, items, "testjson.another")
+	assert.Contains(t, items, "testjson")
+	assert.Contains(t, items, "__Type__")
+	assert.Contains(t, items, "__Len__")
+}
+
+func TestProcessJsonBlobWithArray(t *testing.T) {
+	InitTestLogger(t)
+	items, err := UnmarshalJSON("test", "{\"key\":123,\"anArray\":[\"one\",\"two\"]}")
+	assert.NoError(t, err)
+	assert.NotNil(t, items)
+	assert.Contains(t, items, "testjson.key")
+	assert.Contains(t, items, "testjson.anArray")
+	assert.Contains(t, items, "testjson")
+	assert.Equal(t, items["anArray.__Type__"], "Array")
+	assert.Equal(t, items["anArray.__Len__"], "2")
+	assert.Contains(t, items, "testjson.anArray.anArray0")
+	assert.Equal(t, items["testjson.anArray.anArray0"], "one")
+	assert.Contains(t, items, "testjson.anArray.anArray1")
+	assert.Equal(t, items["testjson.anArray.anArray1"], "two")
+	assert.Contains(t, items, "__Type__")
+	assert.Equal(t, items["__Type__"], "Map")
+	assert.Contains(t, items, "__Len__")
+}

types_executor.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"reflect"
+	"os"
 	"strings"
 
 	"github.com/gosimple/slug"
@@ -164,6 +164,27 @@ func GetExecutorResult(r interface{}) map[string]interface{} {
 	if err != nil {
 		panic(err)
 	}
+	jsonUpdates := os.Getenv(LAZY_JSON_EXPANSION_FLAG)
+	if jsonUpdates == FLAG_ENABLED {
+		for key, value := range d {
+			switch z := value.(type) {
+			case string:
+				m, e := UnmarshalJSON(key, z)
+				if e != nil {
+					panic(e)
+				}
+				if len(m) > 0 {
+					for s, i := range m {
+						//we are skipping this as whenever we use a value we will do DumpString or Dump which means will generate those values
+						if !strings.Contains(strings.ToUpper(s), "__TYPE__") && !strings.Contains(strings.ToUpper(s), "__LEN__") {
+							d[s] = i
+						}
+					}
+				}
+			}
+		}
+	}
+
 	return d
 }
 
@@ -289,7 +310,7 @@ func (v *Venom) RunUserExecutor(ctx context.Context, runner ExecutorRunner, tcIn
 	// here, we have the user executor results.
 	// and for each key in output, we try to add the json version
 	// this will allow user to use json version of output (map, etc...)
-	// because, it's not possible to to that:
+	// because, it's not possible to do that:
 	// output:
 	//   therawout: {{.result.systemout}}
 	//
@@ -300,32 +321,20 @@ func (v *Venom) RunUserExecutor(ctx context.Context, runner ExecutorRunner, tcIn
 		return nil, errors.Wrapf(err, "unable to compute result")
 	}
 
-	for k, v := range result {
-		switch z := v.(type) {
-		case string:
-			var outJSON interface{}
-			if err := JSONUnmarshal([]byte(z), &outJSON); err == nil {
-				result[k+"json"] = outJSON
-				// Now we have to dump this object, but the key will change if this is a array or not
-				if reflect.ValueOf(outJSON).Kind() == reflect.Slice {
-					prefix := k + "json"
-					splitPrefix := strings.Split(prefix, ".")
-					prefix += "." + splitPrefix[len(splitPrefix)-1]
-					outJSONDump, err := Dump(outJSON)
-					if err != nil {
-						return nil, errors.Wrapf(err, "unable to compute result")
-					}
-					for ko, vo := range outJSONDump {
-						result[prefix+ko] = vo
-					}
-				} else {
-					outJSONDump, err := DumpWithPrefix(outJSON, k+"json")
-					if err != nil {
-						return nil, errors.Wrapf(err, "unable to compute result")
-					}
-					for ko, vo := range outJSONDump {
-						result[ko] = vo
-					}
+	lazilyLoadJSON := os.Getenv(LAZY_JSON_EXPANSION_FLAG) == FLAG_ENABLED
+	if lazilyLoadJSON {
+		//ignore the json
+		Debug(ctx, "%s", "ignore the json unmarshalling till we need it")
+	} else {
+		for k, value := range result {
+			switch z := value.(type) {
+			case string:
+				items, err := UnmarshalJSON(k, z)
+				if err != nil {
+					return nil, err
+				}
+				for key, nv := range items {
+					result[key] = nv
 				}
 			}
 		}

types_executor_test.go

@@ -0,0 +1,20 @@
+package venom
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetExecutorResult(t *testing.T) {
+	payload := map[string]string{}
+	payload["key"] = "{\"name\":\"test\"}"
+	os.Setenv(LAZY_JSON_EXPANSION_FLAG, FLAG_ENABLED)
+	result := GetExecutorResult(payload)
+	assert.NotNil(t, result)
+	assert.Equal(t, "Map", result["__type__"])
+	assert.Equal(t, int64(1), result["__len__"])
+	assert.Equal(t, "test", result["keyjson.name"])
+	os.Unsetenv(LAZY_JSON_EXPANSION_FLAG)
+}

venom.go

@@ -27,6 +27,11 @@ var (
 	IsTest  = ""
 )
 
+const (
+	LAZY_JSON_EXPANSION_FLAG = "VENOM_NO_JSON_EXPANSION"
+	FLAG_ENABLED             = "ON"
+)
+
 func OSExit(exitCode int) {
 	if IsTest != "" {
 		bincover.ExitCode = exitCode