Bug fix: custom query in Count doesn't return correct result (cadence…

…-workflow#6179)

* return error instead of log

* CountWorkflowExecutions

common/persistence/pinot/pinot_visibility_store.go

@@ -537,7 +537,11 @@ func (v *pinotVisibilityStore) ScanWorkflowExecutions(ctx context.Context, reque
 }
 
 func (v *pinotVisibilityStore) CountWorkflowExecutions(ctx context.Context, request *p.CountWorkflowExecutionsRequest) (*p.CountWorkflowExecutionsResponse, error) {
-	query := v.getCountWorkflowExecutionsQuery(v.pinotClient.GetTableName(), request)
+	query, err := v.getCountWorkflowExecutionsQuery(v.pinotClient.GetTableName(), request)
+	if err != nil {
+		v.logger.Error(fmt.Sprintf("failed to build count workflow executions query %v", err))
+		return nil, err
+	}
 
 	resp, err := v.pinotClient.CountByQuery(query)
 	if err != nil {
@@ -822,9 +826,9 @@ func (f *PinotQueryFilter) addTimeRange(obj string, earliest interface{}, latest
 	f.string += fmt.Sprintf("%s BETWEEN %v AND %v\n", obj, earliest, latest)
 }
 
-func (v *pinotVisibilityStore) getCountWorkflowExecutionsQuery(tableName string, request *p.CountWorkflowExecutionsRequest) string {
+func (v *pinotVisibilityStore) getCountWorkflowExecutionsQuery(tableName string, request *p.CountWorkflowExecutionsRequest) (string, error) {
 	if request == nil {
-		return ""
+		return "", nil
 	}
 
 	query := NewPinotCountQuery(tableName)
@@ -837,23 +841,23 @@ func (v *pinotVisibilityStore) getCountWorkflowExecutionsQuery(tableName string,
 
 	// if customized query is empty, directly return
 	if requestQuery == "" {
-		return query.String()
+		return query.String(), nil
 	}
 
 	requestQuery = filterPrefix(requestQuery)
 
 	comparExpr, _ := parseOrderBy(requestQuery)
 	comparExpr, err := v.pinotQueryValidator.ValidateQuery(comparExpr)
 	if err != nil {
-		v.logger.Error(fmt.Sprintf("pinot query validator error: %s", err))
+		return "", fmt.Errorf("pinot query validator error: %w, query: %s", err, request.Query)
 	}
 
 	comparExpr = filterPrefix(comparExpr)
 	if comparExpr != "" {
 		query.filters.addQuery(comparExpr)
 	}
 
-	return query.String()
+	return query.String(), nil
 }
 
 func (v *pinotVisibilityStore) getListWorkflowExecutionsByQueryQuery(tableName string, request *p.ListWorkflowExecutionsByQueryRequest) (string, error) {

common/persistence/pinot/pinot_visibility_store_test.go

@@ -1104,6 +1104,14 @@ func TestCountWorkflowExecutions(t *testing.T) {
 			},
 			expectedError: nil,
 		},
+		"Case3: query error case": {
+			request:      &p.CountWorkflowExecutionsRequest{Domain: testDomain, DomainUUID: testDomainID, Query: "CustomKeywordField = missing"},
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("pinot query validator error: invalid comparison expression, right, query: CustomKeywordField = missing"),
+		},
 	}
 
 	for name, test := range tests {
@@ -1192,6 +1200,15 @@ AND WorkflowID = 'wfid'
 			expectedRes:   expectEmptyQueryResult,
 			expectedError: nil,
 		},
+		"Case3: custom attr is missing case": {
+			request: &p.CountWorkflowExecutionsRequest{
+				DomainUUID: testDomainID,
+				Domain:     testDomain,
+				Query:      "CustomKeywordField = missing",
+			},
+			expectedRes:   "",
+			expectedError: fmt.Errorf("pinot query validator error: invalid comparison expression, right, query: CustomKeywordField = missing"),
+		},
 	}
 
 	for name, test := range tests {
@@ -1205,8 +1222,11 @@ AND WorkflowID = 'wfid'
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			res := visibilityStore.getCountWorkflowExecutionsQuery(testTableName, test.request)
+			res, err := visibilityStore.getCountWorkflowExecutionsQuery(testTableName, test.request)
 			assert.Equal(t, test.expectedRes, res)
+			if test.expectedError != nil {
+				assert.Equal(t, test.expectedError.Error(), err.Error())
+			}
 		})
 	}
 }

common/pinot/pinotQueryValidator_test.go

@@ -90,6 +90,10 @@ func TestValidateQuery(t *testing.T) {
 			query:     "CloseTime != missing",
 			validated: "CloseTime != -1",
 		},
+		"Case8-3: query with custom attr": {
+			query: "CustomKeywordField = missing",
+			err:   "invalid comparison expression, right",
+		},
 		"Case9: invalid where expression": {
 			query: "InvalidWhereExpr",
 			err:   "invalid where clause",