Adds new query type field and query type multiplexer (mux) (#152)

Adds new field, QueryType, to QueryDataRequest.
Adds new QueryTypeMux to be able to register a handler 
per query type. This should ease the handling when supporting 
queries of different types.

Closes #75

Co-Authored-By: Kyle Brandt <[email protected]>
Co-authored-by: Arve Knudsen <[email protected]>

Author: 3 people

Diff for: backend/convert_from_protobuf.go

@@ -91,6 +91,7 @@ func (f ConvertFromProtobuf) TimeRange(proto *pluginv2.TimeRange) TimeRange {
 func (f ConvertFromProtobuf) DataQuery(proto *pluginv2.DataQuery) *DataQuery {
 	return &DataQuery{
 		RefID:         proto.RefId,
+		QueryType:     proto.QueryType,
 		MaxDataPoints: proto.MaxDataPoints,
 		TimeRange:     f.TimeRange(proto.TimeRange),
 		Interval:      time.Duration(proto.IntervalMS) * time.Millisecond,

Diff for: backend/convert_from_protobuf_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/davecgh/go-spew/spew"
 	"github.com/grafana/grafana-plugin-sdk-go/genproto/pluginv2"
 	"github.com/mitchellh/reflectwalk"
 	"github.com/stretchr/testify/require"
@@ -30,7 +29,6 @@ func (w *walker) StructField(f reflect.StructField, v reflect.Value) error {
 	if f.PkgPath != "" {
 		return nil
 	}
-	spew.Dump(f)
 	w.FieldCount++
 	if v.IsZero() {
 		w.ZeroValueFieldCount++
@@ -304,6 +302,7 @@ var protoDataQuery = &pluginv2.DataQuery{
 	TimeRange:     protoTimeRange,
 	IntervalMS:    60 * 1000,
 	Json:          []byte(`{ "query": "SELECT * from FUN"`),
+	QueryType:     "qt",
 }
 
 func TestConvertFromProtobufDataQuery(t *testing.T) {
@@ -334,6 +333,8 @@ func TestConvertFromProtobufDataQuery(t *testing.T) {
 
 	requireCounter.Equal(t, protoDQ.RefId, sdkDQ.RefID)
 	requireCounter.Equal(t, protoDQ.MaxDataPoints, sdkDQ.MaxDataPoints)
+	requireCounter.Equal(t, protoDQ.QueryType, sdkDQ.QueryType)
+
 	requireCounter.Equal(t, time.Duration(time.Minute), sdkDQ.Interval)
 	requireCounter.Equal(t, sdkTimeRange.From, sdkDQ.TimeRange.From)
 	requireCounter.Equal(t, sdkTimeRange.To, sdkDQ.TimeRange.To)
@@ -406,6 +407,7 @@ func TestConvertFromProtobufQueryDataRequest(t *testing.T) {
 	// Queries
 	requireCounter.Equal(t, protoQDR.Queries[0].RefId, sdkQDR.Queries[0].RefID)
 	requireCounter.Equal(t, protoQDR.Queries[0].MaxDataPoints, sdkQDR.Queries[0].MaxDataPoints)
+	requireCounter.Equal(t, protoQDR.Queries[0].QueryType, sdkQDR.Queries[0].QueryType)
 	requireCounter.Equal(t, time.Duration(time.Minute), sdkQDR.Queries[0].Interval)
 	requireCounter.Equal(t, sdkTimeRange.From, sdkQDR.Queries[0].TimeRange.From)
 	requireCounter.Equal(t, sdkTimeRange.To, sdkQDR.Queries[0].TimeRange.To)

Diff for: backend/convert_to_protobuf.go

@@ -111,6 +111,7 @@ func (t ConvertToProtobuf) CheckHealthResponse(res *CheckHealthResult) *pluginv2
 func (t ConvertToProtobuf) DataQuery(q DataQuery) *pluginv2.DataQuery {
 	return &pluginv2.DataQuery{
 		RefId:         q.RefID,
+		QueryType:     q.QueryType,
 		MaxDataPoints: q.MaxDataPoints,
 		IntervalMS:    q.Interval.Milliseconds(),
 		TimeRange:     t.TimeRange(q.TimeRange),

Diff for: backend/data.go

@@ -35,6 +35,10 @@ type DataQuery struct {
 	// RefID is the unique identifer of the query, set by the frontend call.
 	RefID string
 
+	// QueryType is an optional identifier for the type of query.
+	// It can be used to distinguish different types of queries.
+	QueryType string
+
 	// MaxDataPoints is the maximum number of datapoints that should be returned from a time series query.
 	MaxDataPoints int64
 

Diff for: backend/datasource/doc.go

@@ -0,0 +1,2 @@
+// Package datasource provides utilities for data source plugins.
+package datasource

Diff for: backend/datasource/query_type_mux.go

@@ -0,0 +1,135 @@
+package datasource
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+// QueryTypeHandlerFunc is an adapter to allow the use of
+// ordinary functions as backend.QueryDataHandler. If f is a function
+// with the appropriate signature, QueryTypeHandlerFunc(f) is a
+// Handler that calls f.
+type QueryTypeHandlerFunc func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error)
+
+// QueryData calls f(ctx, req).
+func (fn QueryTypeHandlerFunc) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	return fn(ctx, req)
+}
+
+// QueryTypeMux is a query type multiplexer.
+type QueryTypeMux struct {
+	m               map[string]backend.QueryDataHandler
+	fallbackHandler backend.QueryDataHandler
+}
+
+// NewQueryTypeMux allocates and returns a new QueryTypeMux.
+func NewQueryTypeMux() *QueryTypeMux {
+	return new(QueryTypeMux)
+}
+
+// Handle registers the handler for the given query type.
+//
+// Providing an empty queryType registers the handler as a fallback handler
+// that will be called when query type doesn't match any registered handlers.
+// If handler is nil, Handle panics.
+// If a handler already exists for queryType, Handle panics.
+func (mux *QueryTypeMux) Handle(queryType string, handler backend.QueryDataHandler) {
+	if handler == nil {
+		panic("datasource: nil handler")
+	}
+
+	if mux.m == nil {
+		mux.m = map[string]backend.QueryDataHandler{}
+		mux.fallbackHandler = QueryTypeHandlerFunc(fallbackHandler)
+	}
+
+	if _, exist := mux.m[queryType]; exist {
+		panic("datasource: multiple registrations for " + queryType)
+	}
+
+	if queryType == "" {
+		log.DefaultLogger.Debug("datasource: registering query type fallback handler")
+		mux.fallbackHandler = handler
+		return
+	}
+
+	log.DefaultLogger.Debug("datasource: registering query type handler", "queryType", queryType)
+	mux.m[queryType] = handler
+}
+
+// HandleFunc registers the handler function for the given query type.
+//
+// Providing an empty queryType registers the handler as a fallback handler
+// that will be called when query type doesn't match any registered handlers.
+// If handler is nil, Handle panics.
+// If a handler already exists for queryType, Handle panics.
+func (mux *QueryTypeMux) HandleFunc(queryType string, handler func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error)) {
+	mux.Handle(queryType, QueryTypeHandlerFunc(handler))
+}
+
+// QueryData dispatches the request to the handler(s) whose
+// query type matches the request queries query type.
+func (mux *QueryTypeMux) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	type requestHandler struct {
+		handler backend.QueryDataHandler
+		request *backend.QueryDataRequest
+	}
+
+	requests := map[string]requestHandler{}
+	for _, q := range req.Queries {
+		qt, handler := mux.getHandler(q.QueryType)
+		if _, exists := requests[qt]; !exists {
+			requests[qt] = requestHandler{
+				handler: handler,
+				request: &backend.QueryDataRequest{
+					PluginContext: req.PluginContext,
+					Headers:       req.Headers,
+					Queries:       []backend.DataQuery{},
+				},
+			}
+		}
+		requests[qt].request.Queries = append(requests[qt].request.Queries, q)
+	}
+
+	responses := backend.Responses{}
+	for _, rh := range requests {
+		qtResponse, err := rh.handler.QueryData(ctx, rh.request)
+		if err != nil {
+			return nil, err
+		}
+
+		for k, v := range qtResponse.Responses {
+			responses[k] = v
+		}
+	}
+
+	return &backend.QueryDataResponse{
+		Responses: responses,
+	}, nil
+}
+
+func (mux *QueryTypeMux) getHandler(queryType string) (string, backend.QueryDataHandler) {
+	handler, exists := mux.m[queryType]
+	if !exists {
+		return "", mux.fallbackHandler
+	}
+
+	return queryType, handler
+}
+
+func fallbackHandler(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	responses := backend.Responses{}
+	for _, q := range req.Queries {
+		responses[q.RefID] = backend.DataResponse{
+			Error: fmt.Errorf("no handler found for query type '%s'", q.QueryType),
+		}
+	}
+
+	return &backend.QueryDataResponse{
+		Responses: responses,
+	}, nil
+}

Diff for: backend/datasource/query_type_mux_test.go

@@ -0,0 +1,104 @@
+package datasource
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/stretchr/testify/require"
+)
+
+func ExampleQueryTypeMux() {
+	mux := NewQueryTypeMux()
+	mux.HandleFunc("queryTypeA", func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+		// handle queryTypeA
+		return nil, nil
+	})
+	mux.HandleFunc("queryTypeB", func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+		// handle queryTypeB
+		return nil, nil
+	})
+
+	_ = backend.ServeOpts{
+		QueryDataHandler: mux,
+	}
+}
+
+func TestQueryTypeMux(t *testing.T) {
+	mux := NewQueryTypeMux()
+	aHandler := &testHandler{}
+	mux.Handle("a", aHandler)
+	bHandler := &testHandler{}
+	mux.Handle("b", bHandler)
+
+	res, err := mux.QueryData(context.Background(), &backend.QueryDataRequest{
+		Queries: []backend.DataQuery{
+			backend.DataQuery{
+				RefID:     "A",
+				QueryType: "a",
+			},
+			backend.DataQuery{
+				RefID:     "B",
+				QueryType: "b",
+			},
+			backend.DataQuery{
+				RefID:     "C",
+				QueryType: "a",
+			},
+			backend.DataQuery{
+				RefID:     "D",
+				QueryType: "d",
+			},
+		},
+	})
+
+	require.NoError(t, err)
+	require.Equal(t, 1, aHandler.callCount)
+	require.Len(t, aHandler.request.Queries, 2)
+	require.Equal(t, "A", aHandler.request.Queries[0].RefID)
+	require.Equal(t, "C", aHandler.request.Queries[1].RefID)
+
+	require.Equal(t, 1, bHandler.callCount)
+	require.Len(t, bHandler.request.Queries, 1)
+	require.Equal(t, "B", bHandler.request.Queries[0].RefID)
+	require.Len(t, res.Responses, 4)
+	require.Equal(t, "no handler found for query type 'd'", res.Responses["D"].Error.Error())
+
+	t.Run("When overriding fallback handler should call fallback handler", func(t *testing.T) {
+		errBoom := errors.New("BOOM")
+		mux.HandleFunc("", func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+			return nil, errBoom
+		})
+		res, err = mux.QueryData(context.Background(), &backend.QueryDataRequest{
+			Queries: []backend.DataQuery{
+				backend.DataQuery{
+					RefID:     "A",
+					QueryType: "unhandled",
+				},
+			},
+		})
+
+		require.Nil(t, res)
+		require.Error(t, err)
+		require.Equal(t, errBoom, err)
+	})
+}
+
+type testHandler struct {
+	callCount int
+	request   *backend.QueryDataRequest
+}
+
+func (th *testHandler) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	th.callCount++
+	th.request = req
+	responses := backend.Responses{}
+	for _, q := range req.Queries {
+		responses[q.RefID] = backend.DataResponse{}
+	}
+
+	return &backend.QueryDataResponse{
+		Responses: responses,
+	}, nil
+}