grafana/pkg/tsdb/elasticsearch/client/search_request_test.go
Ivana Huckova 2a1a5145d0
Elasticsearch: Fix using of individual query time ranges when querying (#84201)
* WIP - proof of concept

* Update pkg/tsdb/elasticsearch/client/client.go

Co-authored-by: Sven Grossmann <sven.grossmann@grafana.com>

* update and add test

* lint

* Fix lint

* Bring back logging when creating client

---------

Co-authored-by: Sven Grossmann <sven.grossmann@grafana.com>
2024-03-13 12:49:35 +02:00

437 lines
14 KiB
Go

package es
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana/pkg/components/simplejson"
)
func TestSearchRequest(t *testing.T) {
timeField := "@timestamp"
from := time.Date(2018, 5, 10, 17, 50, 0, 0, time.UTC)
to := time.Date(2018, 5, 12, 17, 55, 0, 0, time.UTC)
timeRange := backend.TimeRange{
From: from,
To: to,
}
setup := func() *SearchRequestBuilder {
return NewSearchRequestBuilder(15*time.Second, timeRange)
}
t.Run("When building search request", func(t *testing.T) {
b := setup()
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have size of zero", func(t *testing.T) {
require.Equal(t, 0, sr.Size)
})
t.Run("Should have no sorting", func(t *testing.T) {
require.Equal(t, 0, len(sr.Sort))
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
require.Equal(t, 0, json.Get("size").MustInt(500))
require.Nil(t, json.Get("sort").Interface())
require.Nil(t, json.Get("aggs").Interface())
require.Nil(t, json.Get("query").Interface())
})
})
t.Run("When adding size, sort, filters", func(t *testing.T) {
b := setup()
b.Size(200)
b.Sort(SortOrderDesc, timeField, "boolean")
filters := b.Query().Bool().Filter()
filters.AddDateRangeFilter(timeField, 10, 5, DateFormatEpochMS)
filters.AddQueryStringFilter("test", true)
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have correct size", func(t *testing.T) {
require.Equal(t, 200, sr.Size)
})
t.Run("Should have correct sorting", func(t *testing.T) {
sort, ok := sr.Sort[timeField].(map[string]string)
require.True(t, ok)
require.Equal(t, "desc", sort["order"])
require.Equal(t, "boolean", sort["unmapped_type"])
})
t.Run("Should have range filter", func(t *testing.T) {
f, ok := sr.Query.Bool.Filters[0].(*RangeFilter)
require.True(t, ok)
require.Equal(t, int64(5), f.Gte)
require.Equal(t, int64(10), f.Lte)
require.Equal(t, "epoch_millis", f.Format)
})
t.Run("Should have query string filter", func(t *testing.T) {
f, ok := sr.Query.Bool.Filters[1].(*QueryStringFilter)
require.True(t, ok)
require.Equal(t, "test", f.Query)
require.True(t, f.AnalyzeWildcard)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
require.Equal(t, 200, json.Get("size").MustInt(0))
sort := json.GetPath("sort", timeField)
require.Equal(t, "desc", sort.Get("order").MustString())
require.Equal(t, "boolean", sort.Get("unmapped_type").MustString())
timeRangeFilter := json.GetPath("query", "bool", "filter").GetIndex(0).Get("range").Get(timeField)
require.Equal(t, int64(5), timeRangeFilter.Get("gte").MustInt64())
require.Equal(t, int64(10), timeRangeFilter.Get("lte").MustInt64())
require.Equal(t, DateFormatEpochMS, timeRangeFilter.Get("format").MustString(""))
queryStringFilter := json.GetPath("query", "bool", "filter").GetIndex(1).Get("query_string")
require.Equal(t, true, queryStringFilter.Get("analyze_wildcard").MustBool(false))
require.Equal(t, "test", queryStringFilter.Get("query").MustString(""))
})
})
})
t.Run("When adding doc value field", func(t *testing.T) {
b := setup()
b.AddDocValueField(timeField)
t.Run("should set correct props", func(t *testing.T) {
require.Nil(t, b.customProps["fields"])
scriptFields, ok := b.customProps["script_fields"].(map[string]any)
require.True(t, ok)
require.Equal(t, 0, len(scriptFields))
docValueFields, ok := b.customProps["docvalue_fields"].([]string)
require.True(t, ok)
require.Equal(t, 1, len(docValueFields))
require.Equal(t, timeField, docValueFields[0])
})
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
scriptFields, err := json.Get("script_fields").Map()
require.Nil(t, err)
require.Equal(t, 0, len(scriptFields))
_, err = json.Get("fields").StringArray()
require.Error(t, err)
docValueFields, err := json.Get("docvalue_fields").StringArray()
require.Nil(t, err)
require.Equal(t, 1, len(docValueFields))
require.Equal(t, timeField, docValueFields[0])
})
})
})
t.Run("and adding multiple top level aggs", func(t *testing.T) {
b := setup()
aggBuilder := b.Agg()
aggBuilder.Terms("1", "@hostname", nil)
aggBuilder.DateHistogram("2", "@timestamp", nil)
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have 2 top level aggs", func(t *testing.T) {
aggs := sr.Aggs
require.Equal(t, 2, len(aggs))
require.Equal(t, "1", aggs[0].Key)
require.Equal(t, "terms", aggs[0].Aggregation.Type)
require.Equal(t, "2", aggs[1].Key)
require.Equal(t, "date_histogram", aggs[1].Aggregation.Type)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
require.Equal(t, 2, len(json.Get("aggs").MustMap()))
require.Equal(t, "@hostname", json.GetPath("aggs", "1", "terms", "field").MustString())
require.Equal(t, "@timestamp", json.GetPath("aggs", "2", "date_histogram", "field").MustString())
})
})
})
t.Run("and adding top level agg with child agg", func(t *testing.T) {
b := setup()
aggBuilder := b.Agg()
aggBuilder.Terms("1", "@hostname", func(a *TermsAggregation, ib AggBuilder) {
ib.DateHistogram("2", "@timestamp", nil)
})
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have 1 top level agg and one child agg", func(t *testing.T) {
aggs := sr.Aggs
require.Equal(t, 1, len(aggs))
topAgg := aggs[0]
require.Equal(t, "1", topAgg.Key)
require.Equal(t, "terms", topAgg.Aggregation.Type)
require.Equal(t, 1, len(topAgg.Aggregation.Aggs))
childAgg := aggs[0].Aggregation.Aggs[0]
require.Equal(t, "2", childAgg.Key)
require.Equal(t, "date_histogram", childAgg.Aggregation.Type)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
require.Equal(t, 1, len(json.Get("aggs").MustMap()))
firstLevelAgg := json.GetPath("aggs", "1")
secondLevelAgg := firstLevelAgg.GetPath("aggs", "2")
require.Equal(t, "@hostname", firstLevelAgg.GetPath("terms", "field").MustString())
require.Equal(t, "@timestamp", secondLevelAgg.GetPath("date_histogram", "field").MustString())
})
})
})
t.Run("and adding two top level aggs with child agg", func(t *testing.T) {
b := setup()
aggBuilder := b.Agg()
aggBuilder.Histogram("1", "@hostname", func(a *HistogramAgg, ib AggBuilder) {
ib.DateHistogram("2", "@timestamp", nil)
})
aggBuilder.Filters("3", func(a *FiltersAggregation, ib AggBuilder) {
ib.Terms("4", "@test", nil)
})
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have 2 top level aggs with one child agg each", func(t *testing.T) {
aggs := sr.Aggs
require.Equal(t, 2, len(aggs))
topAggOne := aggs[0]
require.Equal(t, "1", topAggOne.Key)
require.Equal(t, "histogram", topAggOne.Aggregation.Type)
require.Equal(t, 1, len(topAggOne.Aggregation.Aggs))
topAggOnechildAgg := topAggOne.Aggregation.Aggs[0]
require.Equal(t, "2", topAggOnechildAgg.Key)
require.Equal(t, "date_histogram", topAggOnechildAgg.Aggregation.Type)
topAggTwo := aggs[1]
require.Equal(t, "3", topAggTwo.Key)
require.Equal(t, "filters", topAggTwo.Aggregation.Type)
require.Equal(t, 1, len(topAggTwo.Aggregation.Aggs))
topAggTwochildAgg := topAggTwo.Aggregation.Aggs[0]
require.Equal(t, "4", topAggTwochildAgg.Key)
require.Equal(t, "terms", topAggTwochildAgg.Aggregation.Type)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
topAggOne := json.GetPath("aggs", "1")
require.Equal(t, "@hostname", topAggOne.GetPath("histogram", "field").MustString())
topAggOnechildAgg := topAggOne.GetPath("aggs", "2")
require.Equal(t, "@timestamp", topAggOnechildAgg.GetPath("date_histogram", "field").MustString())
topAggTwo := json.GetPath("aggs", "3")
topAggTwochildAgg := topAggTwo.GetPath("aggs", "4")
require.Equal(t, 0, len(topAggTwo.GetPath("filters").MustArray()))
require.Equal(t, "@test", topAggTwochildAgg.GetPath("terms", "field").MustString())
})
})
})
t.Run("and adding top level agg with child agg with child agg", func(t *testing.T) {
b := setup()
aggBuilder := b.Agg()
aggBuilder.Terms("1", "@hostname", func(a *TermsAggregation, ib AggBuilder) {
ib.Terms("2", "@app", func(a *TermsAggregation, ib AggBuilder) {
ib.DateHistogram("3", "@timestamp", nil)
})
})
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have 1 top level agg with one child having a child", func(t *testing.T) {
aggs := sr.Aggs
require.Equal(t, 1, len(aggs))
topAgg := aggs[0]
require.Equal(t, "1", topAgg.Key)
require.Equal(t, "terms", topAgg.Aggregation.Type)
require.Equal(t, 1, len(topAgg.Aggregation.Aggs))
childAgg := topAgg.Aggregation.Aggs[0]
require.Equal(t, "2", childAgg.Key)
require.Equal(t, "terms", childAgg.Aggregation.Type)
childChildAgg := childAgg.Aggregation.Aggs[0]
require.Equal(t, "3", childChildAgg.Key)
require.Equal(t, "date_histogram", childChildAgg.Aggregation.Type)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
topAgg := json.GetPath("aggs", "1")
require.Equal(t, "@hostname", topAgg.GetPath("terms", "field").MustString())
childAgg := topAgg.GetPath("aggs", "2")
require.Equal(t, "@app", childAgg.GetPath("terms", "field").MustString())
childChildAgg := childAgg.GetPath("aggs", "3")
require.Equal(t, "@timestamp", childChildAgg.GetPath("date_histogram", "field").MustString())
})
})
})
t.Run("and adding bucket and metric aggs", func(t *testing.T) {
b := setup()
aggBuilder := b.Agg()
aggBuilder.Terms("1", "@hostname", func(a *TermsAggregation, ib AggBuilder) {
ib.Terms("2", "@app", func(a *TermsAggregation, ib AggBuilder) {
ib.Metric("4", "avg", "@value", nil)
ib.DateHistogram("3", "@timestamp", func(a *DateHistogramAgg, ib AggBuilder) {
ib.Metric("4", "avg", "@value", nil)
ib.Metric("5", "max", "@value", nil)
})
})
})
t.Run("When building search request", func(t *testing.T) {
sr, err := b.Build()
require.Nil(t, err)
t.Run("Should have 1 top level agg with one child having a child", func(t *testing.T) {
aggs := sr.Aggs
require.Equal(t, 1, len(aggs))
topAgg := aggs[0]
require.Equal(t, "1", topAgg.Key)
require.Equal(t, "terms", topAgg.Aggregation.Type)
require.Equal(t, 1, len(topAgg.Aggregation.Aggs))
childAgg := topAgg.Aggregation.Aggs[0]
require.Equal(t, "2", childAgg.Key)
require.Equal(t, "terms", childAgg.Aggregation.Type)
childChildOneAgg := childAgg.Aggregation.Aggs[0]
require.Equal(t, "4", childChildOneAgg.Key)
require.Equal(t, "avg", childChildOneAgg.Aggregation.Type)
childChildTwoAgg := childAgg.Aggregation.Aggs[1]
require.Equal(t, "3", childChildTwoAgg.Key)
require.Equal(t, "date_histogram", childChildTwoAgg.Aggregation.Type)
childChildTwoChildOneAgg := childChildTwoAgg.Aggregation.Aggs[0]
require.Equal(t, "4", childChildTwoChildOneAgg.Key)
require.Equal(t, "avg", childChildTwoChildOneAgg.Aggregation.Type)
childChildTwoChildTwoAgg := childChildTwoAgg.Aggregation.Aggs[1]
require.Equal(t, "5", childChildTwoChildTwoAgg.Key)
require.Equal(t, "max", childChildTwoChildTwoAgg.Aggregation.Type)
})
t.Run("When marshal to JSON should generate correct json", func(t *testing.T) {
body, err := json.Marshal(sr)
require.Nil(t, err)
json, err := simplejson.NewJson(body)
require.Nil(t, err)
termsAgg := json.GetPath("aggs", "1")
require.Equal(t, "@hostname", termsAgg.GetPath("terms", "field").MustString())
termsAggTwo := termsAgg.GetPath("aggs", "2")
require.Equal(t, "@app", termsAggTwo.GetPath("terms", "field").MustString())
termsAggTwoAvg := termsAggTwo.GetPath("aggs", "4")
require.Equal(t, "@value", termsAggTwoAvg.GetPath("avg", "field").MustString())
dateHistAgg := termsAggTwo.GetPath("aggs", "3")
require.Equal(t, "@timestamp", dateHistAgg.GetPath("date_histogram", "field").MustString())
avgAgg := dateHistAgg.GetPath("aggs", "4")
require.Equal(t, "@value", avgAgg.GetPath("avg", "field").MustString())
maxAgg := dateHistAgg.GetPath("aggs", "5")
require.Equal(t, "@value", maxAgg.GetPath("max", "field").MustString())
})
})
})
}
func TestMultiSearchRequest(t *testing.T) {
from := time.Date(2018, 5, 10, 17, 50, 0, 0, time.UTC)
to := time.Date(2018, 5, 12, 17, 55, 0, 0, time.UTC)
timeRange := backend.TimeRange{
From: from,
To: to,
}
t.Run("When adding one search request", func(t *testing.T) {
b := NewMultiSearchRequestBuilder()
b.Search(15*time.Second, timeRange)
t.Run("When building search request should contain one search request", func(t *testing.T) {
mr, err := b.Build()
require.Nil(t, err)
require.Equal(t, 1, len(mr.Requests))
})
})
t.Run("When adding two search requests", func(t *testing.T) {
b := NewMultiSearchRequestBuilder()
b.Search(15*time.Second, timeRange)
b.Search(15*time.Second, timeRange)
t.Run("When building search request should contain two search requests", func(t *testing.T) {
mr, err := b.Build()
require.Nil(t, err)
require.Equal(t, 2, len(mr.Requests))
})
})
}