Add support for $gt, $ge, $lt, $le operators for netquery. Update Dat…

…etimeEncoder to support values specified in seconds
safing · Sep 6, 2023 · d96f841 · d96f841
1 parent b35e4b9
commit d96f841
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 7 deletions.
diff --git a/netquery/orm/encoder.go b/netquery/orm/encoder.go
@@ -156,6 +156,8 @@ func DatetimeEncoder(loc *time.Location) EncodeFunc {
 			val = reflect.Indirect(val)
 		}
 
+		normalizedKind := normalizeKind(valType.Kind())
+
 		// we only care about "time.Time" here
 		var t time.Time
 		switch {
@@ -179,6 +181,9 @@ func DatetimeEncoder(loc *time.Location) EncodeFunc {
 				return nil, false, fmt.Errorf("failed to parse time as RFC3339: %w", err)
 			}
 
+		case (normalizedKind == reflect.Int || normalizedKind == reflect.Uint || normalizedKind == reflect.Float32) && colDef.IsTime:
+			t = time.Unix(val.Int(), 0)
+
 		default:
 			// we don't care ...
 			return nil, false, nil

diff --git a/netquery/query.go b/netquery/query.go
@@ -38,11 +38,15 @@ type (
 	Equal interface{}
 
 	Matcher struct {
-		Equal    interface{}   `json:"$eq,omitempty"`
-		NotEqual interface{}   `json:"$ne,omitempty"`
-		In       []interface{} `json:"$in,omitempty"`
-		NotIn    []interface{} `json:"$notIn,omitempty"`
-		Like     string        `json:"$like,omitempty"`
+		Equal          interface{}   `json:"$eq,omitempty"`
+		NotEqual       interface{}   `json:"$ne,omitempty"`
+		In             []interface{} `json:"$in,omitempty"`
+		NotIn          []interface{} `json:"$notIn,omitempty"`
+		Like           string        `json:"$like,omitempty"`
+		Greater        *float64      `json:"$gt,omitempty"`
+		GreaterOrEqual *float64      `json:"$ge,omitempty"`
+		Less           *float64      `json:"$lt,omitempty"`
+		LessOrEqual    *float64      `json:"$le,omitempty"`
 	}
 
 	Count struct {
@@ -258,6 +262,22 @@ func (match Matcher) Validate() error {
 		found++
 	}
 
+	if match.Greater != nil {
+		found++
+	}
+
+	if match.GreaterOrEqual != nil {
+		found++
+	}
+
+	if match.Less != nil {
+		found++
+	}
+
+	if match.LessOrEqual != nil {
+		found++
+	}
+
 	if found == 0 {
 		return fmt.Errorf("no conditions specified")
 	}
@@ -318,10 +338,15 @@ func (match Matcher) toSQLConditionClause(ctx context.Context, suffix string, co
 			params[uniqKey] = encodedValue
 		}
 
+		nameStmt := colDef.Name
+		if colDef.IsTime && colDef.Type == sqlite.TypeText {
+			nameStmt = fmt.Sprintf("strftime('%%s', %s)+0", nameStmt)
+		}
+
 		if len(placeholder) == 1 && !list {
-			queryParts = append(queryParts, fmt.Sprintf("%s %s %s", colDef.Name, operator, placeholder[0]))
+			queryParts = append(queryParts, fmt.Sprintf("%s %s %s", nameStmt, operator, placeholder[0]))
 		} else {
-			queryParts = append(queryParts, fmt.Sprintf("%s %s ( %s )", colDef.Name, operator, strings.Join(placeholder, ", ")))
+			queryParts = append(queryParts, fmt.Sprintf("%s %s ( %s )", nameStmt, operator, strings.Join(placeholder, ", ")))
 		}
 	}
 
@@ -345,6 +370,22 @@ func (match Matcher) toSQLConditionClause(ctx context.Context, suffix string, co
 		add("LIKE", "like", false, match.Like)
 	}
 
+	if match.Greater != nil {
+		add(">", "gt", false, *match.Greater)
+	}
+
+	if match.GreaterOrEqual != nil {
+		add(">=", "ge", false, *match.GreaterOrEqual)
+	}
+
+	if match.Less != nil {
+		add("<", "lt", false, *match.Less)
+	}
+
+	if match.LessOrEqual != nil {
+		add("<=", "le", false, *match.LessOrEqual)
+	}
+
 	if len(queryParts) == 0 {
 		// this is an empty matcher without a single condition.
 		// we convert that to a no-op TRUE value