Hit response refactor (#289)

quesma/model/query.go

@@ -46,7 +46,6 @@ type (
 		NoDBQuery   bool         // true <=> we don't need query to DB here, true in some pipeline aggregations
 		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.
-		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_parser.go

@@ -7,6 +7,7 @@ import (
 	"mitmproxy/quesma/clickhouse"
 	"mitmproxy/quesma/logger"
 	"mitmproxy/quesma/model"
+	"mitmproxy/quesma/model/typical_queries"
 	"mitmproxy/quesma/queryparser/lucene"
 	wc "mitmproxy/quesma/queryparser/where_clause"
 	"mitmproxy/quesma/quesma/types"
@@ -75,6 +76,9 @@ func (cw *ClickhouseQueryTranslator) buildListQueryIfNeeded(
 	}
 	if fullQuery != nil {
 		fullQuery.QueryInfoType = queryInfo.Typ
+		// TODO: pass right arguments
+		queryType := typical_queries.NewHits(cw.Ctx, cw.Table, &highlighter, fullQuery.OrderByFieldNames(), true, false, false)
+		fullQuery.Type = &queryType
 		fullQuery.Highlighter = highlighter
 	}
 

quesma/queryparser/query_translator.go

@@ -99,14 +99,6 @@ func (cw *ClickhouseQueryTranslator) MakeAsyncSearchResponse(ResultSet []model.Q
 	return &response, nil
 }
 
-func (cw *ClickhouseQueryTranslator) MakeAsyncSearchResponseMarshalled(ResultSet []model.QueryResultRow, query *model.Query, asyncRequestIdStr string, isPartial bool) ([]byte, error) {
-	response, err := cw.MakeAsyncSearchResponse(ResultSet, query, asyncRequestIdStr, isPartial)
-	if err != nil {
-		return nil, err
-	}
-	return response.Marshal()
-}
-
 func (cw *ClickhouseQueryTranslator) finishMakeResponse(query *model.Query, ResultSet []model.QueryResultRow, level int) []model.JsonMap {
 	// fmt.Println("FinishMakeResponse", query, ResultSet, level, query.Type.String())
 	if query.Type.IsBucketAggregation() {
@@ -231,33 +223,42 @@ func (cw *ClickhouseQueryTranslator) MakeAggregationPartOfResponse(queries []*mo
 	return aggregations
 }
 
-func (cw *ClickhouseQueryTranslator) MakeSearchResponse(queries []*model.Query, ResultSets [][]model.QueryResultRow) *model.SearchResp {
-	var hitsQuery *model.Query
-	var hitsResultSet []model.QueryResultRow
-
-	// Process hits as last aggregation
-	if len(queries) > 0 && len(ResultSets) > 0 && query_util.IsNonAggregationQuery(queries[len(queries)-1]) {
-		hitsQuery = queries[len(queries)-1]
-		hitsResultSet = ResultSets[len(ResultSets)-1]
+func (cw *ClickhouseQueryTranslator) makeHits(queries []*model.Query, results [][]model.QueryResultRow) (queriesWithoutHits []*model.Query, resultsWithoutHits [][]model.QueryResultRow, hit *model.SearchHits) {
+	hitsIndex := -1
+	for i, query := range queries {
+		if query.QueryInfoType == model.ListAllFields || query.QueryInfoType == model.ListByField {
+			if hitsIndex != -1 {
+				logger.WarnWithCtx(cw.Ctx).Msgf("multiple hits queries found in queries: %v", queries)
+			}
+			hitsIndex = i
+		} else {
+			queriesWithoutHits = append(queriesWithoutHits, query)
+			resultsWithoutHits = append(resultsWithoutHits, results[i])
+		}
+	}
 
-		queries = queries[:len(queries)-1]
-		ResultSets = ResultSets[:len(ResultSets)-1]
-	} else {
-		hitsResultSet = make([]model.QueryResultRow, 0)
+	if hitsIndex == -1 {
+		return queriesWithoutHits, resultsWithoutHits, nil
 	}
 
-	var highlighter *model.Highlighter
-	var orderByFieldNames []string
-	if hitsQuery != nil {
-		highlighter = &hitsQuery.Highlighter
-		orderByFieldNames = hitsQuery.OrderByFieldNames()
+	hitsQuery := queries[hitsIndex]
+	hitsResultSet := results[hitsIndex]
+
+	if hitsQuery.Type == nil {
+		logger.ErrorWithCtx(cw.Ctx).Msgf("hits query type is nil: %v", hitsQuery)
+		return queriesWithoutHits, resultsWithoutHits, nil
 	}
+	hitsPartOfResponse := hitsQuery.Type.TranslateSqlResponseToJson(hitsResultSet, 0)
 
-	// TODO it should be created during parsing, like aggregations
-	hits := typical_queries.NewHits(cw.Ctx, cw.Table, highlighter, orderByFieldNames, true, false, false)
-	hitsPartOfResponse := hits.TranslateSqlResponseToJson(hitsResultSet, 0)
+	hitsResponse := hitsPartOfResponse[0]["hits"].(model.SearchHits)
+	return queriesWithoutHits, resultsWithoutHits, &hitsResponse
+}
 
-	// process count:
+func (cw *ClickhouseQueryTranslator) MakeSearchResponse(queries []*model.Query, ResultSets [][]model.QueryResultRow) *model.SearchResp {
+	var hits *model.SearchHits
+	queries, ResultSets, hits = cw.makeHits(queries, ResultSets) // get hits and remove it from queries
+
+	// TODO: process count:
 	// a) we have count query -> we're done
 	// b) we have hits or facets -> we're done
 	// c) we don't have above: we return len(biggest resultset(all aggregations))
@@ -291,12 +292,14 @@ func (cw *ClickhouseQueryTranslator) MakeSearchResponse(queries []*model.Query,
 			Failed:     0,
 		},
 	}
-	if hitsTyped, ok := hitsPartOfResponse[0]["hits"].(model.SearchHits); ok {
-		response.Hits = hitsTyped
-		response.Hits.Total = &model.Total{
-			Value:    int(totalCount),
-			Relation: "eq",
-		}
+	if hits != nil {
+		response.Hits = *hits
+	} else {
+		response.Hits = model.SearchHits{}
+	}
+	response.Hits.Total = &model.Total{
+		Value:    int(totalCount),
+		Relation: "eq",
 	}
 	return response
 }

quesma/queryparser/query_translator_test.go

@@ -60,8 +60,10 @@ const (
 func TestSearchResponse(t *testing.T) {
 	row := []model.QueryResultRow{{}}
 	cw := ClickhouseQueryTranslator{Table: &clickhouse.Table{Name: "test"}, Ctx: context.Background()}
-	searchRespBuf, err := cw.MakeAsyncSearchResponseMarshalled(row, &model.Query{QueryInfoType: model.ListAllFields, Highlighter: NewEmptyHighlighter()}, asyncRequestIdStr, false)
+	searchResp, err := cw.MakeAsyncSearchResponse(row, &model.Query{QueryInfoType: model.ListAllFields, Highlighter: NewEmptyHighlighter()}, asyncRequestIdStr, false)
 	require.NoError(t, err)
+	searchRespBuf, err2 := searchResp.Marshal()
+	require.NoError(t, err2)
 	var searchResponseResult model.SearchResp
 	err = json.Unmarshal(searchRespBuf, &searchResponseResult)
 	require.NoError(t, err)
@@ -447,10 +449,12 @@ func TestMakeResponseAsyncSearchQuery(t *testing.T) {
 			if i != 0 {
 				t.Skip()
 			}
-			ourResponse, err := cw.MakeAsyncSearchResponseMarshalled(args[i].ourQueryResult, tt.query, asyncRequestIdStr, false)
+			ourResponse, err := cw.MakeAsyncSearchResponse(args[i].ourQueryResult, tt.query, asyncRequestIdStr, false)
 			assert.NoError(t, err)
+			ourResponseBuf, err2 := ourResponse.Marshal()
+			assert.NoError(t, err2)
 
-			actualMinusExpected, expectedMinusActual, err := util.JsonDifference(string(ourResponse), args[i].elasticResponseJson)
+			actualMinusExpected, expectedMinusActual, err := util.JsonDifference(string(ourResponseBuf), args[i].elasticResponseJson)
 			pp.Println(actualMinusExpected, expectedMinusActual)
 			assert.NoError(t, err)