Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

⭐ dict.recurse( .. ) #3082

Merged
merged 2 commits into from
Jan 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions llx/builtin.go
Original file line number Diff line number Diff line change
Expand Up @@ -536,8 +536,9 @@ func init() {
"trim": {f: dictTrimV2, Label: "trim"},
"keys": {f: dictKeysV2, Label: "keys"},
"values": {f: dictValuesV2, Label: "values"},
"where": {f: dictWhereV2, Label: "where"},
"$whereNot": {f: dictWhereNotV2},
"where": {f: dictWhere, Label: "where"},
"recurse": {f: dictRecurse, Label: "recurse"},
"$whereNot": {f: dictWhereNot},
"$all": {f: dictAllV2},
"$none": {f: dictNoneV2},
"$any": {f: dictAnyV2},
Expand Down
136 changes: 131 additions & 5 deletions llx/builtin_map.go
Original file line number Diff line number Diff line change
Expand Up @@ -722,7 +722,13 @@ func _dictArrayWhere(e *blockExecutor, list []interface{}, chunk *Chunk, ref uin
return nil, 0, nil
}

func _dictWhereV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64, invert bool) (*RawData, uint64, error) {
// Query a dict object, cycling over all child objects it contains with the given
// query function. This only works in cases where q query function makes sense,
// like:
// - [a, b, c].query( F ) = [a, b]
// - {a: b, c: d}.query( F ) = {a: b}
// Note: Results get stored in cache for this ref
func _dictWhere(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64, invert bool) (*RawData, uint64, error) {
itemsRef := chunk.Function.Args[0]
items, rref, err := e.resolveValue(itemsRef, ref)
if err != nil || rref > 0 {
Expand Down Expand Up @@ -811,12 +817,132 @@ func _dictWhereV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64, inv
return nil, 0, nil
}

func dictWhereV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
return _dictWhereV2(e, bind, chunk, ref, false)
func dictWhere(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
return _dictWhere(e, bind, chunk, ref, false)
}

func dictWhereNotV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
return _dictWhereV2(e, bind, chunk, ref, true)
func dictWhereNot(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
return _dictWhere(e, bind, chunk, ref, true)
}

func filterList(e *blockExecutor, list []any, chunk *Chunk, ref uint64, invert bool) ([]any, uint64, error) {
arg1 := chunk.Function.Args[1]
fref, ok := arg1.RefV2()
if !ok {
return nil, 0, errors.New("failed to retrieve function reference (in dict list query)")
}

dref, err := e.ensureArgsResolved(chunk.Function.Args[2:], ref)
if dref != 0 || err != nil {
return nil, dref, err
}

argsList := make([][]*RawData, len(list))
for i, value := range list {
argsList[i] = []*RawData{
{
Type: types.Dict,
Value: i,
},
{
Type: types.Dict,
Value: value,
},
}
}

var res []any
err = e.runFunctionBlocks(argsList, fref, func(results []arrayBlockCallResult, errs []error) {
resList := []any{}
for i, res := range results {
if res.isTruthy() == !invert {
key := argsList[i][0].Value.(int)
resList = append(resList, list[key])
}
}
res = resList
})

if err != nil {
return nil, 0, err
}

return res, 0, nil
}

// The recurse function only works on lists and maps. It traverses child structures
// and finds any objects that match the given search function
func _dictRecurse(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64, invert bool) (*RawData, uint64, error) {
itemsRef := chunk.Function.Args[0]
items, rref, err := e.resolveValue(itemsRef, ref)
if err != nil || rref > 0 {
return nil, rref, err
}

if items.Value == nil {
return &RawData{Type: items.Type}, 0, nil
}

remaining := []any{items.Value}
res := []any{}
var candidate any
for len(remaining) != 0 {
candidate = remaining[0]
remaining = remaining[1:]

if candidate == nil {
continue
}

var list []any
if x, ok := candidate.([]any); ok {
list = x
} else if x, ok := candidate.(map[string]any); ok {
list = make([]any, len(x))
i := 0
for _, v := range x {
list[i] = v
i++
}
}

if len(list) == 0 {
continue
}

partial, dref, err := filterList(e, list, chunk, ref, invert)
if dref != 0 || err != nil {
return nil, dref, err
}
res = append(res, partial...)

// we only add items to the remaining, that actually match our type requirements
// (this is trying to keep additional memory juggling to a minimum, instead
// of doing this one level deeper)
for i := range list {
li := list[i]
if x, ok := li.([]any); ok {
remaining = append(remaining, x)
} else if x, ok := li.(map[string]any); ok {
remaining = append(remaining, x)
}
}
}

data := &RawData{
Type: types.Dict,
Value: res,
}
e.cache.Store(ref, &stepCache{
Result: data,
IsStatic: false,
})

return data, 0, nil
}

func dictRecurse(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
return _dictRecurse(e, bind, chunk, ref, false)
}

func dictAllV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
Expand Down
1 change: 1 addition & 0 deletions mqlc/builtin.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ func init() {
"first": {typ: dictType, signature: FunctionSignature{}},
"last": {typ: dictType, signature: FunctionSignature{}},
"where": {compile: compileDictWhere, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
"recurse": {compile: compileDictRecurse, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
"contains": {compile: compileDictContains, typ: boolType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
"in": {typ: boolType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.Array(types.String)}}},
"containsOnly": {compile: compileDictContainsOnly, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
Expand Down
20 changes: 14 additions & 6 deletions mqlc/builtin_map.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (
"go.mondoo.com/cnquery/v10/types"
)

func compileDictWhere(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
func compileDictQuery(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
if call == nil {
return types.Nil, errors.New("missing filter argument for calling '" + id + "'")
}
Expand Down Expand Up @@ -129,8 +129,16 @@ func compileDictWhere(c *compiler, typ types.Type, ref uint64, id string, call *
return typ, nil
}

func compileDictWhere(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
return compileDictQuery(c, typ, ref, id, call)
}

func compileDictRecurse(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
return compileDictQuery(c, typ, ref, "recurse", call)
}

func compileDictContains(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
_, err := compileDictWhere(c, typ, ref, "where", call)
_, err := compileDictQuery(c, typ, ref, "where", call)
if err != nil {
return types.Nil, err
}
Expand Down Expand Up @@ -326,7 +334,7 @@ func compileDictContainsNone(c *compiler, typ types.Type, ref uint64, id string,
}

func compileDictAll(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
_, err := compileDictWhere(c, typ, ref, "$whereNot", call)
_, err := compileDictQuery(c, typ, ref, "$whereNot", call)
if err != nil {
return types.Nil, err
}
Expand All @@ -352,7 +360,7 @@ func compileDictAll(c *compiler, typ types.Type, ref uint64, id string, call *pa
}

func compileDictAny(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
_, err := compileDictWhere(c, typ, ref, "where", call)
_, err := compileDictQuery(c, typ, ref, "where", call)
if err != nil {
return types.Nil, err
}
Expand All @@ -378,7 +386,7 @@ func compileDictAny(c *compiler, typ types.Type, ref uint64, id string, call *pa
}

func compileDictOne(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
_, err := compileDictWhere(c, typ, ref, "where", call)
_, err := compileDictQuery(c, typ, ref, "where", call)
if err != nil {
return types.Nil, err
}
Expand All @@ -404,7 +412,7 @@ func compileDictOne(c *compiler, typ types.Type, ref uint64, id string, call *pa
}

func compileDictNone(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
_, err := compileDictWhere(c, typ, ref, "where", call)
_, err := compileDictQuery(c, typ, ref, "where", call)
if err != nil {
return types.Nil, err
}
Expand Down
2 changes: 1 addition & 1 deletion providers-sdk/v1/testutils/testdata/arch.json
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@
"Fields": {
"content": {
"type": "\u0007",
"value": "{\n \"_\": null,\n \"true\": true,\n \"1\": 1,\n \"1.0\": 1.0,\n \"int-array\": [1,2,3],\n \"dict\": {\n \"ee\": 3,\n \"ej\": 4,\n \"ek\": 5\n },\n \"f\": [{\"ff\": 3}],\n \"string-array\": [\"a\", \"b\", \"c\"],\n \"hello\": \"hello\",\n \"date\": \"2016-01-28T23:02:24Z\",\n \"aoa\": [[1, 2], 3]\n}\n"
"value": "{\n \"_\": null,\n \"true\": true,\n \"1\": 1,\n \"1.0\": 1.0,\n \"int-array\": [1,2,3],\n \"dict\": {\n \"ee\": 3,\n \"ej\": 4,\n \"ek\": 5\n },\n \"f\": [{\"ff\": 3}],\n \"string-array\": [\"a\", \"b\", \"c\"],\n \"hello\": \"hello\",\n \"date\": \"2016-01-28T23:02:24Z\",\n \"aoa\": [[1, 2], 3],\n \"users\": [\n {\n \"name\": \"yor\",\n \"children\": [\n {\"name\": \"anya\"}\n ]\n },\n {\"name\": \"loid\"}\n ],\n \"zzzlast\": \"🌒\"}\n"
},
"path": {
"type": "\u0007",
Expand Down
14 changes: 9 additions & 5 deletions providers/os/resources/mql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ func TestOS_Vars(t *testing.T) {
x.TestSimple(t, []testutils.SimpleTest{
{
Code: "p = file('/dummy.json'); parse.json(file: p).params.length",
Expectation: int64(11),
Expectation: int64(13),
},
})
}
Expand All @@ -97,15 +97,15 @@ func TestMap(t *testing.T) {
},
{
Code: "parse.json('/dummy.json').params.length",
Expectation: int64(11),
Expectation: int64(13),
},
{
Code: "parse.json('/dummy.json').params.keys.length",
Expectation: int64(11),
Expectation: int64(13),
},
{
Code: "parse.json('/dummy.json').params.values.length",
Expectation: int64(11),
Expectation: int64(13),
},
{
Code: "parse.json('/dummy.json').params { _['Protocol'] != 1 }",
Expand Down Expand Up @@ -383,6 +383,10 @@ func TestDict_Methods_Map(t *testing.T) {
Code: p + "params['string-array'].where(_ == 'a')",
Expectation: []interface{}{"a"},
},
{
Code: p + "params.users.recurse(name != empty).map(name)",
Expectation: []any{"yor", "loid", "anya"},
},
{
Code: p + "params['string-array'].in(['a', 'b', 'c'])",
Expectation: true,
Expand Down Expand Up @@ -457,7 +461,7 @@ func TestDict_Methods_Map(t *testing.T) {
},
{
Code: p + "params.last",
Expectation: true,
Expectation: "🌒",
},
{
Code: p + "params['aoa'].flat",
Expand Down
Loading