Simplify multisearch

Author: ddelemeny

pkg/quickwit/client/client.go

@@ -31,6 +31,7 @@ type DatasourceInfo struct {
 	ShouldInit                 bool
 }
 
+// TODO: Move ConfiguredFields closer to handlers, the client layer doesn't need this stuff
 type ConfiguredFields struct {
 	TimeField        string
 	TimeOutputFormat string
@@ -43,61 +44,23 @@ type Client interface {
 	ExecuteMultisearch(r []*SearchRequest) (*MultiSearchResponse, error)
 }
 
+var logger = log.New()
+
 // NewClient creates a new Quickwit client
 var NewClient = func(ctx context.Context, ds *DatasourceInfo) (Client, error) {
-	logger := log.New()
 	logger.Debug("Creating new client", "index", ds.Database)
 
 	return &baseClientImpl{
-		logger: logger,
-		ctx:    ctx,
-		ds:     ds,
-		index:  ds.Database,
+		ctx:   ctx,
+		ds:    ds,
+		index: ds.Database,
 	}, nil
 }
 
 type baseClientImpl struct {
-	ctx    context.Context
-	ds     *DatasourceInfo
-	index  string
-	logger log.Logger
-}
-
-type multiRequest struct {
-	header   map[string]interface{}
-	body     interface{}
-	interval time.Duration
-}
-
-func (c *baseClientImpl) encodeBatchRequests(requests []*multiRequest) ([]byte, error) {
-	c.logger.Debug("Encoding batch requests to json", "batch requests", len(requests))
-	start := time.Now()
-
-	payload := bytes.Buffer{}
-	for _, r := range requests {
-		reqHeader, err := json.Marshal(r.header)
-		if err != nil {
-			return nil, err
-		}
-		payload.WriteString(string(reqHeader) + "\n")
-
-		reqBody, err := json.Marshal(r.body)
-
-		if err != nil {
-			return nil, err
-		}
-
-		body := string(reqBody)
-		body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.interval.Milliseconds(), 10))
-		body = strings.ReplaceAll(body, "$__interval", r.interval.String())
-
-		payload.WriteString(body + "\n")
-	}
-
-	elapsed := time.Since(start)
-	c.logger.Debug("Encoded batch requests to json", "took", elapsed)
-
-	return payload.Bytes(), nil
+	ctx   context.Context
+	ds    *DatasourceInfo
+	index string
 }
 
 func (c *baseClientImpl) makeRequest(method, uriPath, uriQuery string, body []byte) (*http.Request, error) {
@@ -122,9 +85,7 @@ func (c *baseClientImpl) makeRequest(method, uriPath, uriQuery string, body []by
 }
 
 func (c *baseClientImpl) ExecuteMultisearch(requests []*SearchRequest) (*MultiSearchResponse, error) {
-	c.logger.Debug("Executing multisearch", "search requests", requests)
-
-	req, err := c.createMultiSearchRequests(requests)
+	req, err := c.createMultiSearchRequest(requests, c.index)
 	if err != nil {
 		return nil, err
 	}
@@ -135,11 +96,11 @@ func (c *baseClientImpl) ExecuteMultisearch(requests []*SearchRequest) (*MultiSe
 	}
 	defer func() {
 		if err := res.Body.Close(); err != nil {
-			c.logger.Warn("Failed to close response body", "err", err)
+			logger.Warn("Failed to close response body", "err", err)
 		}
 	}()
 
-	c.logger.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)
+	logger.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)
 
 	if res.StatusCode >= 400 {
 		qe := QuickwitQueryError{
@@ -151,12 +112,12 @@ func (c *baseClientImpl) ExecuteMultisearch(requests []*SearchRequest) (*MultiSe
 		}
 
 		errorPayload, _ := json.Marshal(qe)
-		c.logger.Error(string(errorPayload))
+		logger.Error(string(errorPayload))
 		return nil, fmt.Errorf(string(errorPayload))
 	}
 
 	start := time.Now()
-	c.logger.Debug("Decoding multisearch json response")
+	logger.Debug("Decoding multisearch json response")
 
 	var msr MultiSearchResponse
 	dec := json.NewDecoder(res.Body)
@@ -166,44 +127,53 @@ func (c *baseClientImpl) ExecuteMultisearch(requests []*SearchRequest) (*MultiSe
 	}
 
 	elapsed := time.Since(start)
-	c.logger.Debug("Decoded multisearch json response", "took", elapsed)
+	logger.Debug("Decoded multisearch json response", "took", elapsed)
 
 	return &msr, nil
 }
 
-func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) (*http.Request, error) {
-	multiRequests := []*multiRequest{}
-
-	for _, searchReq := range searchRequests {
-		mr := multiRequest{
-			header: map[string]interface{}{
-				"ignore_unavailable": true,
-				"index":              strings.Split(c.index, ","),
-			},
-			body:     searchReq,
-			interval: searchReq.Interval,
+func (c *baseClientImpl) makeMultiSearchPayload(searchRequests []*SearchRequest, index string) ([]byte, error) {
+	// Format, marshall and interpolate
+	payload := bytes.Buffer{}
+	for _, r := range searchRequests {
+		header := map[string]interface{}{
+			"ignore_unavailable": true,
+			"index":              strings.Split(index, ","),
+		}
+		reqHeader, err := json.Marshal(header)
+		if err != nil {
+			return nil, err
 		}
+		payload.WriteString(string(reqHeader) + "\n")
+
+		reqBody, err := json.Marshal(r)
 
-		multiRequests = append(multiRequests, &mr)
+		if err != nil {
+			return nil, err
+		}
+
+		body := string(reqBody)
+		body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.Interval.Milliseconds(), 10))
+		body = strings.ReplaceAll(body, "$__interval", r.Interval.String())
+
+		payload.WriteString(body + "\n")
 	}
+	return payload.Bytes(), nil
+}
 
-	bytes, err := c.encodeBatchRequests(multiRequests)
+func (c *baseClientImpl) createMultiSearchRequest(requests []*SearchRequest, index string) (*http.Request, error) {
+	body, err := c.makeMultiSearchPayload(requests, index)
 	if err != nil {
 		return nil, err
 	}
 
-	queryParams := c.getMultiSearchQueryParameters()
-
-	return c.makeRequest(http.MethodPost, "_elastic/_msearch", queryParams, bytes)
-}
-
-func (c *baseClientImpl) getMultiSearchQueryParameters() string {
 	var qs []string
-
 	maxConcurrentShardRequests := c.ds.MaxConcurrentShardRequests
 	if maxConcurrentShardRequests == 0 {
 		maxConcurrentShardRequests = 5
 	}
 	qs = append(qs, fmt.Sprintf("max_concurrent_shard_requests=%d", maxConcurrentShardRequests))
-	return strings.Join(qs, "&")
+	queryParams := strings.Join(qs, "&")
+
+	return c.makeRequest(http.MethodPost, "_elastic/_msearch", queryParams, body)
 }