[v9.5.x] Prometheus: Handle the response with different field key order (#74734)

Prometheus: Handle the response with different field key order (#74567) * Handle the response with different field key order * More unit tests to cover edge cases * Cover more edge cases * make it simpler * Better test inputs (cherry picked from commit 3107459e57)
2023-09-12 16:00:57 +02:00
parent 25a403cf78
commit eb4b6e45b7
3 changed files with 169 additions and 24 deletions
@@ -0,0 +1,57 @@
+package querydata
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/grafana/grafana/pkg/tsdb/prometheus/models"
+	"github.com/grafana/grafana/pkg/tsdb/prometheus/querydata/exemplar"
+)
+
+func TestQueryData_parseResponse(t *testing.T) {
+	qd := QueryData{exemplarSampler: exemplar.NewStandardDeviationSampler}
+
+	t.Run("resultType is before result the field must parsed normally", func(t *testing.T) {
+		resBody := `{"data":{"resultType":"vector", "result":[{"metric":{"__name__":"some_name","environment":"some_env","id":"some_id","instance":"some_instance:1234","job":"some_job","name":"another_name","region":"some_region"},"value":[1.1,"2"]}]},"status":"success"}`
+		res := &http.Response{Body: io.NopCloser(bytes.NewBufferString(resBody))}
+		result := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Nil(t, result.Error)
+		assert.Len(t, result.Frames, 1)
+	})
+
+	t.Run("resultType is after the result field must parsed normally", func(t *testing.T) {
+		resBody := `{"data":{"result":[{"metric":{"__name__":"some_name","environment":"some_env","id":"some_id","instance":"some_instance:1234","job":"some_job","name":"another_name","region":"some_region"},"value":[1.1,"2"]}],"resultType":"vector"},"status":"success"}`
+		res := &http.Response{Body: io.NopCloser(bytes.NewBufferString(resBody))}
+		result := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Nil(t, result.Error)
+		assert.Len(t, result.Frames, 1)
+	})
+
+	t.Run("no resultType is existed in the data", func(t *testing.T) {
+		resBody := `{"data":{"result":[{"metric":{"__name__":"some_name","environment":"some_env","id":"some_id","instance":"some_instance:1234","job":"some_job","name":"another_name","region":"some_region"},"value":[1.1,"2"]}]},"status":"success"}`
+		res := &http.Response{Body: io.NopCloser(bytes.NewBufferString(resBody))}
+		result := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Nil(t, result.Frames[0].Fields)
+	})
+
+	t.Run("resultType is set as empty string before result", func(t *testing.T) {
+		resBody := `{"data":{"resultType":"", "result":[{"metric":{"__name__":"some_name","environment":"some_env","id":"some_id","instance":"some_instance:1234","job":"some_job","name":"another_name","region":"some_region"},"value":[1.1,"2"]}]},"status":"success"}`
+		res := &http.Response{Body: io.NopCloser(bytes.NewBufferString(resBody))}
+		result := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Error(t, result.Error)
+		assert.Equal(t, result.Error.Error(), "unknown result type: ")
+	})
+
+	t.Run("resultType is set as empty string after result", func(t *testing.T) {
+		resBody := `{"data":{"result":[{"metric":{"__name__":"some_name","environment":"some_env","id":"some_id","instance":"some_instance:1234","job":"some_job","name":"another_name","region":"some_region"},"value":[1.1,"2"]}],"resultType":""},"status":"success"}`
+		res := &http.Response{Body: io.NopCloser(bytes.NewBufferString(resBody))}
+		result := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Error(t, result.Error)
+		assert.Equal(t, result.Error.Error(), "unknown result type: ")
+	})
+}
@@ -0,0 +1,66 @@
+// Package jsonitere wraps json-iterator/go's Iterator methods with error returns
+// so linting can catch unchecked errors.
+// The underlying iterator's Error property is returned and not reset.
+// See json-iterator/go for method documentation and additional methods that
+// can be added to this library.
+package jsonitere
+
+import (
+	j "github.com/json-iterator/go"
+)
+
+type Iterator struct {
+	// named property instead of embedded so there is no
+	// confusion about which method or property is called
+	i *j.Iterator
+}
+
+func NewIterator(i *j.Iterator) *Iterator {
+	return &Iterator{i}
+}
+
+func (iter *Iterator) Read() (any, error) {
+	return iter.i.Read(), iter.i.Error
+}
+
+func (iter *Iterator) ReadAny() (j.Any, error) {
+	return iter.i.ReadAny(), iter.i.Error
+}
+
+func (iter *Iterator) ReadArray() (bool, error) {
+	return iter.i.ReadArray(), iter.i.Error
+}
+
+func (iter *Iterator) ReadObject() (string, error) {
+	return iter.i.ReadObject(), iter.i.Error
+}
+
+func (iter *Iterator) ReadString() (string, error) {
+	return iter.i.ReadString(), iter.i.Error
+}
+
+func (iter *Iterator) WhatIsNext() (j.ValueType, error) {
+	return iter.i.WhatIsNext(), iter.i.Error
+}
+
+func (iter *Iterator) Skip() error {
+	iter.i.Skip()
+	return iter.i.Error
+}
+
+func (iter *Iterator) SkipAndReturnBytes() []byte {
+	return iter.i.SkipAndReturnBytes()
+}
+
+func (iter *Iterator) ReadVal(obj any) error {
+	iter.i.ReadVal(obj)
+	return iter.i.Error
+}
+
+func (iter *Iterator) ReadFloat64() (float64, error) {
+	return iter.i.ReadFloat64(), iter.i.Error
+}
+
+func (iter *Iterator) ReadInt8() (int8, error) {
+	return iter.i.ReadInt8(), iter.i.Error
+}
@@ -106,37 +106,30 @@ func readPrometheusData(iter *jsoniter.Iterator, opt Options) backend.DataRespon

 	resultType := ""
 	var rsp backend.DataResponse
+	resultTypeFound := false
+	var resultBytes []byte

 	for l1Field := iter.ReadObject(); l1Field != ""; l1Field = iter.ReadObject() {
 		switch l1Field {
 		case "resultType":
 			resultType = iter.ReadString()
+			resultTypeFound = true
+
+			// if we have saved resultBytes we will parse them here
+			// we saved them because when we had them we don't know the resultType
+			if len(resultBytes) > 0 {
+				ji := jsoniter.ParseBytes(jsoniter.ConfigDefault, resultBytes)
+				rsp = readResult(resultType, rsp, ji, opt)
+			}

 		case "result":
-			switch resultType {
-			case "matrix":
-				if opt.MatrixWideSeries {
-					rsp = readMatrixOrVectorWide(iter, resultType, opt)
-				} else {
-					rsp = readMatrixOrVectorMulti(iter, resultType, opt)
-				}
-			case "vector":
-				if opt.VectorWideSeries {
-					rsp = readMatrixOrVectorWide(iter, resultType, opt)
-				} else {
-					rsp = readMatrixOrVectorMulti(iter, resultType, opt)
-				}
-			case "streams":
-				rsp = readStream(iter)
-			case "string":
-				rsp = readString(iter)
-			case "scalar":
-				rsp = readScalar(iter)
-			default:
-				iter.Skip()
-				rsp = backend.DataResponse{
-					Error: fmt.Errorf("unknown result type: %s", resultType),
-				}
+			// for some rare cases resultType is coming after the result.
+			// when that happens we save the bytes and parse them after reading resultType
+			// see: https://github.com/grafana/grafana/issues/64693
+			if resultTypeFound {
+				rsp = readResult(resultType, rsp, iter, opt)
+			} else {
+				resultBytes = iter.SkipAndReturnBytes()
 			}

 		case "stats":
@@ -157,7 +150,36 @@ func readPrometheusData(iter *jsoniter.Iterator, opt Options) backend.DataRespon
 			logf("[data] TODO, support key: %s / %v\n", l1Field, v)
 		}
 	}
+	return rsp
+}

+// will read the result object based on the resultType and return a DataResponse
+func readResult(resultType string, rsp backend.DataResponse, iter *jsoniter.Iterator, opt Options) backend.DataResponse {
+	switch resultType {
+	case "matrix":
+		if opt.MatrixWideSeries {
+			rsp = readMatrixOrVectorWide(iter, resultType, opt)
+		} else {
+			rsp = readMatrixOrVectorMulti(iter, resultType, opt)
+		}
+	case "vector":
+		if opt.VectorWideSeries {
+			rsp = readMatrixOrVectorWide(iter, resultType, opt)
+		} else {
+			rsp = readMatrixOrVectorMulti(iter, resultType, opt)
+		}
+	case "streams":
+		rsp = readStream(iter)
+	case "string":
+		rsp = readString(iter)
+	case "scalar":
+		rsp = readScalar(iter)
+	default:
+		iter.Skip()
+		rsp = backend.DataResponse{
+			Error: fmt.Errorf("unknown result type: %s", resultType),
+		}
+	}
 	return rsp
 }