Add JMESPath function for dynamic object/array lookup (#7136)

* Fix JMESPath functions error message

JMESPath functions `parse_yaml`, `items` and `object_from_lists` use
wrong format string arguments for an error message and count the
argument from 0 instead of 1.

Fix the format string args and add 1 to the argument index.

Also improve the error message itself.

Signed-off-by: Andreas Brehmer <andreas.brehmer@sap.com>

* Add JMESPath function `lookup`

`lookup` allows for dynamic lookups of objects and arrays, i.e. where
the key/index to look up is determined during the JMESPath query and
thus cannot be injected upfront.

Signed-off-by: Andreas Brehmer <andreas.brehmer@sap.com>

---------

Signed-off-by: Andreas Brehmer <andreas.brehmer@sap.com>
Co-authored-by: shuting <shuting@nirmata.com>
Co-authored-by: Charles-Edouard Brétéché <charles.edouard@nirmata.com>

pkg/engine/jmespath/error.go

@@ -6,7 +6,7 @@ import (
 
 const (
 	errorPrefix              = "JMESPath function '%s': "
-	invalidArgumentTypeError = errorPrefix + "%d argument is expected of %s type"
+	invalidArgumentTypeError = errorPrefix + "argument #%d is not of type %s"
 	genericError             = errorPrefix + "%s"
 	argOutOfBoundsError      = errorPrefix + "%d argument is out of bounds (%d)"
 	zeroDivisionError        = errorPrefix + "Zero divisor passed"

pkg/engine/jmespath/functions.go

@@ -65,6 +65,7 @@ var (
 	semverCompare          = "semver_compare"
 	parseJson              = "parse_json"
 	parseYAML              = "parse_yaml"
+	lookup                 = "lookup"
 	items                  = "items"
 	objectFromLists        = "object_from_lists"
 	random                 = "random"
@@ -404,6 +405,17 @@ func GetFunctions(configuration config.Configuration) []FunctionEntry {
 		},
 		ReturnType: []jpType{jpAny},
 		Note:       "decodes a valid YAML encoded string to the appropriate type provided it can be represented as JSON",
+	}, {
+		FunctionEntry: gojmespath.FunctionEntry{
+			Name: lookup,
+			Arguments: []argSpec{
+				{Types: []jpType{jpObject, jpArray}},
+				{Types: []jpType{jpString, jpNumber}},
+			},
+			Handler: jpLookup,
+		},
+		ReturnType: []jpType{jpAny},
+		Note:       "returns the value corresponding to the given key/index in the given object/array",
 	}, {
 		FunctionEntry: gojmespath.FunctionEntry{
 			Name: items,
@@ -957,14 +969,43 @@ func jpParseYAML(arguments []interface{}) (interface{}, error) {
 	return output, err
 }
 
+func jpLookup(arguments []interface{}) (interface{}, error) {
+	switch input := arguments[0].(type) {
+	case map[string]interface{}:
+		key, ok := arguments[1].(string)
+		if !ok {
+			return nil, formatError(invalidArgumentTypeError, lookup, 2, "String")
+		}
+		return input[key], nil
+	case []interface{}:
+		key, ok := arguments[1].(float64)
+		if !ok {
+			return nil, formatError(invalidArgumentTypeError, lookup, 2, "Number")
+		}
+		keyInt, err := intNumber(key)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"JMESPath function '%s': argument #2: %s",
+				lookup, err.Error(),
+			)
+		}
+		if keyInt < 0 || keyInt > len(input)-1 {
+			return nil, nil
+		}
+		return input[keyInt], nil
+	default:
+		return nil, formatError(invalidArgumentTypeError, lookup, 1, "Object or Array")
+	}
+}
+
 func jpItems(arguments []interface{}) (interface{}, error) {
 	keyName, ok := arguments[1].(string)
 	if !ok {
-		return nil, formatError(invalidArgumentTypeError, items, arguments, 1, "String")
+		return nil, formatError(invalidArgumentTypeError, items, 2, "String")
 	}
 	valName, ok := arguments[2].(string)
 	if !ok {
-		return nil, formatError(invalidArgumentTypeError, items, arguments, 2, "String")
+		return nil, formatError(invalidArgumentTypeError, items, 3, "String")
 	}
 	switch input := arguments[0].(type) {
 	case map[string]interface{}:
@@ -992,18 +1033,18 @@ func jpItems(arguments []interface{}) (interface{}, error) {
 		}
 		return arrayOfObj, nil
 	default:
-		return nil, formatError(invalidArgumentTypeError, items, arguments, 0, "Object or Array")
+		return nil, formatError(invalidArgumentTypeError, items, 1, "Object or Array")
 	}
 }
 
 func jpObjectFromLists(arguments []interface{}) (interface{}, error) {
 	keys, ok := arguments[0].([]interface{})
 	if !ok {
-		return nil, formatError(invalidArgumentTypeError, objectFromLists, arguments, 0, "Array")
+		return nil, formatError(invalidArgumentTypeError, objectFromLists, 1, "Array")
 	}
 	values, ok := arguments[1].([]interface{})
 	if !ok {
-		return nil, formatError(invalidArgumentTypeError, objectFromLists, arguments, 1, "Array")
+		return nil, formatError(invalidArgumentTypeError, objectFromLists, 2, "Array")
 	}
 
 	output := map[string]interface{}{}
@@ -1011,7 +1052,7 @@ func jpObjectFromLists(arguments []interface{}) (interface{}, error) {
 	for i, ikey := range keys {
 		key, err := ifaceToString(ikey)
 		if err != nil {
-			return nil, formatError(invalidArgumentTypeError, objectFromLists, arguments, 0, "StringArray")
+			return nil, formatError(invalidArgumentTypeError, objectFromLists, 1, "StringArray")
 		}
 		if i < len(values) {
 			output[key] = values[i]

pkg/engine/jmespath/functions_test.go

@@ -846,6 +846,203 @@ func Test_SemverCompare(t *testing.T) {
 	}
 }
 
+func Test_Lookup(t *testing.T) {
+	testCases := []struct {
+		collection     string
+		key            string
+		expectedResult string
+	}{
+		// objects
+		/////////////////
+
+		// not found
+		{
+			collection:     `{}`,
+			key:            `"key1"`,
+			expectedResult: `null`,
+		},
+		{
+			collection:     `{"key1": "value1"}`,
+			key:            `"key2"`,
+			expectedResult: `null`,
+		},
+
+		// found
+		{
+			collection:     `{"key1": "value1"}`,
+			key:            `"key1"`,
+			expectedResult: `"value1"`,
+		},
+		{
+			collection:     `{"": "value1"}`,
+			key:            `""`,
+			expectedResult: `"value1"`,
+		},
+
+		// result types
+		{
+			collection:     `{"k": 123}`,
+			key:            `"k"`,
+			expectedResult: `123`,
+		},
+		{
+			collection:     `{"k": 12.34}`,
+			key:            `"k"`,
+			expectedResult: `12.34`,
+		},
+		{
+			collection:     `{"k": true}`,
+			key:            `"k"`,
+			expectedResult: `true`,
+		},
+		{
+			collection:     `{"k": false}`,
+			key:            `"k"`,
+			expectedResult: `false`,
+		},
+		{
+			collection:     `{"k": null}`,
+			key:            `"k"`,
+			expectedResult: `null`,
+		},
+		{
+			collection:     `{"k":  [7, "x", true] }`,
+			key:            `"k"`,
+			expectedResult: `[7, "x", true]`,
+		},
+		{
+			collection:     `{"k": {"key1":true}}`,
+			key:            `"k"`,
+			expectedResult: `{"key1":true}`,
+		},
+
+		// arrays
+		/////////////////
+
+		// not found
+		{
+			collection:     `[]`,
+			key:            `0`,
+			expectedResult: `null`,
+		},
+		{
+			collection:     `["item0"]`,
+			key:            `-1`,
+			expectedResult: `null`,
+		},
+		{
+			collection:     `["item0"]`,
+			key:            `1`,
+			expectedResult: `null`,
+		},
+
+		// found
+		{
+			collection:     `["item0"]`,
+			key:            `0`,
+			expectedResult: `"item0"`,
+		},
+		{
+			collection:     `["item0", "item1", "item2", "item3"]`,
+			key:            `2`,
+			expectedResult: `"item2"`,
+		},
+		{
+			collection:     `["item0", "item1"]`,
+			key:            `0.99999999999999999999999999999999999999999999`,
+			expectedResult: `"item1"`,
+		},
+
+		// result types
+		{
+			collection:     `[123]`,
+			key:            `0`,
+			expectedResult: `123`,
+		},
+		{
+			collection:     `[12.34]`,
+			key:            `0`,
+			expectedResult: `12.34`,
+		},
+		{
+			collection:     `[true]`,
+			key:            `0`,
+			expectedResult: `true`,
+		},
+		{
+			collection:     `[false]`,
+			key:            `0`,
+			expectedResult: `false`,
+		},
+		{
+			collection:     `[null]`,
+			key:            `0`,
+			expectedResult: `null`,
+		},
+		{
+			collection:     `[ [7, "x", true] ]`,
+			key:            `0`,
+			expectedResult: `[7, "x", true]`,
+		},
+		{
+			collection:     `[{"key1":true}]`,
+			key:            `0`,
+			expectedResult: `{"key1":true}`,
+		},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			query, err := newJMESPath(cfg, "lookup(`"+tc.collection+"`,`"+tc.key+"`)")
+			assert.NilError(t, err)
+
+			result, err := query.Search("")
+			assert.NilError(t, err)
+
+			var expectedResult interface{}
+			err = json.Unmarshal([]byte(tc.expectedResult), &expectedResult)
+			assert.NilError(t, err)
+
+			assert.DeepEqual(t, result, expectedResult)
+		})
+	}
+}
+
+func Test_Lookup_InvalidArgs(t *testing.T) {
+	testCases := []struct {
+		collection  string
+		key         string
+		expectedMsg string
+	}{
+		// invalid key type
+		{
+			collection:  `{}`,
+			key:         `123`,
+			expectedMsg: `argument #2 is not of type String`,
+		},
+		{
+			collection:  `[]`,
+			key:         `"abc"`,
+			expectedMsg: `argument #2 is not of type Number`,
+		},
+
+		// invalid value
+		{
+			collection:  `[]`,
+			key:         `1.5`,
+			expectedMsg: `JMESPath function 'lookup': argument #2: expected an integer number but got: 1.5`,
+		},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			query, err := newJMESPath(cfg, "lookup(`"+tc.collection+"`,`"+tc.key+"`)")
+			assert.NilError(t, err)
+
+			_, err = query.Search("")
+			assert.ErrorContains(t, err, tc.expectedMsg)
+		})
+	}
+}
+
 func Test_Items(t *testing.T) {
 	testCases := []struct {
 		object         string

pkg/engine/jmespath/utils.go

@@ -1,6 +1,8 @@
 package jmespath
 
 import (
+	"fmt"
+	"math"
 	"reflect"
 )
 
@@ -17,3 +19,14 @@ func validateArg(f string, arguments []interface{}, index int, expectedType refl
 	}
 	return arg, nil
 }
+
+func intNumber(number float64) (int, error) {
+	if math.IsInf(number, 0) || math.IsNaN(number) || math.Trunc(number) != number {
+		return 0, fmt.Errorf("expected an integer number but got: %g", number)
+	}
+	intNumber := int(number)
+	if float64(intNumber) != number {
+		return 0, fmt.Errorf("number is outside the range of integer numbers: %g", number)
+	}
+	return intNumber, nil
+}

pkg/engine/jmespath/utils_test.go

@@ -0,0 +1,80 @@
+package jmespath
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	"gotest.tools/assert"
+)
+
+func Test_intNumber(t *testing.T) {
+	testCases := []struct {
+		number         float64
+		expectedResult int
+	}{
+		{
+			number:         0.0,
+			expectedResult: 0,
+		},
+		{
+			number:         1.0,
+			expectedResult: 1,
+		},
+		{
+			number:         -1.0,
+			expectedResult: -1,
+		},
+		{
+			number:         math.MaxInt32,
+			expectedResult: math.MaxInt32,
+		},
+		{
+			number:         math.MinInt32,
+			expectedResult: math.MinInt32,
+		},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			result, resultErr := intNumber(tc.number)
+
+			assert.NilError(t, resultErr)
+			assert.Equal(t, result, tc.expectedResult)
+		})
+	}
+}
+
+func Test_intNumber_Error(t *testing.T) {
+	testCases := []struct {
+		number      float64
+		expectedMsg string
+	}{
+		{
+			number:      1.5,
+			expectedMsg: `expected an integer number but got: 1.5`,
+		},
+		{
+			number:      math.NaN(),
+			expectedMsg: `expected an integer number but got: NaN`,
+		},
+		{
+			number:      math.Inf(1),
+			expectedMsg: `expected an integer number but got: +Inf`,
+		},
+		{
+			number:      math.Inf(-1),
+			expectedMsg: `expected an integer number but got: -Inf`,
+		},
+		{
+			number:      math.MaxFloat64,
+			expectedMsg: `number is outside the range of integer numbers: 1.7976931348623157e+308`,
+		},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			_, resultErr := intNumber(tc.number)
+
+			assert.Error(t, resultErr, tc.expectedMsg)
+		})
+	}
+}