Update pinot query validator to handle raw time string (#5656)

common/pinot/pinotQueryValidator.go

@@ -27,6 +27,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/xwb1989/sqlparser"
 
@@ -378,24 +379,42 @@ func processCustomString(comparisonExpr *sqlparser.ComparisonExpr, colNameStr st
 func trimTimeFieldValueFromNanoToMilliSeconds(original *sqlparser.SQLVal) (*sqlparser.SQLVal, error) {
 	// Convert the SQLVal to a string
 	valStr := string(original.Val)
-	// Convert to an integer
-	valInt, err := strconv.ParseInt(valStr, 10, 64)
+	newVal, err := parseTime(valStr)
 	if err != nil {
 		return original, fmt.Errorf("error: failed to parse int from SQLVal %s", valStr)
 	}
 
-	var newVal int64
-	if valInt < 0 { //exclude open workflow which time field will be -1
-		newVal = valInt
-	} else if len(valStr) > 13 { // Assuming nanoseconds if more than 13 digits
-		newVal = valInt / 1000000 // Convert time to milliseconds
-	} else {
-		newVal = valInt
-	}
-
 	// Convert the new value back to SQLVal
 	return &sqlparser.SQLVal{
 		Type: sqlparser.IntVal,
 		Val:  []byte(strconv.FormatInt(newVal, 10)),
 	}, nil
 }
+
+func parseTime(timeStr string) (int64, error) {
+	if len(timeStr) == 0 {
+		return 0, errors.New("invalid time string")
+	}
+
+	// try to parse
+	parsedTime, err := time.Parse(time.RFC3339, timeStr)
+	if err == nil {
+		return parsedTime.UnixMilli(), nil
+	}
+
+	// treat as raw time
+	valInt, err := strconv.ParseInt(timeStr, 10, 64)
+	if err == nil {
+		var newVal int64
+		if valInt < 0 { //exclude open workflow which time field will be -1
+			newVal = valInt
+		} else if len(timeStr) > 13 { // Assuming nanoseconds if more than 13 digits
+			newVal = valInt / 1000000 // Convert time to milliseconds
+		} else {
+			newVal = valInt
+		}
+		return newVal, nil
+	}
+
+	return 0, errors.New("invalid time string")
+}

common/pinot/pinotQueryValidator_test.go

@@ -173,6 +173,18 @@ func TestValidateQuery(t *testing.T) {
 			query:     "StartTime = 170731995093",
 			validated: "StartTime = 170731995093",
 		},
+		"Case15-13: value in raw string for equal statement": {
+			query:     "StartTime = '2024-02-07T15:32:30Z'",
+			validated: "StartTime = 1707319950000",
+		},
+		"Case15-14: value in raw string for not equal statement": {
+			query:     "StartTime > '2024-02-07T15:32:30Z'",
+			validated: "StartTime > 1707319950000",
+		},
+		"Case15-15: value in raw string for range statement": {
+			query:     "StartTime between '2024-02-07T15:32:30Z' and '2024-02-07T15:33:30Z'",
+			validated: "StartTime between 1707319950000 and 1707320010000",
+		},
 		"Case16-1: custom int attribute greater than or equal to": {
 			query:     "CustomIntField >= 0",
 			validated: "(JSON_MATCH(Attr, '\"$.CustomIntField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) >= 0)",
@@ -205,3 +217,28 @@ func TestValidateQuery(t *testing.T) {
 		})
 	}
 }
+
+func TestParseTime(t *testing.T) {
+	var tests = []struct {
+		name     string
+		timeStr  string
+		expected int64
+		hasErr   bool
+	}{
+		{"empty string", "", 0, true},
+		{"valid RFC3339", "2024-02-07T15:32:30Z", 1707319950000, false},
+		{"valid unix milli string", "1707319950000", 1707319950000, false},
+		{"valid unix nano string", "1707319950000000000", 1707319950000, false},
+		{"invalid string", "invalid", 0, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			parsed, err := parseTime(tt.timeStr)
+			assert.Equal(t, parsed, tt.expected)
+			if tt.hasErr {
+				assert.Error(t, err)
+			}
+		})
+	}
+}