feat: Starlark package arguments will be parsed as a deep Struct when "_kurtosis_parser": "struct" is passed in the arguments JSON

core/server/api_container/server/startosis_engine/package_io/package_io.go

@@ -5,20 +5,26 @@ import (
 	"github.com/sirupsen/logrus"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkjson"
+	"go.starlark.net/starlarkstruct"
+	"reflect"
 )
 
 const (
 	decoderKey  = "decode"
 	encoderKey  = "encode"
 	indenterKey = "indent"
+
+	kurtosisParserKey    = "_kurtosis_parser"
+	kurtosisParserStruct = "struct"
+	kurtosisParserDict   = "dict"
 )
 
 var (
 	noKwargs []starlark.Tuple
 )
 
 // DeserializeArgs deserializes the Kurtosis package args, which should be serialized JSON, into a *starlark.Dict type.
-func DeserializeArgs(thread *starlark.Thread, serializedJsonArgs string) (*starlark.Dict, *startosis_errors.InterpretationError) {
+func DeserializeArgs(thread *starlark.Thread, serializedJsonArgs string) (starlark.Value, *startosis_errors.InterpretationError) {
 	if !starlarkjson.Module.Members.Has(decoderKey) {
 		return nil, startosis_errors.NewInterpretationError("Unable to deserialize package input because Starlark deserializer was not found.")
 	}
@@ -39,7 +45,28 @@ func DeserializeArgs(thread *starlark.Thread, serializedJsonArgs string) (*starl
 		// TODO: we could easily support any kind of starlark.Value here
 		return nil, startosis_errors.NewInterpretationError("Unable to parse package input '%v' into a dictionary. JSON other than dictionaries aren't support right now.", deserializedInputValue)
 	}
-	return parsedDeserializedInputValue, nil
+
+	kurtosisParserValueMaybe, found, err := parsedDeserializedInputValue.Get(starlark.String(kurtosisParserKey))
+	if err != nil {
+		return nil, startosis_errors.WrapWithInterpretationError(err, "Unable to parse package input '%v' into a valid dictionary. JSON other than dictionaries aren't support right now.", deserializedInputValue)
+	}
+	if !found || kurtosisParserValueMaybe == starlark.String(kurtosisParserDict) {
+		// we default to dict to not break back-compat
+		return parsedDeserializedInputValue, nil
+	}
+	if kurtosisParserValueMaybe != starlark.String(kurtosisParserStruct) {
+		return nil, startosis_errors.NewInterpretationError("Kurtosis parser instruction not recognized in package JSON input. The only valid values right now are '%s' and '%s', found '%s'", kurtosisParserDict, kurtosisParserStruct, kurtosisParserValueMaybe)
+	}
+
+	convertedDeserializedInputValue, interpretationErr := convertValueToStructIfPossible(parsedDeserializedInputValue)
+	if err != nil {
+		return nil, interpretationErr
+	}
+	deserializedInputValueAsStruct, ok := convertedDeserializedInputValue.(*starlarkstruct.Struct)
+	if !ok {
+		return nil, startosis_errors.NewInterpretationError("Unable to parse the package JSON input as a Starlark struct as indicated by the '%s:%s' key-value in the JSON. It has been parsed as a '%s'", kurtosisParserKey, kurtosisParserStruct, reflect.TypeOf(deserializedInputValueAsStruct))
+	}
+	return deserializedInputValueAsStruct, nil
 }
 
 func SerializeOutputObject(thread *starlark.Thread, outputObject starlark.Value) (string, *startosis_errors.InterpretationError) {
@@ -91,3 +118,46 @@ func tryIndentJson(thread *starlark.Thread, unindentedJson starlark.Value) starl
 	}
 	return indentedJson
 }
+
+// convertValueToStructIfPossible tries to convert a starlark.Value type to a starlarkstruct.Struct.
+// It no-ops for most Starlark "simple" types, like string, integer, even iterables.
+// It is expected to successfully convert starlark.Dict to starlarkstruct.Struct. Note however that there are certain
+// cases where this is not possible, like when the dictionary contain non-string keys for example. In this case, this
+// function throws an error. This error should never be hit here because what is being passed comes from serialized
+// JSON, which cannot contain non-string keys.
+func convertValueToStructIfPossible(genericValue starlark.Value) (starlark.Value, *startosis_errors.InterpretationError) {
+	switch value := genericValue.(type) {
+	case starlark.NoneType, starlark.Bool, starlark.String, starlark.Bytes, starlark.Int, starlark.Float:
+		return value, nil
+	case *starlark.List, *starlark.Set, starlark.Tuple:
+		return value, nil
+	case *starlark.Dict:
+		// Dictionaries returned by JSON deserialization should have strings as keys. We therefore convert them to struct to facilitate reading from them in Starlark
+		stringDict := starlark.StringDict{}
+		for _, key := range value.Keys() {
+			stringKey, ok := key.(starlark.String)
+			if !ok {
+				return nil, startosis_errors.NewInterpretationError("JSON input was deserialized in an unexpected manner. It seems some JSON keys were not string, which is currently not supported in Kurtosis (key: '%s', type: '%s')", key, reflect.TypeOf(key))
+			}
+			genericDictValue, found, err := value.Get(key)
+			if !found {
+				return nil, startosis_errors.NewInterpretationError("Unexpected error postprocessing JSON input. No value associated with key '%s'", key)
+			}
+			if err != nil {
+				return nil, startosis_errors.NewInterpretationError("Unexpected error postprocessing JSON input (key: '%s')", key)
+
+			}
+			postProcessedValue, interpretationError := convertValueToStructIfPossible(genericDictValue)
+			if err != nil {
+				// do not wrap the interpretation error here as it's coming from a recursive call.
+				return nil, interpretationError
+			}
+			stringDict[stringKey.GoString()] = postProcessedValue
+		}
+		return starlarkstruct.FromStringDict(starlarkstruct.Default, stringDict), nil
+	case *starlarkstruct.Struct:
+		return value, nil
+	default:
+		return nil, startosis_errors.NewInterpretationError("Unexpected type when trying to deserialize module input. Data will be passed to the module with no processing (unsupported type was: '%s').", reflect.TypeOf(genericValue))
+	}
+}

core/server/api_container/server/startosis_engine/package_io/package_io_test.go

@@ -1,8 +1,10 @@
 package package_io
 
 import (
+	"fmt"
 	"github.com/stretchr/testify/require"
 	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
 	"testing"
 )
 
@@ -29,9 +31,43 @@ func createDict(t *testing.T) *starlark.Dict {
 }
 
 func TestPackageIo_DeserializeArgs(t *testing.T) {
+	expectedResultDict := createDict(t)
+
 	result, interpretationErr := DeserializeArgs(&starlark.Thread{}, complexInputJson) //nolint:exhaustruct
 	require.Nil(t, interpretationErr)
-	equal, err := starlark.Equal(createDict(t), result)
+	equal, err := starlark.Equal(expectedResultDict, result)
+	require.Nil(t, err)
+	require.True(t, equal)
+}
+
+func TestPackageIo_DeserializeArgsToStruct(t *testing.T) {
+	complexInputJsonAsStruct := fmt.Sprintf(`{
+	%q: %q,
+	"struct": {},
+	"float": 3.4,
+	"int": 1,
+	"list": [
+		"a",
+		1,
+		{}
+	],
+	"string": "Hello World!"
+}`, kurtosisParserKey, kurtosisParserStruct)
+
+	expectedResultStringDict := starlark.StringDict{
+		kurtosisParserKey: starlark.String(kurtosisParserStruct),
+		"string":          starlark.String("Hello World!"),
+		"int":             starlark.MakeInt(1),
+		"float":           starlark.Float(3.4),
+		"struct":          starlarkstruct.FromStringDict(starlarkstruct.Default, starlark.StringDict{}),
+		"list":            starlark.NewList([]starlark.Value{starlark.String("a"), starlark.MakeInt(1), starlark.NewDict(1)}),
+	}
+	expectedResultStruct := starlarkstruct.FromStringDict(starlarkstruct.Default, expectedResultStringDict)
+
+	result, interpretationErr := DeserializeArgs(&starlark.Thread{}, complexInputJsonAsStruct) //nolint:exhaustruct
+	require.Nil(t, interpretationErr)
+
+	equal, err := starlark.Equal(expectedResultStruct, result)
 	require.Nil(t, err)
 	require.True(t, equal)
 }