From 2926d7f998279be9b56c5babfb89496194796e51 Mon Sep 17 00:00:00 2001 From: Grzegorz Piwowarek Date: Fri, 24 May 2024 08:13:58 +0200 Subject: [PATCH] Propagate sort properties to the response (#191) In our search responses, we were hardcoding two values: ``` hits[i].Sort = []any{"2024-01-30T19:38:54.607Z", 2944} ``` Replaced this with actual values resolved based on provided sort fields. --- quesma/elasticsearch/format.go | 13 +++++++++ quesma/eql/query_translator.go | 11 +++++--- quesma/model/query.go | 19 ++++++++++--- quesma/queryparser/query_translator.go | 31 +++++++++++++++------ quesma/queryparser/query_translator_test.go | 14 ++++++---- quesma/quesma/query_translator.go | 6 ++-- 6 files changed, 68 insertions(+), 26 deletions(-) create mode 100644 quesma/elasticsearch/format.go diff --git a/quesma/elasticsearch/format.go b/quesma/elasticsearch/format.go new file mode 100644 index 000000000..1e1ac3734 --- /dev/null +++ b/quesma/elasticsearch/format.go @@ -0,0 +1,13 @@ +package elasticsearch + +import "time" + +func FormatSortValue(i any) any { + switch v := i.(type) { + case time.Time: + // When returned as part of `sort`, timestamps are always returned as millis + return v.UnixMilli() + default: + return i + } +} diff --git a/quesma/eql/query_translator.go b/quesma/eql/query_translator.go index de62f1fe2..822aebc91 100644 --- a/quesma/eql/query_translator.go +++ b/quesma/eql/query_translator.go @@ -3,6 +3,7 @@ package eql import ( "context" "mitmproxy/quesma/clickhouse" + "mitmproxy/quesma/elasticsearch" "mitmproxy/quesma/eql/transform" "mitmproxy/quesma/logger" "mitmproxy/quesma/model" @@ -25,7 +26,6 @@ func (cw *ClickhouseEQLQueryTranslator) MakeSearchResponse(ResultSet []model.Que // This shares a lot of code with the ClickhouseQueryTranslator // - hits := make([]model.SearchHit, len(ResultSet)) for i := range ResultSet { resultRow := ResultSet[i] @@ -38,9 +38,12 @@ func (cw *ClickhouseEQLQueryTranslator) MakeSearchResponse(ResultSet []model.Que hits[i].Index = cw.Table.Name hits[i].Score = 1 hits[i].Version = 1 - hits[i].Sort = []any{ - "2024-01-30T19:38:54.607Z", - 2944, + } + for _, property := range query.SortFields.Properties() { + if val, ok := hits[i].Fields[property]; ok { + hits[i].Sort = append(hits[i].Sort, elasticsearch.FormatSortValue(val[0])) + } else { + logger.WarnWithCtx(cw.Ctx).Msgf("property %s not found in fields", property) } } } diff --git a/quesma/model/query.go b/quesma/model/query.go index 2ee851a95..22d8aba60 100644 --- a/quesma/model/query.go +++ b/quesma/model/query.go @@ -11,9 +11,20 @@ import ( const RowNumberColumnName = "row_number" const EmptyFieldSelection = "''" // we can query SELECT '', that's why such quotes -type SortField struct { - Field string - Desc bool +type ( + SortFields []SortField + SortField struct { + Field string + Desc bool + } +) + +func (sf SortFields) Properties() []string { + properties := make([]string, 0) + for _, sortField := range sf { + properties = append(properties, sortField.Field) + } + return properties } type Highlighter struct { @@ -40,7 +51,7 @@ type Query struct { Parent string // parent aggregation name, used in some pipeline aggregations Aggregators []Aggregator // keeps names of aggregators, e.g. "0", "1", "2", "suggestions". Needed for JSON response. Type QueryType - SortFields []SortField // fields to sort by + SortFields SortFields // fields to sort by SubSelect string // dictionary to add as 'meta' field in the response. // WARNING: it's probably not passed everywhere where it's needed, just in one place. diff --git a/quesma/queryparser/query_translator.go b/quesma/queryparser/query_translator.go index 952f11a73..7b5883d4f 100644 --- a/quesma/queryparser/query_translator.go +++ b/quesma/queryparser/query_translator.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "mitmproxy/quesma/clickhouse" + "mitmproxy/quesma/elasticsearch" "mitmproxy/quesma/kibana" "mitmproxy/quesma/logger" "mitmproxy/quesma/model" @@ -89,7 +90,7 @@ func (cw *ClickhouseQueryTranslator) highlightHit(hit *model.SearchHit, highligh } } -func (cw *ClickhouseQueryTranslator) makeSearchResponseNormal(ResultSet []model.QueryResultRow, highlighter model.Highlighter) *model.SearchResp { +func (cw *ClickhouseQueryTranslator) makeSearchResponseNormal(ResultSet []model.QueryResultRow, highlighter model.Highlighter, sortProperties []string) *model.SearchResp { hits := make([]model.SearchHit, len(ResultSet)) for i, row := range ResultSet { hits[i] = model.SearchHit{ @@ -97,9 +98,17 @@ func (cw *ClickhouseQueryTranslator) makeSearchResponseNormal(ResultSet []model. Source: []byte(row.String(cw.Ctx)), Fields: make(map[string][]interface{}), Highlight: make(map[string][]string), + Sort: []any{}, } cw.highlightHit(&hits[i], highlighter, ResultSet[i]) hits[i].ID = cw.computeIdForDocument(hits[i], strconv.Itoa(i+1)) + for _, property := range sortProperties { + if val, ok := hits[i].Fields[property]; ok { + hits[i].Sort = append(hits[i].Sort, elasticsearch.FormatSortValue(val[0])) + } else { + logger.WarnWithCtx(cw.Ctx).Msgf("property %s not found in fields", property) + } + } } return &model.SearchResp{ @@ -149,11 +158,12 @@ func EmptyAsyncSearchResponse(id string, isPartial bool, completionStatus int) ( func (cw *ClickhouseQueryTranslator) MakeSearchResponse(ResultSet []model.QueryResultRow, query model.Query) (*model.SearchResp, error) { switch query.QueryInfo.Typ { case model.Normal: - return cw.makeSearchResponseNormal(ResultSet, query.Highlighter), nil + return cw.makeSearchResponseNormal(ResultSet, query.Highlighter, query.SortFields.Properties()), nil case model.Facets, model.FacetsNumeric: return cw.makeSearchResponseFacets(ResultSet, query.QueryInfo.Typ), nil case model.ListByField, model.ListAllFields: - return cw.makeSearchResponseList(ResultSet, query.QueryInfo.Typ, query.Highlighter), nil + + return cw.makeSearchResponseList(ResultSet, query.QueryInfo.Typ, query.Highlighter, query.SortFields.Properties()), nil default: return nil, fmt.Errorf("unknown SearchQueryType: %v", query.QueryInfo.Typ) } @@ -303,7 +313,7 @@ func (cw *ClickhouseQueryTranslator) computeIdForDocument(doc model.SearchHit, d return pseudoUniqueId } -func (cw *ClickhouseQueryTranslator) makeSearchResponseList(ResultSet []model.QueryResultRow, typ model.SearchQueryType, highlighter model.Highlighter) *model.SearchResp { +func (cw *ClickhouseQueryTranslator) makeSearchResponseList(ResultSet []model.QueryResultRow, typ model.SearchQueryType, highlighter model.Highlighter, sortProperties []string) *model.SearchResp { hits := make([]model.SearchHit, len(ResultSet)) for i := range ResultSet { hits[i].Fields = make(map[string][]interface{}) @@ -313,13 +323,16 @@ func (cw *ClickhouseQueryTranslator) makeSearchResponseList(ResultSet []model.Qu hits[i].Index = cw.Table.Name hits[i].Score = 1 hits[i].Version = 1 - hits[i].Sort = []any{ - "2024-01-30T19:38:54.607Z", - 2944, - } } cw.highlightHit(&hits[i], highlighter, ResultSet[i]) hits[i].ID = cw.computeIdForDocument(hits[i], strconv.Itoa(i+1)) + for _, property := range sortProperties { + if val, ok := hits[i].Fields[property]; ok { + hits[i].Sort = append(hits[i].Sort, elasticsearch.FormatSortValue(val[0])) + } else { + logger.WarnWithCtx(cw.Ctx).Msgf("property %s not found in fields", property) + } + } } return &model.SearchResp{ @@ -466,7 +479,7 @@ func (cw *ClickhouseQueryTranslator) MakeResponseAggregation(queries []model.Que hits := []model.SearchHit{} // Process hits as last aggregation if len(queries) > 0 && len(ResultSets) > 0 && queries[len(queries)-1].IsWildcard() { - response := cw.makeSearchResponseNormal(ResultSets[len(ResultSets)-1], queries[len(queries)-1].Highlighter) + response := cw.makeSearchResponseNormal(ResultSets[len(ResultSets)-1], queries[len(queries)-1].Highlighter, queries[len(queries)-1].SortFields.Properties()) hits = response.Hits.Hits queries = queries[:len(queries)-1] ResultSets = ResultSets[:len(ResultSets)-1] diff --git a/quesma/queryparser/query_translator_test.go b/quesma/queryparser/query_translator_test.go index 41f0ac613..0e4a1e23c 100644 --- a/quesma/queryparser/query_translator_test.go +++ b/quesma/queryparser/query_translator_test.go @@ -361,7 +361,7 @@ func TestMakeResponseAsyncSearchQuery(t *testing.T) { }, "sort": [ "2024-01-30T19:39:35.767Z", - 16 + "apollo" ] }, { @@ -400,7 +400,7 @@ func TestMakeResponseAsyncSearchQuery(t *testing.T) { }, "sort": [ "2024-01-30T19:38:54.607Z", - 2944 + "apollo" ] } ], @@ -449,13 +449,17 @@ func TestMakeResponseAsyncSearchQuery(t *testing.T) { cw := ClickhouseQueryTranslator{Table: &clickhouse.Table{Name: "test"}, Ctx: context.Background()} for i, tt := range args { t.Run(tt.queryType.String(), func(t *testing.T) { - ourResponse, err := cw.MakeAsyncSearchResponseMarshalled(args[i].ourQueryResult, model.Query{QueryInfo: model.SearchQueryInfo{Typ: args[i].queryType}, Highlighter: NewEmptyHighlighter()}, asyncRequestIdStr, false) + ourResponse, err := cw.MakeAsyncSearchResponseMarshalled(args[i].ourQueryResult, model.Query{ + QueryInfo: model.SearchQueryInfo{Typ: args[i].queryType}, + Highlighter: NewEmptyHighlighter(), + SortFields: []model.SortField{{Field: "@timestamp", Desc: true}, {Field: "host.name", Desc: true}}, + }, asyncRequestIdStr, false) assert.NoError(t, err) actualMinusExpected, expectedMinusActual, err := util.JsonDifference(string(ourResponse), args[i].elasticResponseJson) assert.NoError(t, err) - assert.Empty(t, actualMinusExpected) - assert.Empty(t, expectedMinusActual) + assert.Empty(t, actualMinusExpected, "actualMinusExpected: %s", actualMinusExpected) + assert.Empty(t, expectedMinusActual, "expectedMinusActual: %s", expectedMinusActual) }) } } diff --git a/quesma/quesma/query_translator.go b/quesma/quesma/query_translator.go index 9b5a48ad3..adb2fa294 100644 --- a/quesma/quesma/query_translator.go +++ b/quesma/quesma/query_translator.go @@ -33,10 +33,8 @@ const ( func NewQueryTranslator(ctx context.Context, language QueryLanguage, table *clickhouse.Table, logManager *clickhouse.LogManager, dateMathRenderer string) (queryTranslator IQueryTranslator) { switch language { case QueryLanguageEQL: - queryTranslator = &eql.ClickhouseEQLQueryTranslator{ClickhouseLM: logManager, Table: table, Ctx: ctx} + return &eql.ClickhouseEQLQueryTranslator{ClickhouseLM: logManager, Table: table, Ctx: ctx} default: - queryTranslator = &queryparser.ClickhouseQueryTranslator{ClickhouseLM: logManager, Table: table, Ctx: ctx, DateMathRenderer: dateMathRenderer} + return &queryparser.ClickhouseQueryTranslator{ClickhouseLM: logManager, Table: table, Ctx: ctx, DateMathRenderer: dateMathRenderer} } - - return queryTranslator }