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

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-11 13:41:29 +02:00
parent 82ac2432d4
commit aa758499ea
2 changed files with 106 additions and 24 deletions
@@ -0,0 +1,61 @@
+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, err := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Nil(t, err)
+		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, err := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Nil(t, err)
+		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.Equal(t, 0, len(result.Frames))
+	})
+
+	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, err := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Equal(t, err.Error(), "unknown result type: ")
+		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, err := qd.parseResponse(context.Background(), &models.Query{}, res)
+		assert.Equal(t, err.Error(), "unknown result type: ")
+		assert.Error(t, result.Error)
+		assert.Equal(t, result.Error.Error(), "unknown result type: ")
+	})
+}
@@ -104,38 +104,30 @@ func readPrometheusData(iter *jsoniter.Iterator, opt Options) backend.DataRespon
 	}

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

 	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)
-				} else {
-					rsp = readMatrixOrVectorMulti(iter, resultType)
-				}
-			case "vector":
-				if opt.VectorWideSeries {
-					rsp = readMatrixOrVectorWide(iter, resultType)
-				} else {
-					rsp = readMatrixOrVectorMulti(iter, resultType)
-				}
-			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":
@@ -156,7 +148,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)
+		} else {
+			rsp = readMatrixOrVectorMulti(iter, resultType)
+		}
+	case "vector":
+		if opt.VectorWideSeries {
+			rsp = readMatrixOrVectorWide(iter, resultType)
+		} else {
+			rsp = readMatrixOrVectorMulti(iter, resultType)
+		}
+	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
 }