[v10.4.x] Elasticsearch: Fix legend for alerting, expressions and pre…

…viously frontend queries (#84685)

* Elasticsearch: Fix legend for alerting, expressions and previously frontend queries (#84485)

* Elasticsearch: Fix legend for alerting, expressions and previously frontend queries

* Add comment

* Update comment

(cherry picked from commit 494d169)

* Pass correct param to NewClient (this is not needed on main)

---------

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
Co-authored-by: Ivana Huckova <ivana.huckova@gmail.com>

pkg/tsdb/elasticsearch/data_query.go

@@ -23,20 +23,27 @@ const (
 )
 
 type elasticsearchDataQuery struct {
-	client      es.Client
-	dataQueries []backend.DataQuery
-	logger      log.Logger
-	ctx         context.Context
-	tracer      tracing.Tracer
+	client               es.Client
+	dataQueries          []backend.DataQuery
+	logger               log.Logger
+	ctx                  context.Context
+	tracer               tracing.Tracer
+	keepLabelsInResponse bool
 }
 
-var newElasticsearchDataQuery = func(ctx context.Context, client es.Client, dataQuery []backend.DataQuery, logger log.Logger, tracer tracing.Tracer) *elasticsearchDataQuery {
+var newElasticsearchDataQuery = func(ctx context.Context, client es.Client, req *backend.QueryDataRequest, logger log.Logger, tracer tracing.Tracer) *elasticsearchDataQuery {
+	_, fromAlert := req.Headers[headerFromAlert]
+	fromExpression := req.GetHTTPHeader(headerFromExpression) != ""
+
 	return &elasticsearchDataQuery{
 		client:      client,
-		dataQueries: dataQuery,
+		dataQueries: req.Queries,
 		logger:      logger,
 		ctx:         ctx,
 		tracer:      tracer,
+		// To maintain backward compatibility, it is necessary to keep labels in responses for alerting and expressions queries.
+		// Historically, these labels have been used in alerting rules and transformations.
+		keepLabelsInResponse: fromAlert || fromExpression,
 	}
 }
 
@@ -77,7 +84,7 @@ func (e *elasticsearchDataQuery) execute() (*backend.QueryDataResponse, error) {
 		return errorsource.AddErrorToResponse(e.dataQueries[0].RefID, response, err), nil
 	}
 
-	return parseResponse(e.ctx, res.Responses, queries, e.client.GetConfiguredFields(), e.logger, e.tracer)
+	return parseResponse(e.ctx, res.Responses, queries, e.client.GetConfiguredFields(), e.keepLabelsInResponse, e.logger, e.tracer)
 }
 
 func (e *elasticsearchDataQuery) processQuery(q *Query, ms *es.MultiSearchRequestBuilder, from, to int64) error {

pkg/tsdb/elasticsearch/data_query_test.go

@@ -1862,6 +1862,6 @@ func executeElasticsearchDataQuery(c es.Client, body string, from, to time.Time)
 			},
 		},
 	}
-	query := newElasticsearchDataQuery(context.Background(), c, dataRequest.Queries, log.New("test.logger"), tracing.InitializeTracerForTest())
+	query := newElasticsearchDataQuery(context.Background(), c, &dataRequest, log.New("test.logger"), tracing.InitializeTracerForTest())
 	return query.execute()
 }

pkg/tsdb/elasticsearch/elasticsearch.go

@@ -24,12 +24,18 @@ import (
 	"github.com/grafana/grafana/pkg/infra/httpclient"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
-	ngalertmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
 	es "github.com/grafana/grafana/pkg/tsdb/elasticsearch/client"
 )
 
 var eslog = log.New("tsdb.elasticsearch")
 
+const (
+	// headerFromExpression is used by data sources to identify expression queries
+	headerFromExpression = "X-Grafana-From-Expr"
+	// headerFromAlert is used by datasources to identify alert queries
+	headerFromAlert = "FromAlert"
+)
+
 type Service struct {
 	httpClientProvider httpclient.Provider
 	im                 instancemgmt.InstanceManager
@@ -48,28 +54,28 @@ func ProvideService(httpClientProvider httpclient.Provider, tracer tracing.Trace
 
 func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
 	dsInfo, err := s.getDSInfo(ctx, req.PluginContext)
-	_, fromAlert := req.Headers[ngalertmodels.FromAlertHeaderName]
+	_, fromAlert := req.Headers[headerFromAlert]
 	logger := s.logger.FromContext(ctx).New("fromAlert", fromAlert)
 
 	if err != nil {
 		logger.Error("Failed to get data source info", "error", err)
 		return &backend.QueryDataResponse{}, err
 	}
 
-	return queryData(ctx, req.Queries, dsInfo, logger, s.tracer)
+	return queryData(ctx, req, dsInfo, logger, s.tracer)
 }
 
 // separate function to allow testing the whole transformation and query flow
-func queryData(ctx context.Context, queries []backend.DataQuery, dsInfo *es.DatasourceInfo, logger log.Logger, tracer tracing.Tracer) (*backend.QueryDataResponse, error) {
-	if len(queries) == 0 {
+func queryData(ctx context.Context, req *backend.QueryDataRequest, dsInfo *es.DatasourceInfo, logger log.Logger, tracer tracing.Tracer) (*backend.QueryDataResponse, error) {
+	if len(req.Queries) == 0 {
 		return &backend.QueryDataResponse{}, fmt.Errorf("query contains no queries")
 	}
 
-	client, err := es.NewClient(ctx, dsInfo, queries[0].TimeRange, logger, tracer)
+	client, err := es.NewClient(ctx, dsInfo, req.Queries[0].TimeRange, logger, tracer)
 	if err != nil {
 		return &backend.QueryDataResponse{}, err
 	}
-	query := newElasticsearchDataQuery(ctx, client, queries, logger, tracer)
+	query := newElasticsearchDataQuery(ctx, client, req, logger, tracer)
 	return query.execute()
 }
 

pkg/tsdb/elasticsearch/querydata_test.go

@@ -114,6 +114,9 @@ type queryDataTestResult struct {
 
 func queryDataTestWithResponseCode(queriesBytes []byte, responseStatusCode int, responseBytes []byte) (queryDataTestResult, error) {
 	queries, err := newFlowTestQueries(queriesBytes)
+	req := backend.QueryDataRequest{
+		Queries: queries,
+	}
 	if err != nil {
 		return queryDataTestResult{}, err
 	}
@@ -138,7 +141,7 @@ func queryDataTestWithResponseCode(queriesBytes []byte, responseStatusCode int,
 		return nil
 	})
 
-	result, err := queryData(context.Background(), queries, dsInfo, log.New("test.logger"), tracing.InitializeTracerForTest())
+	result, err := queryData(context.Background(), &req, dsInfo, log.New("test.logger"), tracing.InitializeTracerForTest())
 	if err != nil {
 		return queryDataTestResult{}, err
 	}

pkg/tsdb/elasticsearch/response_parser.go

@@ -47,7 +47,7 @@ const (
 
 var searchWordsRegex = regexp.MustCompile(regexp.QuoteMeta(es.HighlightPreTagsString) + `(.*?)` + regexp.QuoteMeta(es.HighlightPostTagsString))
 
-func parseResponse(ctx context.Context, responses []*es.SearchResponse, targets []*Query, configuredFields es.ConfiguredFields, logger log.Logger, tracer tracing.Tracer) (*backend.QueryDataResponse, error) {
+func parseResponse(ctx context.Context, responses []*es.SearchResponse, targets []*Query, configuredFields es.ConfiguredFields, keepLabelsInResponse bool, logger log.Logger, tracer tracing.Tracer) (*backend.QueryDataResponse, error) {
 	result := backend.QueryDataResponse{
 		Responses: backend.Responses{},
 	}
@@ -117,7 +117,7 @@ func parseResponse(ctx context.Context, responses []*es.SearchResponse, targets
 				resSpan.End()
 				return &backend.QueryDataResponse{}, err
 			}
-			nameFields(queryRes, target)
+			nameFields(queryRes, target, keepLabelsInResponse)
 			trimDatapoints(queryRes, target)
 
 			result.Responses[target.RefID] = queryRes
@@ -888,7 +888,7 @@ func getSortedLabelValues(labels data.Labels) []string {
 	return values
 }
 
-func nameFields(queryResult backend.DataResponse, target *Query) {
+func nameFields(queryResult backend.DataResponse, target *Query, keepLabelsInResponse bool) {
 	set := make(map[string]struct{})
 	frames := queryResult.Frames
 	for _, v := range frames {
@@ -907,10 +907,18 @@ func nameFields(queryResult backend.DataResponse, target *Query) {
 			// another is "number"
 			valueField := frame.Fields[1]
 			fieldName := getFieldName(*valueField, target, metricTypeCount)
-			// We need to remove labels so they are not added to legend as duplicates
-			// ensures backward compatibility with "frontend" version of the plugin
-			valueField.Labels = nil
-			frame.Name = fieldName
+			// If we  need to keep the labels in the response, to prevent duplication in names and to keep
+			// backward compatibility with alerting and expressions we use DisplayNameFromDS
+			if keepLabelsInResponse {
+				if valueField.Config == nil {
+					valueField.Config = &data.FieldConfig{}
+				}
+				valueField.Config.DisplayNameFromDS = fieldName
+				// If we don't need to keep labels (how frontend mode worked), we use frame.Name and remove labels
+			} else {
+				valueField.Labels = nil
+				frame.Name = fieldName
+			}
 		}
 	}
 }