From b71b27ab5d896d6968d84b8c5631814e1a2e4da6 Mon Sep 17 00:00:00 2001
From: Idriss Neumann <idriss.neumann@comwork.io>
Date: Tue, 28 Nov 2023 08:20:25 +0100
Subject: [PATCH] Issue #18: searching keyword to highlight from the backend
 side

---
 pkg/quickwit/response_parser.go      | 41 ++++++++++++++++------------
 pkg/quickwit/response_parser_test.go | 30 ++++++++++++++++++++
 2 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/pkg/quickwit/response_parser.go b/pkg/quickwit/response_parser.go
index 818d0bc..4d83090 100644
--- a/pkg/quickwit/response_parser.go
+++ b/pkg/quickwit/response_parser.go
@@ -95,11 +95,30 @@ func parseResponse(responses []*es.SearchResponse, targets []*Query, configuredF
 	return &result, nil
 }
 
+func parseLuceneQuery(query string) []string {
+	var keywords []string
+
+	termRegex := regexp.MustCompile(`("[^"]+"|\S+)`)
+	matches := termRegex.FindAllString(query, -1)
+
+	for _, match := range matches {
+		if match[0] == '"' && match[len(match)-1] == '"' {
+			match = match[1 : len(match)-1]
+		}
+
+		keywords = append(keywords, strings.ReplaceAll(match, "*", ""))
+	}
+
+	return keywords
+}
+
 func processLogsResponse(res *es.SearchResponse, target *Query, configuredFields es.ConfiguredFields, queryRes *backend.DataResponse) error {
 	propNames := make(map[string]bool)
 	docs := make([]map[string]interface{}, len(res.Hits.Hits))
 	searchWords := make(map[string]bool)
 
+	highlights := parseLuceneQuery(target.RawQuery)
+
 	for hitIdx, hit := range res.Hits.Hits {
 		var flattened map[string]interface{}
 		if hit["_source"] != nil {
@@ -132,23 +151,6 @@ func processLogsResponse(res *es.SearchResponse, target *Query, configuredFields
 			propNames[key] = true
 		}
 
-		// FIXME: Quickwit does not support highlight. Should we replace this by a custom highlighter?
-		// Process highlight to searchWords
-		if highlights, ok := doc["highlight"].(map[string]interface{}); ok {
-			for _, highlight := range highlights {
-				if highlightList, ok := highlight.([]interface{}); ok {
-					for _, highlightValue := range highlightList {
-						str := fmt.Sprintf("%v", highlightValue)
-						matches := searchWordsRegex.FindAllStringSubmatch(str, -1)
-
-						for _, v := range matches {
-							searchWords[v[1]] = true
-						}
-					}
-				}
-			}
-		}
-
 		docs[hitIdx] = doc
 	}
 
@@ -158,6 +160,11 @@ func processLogsResponse(res *es.SearchResponse, target *Query, configuredFields
 	frames := data.Frames{}
 	frame := data.NewFrame("", fields...)
 	setPreferredVisType(frame, data.VisTypeLogs)
+
+	for _, keyword := range highlights {
+		searchWords[keyword] = true
+	}
+
 	setLogsCustomMeta(frame, searchWords, stringToIntWithDefaultValue(target.Metrics[0].Settings.Get("limit").MustString(), defaultSize))
 	frames = append(frames, frame)
 
diff --git a/pkg/quickwit/response_parser_test.go b/pkg/quickwit/response_parser_test.go
index e7f111d..5bec429 100644
--- a/pkg/quickwit/response_parser_test.go
+++ b/pkg/quickwit/response_parser_test.go
@@ -3178,6 +3178,36 @@ func TestLabelOrderInFieldName(t *testing.T) {
 	requireTimeSeriesName(t, "val1 error", frames[5])
 }
 
+func TestParseLuceneQuery(t *testing.T) {
+	t.Run("Empty term query", func(t *testing.T) {
+		query := ""
+		highlights := parseLuceneQuery(query)
+		require.Len(t, highlights, 0)
+	})
+
+	t.Run("Simple term query", func(t *testing.T) {
+		query := "foo"
+		highlights := parseLuceneQuery(query)
+		require.Len(t, highlights, 1)
+		require.Equal(t, "foo", highlights[0])
+	})
+
+	t.Run("Multi term query", func(t *testing.T) {
+		query := "foo bar"
+		highlights := parseLuceneQuery(query)
+		require.Len(t, highlights, 2)
+		require.Equal(t, "foo", highlights[0])
+		require.Equal(t, "bar", highlights[1])
+	})
+
+	t.Run("Wildcard query", func(t *testing.T) {
+		query := "foo*"
+		highlights := parseLuceneQuery(query)
+		require.Len(t, highlights, 1)
+		require.Equal(t, "foo", highlights[0])
+	})
+}
+
 func TestFlatten(t *testing.T) {
 	t.Run("Flattens simple object", func(t *testing.T) {
 		obj := map[string]interface{}{