Range query for custom attribute

common/persistence/pinot/pinotVisibilityStore_test.go

@@ -185,7 +185,7 @@ LIMIT 0, 10
 FROM %s
 WHERE DomainID = 'bfd5c907-f899-4baf-a7b2-2ab85e623ebd'
 AND IsDeleted = false
-AND WorkflowID = 'wid' and ((JSON_MATCH(Attr, '"$.CustomStringField" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or CustomIntField between 1 and 10)
+AND WorkflowID = 'wid' and ((JSON_MATCH(Attr, '"$.CustomStringField" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or (JSON_MATCH(Attr, '"$.CustomIntField" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) >= 1 AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) <= 10))
 Order BY StartTime DESC
 LIMIT 0, 10
 `, testTableName),
@@ -239,7 +239,7 @@ LIMIT 0, 10
 FROM %s
 WHERE DomainID = 'bfd5c907-f899-4baf-a7b2-2ab85e623ebd'
 AND IsDeleted = false
-AND CloseStatus < 0 and (JSON_MATCH(Attr, '"$.CustomKeywordField"=''keywordCustomized''') or JSON_MATCH(Attr, '"$.CustomKeywordField[*]"=''keywordCustomized''')) and JSON_MATCH(Attr, '"$.CustomIntField"=''10''') and (JSON_MATCH(Attr, '"$.CustomStringField" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'String field is for text*'))
+AND CloseStatus < 0 and (JSON_MATCH(Attr, '"$.CustomKeywordField"=''keywordCustomized''') or JSON_MATCH(Attr, '"$.CustomKeywordField[*]"=''keywordCustomized''')) and (JSON_MATCH(Attr, '"$.CustomIntField" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) <= 10) and (JSON_MATCH(Attr, '"$.CustomStringField" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'String field is for text*'))
 Order by DomainID Desc
 LIMIT 11, 10
 `, testTableName),

common/pinot/pinotQueryValidator.go

@@ -91,24 +91,58 @@ func (qv *VisibilityQueryValidator) validateWhereExpr(expr sqlparser.Expr) (stri
 	if expr == nil {
 		return "", nil
 	}
-	buf := sqlparser.NewTrackedBuffer(nil)
-
 	switch expr := expr.(type) {
 	case *sqlparser.AndExpr, *sqlparser.OrExpr:
 		return qv.validateAndOrExpr(expr)
 	case *sqlparser.ComparisonExpr:
 		return qv.validateComparisonExpr(expr)
 	case *sqlparser.RangeCond:
-		expr.Format(buf)
-		return buf.String(), nil
-		//return qv.validateRangeExpr(expr)
+		return qv.validateRangeExpr(expr)
 	case *sqlparser.ParenExpr:
 		return qv.validateWhereExpr(expr.Expr)
 	default:
 		return "", errors.New("invalid where clause")
 	}
 }
 
+// for "between...and..." only
+// <, >, >=, <= are included in validateComparisonExpr()
+func (qv *VisibilityQueryValidator) validateRangeExpr(expr sqlparser.Expr) (string, error) {
+	buf := sqlparser.NewTrackedBuffer(nil)
+	rangeCond := expr.(*sqlparser.RangeCond)
+	colName, ok := rangeCond.Left.(*sqlparser.ColName)
+	if !ok {
+		return "", errors.New("invalid range expression: fail to get colname")
+	}
+	colNameStr := colName.Name.String()
+
+	if !qv.isValidSearchAttributes(colNameStr) {
+		return "", fmt.Errorf("invalid search attribute %q", colNameStr)
+	}
+
+	if definition.IsSystemIndexedKey(colNameStr) {
+		expr.Format(buf)
+		return buf.String(), nil
+	}
+
+	//lowerBound, ok := rangeCond.From.(*sqlparser.ColName)
+	lowerBound, ok := rangeCond.From.(*sqlparser.SQLVal)
+	if !ok {
+		return "", errors.New("invalid range expression: fail to get lowerbound")
+	}
+	lowerBoundString := string(lowerBound.Val)
+
+	upperBound, ok := rangeCond.To.(*sqlparser.SQLVal)
+	if !ok {
+		return "", errors.New("invalid range expression: fail to get upperbound")
+	}
+	upperBoundString := string(upperBound.Val)
+
+	return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\" is not null') "+
+		"AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.%s') AS INT) >= %s "+
+		"AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.%s') AS INT) <= %s)", colNameStr, colNameStr, lowerBoundString, colNameStr, upperBoundString), nil
+}
+
 func (qv *VisibilityQueryValidator) validateAndOrExpr(expr sqlparser.Expr) (string, error) {
 	var leftExpr sqlparser.Expr
 	var rightExpr sqlparser.Expr
@@ -248,24 +282,48 @@ func (qv *VisibilityQueryValidator) processCustomKey(expr sqlparser.Expr) (strin
 	// get the value type
 	indexValType := common.ConvertIndexedValueTypeToInternalType(valType, log.NewNoop())
 
-	// Case2-1: when it is string, need partial match
-	if indexValType == types.IndexedValueTypeString {
-		// change to like statement for partial match
-		comparisonExpr.Operator = sqlparser.LikeStr
-		comparisonExpr.Right = &sqlparser.SQLVal{
-			Type: sqlparser.StrVal,
-			Val:  []byte("%" + colValStr + "%"),
-		}
-		//return fmt.Sprintf("JSON_EXTRACT_SCALAR(Attr, '$.%s', 'STRING') LIKE '%%%s%%'", colNameStr, colValStr), nil
-		return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\" is not null') "+
-			"AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.%s', 'string'), '%s*'))", colNameStr, colNameStr, colValStr), nil
+	operator := comparisonExpr.Operator
+
+	switch indexValType {
+	case types.IndexedValueTypeString:
+		return processCustomString(comparisonExpr, colNameStr, colValStr), nil
+	case types.IndexedValueTypeKeyword:
+		return processCustomKeyword(operator, colNameStr, colValStr), nil
+	case types.IndexedValueTypeDatetime:
+		return processCustomNum(operator, colNameStr, colValStr, "BIGINT"), nil
+	case types.IndexedValueTypeDouble:
+		return processCustomNum(operator, colNameStr, colValStr, "DOUBLE"), nil
+	case types.IndexedValueTypeInt:
+		return processCustomNum(operator, colNameStr, colValStr, "INT"), nil
+	default:
+		return processEqual(colNameStr, colValStr), nil
 	}
-	// case2-2: otherwise, exact match
-	// case2-2-1: if it is keyword, need to deal with a situation when value is an array
-	if indexValType == types.IndexedValueTypeKeyword {
-		return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\"=''%s''') or JSON_MATCH(Attr, '\"$.%s[*]\"=''%s'''))",
-			colNameStr, colValStr, colNameStr, colValStr), nil
+}
+
+func processCustomNum(operator string, colNameStr string, colValStr string, valType string) string {
+	if operator == sqlparser.EqualStr {
+		return processEqual(colNameStr, colValStr)
+	}
+	return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\" is not null') "+
+		"AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.%s') AS %s) %s %s)", colNameStr, colNameStr, valType, operator, colValStr)
+}
+
+func processEqual(colNameStr string, colValStr string) string {
+	return fmt.Sprintf("JSON_MATCH(Attr, '\"$.%s\"=''%s''')", colNameStr, colValStr)
+}
+
+func processCustomKeyword(operator string, colNameStr string, colValStr string) string {
+	return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\"%s''%s''') or JSON_MATCH(Attr, '\"$.%s[*]\"%s''%s'''))",
+		colNameStr, operator, colValStr, colNameStr, operator, colValStr)
+}
+
+func processCustomString(comparisonExpr *sqlparser.ComparisonExpr, colNameStr string, colValStr string) string {
+	// change to like statement for partial match
+	comparisonExpr.Operator = sqlparser.LikeStr
+	comparisonExpr.Right = &sqlparser.SQLVal{
+		Type: sqlparser.StrVal,
+		Val:  []byte("%" + colValStr + "%"),
 	}
-	// case2-2-2: other cases:
-	return fmt.Sprintf("JSON_MATCH(Attr, '\"$.%s\"=''%s''')", colNameStr, colValStr), nil
+	return fmt.Sprintf("(JSON_MATCH(Attr, '\"$.%s\" is not null') "+
+		"AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.%s', 'string'), '%s*'))", colNameStr, colNameStr, colValStr)
 }

common/pinot/pinotQueryValidator_test.go

@@ -59,7 +59,7 @@ func TestValidateQuery(t *testing.T) {
 		},
 		"Case6-1: complex query I: with parenthesis": {
 			query:     "(CustomStringField = 'custom and custom2 or custom3 order by') or CustomIntField between 1 and 10",
-			validated: "((JSON_MATCH(Attr, '\"$.CustomStringField\" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or CustomIntField between 1 and 10)",
+			validated: "((JSON_MATCH(Attr, '\"$.CustomStringField\" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or (JSON_MATCH(Attr, '\"$.CustomIntField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) >= 1 AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) <= 10))",
 		},
 		"Case6-2: complex query II: with only system keys": {
 			query:     "DomainID = 'd-id' and (RunID = 'run-id' or WorkflowID = 'wid')",
@@ -71,7 +71,7 @@ func TestValidateQuery(t *testing.T) {
 		},
 		"Case6-4: complex query IV": {
 			query:     "WorkflowID = 'wid' and (CustomStringField = 'custom and custom2 or custom3 order by' or CustomIntField between 1 and 10)",
-			validated: "WorkflowID = 'wid' and ((JSON_MATCH(Attr, '\"$.CustomStringField\" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or CustomIntField between 1 and 10)",
+			validated: "WorkflowID = 'wid' and ((JSON_MATCH(Attr, '\"$.CustomStringField\" is not null') AND REGEXP_LIKE(JSON_EXTRACT_SCALAR(Attr, '$.CustomStringField', 'string'), 'custom and custom2 or custom3 order by*')) or (JSON_MATCH(Attr, '\"$.CustomIntField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) >= 1 AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) <= 10))",
 		},
 		"Case7: invalid sql query": {
 			query: "Invalid SQL",
@@ -113,6 +113,43 @@ func TestValidateQuery(t *testing.T) {
 			query:     "CustomIntField = 1 or CustomIntField = 2",
 			validated: "(JSON_MATCH(Attr, '\"$.CustomIntField\"=''1''') or JSON_MATCH(Attr, '\"$.CustomIntField\"=''2'''))",
 		},
+		"Case14-1: range query: custom filed": {
+			query:     "CustomIntField BETWEEN 1 AND 2",
+			validated: "(JSON_MATCH(Attr, '\"$.CustomIntField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) >= 1 AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomIntField') AS INT) <= 2)",
+		},
+		"Case14-2: range query: system filed": {
+			query:     "NumClusters BETWEEN 1 AND 2",
+			validated: "NumClusters between 1 and 2",
+		},
+		"Case15-1: custom date attribute less than": {
+			query:     "CustomDatetimeField < 1697754674",
+			validated: "(JSON_MATCH(Attr, '\"$.CustomDatetimeField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomDatetimeField') AS BIGINT) < 1697754674)",
+		},
+		"Case15-2: custom date attribute greater than or equal to": {
+			query:     "CustomDatetimeField >= 1697754674",
+			validated: "(JSON_MATCH(Attr, '\"$.CustomDatetimeField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomDatetimeField') AS BIGINT) >= 1697754674)",
+		},
+		"Case15-3: system date attribute greater than or equal to": {
+			query:     "StartTime >= 1697754674",
+			validated: "StartTime >= 1697754674",
+		},
+		"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)",
+		},
+		"Case16-2: custom double attribute greater than or equal to": {
+			query:     "CustomDoubleField >= 0",
+			validated: "(JSON_MATCH(Attr, '\"$.CustomDoubleField\" is not null') AND CAST(JSON_EXTRACT_SCALAR(Attr, '$.CustomDoubleField') AS DOUBLE) >= 0)",
+		},
+		"Case17: custom keyword attribute greater than or equal to. Will return error run time": {
+			query:     "CustomKeywordField < 0",
+			validated: "(JSON_MATCH(Attr, '\"$.CustomKeywordField\"<''0''') or JSON_MATCH(Attr, '\"$.CustomKeywordField[*]\"<''0'''))",
+		},
+		// TODO
+		"Case18: custom int order by. Will have errors at run time. Doesn't support for now": {
+			query:     "CustomIntField = 0 order by CustomIntField desc",
+			validated: "JSON_MATCH(Attr, '\"$.CustomIntField\"=''0''') order by CustomIntField desc",
+		},
 	}
 
 	for name, test := range tests {