Add close status parse method in pinot query validator (#5680)

* Add close status parse method in pinot query validator

common/pinot/pinotQueryValidator.go

@@ -295,6 +295,16 @@ func (qv *VisibilityQueryValidator) processSystemKey(expr sqlparser.Expr) (strin
 				return "", fmt.Errorf("trim time field %s got error: %w", colNameStr, err)
 			}
 			comparisonExpr.Right = trimmed
+		} else if colNameStr == "CloseStatus" {
+			sqlVal, ok := comparisonExpr.Right.(*sqlparser.SQLVal)
+			if !ok {
+				return "", fmt.Errorf("right comparison is invalid: %v", comparisonExpr.Right)
+			}
+			closeStatus, err := parseCloseStatus(sqlVal)
+			if err != nil {
+				return "", fmt.Errorf("parse CloseStatus field got error: %w", err)
+			}
+			comparisonExpr.Right = closeStatus
 		}
 	}
 
@@ -417,3 +427,24 @@ func parseTime(timeStr string) (int64, error) {
 
 	return 0, errors.New("invalid time string")
 }
+
+func parseCloseStatus(original *sqlparser.SQLVal) (*sqlparser.SQLVal, error) {
+	statusStr := string(original.Val)
+
+	// first check if already in int64 format
+	if _, err := strconv.ParseInt(statusStr, 10, 64); err == nil {
+		return original, nil
+	}
+
+	// try to parse close status string
+	var parsedStatus types.WorkflowExecutionCloseStatus
+	err := parsedStatus.UnmarshalText([]byte(statusStr))
+	if err != nil {
+		return nil, err
+	}
+
+	return &sqlparser.SQLVal{
+		Type: sqlparser.IntVal,
+		Val:  []byte(strconv.FormatInt(int64(parsedStatus), 10)),
+	}, nil
+}

common/pinot/pinotQueryValidator_test.go

@@ -26,6 +26,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/xwb1989/sqlparser"
 
 	"github.com/uber/cadence/common/definition"
 	"github.com/uber/cadence/common/dynamicconfig"
@@ -210,6 +211,10 @@ func TestValidateQuery(t *testing.T) {
 			query:     "CustomIntField = 0 order by CustomIntField desc",
 			validated: "JSON_MATCH(Attr, '\"$.CustomIntField\"=''0''') order by CustomIntField desc",
 		},
+		"case 19: close status parse": {
+			query:     "CloseStatus = 'CONTINUED_AS_NEW'",
+			validated: "CloseStatus = 4",
+		},
 	}
 
 	for name, test := range tests {
@@ -250,3 +255,41 @@ func TestParseTime(t *testing.T) {
 		})
 	}
 }
+
+func TestParseCloseStatus(t *testing.T) {
+	tests := []struct {
+		input       string
+		expected    *sqlparser.SQLVal
+		expectedErr bool
+	}{
+		{
+			input:       "4",
+			expected:    &sqlparser.SQLVal{Type: sqlparser.IntVal, Val: []byte("4")},
+			expectedErr: false,
+		},
+		{
+			input:       "CANCELED",
+			expected:    &sqlparser.SQLVal{Type: sqlparser.IntVal, Val: []byte("2")},
+			expectedErr: false,
+		},
+		{
+			input:       "invalid",
+			expected:    nil,
+			expectedErr: true, // expected error,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.input, func(t *testing.T) {
+			original := &sqlparser.SQLVal{Type: sqlparser.IntVal, Val: []byte(test.input)}
+			result, err := parseCloseStatus(original)
+
+			assert.Equal(t, test.expected, result)
+			if test.expectedErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}