Prometheus: Add Exemplar sampling for streaming parser (#56049)

pkg/tsdb/prometheus/models/query.go

@@ -133,8 +133,8 @@ func (query *Query) TimeRange() TimeRange {
 	return TimeRange{
 		Step: query.Step,
 		// Align query range to step. It rounds start and end down to a multiple of step.
-		Start: alignTimeRange(query.Start, query.Step, query.UtcOffsetSec),
-		End:   alignTimeRange(query.End, query.Step, query.UtcOffsetSec),
+		Start: AlignTimeRange(query.Start, query.Step, query.UtcOffsetSec),
+		End:   AlignTimeRange(query.End, query.Step, query.UtcOffsetSec),
 	}
 }
 
@@ -225,6 +225,8 @@ func isVariableInterval(interval string) bool {
 	return false
 }
 
-func alignTimeRange(t time.Time, step time.Duration, offset int64) time.Time {
-	return time.Unix(int64(math.Floor((float64(t.Unix()+offset)/step.Seconds()))*step.Seconds()-float64(offset)), 0)
+func AlignTimeRange(t time.Time, step time.Duration, offset int64) time.Time {
+	offsetNano := float64(offset * 1e9)
+	stepNano := float64(step.Nanoseconds())
+	return time.Unix(0, int64(math.Floor((float64(t.UnixNano())+offsetNano)/stepNano)*stepNano-offsetNano)).UTC()
 }

pkg/tsdb/prometheus/models/query_test.go

@@ -1,6 +1,7 @@
 package models_test
 
 import (
+	"reflect"
 	"testing"
 	"time"
 
@@ -15,7 +16,7 @@ var (
 	intervalCalculator = intervalv2.NewCalculator()
 )
 
-func TestPrometheus_timeSeriesQuery_parseTimeSeriesQuery(t *testing.T) {
+func TestParse(t *testing.T) {
 	t.Run("parsing query from unified alerting", func(t *testing.T) {
 		timeRange := backend.TimeRange{
 			From: now,
@@ -449,3 +450,27 @@ func queryContext(json string, timeRange backend.TimeRange) backend.DataQuery {
 		RefID:     "A",
 	}
 }
+
+func TestAlignTimeRange(t *testing.T) {
+	type args struct {
+		t      time.Time
+		step   time.Duration
+		offset int64
+	}
+	tests := []struct {
+		name string
+		args args
+		want time.Time
+	}{
+		{name: "second step", args: args{t: time.Unix(1664816826, 0), step: 10 * time.Second, offset: 0}, want: time.Unix(1664816820, 0).UTC()},
+		{name: "millisecond step", args: args{t: time.Unix(1664816825, 5*int64(time.Millisecond)), step: 10 * time.Millisecond, offset: 0}, want: time.Unix(1664816825, 0).UTC()},
+		{name: "second step with offset", args: args{t: time.Unix(1664816825, 5*int64(time.Millisecond)), step: 2 * time.Second, offset: -3}, want: time.Unix(1664816825, 0).UTC()},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := models.AlignTimeRange(tt.args.t, tt.args.step, tt.args.offset); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("AlignTimeRange() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

pkg/tsdb/prometheus/models/result.go

@@ -0,0 +1,42 @@
+package models
+
+import "github.com/grafana/grafana-plugin-sdk-go/data"
+
+type ResultType string
+
+const (
+	ResultTypeMatrix   ResultType = "matrix"
+	ResultTypeExemplar ResultType = "exemplar"
+	ResultTypeVector   ResultType = "vector"
+	ResultTypeUnknown  ResultType = ""
+)
+
+func ResultTypeFromFrame(frame *data.Frame) ResultType {
+	if frame.Meta.Custom == nil {
+		return ResultTypeUnknown
+	}
+	custom, ok := frame.Meta.Custom.(map[string]string)
+	if !ok {
+		return ResultTypeUnknown
+	}
+
+	rt, ok := custom["resultType"]
+	if !ok {
+		return ResultTypeUnknown
+	}
+
+	switch rt {
+	case ResultTypeMatrix.String():
+		return ResultTypeMatrix
+	case ResultTypeExemplar.String():
+		return ResultTypeExemplar
+	case ResultTypeVector.String():
+		return ResultTypeVector
+	}
+
+	return ResultTypeUnknown
+}
+
+func (r ResultType) String() string {
+	return string(r)
+}