Bugfix: Pass docvalue_fields for elasticsearch

pkg/opensearch/client/search_request.go

@@ -93,26 +93,35 @@ func (b *SearchRequestBuilder) Sort(order, field, unmappedType string) *SearchRe
 	return b
 }
 
-// AddDocValueField adds a doc value field to the search request
-func (b *SearchRequestBuilder) AddDocValueField(field string) *SearchRequestBuilder {
-	b.customProps["script_fields"] = make(map[string]interface{})
+const timeFormat = "strict_date_optional_time_nanos"
 
-	if b.version.Major() < 5 && b.flavor == Elasticsearch {
-		b.customProps["fielddata_fields"] = []string{field}
-		b.customProps["fields"] = []string{"*", "_source"}
-	} else {
-		b.customProps["docvalue_fields"] = []string{field}
-	}
-
-	return b
-}
-
-// AddTimeFieldWithStandardizedFormat adds timeField as field with standardized time format to not receive
+// SetCustomProps adds timeField as field with standardized time format to not receive
 // invalid formats that Elasticsearch/OpenSearch can parse, but our frontend can't (e.g. yyyy_MM_dd_HH_mm_ss)
 // https://opensearch.org/docs/latest/api-reference/search/#request-body
 // https://opensearch.org/docs/latest/field-types/supported-field-types/date/#full-date-formats
-func (b *SearchRequestBuilder) AddTimeFieldWithStandardizedFormat(timeField string) {
-	b.customProps["fields"] = []map[string]string{{"field": timeField, "format": "strict_date_optional_time_nanos"}}
+// This is added to different keys in customProps depending on the flavor and version.
+//
+// This also adds the required {"*", "_source"} value to "fields" for very old versions of Elasticsearch
+// for log queries.
+func (b *SearchRequestBuilder) SetCustomProps(timeField string, luceneQueryType string) {
+	// defaults - OpenSearch or Elasticsearch > 7
+	var key = "fields"
+	var value any = []any{map[string]string{"field": timeField, "format": timeFormat}}
+	if b.flavor == OpenSearch && luceneQueryType == "logs" {
+		b.customProps["docvalue_fields"] = []any{timeField}
+	}
+	if b.flavor == Elasticsearch {
+		if b.version.Major() >= 5 && b.version.Major() <= 7 {
+			key = "docvalue_fields"
+		} else {
+			if b.version.Major() < 5 && luceneQueryType == "logs" {
+				b.customProps["fielddata_fields"] = []any{timeField}
+				b.customProps["fields"] = []any{"*", "_source", value}
+				return
+			}
+		}
+	}
+	b.customProps[key] = value
 }
 
 // Query creates and return a query builder

pkg/opensearch/client/search_request_test.go

@@ -146,22 +146,44 @@ func TestSearchRequest(t *testing.T) {
 		})
 
 		t.Run("When adding doc value field", func(t *testing.T) {
-			b.AddDocValueField(timeField)
+			t.Run("should set correct props for logs", func(t *testing.T) {
+				b = NewSearchRequestBuilder(OpenSearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
 
-			t.Run("should set correct props", func(t *testing.T) {
-				assert.Nil(t, b.customProps["fields"])
-
-				scriptFields, ok := b.customProps["script_fields"].(map[string]interface{})
-				assert.True(t, ok)
-				assert.Len(t, scriptFields, 0)
+				b.SetCustomProps(timeField, "logs")
 
-				docValueFields, ok := b.customProps["docvalue_fields"].([]string)
+				docValueFields, ok := b.customProps["docvalue_fields"].([]any)
 				assert.True(t, ok)
 				assert.Len(t, docValueFields, 1)
 				assert.Equal(t, timeField, docValueFields[0])
+
+				fields, ok := b.customProps["fields"].([]any)
+				assert.True(t, ok)
+				assert.Len(t, fields, 1)
+				resultTimeField, ok := fields[0].(map[string]string)["field"]
+				assert.True(t, ok)
+				assert.Equal(t, timeField, resultTimeField)
+
 			})
 
-			t.Run("When building search request", func(t *testing.T) {
+			t.Run("should set correct props for raw_document", func(t *testing.T) {
+				b = NewSearchRequestBuilder(OpenSearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+
+				b.SetCustomProps(timeField, "raw_document")
+
+				assert.Nil(t, b.customProps["docvalue_fields"])
+
+				fields, ok := b.customProps["fields"].([]any)
+				assert.True(t, ok)
+				assert.Len(t, fields, 1)
+				resultTimeField, ok := fields[0].(map[string]string)["field"]
+				assert.True(t, ok)
+				assert.Equal(t, timeField, resultTimeField)
+
+			})
+
+			t.Run("When building search request for logs", func(t *testing.T) {
+				b = NewSearchRequestBuilder(OpenSearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+				b.SetCustomProps(timeField, "logs")
 				sr, err := b.Build()
 				assert.NoError(t, err)
 
@@ -171,47 +193,94 @@ func TestSearchRequest(t *testing.T) {
 					json, err := simplejson.NewJson(body)
 					assert.NoError(t, err)
 
-					scriptFields, err := json.Get("script_fields").Map()
-					assert.NoError(t, err)
-					assert.Len(t, scriptFields, 0)
-
-					_, err = json.Get("fields").StringArray()
-					assert.Error(t, err)
-
 					docValueFields, err := json.Get("docvalue_fields").StringArray()
 					assert.NoError(t, err)
 					assert.Len(t, docValueFields, 1)
 					assert.Equal(t, docValueFields[0], timeField)
+
+					fields, err := json.Get("fields").Array()
+					assert.NoError(t, err)
+					assert.Len(t, fields, 1)
+					assert.Equal(t, fields[0].(map[string]interface{})["field"], timeField)
+				})
+			})
+
+			t.Run("When building search request for raw_document", func(t *testing.T) {
+				b = NewSearchRequestBuilder(OpenSearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+				b.SetCustomProps(timeField, "raw_document")
+				sr, err := b.Build()
+				assert.NoError(t, err)
+
+				t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
+					body, err := json.Marshal(sr)
+					assert.NoError(t, err)
+					json, err := simplejson.NewJson(body)
+					assert.NoError(t, err)
+
+					docvalueFields := json.Get("docvalue_fields").Interface()
+					assert.Nil(t, docvalueFields)
+
+					fields, err := json.Get("fields").Array()
+					assert.NoError(t, err)
+					assert.Len(t, fields, 1)
+					assert.Equal(t, fields[0].(map[string]interface{})["field"], timeField)
 				})
 			})
 		})
 	})
 
 	t.Run("Given new search request builder for Elasticsearch 2.0.0", func(t *testing.T) {
 		version, _ := semver.NewVersion("2.0.0")
-		b := NewSearchRequestBuilder(Elasticsearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
-
-		t.Run("When adding doc value field", func(t *testing.T) {
-			b.AddDocValueField(timeField)
+		t.Run("When adding doc value field for logs", func(t *testing.T) {
+			b := NewSearchRequestBuilder(Elasticsearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+			b.SetCustomProps(timeField, "logs")
 
 			t.Run("should set correct props", func(t *testing.T) {
-				fields, ok := b.customProps["fields"].([]string)
+				fields, ok := b.customProps["fields"].([]any)
 				assert.True(t, ok)
-				assert.Len(t, fields, 2)
+				assert.Len(t, fields, 3)
 				assert.Equal(t, "*", fields[0])
 				assert.Equal(t, "_source", fields[1])
+				resultTimeField, ok := fields[2].([]any)[0].(map[string]string)["field"]
+				assert.True(t, ok)
+				assert.Equal(t, timeField, resultTimeField)
 
-				scriptFields, ok := b.customProps["script_fields"].(map[string]interface{})
+				fieldDataField, ok := b.customProps["fielddata_fields"].([]any)
 				assert.True(t, ok)
-				assert.Len(t, scriptFields, 0)
+				assert.Equal(t, fieldDataField[0], timeField)
+
+			})
+		})
+		t.Run("When adding doc value field for raw_document", func(t *testing.T) {
+			b := NewSearchRequestBuilder(Elasticsearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+			b.SetCustomProps(timeField, "raw_document")
 
-				fieldDataFields, ok := b.customProps["fielddata_fields"].([]string)
+			t.Run("should set correct props", func(t *testing.T) {
+				fields, ok := b.customProps["fields"].([]any)
 				assert.True(t, ok)
-				assert.Len(t, fieldDataFields, 1)
-				assert.Equal(t, timeField, fieldDataFields[0])
+				assert.Len(t, fields, 1)
+				resultTimeField, ok := fields[0].(map[string]string)["field"]
+				assert.True(t, ok)
+				assert.Equal(t, timeField, resultTimeField)
+
+				assert.Nil(t, b.customProps["fielddata_fields"])
 			})
 		})
 	})
+
+	t.Run("Given new search request builder for Elasticsearch 7.0.0", func(t *testing.T) {
+		version, _ := semver.NewVersion("7.0.0")
+		t.Run("When adding timestamp format should add new time format field", func(t *testing.T) {
+			b := NewSearchRequestBuilder(Elasticsearch, version, tsdb.Interval{Value: 15 * time.Second, Text: "15s"})
+			b.SetCustomProps(timeField, "logs")
+
+			fields, ok := b.customProps["docvalue_fields"].([]any)
+			assert.True(t, ok)
+			assert.Len(t, fields, 1)
+			assert.Equal(t, map[string]string{"field": timeField, "format": "strict_date_optional_time_nanos"}, fields[0])
+
+		})
+	})
 }
 
 func TestMultiSearchRequest(t *testing.T) {

pkg/opensearch/lucene_handler.go

@@ -119,8 +119,8 @@ func getTraceId(rawQuery string) string {
 func processLogsQuery(q *Query, b *client.SearchRequestBuilder, from, to int64, defaultTimeField string) {
 	metric := q.Metrics[0]
 	b.Sort(descending, defaultTimeField, "boolean")
-	b.AddDocValueField(defaultTimeField)
-	b.AddTimeFieldWithStandardizedFormat(defaultTimeField)
+	b.SetCustomProps(defaultTimeField, "logs")
+
 	sizeString := metric.Settings.Get("size").MustString()
 	size, err := strconv.Atoi(sizeString)
 	if err != nil {
@@ -164,7 +164,7 @@ func processDocumentQuery(q *Query, b *client.SearchRequestBuilder, defaultTimeF
 	order := metric.Settings.Get("order").MustString()
 	b.Sort(order, defaultTimeField, "boolean")
 	b.Sort(order, "_doc", "")
-	b.AddTimeFieldWithStandardizedFormat(defaultTimeField)
+	b.SetCustomProps(defaultTimeField, "raw_document")
 	sizeString := metric.Settings.Get("size").MustString()
 	size, err := strconv.Atoi(sizeString)
 	if err != nil {
@@ -322,7 +322,7 @@ func getParametersFromServiceMapResult(smResult *client.SearchResponse) ([]strin
 	// ensure consistent order for the snapshot tests in lucene_service_map_test.go
 	sort.Strings(services)
 	sort.Strings(operations)
-	
+
 	return services, operations
 }
 

pkg/opensearch/opensearch.go

@@ -88,7 +88,7 @@ func handleServiceMapPrefetch(ctx context.Context, osClient client.Client, req *
 		queryType := model.Get("queryType").MustString()
 		luceneQueryType := model.Get("luceneQueryType").MustString()
 		serviceMapRequested := model.Get("serviceMap").MustBool(false)
-		if queryType == Lucene && luceneQueryType == "Traces" && serviceMapRequested {
+		if queryType == Lucene && luceneQueryType == luceneQueryTypeTraces && serviceMapRequested {
 			prefetchQuery := createServiceMapPrefetchQuery(query)
 			q := newQueryRequest(osClient, []backend.DataQuery{prefetchQuery}, req.PluginContext.DataSourceInstanceSettings, intervalCalculator)
 			response, err := q.execute(ctx)

pkg/opensearch/query_request.go

@@ -94,7 +94,7 @@ func parse(reqQueries []backend.DataQuery) ([]*Query, error) {
 		// For queries requesting the service map, we inject extra queries to handle retrieving
 		// the required information
 		hasServiceMap := model.Get("serviceMap").MustBool(false)
-		if luceneQueryType == "Traces" && hasServiceMap {
+		if luceneQueryType == luceneQueryTypeTraces && hasServiceMap {
 			// The Prefetch request is used by itself for internal use, to get the parameters
 			// necessary for the Stats request. In this case there's no original query to
 			// pass along, so we continue below.

pkg/opensearch/snapshot_tests/lucene_logs_test.go

@@ -41,7 +41,7 @@ func Test_logs_request(t *testing.T) {
 
 	// assert request's header and query
 	expectedRequest := `{"ignore_unavailable":true,"index":"","search_type":"query_then_fetch"}
-{"aggs":{"1":{"date_histogram":{"field":"timestamp","interval":"100ms","min_doc_count":0,"extended_bounds":{"min":1668422437218,"max":1668422625668},"format":"epoch_millis"}}},"docvalue_fields":["timestamp"],"fields":[{"field":"timestamp","format":"strict_date_optional_time_nanos"}],"query":{"bool":{"filter":[{"range":{"timestamp":{"format":"epoch_millis","gte":1668422437218,"lte":1668422625668}}},{"query_string":{"analyze_wildcard":true,"query":"FlightDelayType:\"Carrier Delay\" AND Carrier:Open*"}}]}},"script_fields":{},"size":500,"sort":[{"timestamp":{"order":"desc","unmapped_type":"boolean"}}]}
+{"aggs":{"1":{"date_histogram":{"field":"timestamp","interval":"100ms","min_doc_count":0,"extended_bounds":{"min":1668422437218,"max":1668422625668},"format":"epoch_millis"}}},"docvalue_fields":["timestamp"],"fields":[{"field":"timestamp","format":"strict_date_optional_time_nanos"}],"query":{"bool":{"filter":[{"range":{"timestamp":{"format":"epoch_millis","gte":1668422437218,"lte":1668422625668}}},{"query_string":{"analyze_wildcard":true,"query":"FlightDelayType:\"Carrier Delay\" AND Carrier:Open*"}}]}},"size":500,"sort":[{"timestamp":{"order":"desc","unmapped_type":"boolean"}}]}
 `
 	assert.Equal(t, expectedRequest, string(interceptedRequest))
 }