Support _msearch queries targetting mixed sources (#1278)

This is a very basic support for `_msearch` targeting Elasticsearch, not
ClickHouse.

The main goal is to provide better Grafana experience which almost
exclusively relies on `_msearch` endpoint.

In the future we might expand on this, do a bigger refactor or support
parrallel execution of multi search queries.

Author: mieciu

quesma/processors/es_to_ch_ingest/elasticsearch_to_clickhouse_ingest_processor.go

@@ -4,13 +4,13 @@
 package es_to_ch_ingest
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"github.com/QuesmaOrg/quesma/quesma/backend_connectors"
 	"github.com/QuesmaOrg/quesma/quesma/frontend_connectors"
 	"github.com/QuesmaOrg/quesma/quesma/ingest"
 	"github.com/QuesmaOrg/quesma/quesma/logger"
+	"github.com/QuesmaOrg/quesma/quesma/util"
 	quesma_api "github.com/QuesmaOrg/quesma/quesma/v2/core"
 
 	"github.com/QuesmaOrg/quesma/quesma/processors"
@@ -20,7 +20,6 @@ import (
 	"github.com/QuesmaOrg/quesma/quesma/quesma/types"
 	"github.com/QuesmaOrg/quesma/quesma/schema"
 	"github.com/rs/zerolog/log"
-	"io"
 	"net/http"
 )
 
@@ -165,15 +164,6 @@ func (p *ElasticsearchToClickHouseIngestProcessor) GetSupportedBackendConnectors
 	return []quesma_api.BackendConnectorType{quesma_api.ClickHouseSQLBackend, quesma_api.ElasticsearchBackend}
 }
 
-func ReadResponseBody(resp *http.Response) ([]byte, error) {
-	respBody, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-	resp.Body = io.NopCloser(bytes.NewBuffer(respBody))
-	return respBody, nil
-}
-
 func (p *ElasticsearchToClickHouseIngestProcessor) routeToElasticsearch(metadata map[string]interface{}, req *http.Request) (map[string]interface{}, *quesma_api.Result, error) {
 	metadata[es_to_ch_common.RealSourceHeader] = es_to_ch_common.RealSourceElasticsearch
 	esConn, err := p.getElasticsearchBackendConnector()
@@ -184,7 +174,7 @@ func (p *ElasticsearchToClickHouseIngestProcessor) routeToElasticsearch(metadata
 	if err != nil {
 		return metadata, nil, err
 	}
-	respBody, err := ReadResponseBody(resp)
+	respBody, err := util.ReadResponseBody(resp)
 	if err != nil {
 		return metadata, nil, fmt.Errorf("failed to read response body from Elastic")
 	}

quesma/processors/es_to_ch_query/elasticsearch_to_clickhouse_query_processor.go

@@ -4,8 +4,6 @@
 package es_to_ch_query
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"fmt"
 	"github.com/QuesmaOrg/quesma/quesma/backend_connectors"
@@ -17,6 +15,7 @@ import (
 	quesm "github.com/QuesmaOrg/quesma/quesma/quesma"
 	"github.com/QuesmaOrg/quesma/quesma/quesma/config"
 	"github.com/QuesmaOrg/quesma/quesma/quesma/types"
+	"github.com/QuesmaOrg/quesma/quesma/util"
 	quesma_api "github.com/QuesmaOrg/quesma/quesma/v2/core"
 	"github.com/QuesmaOrg/quesma/quesma/v2/core/tracing"
 	"io"
@@ -251,7 +250,7 @@ func (p *ElasticsearchToClickHouseQueryProcessor) routeToElasticsearch(metadata
 	if err != nil {
 		return metadata, nil, fmt.Errorf("failed sending request to Elastic")
 	}
-	respBody, err := ReadResponseBody(resp)
+	respBody, err := util.ReadResponseBody(resp)
 	if err != nil {
 		return metadata, nil, fmt.Errorf("failed to read response body from Elastic")
 	}
@@ -276,26 +275,6 @@ func findQueryTarget(index string, processorConfig config.QuesmaProcessorConfig)
 	}
 }
 
-func ReadResponseBody(resp *http.Response) ([]byte, error) {
-	var reader io.Reader
-	if resp.Header.Get("Content-Encoding") == "gzip" {
-		gzipReader, err := gzip.NewReader(resp.Body)
-		if err != nil {
-			return nil, err
-		}
-		defer gzipReader.Close()
-		reader = gzipReader
-	} else {
-		reader = resp.Body
-	}
-	respBody, err := io.ReadAll(reader)
-	if err != nil {
-		return nil, err
-	}
-	resp.Body = io.NopCloser(bytes.NewBuffer(respBody))
-	return respBody, nil
-}
-
 func GetQueryFromRequest(req *http.Request) (types.JSON, error) {
 	reqBody, err := io.ReadAll(req.Body)
 	if err != nil {

quesma/quesma/search.go

@@ -235,10 +235,13 @@ func (q *QueryRunner) HandleMultiSearch(ctx context.Context, defaultIndexName st
 	}
 	logger.DebugWithCtx(ctx).Msgf("handling multisearch: queries=%d, indices=[%s], defaultIndex=[%s]", len(queries), queriedIndices, defaultIndexName)
 	for _, query := range queries {
-
-		// TODO ask table resolver here and go to the right connector or connectors
-
-		responseBody, err := q.HandleSearch(ctx, query.indexName, query.query)
+		var responseBody []byte
+		if q.shouldRouteQueryToElasticsearch(query) { // this branch is here to get response from multi-search query targeted an index not stored in Clickhouse
+			// this is also a shortcut that we took to delay a bigger refactor, eventually HandleMultiSearch should dispatch all individual queries to proper connector, similarly to `_bulk` endpoint
+			responseBody, err = q.forwardToElasticsearch(ctx, query.indexName, query.query)
+		} else {
+			responseBody, err = q.HandleSearch(ctx, query.indexName, query.query)
+		}
 
 		if err != nil {
 
@@ -290,6 +293,27 @@ func (q *QueryRunner) HandleMultiSearch(ctx context.Context, defaultIndexName st
 	return responseBody, nil
 }
 
+func (q *QueryRunner) forwardToElasticsearch(ctx context.Context, indexName string, query types.JSON) ([]byte, error) {
+	logger.DebugWithCtx(ctx).Msgf("_msearch query to index=%s forwarded to Elasticsearch", indexName)
+	esClient := elasticsearch.NewSimpleClient(&q.cfg.Elasticsearch)
+	queryBody, _ := query.Bytes()
+	if resp, err := esClient.Request(ctx, "POST", "/_search", queryBody); err != nil {
+		return []byte{}, err
+	} else {
+		respBody, errRead := util.ReadResponseBody(resp)
+		return respBody, errRead
+	}
+}
+
+func (q *QueryRunner) shouldRouteQueryToElasticsearch(query msearchQuery) bool {
+	decision := q.tableResolver.Resolve(quesma_api.QueryPipeline, query.indexName)
+	if len(decision.UseConnectors) == 1 {
+		_, useElastic := decision.UseConnectors[0].(*quesma_api.ConnectorDecisionElastic)
+		return useElastic
+	}
+	return false
+}
+
 func (q *QueryRunner) HandleSearch(ctx context.Context, indexPattern string, body types.JSON) ([]byte, error) {
 	return q.handleSearchCommon(ctx, indexPattern, body, nil)
 }
@@ -457,15 +481,16 @@ func (q *QueryRunner) handleSearchCommon(ctx context.Context, indexPattern strin
 			clickhouseConnector = c
 
 		case *quesma_api.ConnectorDecisionElastic:
-			// NOP
+			// After https://github.com/QuesmaOrg/quesma/pull/1278 we should never land in this situation,
+			// previously this was an escape hatch for `_msearch` payload containing Elasticsearch-targetted query
+			// This code lives only to postpone bigger refactor of `handleSearchCommon` which also supports async and A/B testing
 
 		default:
 			return nil, fmt.Errorf("unknown connector type: %T", c)
 		}
 	}
 
 	if clickhouseConnector == nil {
-		// TODO: at this moment it's possible to land in this situation if `_msearch` payload contains Elasticsearch-targetted query
 		logger.Warn().Msgf("multi-search payload contains Elasticsearch-targetted query")
 		return nil, fmt.Errorf("quesma-processed _msearch payload contains Elasticsearch-targetted query")
 	}

quesma/util/utils.go

@@ -4,6 +4,7 @@ package util
 
 import (
 	"bytes"
+	"compress/gzip"
 	"database/sql"
 	"encoding/base64"
 	"fmt"
@@ -996,3 +997,23 @@ func TableNamePatternRegexp(indexPattern string) *regexp.Regexp {
 	patternCache[indexPattern] = result
 	return result
 }
+
+func ReadResponseBody(resp *http.Response) ([]byte, error) {
+	var reader io.Reader
+	if resp.Header.Get("Content-Encoding") == "gzip" {
+		gzipReader, err := gzip.NewReader(resp.Body)
+		if err != nil {
+			return nil, err
+		}
+		defer gzipReader.Close()
+		reader = gzipReader
+	} else {
+		reader = resp.Body
+	}
+	respBody, err := io.ReadAll(reader)
+	if err != nil {
+		return nil, err
+	}
+	resp.Body = io.NopCloser(bytes.NewBuffer(respBody))
+	return respBody, nil
+}