Fix Kibana's surrounding documents view (related to search_after query) (#1432)

Fixes "surrounding documents" view. It's imperfect (since based on the
"timestamp being id" and lack of PK being a unique constraint in the
ClickHouse/Hydrolix realm) , yet works 😉

It actually took some non-trivial hacks to land for this to even be
feasible, I tried to comment any non-obvious action taken. Most notable:
* Append document id (`_id`) as a drop-in replacement for Lucene's
`_doc` field
* Align all parsing logic accordingly
* In fact ignore that tiebreaker later on since it's not available and
would need recomputing again on transformation (good thing in that even
in Elasticsearch that's non deterministic (ref:
[docs](https://www.elastic.co/docs/reference/elasticsearch/rest-apis/paginate-search-results))
* Fix the timestamp casting issues which caused errors for
non-millisecond timestamp fields

### Screenshots showing that functionality 
<img width="300" alt="image"
src="https://github.com/user-attachments/assets/90c632d3-8365-456c-a9de-8f3c76893a79"
/>

<img width="300" alt="image"
src="https://github.com/user-attachments/assets/d4cb8633-f736-4ff3-9346-424c9cfeb2c6"
/>

Author: mieciu

platform/frontend_connectors/schema_search_after_transformer.go

@@ -8,6 +8,7 @@ import (
 	"github.com/QuesmaOrg/quesma/platform/model"
 	"github.com/QuesmaOrg/quesma/platform/schema"
 	"github.com/QuesmaOrg/quesma/platform/util"
+	"slices"
 )
 
 type (
@@ -104,12 +105,21 @@ func (s searchAfterStrategyBasicAndFast) validateAndParse(query *model.Query, in
 	if !ok {
 		return nil, fmt.Errorf("search_after must be an array, got: %v", query.SearchAfter)
 	}
-	if len(asArray) != len(query.SelectCommand.OrderBy) {
+	// Kibana uses search_after API for pagination in the following way:
+	// "sort": [
+	//  { "@timestamp": { "order": "asc" } },   // this is the main sorting criteria -> timestamp field from the Data View setting
+	//  { "_doc": { "order": "asc" } }			// this is the tiebreaker field, referencing the Lucene document ID
+	//
+	// Quesma ignores the _doc field during query parsing, therefore there's no related `query.SelectCommand.OrderBy` clause.
+	// So the condition below is to handle this Kibana-specific case, in which we ignore the _doc and just sort by the first field only.
+	// In any other case, we expect the search_after to have the same number of elements as the order by fields.
+	if len(asArray) == 2 && slices.Contains(query.SearchAfterFieldNames, "_doc") {
+		asArray = asArray[:len(asArray)-1]
+	} else if len(asArray) != len(query.SelectCommand.OrderBy) {
 		return nil, fmt.Errorf("len(search_after) != len(sortFields), search_after: %v, sortFields: %v", asArray, query.SelectCommand.OrderBy)
 	}
 
-	sortFieldsNr := len(asArray)
-	searchAfterParsed = make([]model.Expr, sortFieldsNr)
+	searchAfterParsed = make([]model.Expr, len(asArray))
 	for i, searchAfterValue := range asArray {
 		column, ok := query.SelectCommand.OrderBy[i].Expr.(model.ColumnRef)
 		if !ok {
@@ -124,8 +134,16 @@ func (s searchAfterStrategyBasicAndFast) validateAndParse(query *model.Query, in
 		if field.Type.Name == "date" || field.Type.Name == "timestamp" {
 			if number, isNumber := util.ExtractNumeric64Maybe(searchAfterValue); isNumber {
 				if number >= 0 && util.IsFloat64AnInt64(number) {
-					// this param will always be timestamp in milliseconds, as we create it like this while rendering hits
-					searchAfterParsed[i] = model.NewFunction("fromUnixTimestamp64Milli", model.NewLiteral(int64(number)))
+					milliTimestamp := model.NewFunction("fromUnixTimestamp64Milli", model.NewLiteral(int64(number)))
+					if field.InternalPropertyType == "DateTime" {
+						// Here we have to make sure that the statement landing in the where clause will match the field type
+						searchAfterParsed[i] = model.NewFunction("toDateTime", milliTimestamp)
+					} else {
+						// Default case of (field.InternalPropertyType == "DateTime64")
+						// which most often is DateTime64(3) allowing millisecond precision
+						searchAfterParsed[i] = milliTimestamp
+					}
+
 				} else {
 					return nil, fmt.Errorf("for basicAndFast strategy, search_after must be a unix timestamp in milliseconds")
 				}

platform/frontend_connectors/schema_transformer.go

@@ -1031,7 +1031,7 @@ func (s *SchemaCheckPass) Transform(queries []*model.Query) ([]*model.Query, err
 
 			query, err = transformation.Transformation(query.Schema, query)
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("error applying transformation %s: %w", transformation.TransformationName, err)
 			}
 
 			if s.cfg.Logging.EnableSQLTracing {

platform/frontend_connectors/search_test.go

@@ -1003,12 +1003,14 @@ func TestSearchAfterParameter_sortByJustTimestamp(t *testing.T) {
 		expectedSQL                 string
 		resultRowsFromDB            [][]any
 		basicAndFastSortFieldPerHit []int64
+		expectedSortValuesCount     int
 	}{
 		{
 			request:                     `{"size": 3, "track_total_hits": false, "sort": [{"@timestamp": {"order": "desc"}}]}`,
 			expectedSQL:                 `SELECT "@timestamp", "message" FROM __quesma_table_name ORDER BY "@timestamp" DESC LIMIT 3`,
 			resultRowsFromDB:            [][]any{{someTime, "m1"}, {someTime, "m2"}, {someTime, "m3"}},
 			basicAndFastSortFieldPerHit: []int64{someTime.UnixMilli(), someTime.UnixMilli(), someTime.UnixMilli()},
+			expectedSortValuesCount:     1,
 		},
 		{
 			request: `
@@ -1024,18 +1026,21 @@ func TestSearchAfterParameter_sortByJustTimestamp(t *testing.T) {
 			expectedSQL:                 `SELECT "@timestamp", "message" FROM __quesma_table_name WHERE fromUnixTimestamp64Milli(1706551896491)>"@timestamp" ORDER BY "@timestamp" DESC LIMIT 3`,
 			resultRowsFromDB:            [][]any{{sub(1), "m8"}, {sub(2), "m9"}, {sub(3), "m10"}},
 			basicAndFastSortFieldPerHit: []int64{sub(1).UnixMilli(), sub(2).UnixMilli(), sub(3).UnixMilli()},
+			expectedSortValuesCount:     2,
 		},
 		{
 			request:                     `{"search_after": [1706551896488], "size": 3, "track_total_hits": false, "sort": [{"@timestamp": {"order": "desc"}}]}`,
 			expectedSQL:                 `SELECT "@timestamp", "message" FROM __quesma_table_name WHERE fromUnixTimestamp64Milli(1706551896488)>"@timestamp" ORDER BY "@timestamp" DESC LIMIT 3`,
 			resultRowsFromDB:            [][]any{{sub(4), "m11"}, {sub(5), "m12"}, {sub(6), "m13"}},
 			basicAndFastSortFieldPerHit: []int64{sub(4).UnixMilli(), sub(5).UnixMilli(), sub(6).UnixMilli()},
+			expectedSortValuesCount:     1,
 		},
 		{
 			request:                     `{"search_after": [1706551896485], "size": 3, "track_total_hits": false, "sort": [{"@timestamp": {"order": "desc"}}]}`,
 			expectedSQL:                 `SELECT "@timestamp", "message" FROM __quesma_table_name WHERE fromUnixTimestamp64Milli(1706551896485)>"@timestamp" ORDER BY "@timestamp" DESC LIMIT 3`,
 			resultRowsFromDB:            [][]any{{sub(7), "m14"}, {sub(8), "m15"}, {sub(9), "m16"}},
 			basicAndFastSortFieldPerHit: []int64{sub(7).UnixMilli(), sub(8).UnixMilli(), sub(9).UnixMilli()},
+			expectedSortValuesCount:     1,
 		},
 	}
 
@@ -1078,7 +1083,7 @@ func TestSearchAfterParameter_sortByJustTimestamp(t *testing.T) {
 			assert.Len(t, hits, len(iteration.resultRowsFromDB))
 			for i, hit := range hits {
 				sortField := hit.(model.JsonMap)["sort"].([]any)
-				assert.Len(t, sortField, 1)
+				assert.Len(t, sortField, iteration.expectedSortValuesCount)
 				assert.Equal(t, float64(iteration.basicAndFastSortFieldPerHit[i]), sortField[0].(float64))
 			}
 		}

platform/model/query.go

@@ -63,8 +63,9 @@ type (
 		// this is schema for current query, this schema should be used in pipeline processing
 		Schema schema.Schema
 
-		Highlighter Highlighter
-		SearchAfter any // Value of query's "search_after" param. Used for pagination of hits. SearchAfterEmpty means no pagination
+		Highlighter           Highlighter
+		SearchAfter           any      // Value of query's "search_after" param. Used for pagination of hits. SearchAfterEmpty means no pagination
+		SearchAfterFieldNames []string // Names of fields used in search_after. These can be different from the order by fields,
 
 		RuntimeMappings map[string]RuntimeMapping
 

platform/model/simple_query.go

@@ -3,9 +3,12 @@
 package model
 
 type SimpleQuery struct {
-	WhereClause Expr
-	OrderBy     []OrderByExpr
-	CanParse    bool
+	WhereClause    Expr
+	OrderBy        []OrderByExpr
+	SortFieldNames []string // SortFieldNames is used to preserve fields listed in the `sort` part of the query.
+	// This can be different from the OrderBy clause, as it may contain Elasticsearch-internal fields like `_doc`.
+	// In that case, it is not reflected in the OrderBy clause, but is still used for assembling the response.
+	CanParse bool
 	// NeedCountWithLimit > 0 means we need count(*) LIMIT NeedCountWithLimit
 	// NeedCountWithLimit 0 (WeNeedUnlimitedCount) means we need count(*) (unlimited)
 	// NeedCountWithLimit -1 (WeDontNeedCount) means we don't need a count(*) query

platform/model/typical_queries/hits.go

@@ -94,6 +94,8 @@ func (query Hits) TranslateSqlResponseToJson(rows []model.QueryResultRow) model.
 		for _, fieldName := range query.sortFieldNames {
 			if val, ok := hit.Fields[fieldName]; ok {
 				hit.Sort = append(hit.Sort, elasticsearch.FormatSortValue(val[0]))
+			} else if fieldName == "_doc" { // Kibana adds _doc as a tiebreaker field for sorting
+				hit.Sort = append(hit.Sort, hit.ID)
 			} else {
 				logger.WarnWithCtx(query.ctx).Msgf("field %s not found in fields", fieldName)
 			}

platform/parsers/elastic_query_dsl/query_parser.go

@@ -119,10 +119,10 @@ func (cw *ClickhouseQueryTranslator) buildListQueryIfNeeded(
 	}
 	if fullQuery != nil {
 		highlighter.SetTokensToHighlight(fullQuery.SelectCommand)
-		// TODO: pass right arguments
-		queryType := typical_queries.NewHits(cw.Ctx, cw.Table, &highlighter, fullQuery.SelectCommand.OrderByFieldNames(), true, false, false, cw.Indexes)
+		queryType := typical_queries.NewHits(cw.Ctx, cw.Table, &highlighter, simpleQuery.SortFieldNames, true, false, false, cw.Indexes)
 		fullQuery.Type = &queryType
 		fullQuery.Highlighter = highlighter
+		fullQuery.SearchAfterFieldNames = simpleQuery.SortFieldNames
 	}
 
 	return fullQuery
@@ -155,7 +155,7 @@ func (cw *ClickhouseQueryTranslator) parseQueryInternal(body types.JSON) (*model
 	}
 
 	if sortPart, ok := queryAsMap["sort"]; ok {
-		parsedQuery.OrderBy = cw.parseSortFields(sortPart)
+		parsedQuery.OrderBy, parsedQuery.SortFieldNames = cw.parseSortFields(sortPart)
 	}
 	size := cw.parseSize(queryAsMap, defaultQueryResultSize)
 
@@ -1080,8 +1080,9 @@ func (cw *ClickhouseQueryTranslator) extractInterval(queryMap QueryMap) (interva
 
 // parseSortFields parses sort fields from the query
 // We're skipping ELK internal fields, like "_doc", "_id", etc. (we only accept field starting with "_" if it exists in our table)
-func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns []model.OrderByExpr) {
+func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns []model.OrderByExpr, sortFieldNames []string) {
 	sortColumns = make([]model.OrderByExpr, 0)
+	sortFieldNames = make([]string, 0)
 	switch sortMaps := sortMaps.(type) {
 	case []any:
 		for _, sortMapAsAny := range sortMaps {
@@ -1093,6 +1094,7 @@ func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns
 
 			// sortMap has only 1 key, so we can just iterate over it
 			for k, v := range sortMap {
+				sortFieldNames = append(sortFieldNames, k)
 				// TODO replace cw.Table.GetFieldInfo with schema.Field[]
 				if strings.HasPrefix(k, "_") && cw.Table.GetFieldInfo(cw.Ctx, ResolveField(cw.Ctx, k, cw.Schema)) == clickhouse.NotExists {
 					// we're skipping ELK internal fields, like "_doc", "_id", etc.
@@ -1125,9 +1127,10 @@ func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns
 				}
 			}
 		}
-		return sortColumns
+		return sortColumns, sortFieldNames
 	case map[string]interface{}:
 		for fieldName, fieldValue := range sortMaps {
+			sortFieldNames = append(sortFieldNames, fieldName)
 			if strings.HasPrefix(fieldName, "_") && cw.Table.GetFieldInfo(cw.Ctx, ResolveField(cw.Ctx, fieldName, cw.Schema)) == clickhouse.NotExists {
 				// TODO Elastic internal fields will need to be supported in the future
 				continue
@@ -1141,10 +1144,11 @@ func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns
 			}
 		}
 
-		return sortColumns
+		return sortColumns, sortFieldNames
 
 	case map[string]string:
 		for fieldName, fieldValue := range sortMaps {
+			sortFieldNames = append(sortFieldNames, fieldName)
 			if strings.HasPrefix(fieldName, "_") && cw.Table.GetFieldInfo(cw.Ctx, ResolveField(cw.Ctx, fieldName, cw.Schema)) == clickhouse.NotExists {
 				// TODO Elastic internal fields will need to be supported in the future
 				continue
@@ -1156,10 +1160,10 @@ func (cw *ClickhouseQueryTranslator) parseSortFields(sortMaps any) (sortColumns
 			}
 		}
 
-		return sortColumns
+		return sortColumns, sortFieldNames
 	default:
 		logger.ErrorWithCtx(cw.Ctx).Msgf("unexpected type of sortMaps: %T, value: %v", sortMaps, sortMaps)
-		return []model.OrderByExpr{}
+		return []model.OrderByExpr{}, sortFieldNames
 	}
 }
 

platform/parsers/elastic_query_dsl/query_parser_test.go

@@ -215,9 +215,10 @@ func TestQueryParserNoAttrsConfig(t *testing.T) {
 
 func Test_parseSortFields(t *testing.T) {
 	tests := []struct {
-		name        string
-		sortMap     any
-		sortColumns []model.OrderByExpr
+		name           string
+		sortMap        any
+		sortColumns    []model.OrderByExpr
+		sortFieldNames []string
 	}{
 		{
 			name: "compound",
@@ -234,34 +235,39 @@ func Test_parseSortFields(t *testing.T) {
 				model.NewSortColumn("no_order_field", model.AscOrder),
 				model.NewSortColumn("_table_field_with_underscore", model.AscOrder),
 			},
+			sortFieldNames: []string{"@timestamp", "service.name", "no_order_field", "_table_field_with_underscore", "_doc"},
 		},
 		{
-			name:        "empty",
-			sortMap:     []any{},
-			sortColumns: []model.OrderByExpr{},
+			name:           "empty",
+			sortMap:        []any{},
+			sortColumns:    []model.OrderByExpr{},
+			sortFieldNames: []string{},
 		},
 		{
 			name: "map[string]string",
 			sortMap: map[string]string{
 				"timestamp": "desc",
 				"_doc":      "desc",
 			},
-			sortColumns: []model.OrderByExpr{model.NewSortColumn("timestamp", model.DescOrder)},
+			sortColumns:    []model.OrderByExpr{model.NewSortColumn("timestamp", model.DescOrder)},
+			sortFieldNames: []string{"timestamp", "_doc"},
 		},
 		{
 			name: "map[string]interface{}",
 			sortMap: map[string]interface{}{
 				"timestamp": "desc",
 				"_doc":      "desc",
 			},
-			sortColumns: []model.OrderByExpr{model.NewSortColumn("timestamp", model.DescOrder)},
+			sortColumns:    []model.OrderByExpr{model.NewSortColumn("timestamp", model.DescOrder)},
+			sortFieldNames: []string{"timestamp", "_doc"},
 		}, {
 			name: "[]map[string]string",
 			sortMap: []any{
 				QueryMap{"@timestamp": "asc"},
 				QueryMap{"_doc": "asc"},
 			},
-			sortColumns: []model.OrderByExpr{model.NewSortColumn("@timestamp", model.AscOrder)},
+			sortColumns:    []model.OrderByExpr{model.NewSortColumn("@timestamp", model.AscOrder)},
+			sortFieldNames: []string{"@timestamp", "_doc"},
 		},
 	}
 	table, _ := clickhouse.NewTable(`CREATE TABLE `+tableName+`
@@ -272,7 +278,9 @@ func Test_parseSortFields(t *testing.T) {
 	cw := ClickhouseQueryTranslator{Table: table, Ctx: context.Background()}
 	for i, tt := range tests {
 		t.Run(util.PrettyTestName(tt.name, i), func(t *testing.T) {
-			assert.Equal(t, tt.sortColumns, cw.parseSortFields(tt.sortMap))
+			orderBy, sortFieldNames := cw.parseSortFields(tt.sortMap)
+			assert.Equal(t, tt.sortColumns, orderBy)
+			assert.ElementsMatch(t, tt.sortFieldNames, sortFieldNames)
 		})
 	}
 }