✨ time.inRange(min, max) (#2959)

Add support for the inRange operator for the time type. Makes it easier to check if a time is between two timestamps. Signed-off-by: Dominik Richter <[email protected]>
mondoohq · Jan 6, 2024 · 59f7bb8 · 59f7bb8
1 parent db91d3c
commit 59f7bb8
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 5 deletions.
diff --git a/llx/builtin.go b/llx/builtin.go
@@ -434,6 +434,7 @@ func init() {
 			string("hours"):   {f: timeHoursV2, Label: "hours"},
 			string("days"):    {f: timeDaysV2, Label: "days"},
 			string("unix"):    {f: timeUnixV2, Label: "unix"},
+			string("inRange"): {f: timeInRange, Label: "inRange"},
 		},
 		types.Dict: {
 			string("==" + types.Nil):                 {f: dictCmpNilV2, Label: "=="},

diff --git a/llx/builtin_simple.go b/llx/builtin_simple.go
@@ -2495,13 +2495,48 @@ func timeDaysV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*Raw
 func timeUnixV2(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
 	t := bind.Value.(*time.Time)
 	if t == nil {
-		return &RawData{Type: types.Array(types.Time)}, 0, nil
+		return &RawData{Type: types.Int}, 0, nil
 	}
 
 	raw := t.Unix()
 	return IntData(int64(raw)), 0, nil
 }
 
+func timeInRange(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {
+	val := bind.Value.(*time.Time)
+	if val == nil {
+		return &RawData{Type: types.Array(types.Time)}, 0, nil
+	}
+
+	minRef := chunk.Function.Args[0]
+	min, rref, err := e.resolveValue(minRef, ref)
+	if err != nil || rref > 0 {
+		return nil, rref, err
+	}
+
+	switch minval := min.Value.(type) {
+	case *time.Time:
+		if val.Before(*minval) {
+			return BoolFalse, 0, nil
+		}
+	}
+
+	maxRef := chunk.Function.Args[1]
+	max, rref, err := e.resolveValue(maxRef, ref)
+	if err != nil || rref > 0 {
+		return nil, rref, err
+	}
+
+	switch maxval := max.Value.(type) {
+	case *time.Time:
+		if val.After(*maxval) {
+			return BoolFalse, 0, nil
+		}
+	}
+
+	return BoolTrue, 0, nil
+}
+
 // stringslice methods
 
 func stringsliceEqString(e *blockExecutor, bind *RawData, chunk *Chunk, ref uint64) (*RawData, uint64, error) {

diff --git a/mqlc/builtin.go b/mqlc/builtin.go
@@ -35,10 +35,10 @@ var builtinFunctions map[types.Type]map[string]compileHandler
 func init() {
 	builtinFunctions = map[types.Type]map[string]compileHandler{
 		types.Int: {
-			"inRange": {typ: boolType, compile: compileInRange},
+			"inRange": {typ: boolType, compile: compileNumberInRange},
 		},
 		types.Float: {
-			"inRange": {typ: boolType, compile: compileInRange},
+			"inRange": {typ: boolType, compile: compileNumberInRange},
 		},
 		types.String: {
 			"contains":  {compile: compileStringContains, typ: boolType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.String}}},
@@ -58,12 +58,13 @@ func init() {
 			"hours":   {typ: intType, signature: FunctionSignature{}},
 			"days":    {typ: intType, signature: FunctionSignature{}},
 			"unix":    {typ: intType, signature: FunctionSignature{}},
+			"inRange": {typ: boolType, compile: compileTimeInRange},
 		},
 		types.Dict: {
 			"[]": {typ: dictType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.Any}}},
 			"{}": {typ: blockType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.FunctionLike}}},
 			// number-ish
-			"inRange": {typ: boolType, compile: compileInRange},
+			"inRange": {typ: boolType, compile: compileNumberInRange},
 			// string-ish
 			"find":      {typ: stringArrayType, signature: FunctionSignature{Required: 1, Args: []types.Type{types.Regex}}},
 			"length":    {typ: intType, signature: FunctionSignature{}},

diff --git a/mqlc/builtin_simple.go b/mqlc/builtin_simple.go
@@ -159,7 +159,7 @@ func compileStringIn(c *compiler, typ types.Type, ref uint64, id string, call *p
 	return types.Bool, nil
 }
 
-func compileInRange(c *compiler, typ types.Type, ref uint64, id string, call *parser.Call) (types.Type, error) {
+func compileNumberInRange(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")
 	}
@@ -184,3 +184,29 @@ func compileInRange(c *compiler, typ types.Type, ref uint64, id string, call *pa
 	})
 	return types.Bool, nil
 }
+
+func compileTimeInRange(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")
+	}
+
+	min, err := callArgTypeIs(c, call, id, "min", 0, types.Time)
+	if err != nil {
+		return types.Nil, err
+	}
+	max, err := callArgTypeIs(c, call, id, "max", 1, types.Time)
+	if err != nil {
+		return types.Nil, err
+	}
+
+	c.addChunk(&llx.Chunk{
+		Call: llx.Chunk_FUNCTION,
+		Id:   "inRange",
+		Function: &llx.Function{
+			Type:    string(types.Bool),
+			Binding: ref,
+			Args:    []*llx.Primitive{min, max},
+		},
+	})
+	return types.Bool, nil
+}
diff --git a/providers/core/resources/mql_test.go b/providers/core/resources/mql_test.go
@@ -775,6 +775,19 @@ func TestMap(t *testing.T) {
 	})
 }
 
+func TestTime(t *testing.T) {
+	x.TestSimple(t, []testutils.SimpleTest{
+		{
+			Code:        "time.now.inRange(time.now, time.tomorrow)",
+			ResultIndex: 1, Expectation: true,
+		},
+		{
+			Code:        "time.now.inRange(time.tomorrow, time.tomorrow)",
+			ResultIndex: 1, Expectation: false,
+		},
+	})
+}
+
 func TestResource_Default(t *testing.T) {
 	x := testutils.InitTester(testutils.LinuxMock())
 	res := x.TestQuery(t, "mondoo")