[esutil.Client] fix routing support for MultiSearch to specify routing for each individual search (#95)

Author: mrparkers

go/v1beta1/storage/esutil/client.go

@@ -63,7 +63,6 @@ type BulkRequestItem struct {
 
 type MultiSearchRequest struct {
 	Index    string
-	Routing  string
 	Searches []*EsSearch
 }
 
@@ -297,6 +296,10 @@ func (c *client) Search(ctx context.Context, request *SearchRequest) (*SearchRes
 		c.esClient.Search.WithContext(ctx),
 	}
 
+	if body.Routing != "" {
+		searchOptions = append(searchOptions, c.esClient.Search.WithRouting(body.Routing))
+	}
+
 	var (
 		searchFrom int
 		pitId      string
@@ -388,19 +391,30 @@ func (c *client) Search(ctx context.Context, request *SearchRequest) (*SearchRes
 func (c *client) MultiSearch(ctx context.Context, request *MultiSearchRequest) (*EsMultiSearchResponse, error) {
 	log := c.logger.Named("MultiSearch")
 
-	searchMetadata, _ := json.Marshal(&EsMultiSearchQueryFragment{
-		Index:   request.Index,
-		Routing: request.Routing,
+	searchMetadataWithoutRouting, _ := json.Marshal(&EsMultiSearchQueryFragment{
+		Index: request.Index,
 	})
-	searchMetadata = append(searchMetadata, '\n')
+	searchMetadataWithoutRouting = append(searchMetadataWithoutRouting, '\n')
 
 	var searchRequestBody bytes.Buffer
 	for _, search := range request.Searches {
 		data, _ := json.Marshal(search)
 		dataBytes := append(data, '\n')
 
-		searchRequestBody.Grow(len(searchMetadata) + len(dataBytes))
-		searchRequestBody.Write(searchMetadata)
+		if search.Routing != "" {
+			searchMetadataWithRouting, _ := json.Marshal(&EsMultiSearchQueryFragment{
+				Index:   request.Index,
+				Routing: search.Routing,
+			})
+			searchMetadataWithRouting = append(searchMetadataWithRouting, '\n')
+
+			searchRequestBody.Grow(len(searchMetadataWithRouting) + len(dataBytes))
+			searchRequestBody.Write(searchMetadataWithRouting)
+		} else {
+			searchRequestBody.Grow(len(searchMetadataWithoutRouting) + len(dataBytes))
+			searchRequestBody.Write(searchMetadataWithoutRouting)
+		}
+
 		searchRequestBody.Write(dataBytes)
 	}
 

go/v1beta1/storage/esutil/client_test.go

@@ -566,6 +566,19 @@ var _ = Describe("elasticsearch client", func() {
 
 				Expect(searchRequest).To(BeEquivalentTo(expectedSearch))
 			})
+
+			When("routing is specified", func() {
+				var expectedRouting string
+
+				BeforeEach(func() {
+					expectedRouting = fake.LetterN(10)
+					expectedSearchRequest.Search.Routing = expectedRouting
+				})
+
+				It("should specify the routing value during the search", func() {
+					Expect(transport.ReceivedHttpRequests[0].URL.Query().Get("routing")).To(Equal(expectedRouting))
+				})
+			})
 		})
 
 		When("the search operation fails", func() {
@@ -802,15 +815,20 @@ var _ = Describe("elasticsearch client", func() {
 			})
 		})
 
-		When("routing is specified", func() {
-			var expectedRouting string
+		When("routing is specified for a search", func() {
+			var (
+				expectedRouting   string
+				randomSearchIndex int
+			)
 
 			BeforeEach(func() {
 				expectedRouting = fake.LetterN(10)
-				expectedMultiSearchRequest.Routing = expectedRouting
+				randomSearchIndex = fake.Number(0, len(expectedSearches)-1)
+
+				expectedMultiSearchRequest.Searches[randomSearchIndex].Routing = expectedRouting
 			})
 
-			It("should include the routing value in each search header", func() {
+			It("should include the routing value in the specified search header", func() {
 				var expectedPayloads []interface{}
 
 				for i := 0; i < len(expectedSearches); i++ {
@@ -820,13 +838,22 @@ var _ = Describe("elasticsearch client", func() {
 				parseNDJSONRequestBody(transport.ReceivedHttpRequests[0].Body, expectedPayloads)
 
 				for i, payload := range expectedPayloads {
-					if i%2 == 0 { // search metadata
+					modulus := i % 2
+					index := (i - modulus) / 2
+
+					if modulus == 0 { // search metadata
 						metadata := payload.(*EsMultiSearchQueryFragment)
 						Expect(metadata.Index).To(Equal(expectedIndex))
-						Expect(metadata.Routing).To(Equal(expectedRouting))
+
+						if index == randomSearchIndex {
+							Expect(metadata.Routing).To(Equal(expectedRouting))
+						} else {
+							Expect(metadata.Routing).To(BeEmpty())
+						}
 					} else { // search
 						search := payload.(*EsSearch)
-						expectedSearch := expectedSearches[(i-1)/2]
+						expectedSearch := expectedSearches[index]
+						expectedSearch.Routing = ""
 
 						Expect(search).To(Equal(expectedSearch))
 					}

go/v1beta1/storage/esutil/types.go

@@ -54,6 +54,7 @@ type EsSearch struct {
 	Sort     map[string]EsSortOrder `json:"sort,omitempty"`
 	Collapse *EsSearchCollapse      `json:"collapse,omitempty"`
 	Pit      *EsSearchPit           `json:"pit,omitempty"`
+	Routing  string                 `json:"-"`
 }
 
 type EsSortOrder string