✨ string.in(list of strings)

Support the inverse of `list.contains(string)` by adding support for `string.in(list)`. This makes queries much more flexible. Example: ```coffee hi.in([a,hi,b]) ``` Signed-off-by: Dominik Richter <dominik.richter@gmail.com>
mondoohq · Jan 6, 2024 · a203f87 · a203f87
1 parent 88dc0b4
commit a203f87
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 29 deletions.
diff --git a/llx/builtin.go b/llx/builtin.go
@@ -319,6 +319,7 @@ func init() {
 			string("contains" + types.Array(types.Int)):    {f: stringContainsArrayIntV2, Label: "contains"},
 			string("contains" + types.Regex):               {f: stringContainsRegex, Label: "contains"},
 			string("contains" + types.Array(types.Regex)):  {f: stringContainsArrayRegex, Label: "contains"},
+			string("in"):        {f: stringInArray, Label: "in"},
 			string("find"):      {f: stringFindV2, Label: "find"},
 			string("camelcase"): {f: stringCamelcaseV2, Label: "camelcase"},
 			string("downcase"):  {f: stringDowncaseV2, Label: "downcase"},
@@ -539,6 +540,7 @@ func init() {
 			string("contains" + types.Array(types.Int)):    {f: dictContainsArrayIntV2, Label: "contains"},
 			string("contains" + types.Regex):               {f: dictContainsRegex, Label: "contains"},
 			string("contains" + types.Array(types.Regex)):  {f: dictContainsArrayRegex, Label: "contains"},
+			"in":           {f: dictIn, Label: "in"},
 			string("find"): {f: dictFindV2, Label: "find"},
 			// NOTE: the following functions are internal ONLY!
 			// We have not yet decided if and how these may be exposed to users

diff --git a/llx/builtin_map.go b/llx/builtin_map.go
@@ -1279,6 +1279,15 @@ func dictContainsArrayRegex(e *blockExecutor, bind *RawData, chunk *Chunk, ref u
 	}
 }
 
+func dictIn(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
+	switch bind.Value.(type) {
+	case string:
+		return stringInArray(e, bind, chunk, ref)
+	default:
+		return nil, 0, errors.New("dict value does not support field `in`")
+	}
+}
+
 func dictFindV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
 	switch bind.Value.(type) {
 	case string:

diff --git a/llx/builtin_simple.go b/llx/builtin_simple.go
@@ -2220,11 +2220,36 @@ func stringContainsArrayRegex(e *blockExecutor, bind *RawData, chunk *Chunk, ref
 		}
 
 		if re.MatchString(bind.Value.(string)) {
-			return BoolData(true), 0, nil
+			return BoolTrue, 0, nil
 		}
 	}
 
-	return BoolData(false), 0, nil
+	return BoolFalse, 0, nil
+}
+
+func stringInArray(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
+	if bind.Value == nil {
+		return BoolFalse, 0, nil
+	}
+
+	argRef := chunk.Function.Args[0]
+	arg, rref, err := e.resolveValue(argRef, ref)
+	if err != nil || rref > 0 {
+		return nil, rref, err
+	}
+
+	if arg.Value == nil {
+		return BoolFalse, 0, nil
+	}
+
+	arr := arg.Value.([]interface{})
+	for i := range arr {
+		v := arr[i].(string)
+		if bind.Value.(string) == v {
+			return BoolTrue, 0, nil
+		}
+	}
+	return BoolFalse, 0, nil
 }
 
 func stringFindV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {

diff --git a/mqlc/builtin.go b/mqlc/builtin.go
@@ -42,6 +42,7 @@ func init() {
 		},
 		types.String: {
 			"contains":  {compile: compileStringContains, typ: boolType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.String}}},
+			"in":        {typ: boolType, compile: compileStringIn},
 			"find":      {typ: stringArrayType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.Regex}}},
 			"length":    {typ: intType, signature: FunctionSignature{}},
 			"camelcase": {typ: stringType, signature: FunctionSignature{}},
@@ -77,6 +78,7 @@ func init() {
 			"last":         {typ: dictType, signature: FunctionSignature{}},
 			"where":        {compile: compileDictWhere, 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}}},
 			"containsAll":  {compile: compileDictContainsAll, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
 			"containsNone": {compile: compileDictContainsNone, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},

diff --git a/mqlc/builtin_simple.go b/mqlc/builtin_simple.go
@@ -12,6 +12,39 @@ import (
 	"go.mondoo.com/cnquery/v9/types"
 )
 
+func callArgTypeIs(c *compiler, call *parser.Call, id string, argName string, idx int, types ...types.Type) (*llx.Primitive, error) {
+	if len(call.Function) <= idx {
+		return nil, errors.New("function " + id + " is missing a " + argName + " (arg #" + strconv.Itoa(idx+1) + ")")
+	}
+
+	arg := call.Function[idx]
+	if arg.Value == nil || arg.Value.Operand == nil {
+		return nil, errors.New("function " + id + " is missing a " + argName + " (arg #" + strconv.Itoa(idx+1) + " is null)")
+	}
+
+	val, err := c.compileOperand(arg.Value.Operand)
+	if err != nil {
+		return nil, err
+	}
+
+	valType, err := c.dereferenceType(val)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, t := range types {
+		if t == valType {
+			return val, nil
+		}
+	}
+
+	var typesStr string
+	for _, t := range types {
+		typesStr += t.Label() + "/"
+	}
+	return nil, errors.New("function " + id + " type mismatch for " + argName + " (expected: " + typesStr[0:len(typesStr)-1] + ", got: " + valType.Label() + ")")
+}
+
 func compileStringContains(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
 	if call == nil || len(call.Function) != 1 {
 		return types.Nil, errors.New("function " + id + " needs one argument (function missing)")
@@ -104,42 +137,31 @@ func compileStringContains(c *compiler, typ types.Type, ref uint64, id string, c
 	}
 }
 
-func callArgTypeIs(c *compiler, call *parser.Call, id string, argName string, idx int, types ...types.Type) (*llx.Primitive, error) {
-	if len(call.Function) <= idx {
-		return nil, errors.New("function " + id + " is missing a " + argName + " (arg #" + strconv.Itoa(idx+1) + ")")
-	}
-
-	arg := call.Function[idx]
-	if arg.Value == nil || arg.Value.Operand == nil {
-		return nil, errors.New("function " + id + " is missing a " + argName + " (arg #" + strconv.Itoa(idx+1) + " is null)")
-	}
-
-	val, err := c.compileOperand(arg.Value.Operand)
-	if err != nil {
-		return nil, err
+func compileStringIn(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
+	if call == nil || len(call.Function) != 1 {
+		return types.Nil, errors.New("function " + id + " needs one argument")
 	}
 
-	valType, err := c.dereferenceType(val)
+	arr, err := callArgTypeIs(c, call, id, "list", 0, types.Array(types.String), types.Array(types.Unset))
 	if err != nil {
-		return nil, err
-	}
-
-	for _, t := range types {
-		if t == valType {
-			return val, nil
-		}
+		return types.Nil, err
 	}
 
-	var typesStr string
-	for _, t := range types {
-		typesStr += t.Label() + "/"
-	}
-	return nil, errors.New("function " + id + " type mismatch for " + argName + " (expected: " + typesStr[0:len(typesStr)-1] + ", got: " + valType.Label() + ")")
+	c.addChunk(&llx.Chunk{
+		Call: llx.Chunk_FUNCTION,
+		Id:   "in",
+		Function: &llx.Function{
+			Type:    string(types.Bool),
+			Binding: ref,
+			Args:    []*llx.Primitive{arr},
+		},
+	})
+	return types.Bool, nil
 }
 
 func compileInRange(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
 	if call == nil || len(call.Function) != 2 {
-		return types.Nil, errors.New("function " + id + " needs two arguments (function missing)")
+		return types.Nil, errors.New("function " + id + " needs two arguments")
 	}
 
 	min, err := callArgTypeIs(c, call, id, "min", 0, types.Int, types.Float, types.Dict)

diff --git a/providers/core/resources/mql_test.go b/providers/core/resources/mql_test.go
@@ -424,6 +424,18 @@ func TestString_Methods(t *testing.T) {
 			Code:        "'hello'.contains([/z/, /ll/])",
 			Expectation: true,
 		},
+		{
+			Code:        "'hi'.in(['one','hi','five'])",
+			Expectation: true,
+		},
+		{
+			Code:        "'hiya'.in(['one','hi','five'])",
+			Expectation: false,
+		},
+		{
+			Code:        "'hiya'.in([])",
+			Expectation: false,
+		},
 		{
 			Code:        "'oh-hello-world!'.camelcase",
 			Expectation: "ohHelloWorld!",