Add generic ES query building utilities

common/elasticsearch/client_v6.go

@@ -34,6 +34,7 @@ import (
 
 	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/config"
+	"github.com/uber/cadence/common/elasticsearch/query"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
 	p "github.com/uber/cadence/common/persistence"
@@ -53,7 +54,7 @@ type (
 	// searchParametersV6 holds all required and optional parameters for executing a search
 	searchParametersV6 struct {
 		Index       string
-		Query       elastic.Query
+		Query       query.Query
 		From        int
 		PageSize    int
 		Sorter      []elastic.Sorter
@@ -126,9 +127,9 @@ func (c *elasticV6) Search(ctx context.Context, request *SearchRequest) (*p.Inte
 		return nil, err
 	}
 
-	var matchQuery *elastic.MatchQuery
+	var matchQuery *query.MatchQuery
 	if request.MatchQuery != nil {
-		matchQuery = elastic.NewMatchQuery(request.MatchQuery.Name, request.MatchQuery.Text)
+		matchQuery = query.NewMatchQuery(request.MatchQuery.Name, request.MatchQuery.Text)
 	}
 
 	searchResult, err := c.getSearchResult(
@@ -214,10 +215,10 @@ func (c *elasticV6) SearchForOneClosedExecution(
 	request *p.InternalGetClosedWorkflowExecutionRequest,
 ) (*p.InternalGetClosedWorkflowExecutionResponse, error) {
 
-	matchDomainQuery := elastic.NewMatchQuery(DomainID, request.DomainUUID)
-	existClosedStatusQuery := elastic.NewExistsQuery(CloseStatus)
-	matchWorkflowIDQuery := elastic.NewMatchQuery(WorkflowID, request.Execution.GetWorkflowID())
-	boolQuery := elastic.NewBoolQuery().Must(matchDomainQuery).Must(existClosedStatusQuery).Must(matchWorkflowIDQuery)
+	matchDomainQuery := query.NewMatchQuery(DomainID, request.DomainUUID)
+	existClosedStatusQuery := query.NewExistsQuery(CloseStatus)
+	matchWorkflowIDQuery := query.NewMatchQuery(WorkflowID, request.Execution.GetWorkflowID())
+	boolQuery := query.NewBoolQuery().Must(matchDomainQuery).Must(existClosedStatusQuery).Must(matchWorkflowIDQuery)
 	rid := request.Execution.GetRunID()
 	if rid != "" {
 		matchRunIDQuery := elastic.NewMatchQuery(RunID, rid)
@@ -412,13 +413,13 @@ func (c *elasticV6) getSearchResult(
 	ctx context.Context,
 	index string,
 	request *p.InternalListWorkflowExecutionsRequest,
-	matchQuery *elastic.MatchQuery,
+	matchQuery *query.MatchQuery,
 	isOpen bool,
 	token *ElasticVisibilityPageToken,
 ) (*elastic.SearchResult, error) {
 
 	// always match domain id
-	boolQuery := elastic.NewBoolQuery().Must(elastic.NewMatchQuery(DomainID, request.DomainUUID))
+	boolQuery := query.NewBoolQuery().Must(query.NewMatchQuery(DomainID, request.DomainUUID))
 
 	if matchQuery != nil {
 		boolQuery = boolQuery.Must(matchQuery)
@@ -434,7 +435,7 @@ func (c *elasticV6) getSearchResult(
 	}
 
 	var rangeQueryField string
-	existClosedStatusQuery := elastic.NewExistsQuery(CloseStatus)
+	existClosedStatusQuery := query.NewExistsQuery(CloseStatus)
 	if isOpen {
 		rangeQueryField = StartTime
 		boolQuery = boolQuery.MustNot(existClosedStatusQuery)
@@ -445,7 +446,7 @@ func (c *elasticV6) getSearchResult(
 
 	earliestTimeStr := strconv.FormatInt(request.EarliestTime.UnixNano()-oneMicroSecondInNano, 10)
 	latestTimeStr := strconv.FormatInt(request.LatestTime.UnixNano()+oneMicroSecondInNano, 10)
-	rangeQuery := elastic.NewRangeQuery(rangeQueryField).Gte(earliestTimeStr).Lte(latestTimeStr)
+	rangeQuery := query.NewRangeQuery(rangeQueryField).Gte(earliestTimeStr).Lte(latestTimeStr)
 	boolQuery = boolQuery.Filter(rangeQuery)
 
 	params := &searchParametersV6{
@@ -455,8 +456,8 @@ func (c *elasticV6) getSearchResult(
 		PageSize: request.PageSize,
 	}
 
-	params.Sorter = append(params.Sorter, elastic.NewFieldSort(rangeQueryField).Desc())
-	params.Sorter = append(params.Sorter, elastic.NewFieldSort(RunID).Desc())
+	params.Sorter = append(params.Sorter, query.NewFieldSort(rangeQueryField).Desc())
+	params.Sorter = append(params.Sorter, query.NewFieldSort(RunID).Desc())
 
 	if ShouldSearchAfter(token) {
 		params.SearchAfter = []interface{}{token.SortValue, token.TieBreaker}

common/elasticsearch/client_v7.go

@@ -34,6 +34,7 @@ import (
 
 	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/config"
+	"github.com/uber/cadence/common/elasticsearch/query"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
 	p "github.com/uber/cadence/common/persistence"
@@ -53,7 +54,7 @@ type (
 	// searchParametersV7 holds all required and optional parameters for executing a search
 	searchParametersV7 struct {
 		Index       string
-		Query       elastic.Query
+		Query       query.Query
 		From        int
 		PageSize    int
 		Sorter      []elastic.Sorter
@@ -126,9 +127,9 @@ func (c *elasticV7) Search(ctx context.Context, request *SearchRequest) (*p.Inte
 		return nil, err
 	}
 
-	var matchQuery *elastic.MatchQuery
+	var matchQuery *query.MatchQuery
 	if request.MatchQuery != nil {
-		matchQuery = elastic.NewMatchQuery(request.MatchQuery.Name, request.MatchQuery.Text)
+		matchQuery = query.NewMatchQuery(request.MatchQuery.Name, request.MatchQuery.Text)
 	}
 
 	searchResult, err := c.getSearchResult(
@@ -214,13 +215,13 @@ func (c *elasticV7) SearchForOneClosedExecution(
 	request *p.InternalGetClosedWorkflowExecutionRequest,
 ) (*p.InternalGetClosedWorkflowExecutionResponse, error) {
 
-	matchDomainQuery := elastic.NewMatchQuery(DomainID, request.DomainUUID)
-	existClosedStatusQuery := elastic.NewExistsQuery(CloseStatus)
-	matchWorkflowIDQuery := elastic.NewMatchQuery(WorkflowID, request.Execution.GetWorkflowID())
-	boolQuery := elastic.NewBoolQuery().Must(matchDomainQuery).Must(existClosedStatusQuery).Must(matchWorkflowIDQuery)
+	matchDomainQuery := query.NewMatchQuery(DomainID, request.DomainUUID)
+	existClosedStatusQuery := query.NewExistsQuery(CloseStatus)
+	matchWorkflowIDQuery := query.NewMatchQuery(WorkflowID, request.Execution.GetWorkflowID())
+	boolQuery := query.NewBoolQuery().Must(matchDomainQuery).Must(existClosedStatusQuery).Must(matchWorkflowIDQuery)
 	rid := request.Execution.GetRunID()
 	if rid != "" {
-		matchRunIDQuery := elastic.NewMatchQuery(RunID, rid)
+		matchRunIDQuery := query.NewMatchQuery(RunID, rid)
 		boolQuery = boolQuery.Must(matchRunIDQuery)
 	}
 
@@ -415,14 +416,13 @@ func (c *elasticV7) getSearchResult(
 	ctx context.Context,
 	index string,
 	request *p.InternalListWorkflowExecutionsRequest,
-	matchQuery *elastic.MatchQuery,
+	matchQuery *query.MatchQuery,
 	isOpen bool,
 	token *ElasticVisibilityPageToken,
 ) (*elastic.SearchResult, error) {
 
 	// always match domain id
-	boolQuery := elastic.NewBoolQuery().Must(elastic.NewMatchQuery(DomainID, request.DomainUUID))
-
+	boolQuery := query.NewBoolQuery().Must(query.NewMatchQuery(DomainID, request.DomainUUID))
 	if matchQuery != nil {
 		boolQuery = boolQuery.Must(matchQuery)
 	}
@@ -457,9 +457,8 @@ func (c *elasticV7) getSearchResult(
 		From:     token.From,
 		PageSize: request.PageSize,
 	}
-
-	params.Sorter = append(params.Sorter, elastic.NewFieldSort(rangeQueryField).Desc())
-	params.Sorter = append(params.Sorter, elastic.NewFieldSort(RunID).Desc())
+	params.Sorter = append(params.Sorter, query.NewFieldSort(rangeQueryField).Desc())
+	params.Sorter = append(params.Sorter, query.NewFieldSort(RunID).Desc())
 
 	if ShouldSearchAfter(token) {
 		params.SearchAfter = []interface{}{token.SortValue, token.TieBreaker}

common/elasticsearch/query/bool_query.go

@@ -0,0 +1,124 @@
+// The MIT License (MIT)
+
+// Copyright (c) 2017-2020 Uber Technologies Inc.
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package query
+
+type Query interface {
+	// Source returns the JSON-serializable query request.
+	Source() (interface{}, error)
+}
+
+type BoolQuery struct {
+	mustClauses    []Query
+	mustNotClauses []Query
+	filterClauses  []Query
+}
+
+// Creates a new bool query.
+func NewBoolQuery() *BoolQuery {
+	return &BoolQuery{
+		mustClauses:    make([]Query, 0),
+		mustNotClauses: make([]Query, 0),
+		filterClauses:  make([]Query, 0),
+	}
+}
+
+func (q *BoolQuery) Must(queries ...Query) *BoolQuery {
+	q.mustClauses = append(q.mustClauses, queries...)
+	return q
+}
+
+func (q *BoolQuery) MustNot(queries ...Query) *BoolQuery {
+	q.mustNotClauses = append(q.mustNotClauses, queries...)
+	return q
+}
+
+func (q *BoolQuery) Filter(filters ...Query) *BoolQuery {
+	q.filterClauses = append(q.filterClauses, filters...)
+	return q
+}
+
+func (q *BoolQuery) Source() (interface{}, error) {
+	query := make(map[string]interface{})
+
+	boolClause := make(map[string]interface{})
+	query["bool"] = boolClause
+
+	// must
+	if len(q.mustClauses) == 1 {
+		src, err := q.mustClauses[0].Source()
+		if err != nil {
+			return nil, err
+		}
+		boolClause["must"] = src
+	} else if len(q.mustClauses) > 1 {
+		var clauses []interface{}
+		for _, subQuery := range q.mustClauses {
+			src, err := subQuery.Source()
+			if err != nil {
+				return nil, err
+			}
+			clauses = append(clauses, src)
+		}
+		boolClause["must"] = clauses
+	}
+
+	// must_not
+	if len(q.mustNotClauses) == 1 {
+		src, err := q.mustNotClauses[0].Source()
+		if err != nil {
+			return nil, err
+		}
+		boolClause["must_not"] = src
+	} else if len(q.mustNotClauses) > 1 {
+		var clauses []interface{}
+		for _, subQuery := range q.mustNotClauses {
+			src, err := subQuery.Source()
+			if err != nil {
+				return nil, err
+			}
+			clauses = append(clauses, src)
+		}
+		boolClause["must_not"] = clauses
+	}
+
+	// filter
+	if len(q.filterClauses) == 1 {
+		src, err := q.filterClauses[0].Source()
+		if err != nil {
+			return nil, err
+		}
+		boolClause["filter"] = src
+	} else if len(q.filterClauses) > 1 {
+		var clauses []interface{}
+		for _, subQuery := range q.filterClauses {
+			src, err := subQuery.Source()
+			if err != nil {
+				return nil, err
+			}
+			clauses = append(clauses, src)
+		}
+		boolClause["filter"] = clauses
+	}
+
+	return query, nil
+}

common/elasticsearch/query/bool_query_test.go

@@ -0,0 +1,46 @@
+// The MIT License (MIT)
+
+// Copyright (c) 2017-2020 Uber Technologies Inc.
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package query
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestBoolQuery(t *testing.T) {
+	q := NewBoolQuery()
+	q = q.MustNot(NewRangeQuery("age").From(10).To(20))
+	src, err := q.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"bool":{"must_not":{"range":{"age":{"from":10,"include_lower":true,"include_upper":true,"to":20}}}}}`
+	if got != expected {
+		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+	}
+}