feat: add validation for exemplar types in timeseries queries (#4631)

Author: marcsanmi

pkg/querybackend/query_time_series.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/grafana/dskit/runutil"
 	"github.com/parquet-go/parquet-go"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
@@ -31,11 +33,15 @@ func init() {
 }
 
 func queryTimeSeries(q *queryContext, query *queryv1.Query) (r *queryv1.Report, err error) {
+	includeExemplars, err := validateExemplarType(query.TimeSeries.ExemplarType)
+	if err != nil {
+		return nil, err
+	}
+
 	opts := []profileIteratorOption{
 		withFetchPartition(false), // Partition data not needed, as we don't access stacktraces at all
 	}
 
-	includeExemplars := query.TimeSeries.ExemplarType == typesv1.ExemplarType_EXEMPLAR_TYPE_INDIVIDUAL
 	if includeExemplars {
 		opts = append(opts,
 			withAllLabels(),
@@ -166,3 +172,18 @@ func (a *timeSeriesAggregator) build() *queryv1.Report {
 		},
 	}
 }
+
+// validateExemplarType validates the exemplar type and returns whether to include exemplars.
+func validateExemplarType(exemplarType typesv1.ExemplarType) (bool, error) {
+	switch exemplarType {
+	case typesv1.ExemplarType_EXEMPLAR_TYPE_UNSPECIFIED,
+		typesv1.ExemplarType_EXEMPLAR_TYPE_NONE:
+		return false, nil
+	case typesv1.ExemplarType_EXEMPLAR_TYPE_INDIVIDUAL:
+		return true, nil
+	case typesv1.ExemplarType_EXEMPLAR_TYPE_SPAN:
+		return false, status.Error(codes.Unimplemented, "exemplar type span is not implemented")
+	default:
+		return false, status.Errorf(codes.InvalidArgument, "unknown exemplar type: %v", exemplarType)
+	}
+}

pkg/querybackend/query_time_series_test.go

@@ -0,0 +1,68 @@
+package querybackend
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
+)
+
+func TestValidateExemplarType(t *testing.T) {
+	tests := []struct {
+		name             string
+		exemplarType     typesv1.ExemplarType
+		expectedInclude  bool
+		expectedErrorMsg string
+		expectedCode     codes.Code
+	}{
+		{
+			name:            "UNSPECIFIED returns false, no error",
+			exemplarType:    typesv1.ExemplarType_EXEMPLAR_TYPE_UNSPECIFIED,
+			expectedInclude: false,
+		},
+		{
+			name:            "NONE returns false, no error",
+			exemplarType:    typesv1.ExemplarType_EXEMPLAR_TYPE_NONE,
+			expectedInclude: false,
+		},
+		{
+			name:            "INDIVIDUAL returns true, no error",
+			exemplarType:    typesv1.ExemplarType_EXEMPLAR_TYPE_INDIVIDUAL,
+			expectedInclude: true,
+		},
+		{
+			name:             "SPAN returns error with Unimplemented code",
+			exemplarType:     typesv1.ExemplarType_EXEMPLAR_TYPE_SPAN,
+			expectedInclude:  false,
+			expectedErrorMsg: "exemplar type span is not implemented",
+			expectedCode:     codes.Unimplemented,
+		},
+		{
+			name:             "Unknown type returns error with InvalidArgument code",
+			exemplarType:     typesv1.ExemplarType(999),
+			expectedInclude:  false,
+			expectedErrorMsg: "unknown exemplar type",
+			expectedCode:     codes.InvalidArgument,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			include, err := validateExemplarType(tt.exemplarType)
+			if tt.expectedErrorMsg != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.expectedErrorMsg)
+				st, ok := status.FromError(err)
+				require.True(t, ok)
+				assert.Equal(t, tt.expectedCode, st.Code())
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.expectedInclude, include)
+			}
+		})
+	}
+}