feat: support for window based pagination in new trace v4

pkg/query-service/app/querier/querier.go

@@ -52,7 +52,8 @@ type querier struct {
 	returnedSeries []*v3.Series
 	returnedErr    error
 
-	UseLogsNewSchema bool
+	UseLogsNewSchema  bool
+	UseTraceNewSchema bool
 }
 
 type QuerierOptions struct {
@@ -308,15 +309,17 @@ func (q *querier) runClickHouseQueries(ctx context.Context, params *v3.QueryRang
 	return results, errQueriesByName, err
 }
 
-func (q *querier) runLogsListQuery(ctx context.Context, params *v3.QueryRangeParamsV3, tsRanges []utils.LogsListTsRange) ([]*v3.Result, map[string]error, error) {
+func (q *querier) runWindowBasedListQuery(ctx context.Context, params *v3.QueryRangeParamsV3, tsRanges []utils.LogsListTsRange) ([]*v3.Result, map[string]error, error) {
 	res := make([]*v3.Result, 0)
 	qName := ""
 	pageSize := uint64(0)
+	limit := uint64(0)
 
 	// se we are considering only one query
 	for name, v := range params.CompositeQuery.BuilderQueries {
 		qName = name
 		pageSize = v.PageSize
+		limit = v.Limit
 	}
 	data := []*v3.Row{}
 
@@ -330,6 +333,7 @@ func (q *querier) runLogsListQuery(ctx context.Context, params *v3.QueryRangePar
 			return nil, nil, err
 		}
 
+		length := uint64(0)
 		// this will to run only once
 		for name, query := range queries {
 			rowList, err := q.reader.GetListResultV3(ctx, query)
@@ -340,24 +344,50 @@ func (q *querier) runLogsListQuery(ctx context.Context, params *v3.QueryRangePar
 				}
 				return nil, errQuriesByName, fmt.Errorf("encountered multiple errors: %s", multierr.Combine(errs...))
 			}
+			length += uint64(len(rowList))
 			data = append(data, rowList...)
 		}
 
-		// append a filter to the params
-		if len(data) > 0 {
-			params.CompositeQuery.BuilderQueries[qName].Filters.Items = append(params.CompositeQuery.BuilderQueries[qName].Filters.Items, v3.FilterItem{
-				Key: v3.AttributeKey{
-					Key:      "id",
-					IsColumn: true,
-					DataType: "string",
-				},
-				Operator: v3.FilterOperatorLessThan,
-				Value:    data[len(data)-1].Data["id"],
-			})
-		}
+		// appending the filter to get the next set of data
+		if params.CompositeQuery.BuilderQueries[qName].DataSource == v3.DataSourceLogs {
+			if length > 0 {
+				params.CompositeQuery.BuilderQueries[qName].Filters.Items = append(params.CompositeQuery.BuilderQueries[qName].Filters.Items, v3.FilterItem{
+					Key: v3.AttributeKey{
+						Key:      "id",
+						IsColumn: true,
+						DataType: "string",
+					},
+					Operator: v3.FilterOperatorLessThan,
+					Value:    data[len(data)-1].Data["id"],
+				})
+			}
+
+			if uint64(len(data)) >= pageSize {
+				break
+			}
+		} else {
+			// we are updating the offset and limit based on the number of traces we have found in the current timerange
+			// eg -
+			// 1)offset = 0, limit = 100, tsRanges = [t1, t10], [t10, 20], [t20, t30]
+			//
+			// if 100 traces are there in [t1, t10] then 100 will return immediately.
+			// if 10 traces are there in [t1, t10] then we get 10, set offset to 0 and limit to 90, search in the next timerange of [t10, 20]
+			// if we don't find any trace in [t1, t10], then we search in [t10, 20] with offset=0, limit=100
+
+			//
+			// 2) offset = 50, limit = 100, tsRanges = [t1, t10], [t10, 20], [t20, t30]
+			//
+			// If we find 100 traces in [t1, t10] then we return immediately
+			// If we find 50 in [t1, t10] then it will set offset = 0 and limit = 50 and search in the next timerange of [t10, 20]
+			// if we don't find any trace in [t1, t10], then we search in [t10, 20] with offset=50, limit=100
+			if length > 0 {
+				params.CompositeQuery.BuilderQueries[qName].Offset = 0
+				params.CompositeQuery.BuilderQueries[qName].Limit = limit - uint64(len(data))
+			}
 
-		if uint64(len(data)) >= pageSize {
-			break
+			if uint64(len(data)) >= limit {
+				break
+			}
 		}
 	}
 	res = append(res, &v3.Result{
@@ -368,15 +398,25 @@ func (q *querier) runLogsListQuery(ctx context.Context, params *v3.QueryRangePar
 }
 
 func (q *querier) runBuilderListQueries(ctx context.Context, params *v3.QueryRangeParamsV3) ([]*v3.Result, map[string]error, error) {
-	// List query has support for only one query.
-	if q.UseLogsNewSchema && params.CompositeQuery != nil && len(params.CompositeQuery.BuilderQueries) == 1 {
+	// List query has support for only one query
+	// we are skipping for PanelTypeTrace as it has a custom order by regardless of what's in the payload
+	if params.CompositeQuery != nil &&
+		len(params.CompositeQuery.BuilderQueries) == 1 &&
+		params.CompositeQuery.PanelType != v3.PanelTypeTrace {
 		for _, v := range params.CompositeQuery.BuilderQueries {
+			if (v.DataSource == v3.DataSourceLogs && !q.UseLogsNewSchema) ||
+				(v.DataSource == v3.DataSourceTraces && !q.UseTraceNewSchema) {
+				break
+			}
+
 			// only allow of logs queries with timestamp ordering desc
-			if v.DataSource == v3.DataSourceLogs && len(v.OrderBy) == 1 && v.OrderBy[0].ColumnName == "timestamp" && v.OrderBy[0].Order == "desc" {
-				startEndArr := utils.GetLogsListTsRanges(params.Start, params.End)
-				if len(startEndArr) > 0 {
-					return q.runLogsListQuery(ctx, params, startEndArr)
-				}
+			// TODO(nitya): allow for timestamp asc
+			if (v.DataSource == v3.DataSourceLogs || v.DataSource == v3.DataSourceTraces) &&
+				len(v.OrderBy) == 1 &&
+				v.OrderBy[0].ColumnName == "timestamp" &&
+				v.OrderBy[0].Order == "desc" {
+				startEndArr := utils.GetListTsRanges(params.Start, params.End)
+				return q.runWindowBasedListQuery(ctx, params, startEndArr)
 			}
 		}
 	}