From 3742d58692f6db75a3bb4bbb5910f9a347b64285 Mon Sep 17 00:00:00 2001
From: Neil Xie <neil.xie@uber.com>
Date: Wed, 14 Feb 2024 11:15:52 -0800
Subject: [PATCH] Fix bug when user pass in not equal query with value missing

---
 common/pinot/pinotQueryValidator.go      |  9 ++++-----
 common/pinot/pinotQueryValidator_test.go | 12 ++++++++++--
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/common/pinot/pinotQueryValidator.go b/common/pinot/pinotQueryValidator.go
index fb380b7c58a..0405fa709bd 100644
--- a/common/pinot/pinotQueryValidator.go
+++ b/common/pinot/pinotQueryValidator.go
@@ -241,7 +241,7 @@ func (qv *VisibilityQueryValidator) processSystemKey(expr sqlparser.Expr) (strin
 	}
 	colNameStr := colName.Name.String()
 
-	if comparisonExpr.Operator != sqlparser.EqualStr {
+	if comparisonExpr.Operator != sqlparser.EqualStr && comparisonExpr.Operator != sqlparser.NotEqualStr {
 		if _, ok := timeSystemKeys[colNameStr]; ok {
 			sqlVal, ok := comparisonExpr.Right.(*sqlparser.SQLVal)
 			if !ok {
@@ -280,10 +280,9 @@ func (qv *VisibilityQueryValidator) processSystemKey(expr sqlparser.Expr) (strin
 		} else {
 			newColVal = "-1" // -1 is the default value for all Closed workflows related fields
 		}
-		comparisonExpr.Right = &sqlparser.ColName{
-			Metadata:  colName.Metadata,
-			Name:      sqlparser.NewColIdent(newColVal),
-			Qualifier: colName.Qualifier,
+		comparisonExpr.Right = &sqlparser.SQLVal{
+			Type: sqlparser.IntVal, // or sqlparser.StrVal if you need to assign a string
+			Val:  []byte(newColVal),
 		}
 	} else {
 		if _, ok := timeSystemKeys[colNameStr]; ok {
diff --git a/common/pinot/pinotQueryValidator_test.go b/common/pinot/pinotQueryValidator_test.go
index eda79836b0f..865634ace5a 100644
--- a/common/pinot/pinotQueryValidator_test.go
+++ b/common/pinot/pinotQueryValidator_test.go
@@ -77,9 +77,13 @@ func TestValidateQuery(t *testing.T) {
 			query: "Invalid SQL",
 			err:   "Invalid query.",
 		},
-		"Case8: query with missing val": {
+		"Case8-1: query with missing val": {
 			query:     "CloseTime = missing",
-			validated: "CloseTime = `-1`",
+			validated: "CloseTime = -1",
+		},
+		"Case8-2: query with not missing case": {
+			query:     "CloseTime != missing",
+			validated: "CloseTime != -1",
 		},
 		"Case9: invalid where expression": {
 			query: "InvalidWhereExpr",
@@ -185,6 +189,10 @@ func TestValidateQuery(t *testing.T) {
 			query:     "StartTime between '2024-02-07T15:32:30Z' and '2024-02-07T15:33:30Z'",
 			validated: "StartTime between 1707319950000 and 1707320010000",
 		},
+		"Case15-16: combined time and missing case": {
+			query:     "CloseTime != missing and StartTime >= 1707662555754408145",
+			validated: "CloseTime != -1 and StartTime >= 1707662555754",
+		},
 		"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)",