Merge pull request #6225 from fpetkovski/grpc-range-query-response

Add support for vector and scalar types in gRPC query range
thanos-io · Mar 20, 2023 · 36de497 · 36de497
2 parents 143c49a + 813240c
commit 36de497
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 8 deletions.
diff --git a/pkg/api/query/grpc.go b/pkg/api/query/grpc.go
@@ -126,7 +126,7 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
 		}
 	case promql.Vector:
 		for _, sample := range vector {
-			floats, histograms := prompb.SamplesFromPromqlPoints([]promql.Point{sample.Point})
+			floats, histograms := prompb.SamplesFromPromqlPoints(sample.Point)
 			series := &prompb.TimeSeries{
 				Labels:     labelpb.ZLabelsFromPromLabels(sample.Metric),
 				Samples:    floats,
@@ -136,7 +136,6 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
 				return err
 			}
 		}
-
 		return nil
 	}
 
@@ -147,7 +146,8 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
 	ctx := srv.Context()
 	if request.TimeoutSeconds != 0 {
 		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, time.Duration(request.TimeoutSeconds))
+		timeout := time.Duration(request.TimeoutSeconds) * time.Second
+		ctx, cancel = context.WithTimeout(ctx, timeout)
 		defer cancel()
 	}
 
@@ -204,10 +204,10 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
 		}
 	}
 
-	switch matrix := result.Value.(type) {
+	switch value := result.Value.(type) {
 	case promql.Matrix:
-		for _, series := range matrix {
-			floats, histograms := prompb.SamplesFromPromqlPoints(series.Points)
+		for _, series := range value {
+			floats, histograms := prompb.SamplesFromPromqlPoints(series.Points...)
 			series := &prompb.TimeSeries{
 				Labels:     labelpb.ZLabelsFromPromLabels(series.Metric),
 				Samples:    floats,
@@ -217,8 +217,25 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
 				return err
 			}
 		}
-
+	case promql.Vector:
+		for _, sample := range value {
+			point := promql.Point{T: sample.T, V: sample.V, H: sample.H}
+			floats, histograms := prompb.SamplesFromPromqlPoints(point)
+			series := &prompb.TimeSeries{
+				Labels:     labelpb.ZLabelsFromPromLabels(sample.Metric),
+				Samples:    floats,
+				Histograms: histograms,
+			}
+			if err := srv.Send(querypb.NewQueryRangeResponse(series)); err != nil {
+				return err
+			}
+		}
 		return nil
+	case promql.Scalar:
+		series := &prompb.TimeSeries{
+			Samples: []prompb.Sample{{Value: value.V, Timestamp: value.T}},
+		}
+		return srv.Send(querypb.NewQueryRangeResponse(series))
 	}
 
 	return nil

diff --git a/pkg/store/storepb/prompb/samples.go b/pkg/store/storepb/prompb/samples.go
@@ -25,7 +25,7 @@ func SamplesFromSamplePairs(samples []model.SamplePair) []Sample {
 
 // SamplesFromPromqlPoints converts a slice of promql.Point
 // to a slice of Sample.
-func SamplesFromPromqlPoints(samples []promql.Point) ([]Sample, []Histogram) {
+func SamplesFromPromqlPoints(samples ...promql.Point) ([]Sample, []Histogram) {
 	floats := make([]Sample, 0, len(samples))
 	histograms := make([]Histogram, 0, len(samples))
 	for _, s := range samples {