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

@@ -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
+	}
+}

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)
 			}
 		}
 	}

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.

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,17 +90,25 @@ 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{
 			Index:     row.Index,
 			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]

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)
 		})
 	}
 }

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
 }