Datasources: Support mixed datasources in a single query (#56832)

* initial cut at refactor - need to run more tests

* fix unit tests

* change newly unused function to test helper

* create unit tests for parsing query requests that cover a range of cases

* add some comments

* rename function to avoid dev confusion

pkg/api/metrics.go

@@ -54,7 +54,7 @@ func (hs *HTTPServer) QueryMetricsV2(c *models.ReqContext) response.Response {
 
 	reqDTO.HTTPRequest = c.Req
 
-	resp, err := hs.queryDataService.QueryData(c.Req.Context(), c.SignedInUser, c.SkipCache, reqDTO, true)
+	resp, err := hs.queryDataService.QueryData(c.Req.Context(), c.SignedInUser, c.SkipCache, reqDTO)
 	if err != nil {
 		return hs.handleQueryMetricsError(err)
 	}

pkg/services/live/live.go

@@ -590,7 +590,7 @@ func (g *GrafanaLive) handleOnRPC(client *centrifuge.Client, e centrifuge.RPCEve
 	if err != nil {
 		return centrifuge.RPCReply{}, centrifuge.ErrorBadRequest
 	}
-	resp, err := g.queryDataService.QueryData(client.Context(), user, false, req, true)
+	resp, err := g.queryDataService.QueryData(client.Context(), user, false, req)
 	if err != nil {
 		logger.Error("Error query data", "user", client.UserID(), "client", client.ID(), "method", e.Method, "error", err)
 		if errors.Is(err, datasources.ErrDataSourceAccessDenied) {

pkg/services/publicdashboards/queries/dashboard_queries.go

@@ -76,21 +76,6 @@ func HasExpressionQuery(queries []*simplejson.Json) bool {
 	return false
 }
 
-func GroupQueriesByDataSource(queries []*simplejson.Json) (result [][]*simplejson.Json) {
-	byDataSource := make(map[string][]*simplejson.Json)
-
-	for _, query := range queries {
-		uid := GetDataSourceUidFromJson(query)
-		byDataSource[uid] = append(byDataSource[uid], query)
-	}
-
-	for _, queries := range byDataSource {
-		result = append(result, queries)
-	}
-
-	return
-}
-
 func GetDataSourceUidFromJson(query *simplejson.Json) string {
 	uid := query.Get("datasource").Get("uid").MustString()
 

pkg/services/publicdashboards/queries/dashboard_queries_test.go

@@ -361,7 +361,7 @@ func TestGroupQueriesByPanelId(t *testing.T) {
 		queries := GroupQueriesByPanelId(json)
 
 		panelId := int64(2)
-		queriesByDatasource := GroupQueriesByDataSource(queries[panelId])
+		queriesByDatasource := groupQueriesByDataSource(t, queries[panelId])
 		require.Len(t, queriesByDatasource[0], 1)
 	})
 	t.Run("will delete exemplar property from target if exists", func(t *testing.T) {
@@ -370,7 +370,7 @@ func TestGroupQueriesByPanelId(t *testing.T) {
 		queries := GroupQueriesByPanelId(json)
 
 		panelId := int64(2)
-		queriesByDatasource := GroupQueriesByDataSource(queries[panelId])
+		queriesByDatasource := groupQueriesByDataSource(t, queries[panelId])
 		for _, query := range queriesByDatasource[0] {
 			_, ok := query.CheckGet("exemplar")
 			require.False(t, ok)
@@ -382,7 +382,7 @@ func TestGroupQueriesByPanelId(t *testing.T) {
 		queries := GroupQueriesByPanelId(json)
 
 		panelId := int64(2)
-		queriesByDatasource := GroupQueriesByDataSource(queries[panelId])
+		queriesByDatasource := groupQueriesByDataSource(t, queries[panelId])
 		require.Len(t, queriesByDatasource[0], 2)
 	})
 	t.Run("can extract no queries from empty dashboard", func(t *testing.T) {
@@ -487,7 +487,7 @@ func TestGroupQueriesByDataSource(t *testing.T) {
 			}`)),
 		}
 
-		queriesByDatasource := GroupQueriesByDataSource(queries)
+		queriesByDatasource := groupQueriesByDataSource(t, queries)
 		require.Len(t, queriesByDatasource, 2)
 		require.Contains(t, queriesByDatasource, []*simplejson.Json{simplejson.MustJson([]byte(`{
             "datasource": {
@@ -565,3 +565,19 @@ func TestSanitizeMetadataFromQueryData(t *testing.T) {
 		}
 	})
 }
+
+func groupQueriesByDataSource(t *testing.T, queries []*simplejson.Json) (result [][]*simplejson.Json) {
+	t.Helper()
+	byDataSource := make(map[string][]*simplejson.Json)
+
+	for _, query := range queries {
+		uid := GetDataSourceUidFromJson(query)
+		byDataSource[uid] = append(byDataSource[uid], query)
+	}
+
+	for _, queries := range byDataSource {
+		result = append(result, queries)
+	}
+
+	return
+}

pkg/services/publicdashboards/service/service.go

@@ -210,7 +210,7 @@ func (pd *PublicDashboardServiceImpl) GetQueryDataResponse(ctx context.Context,
 		return nil, err
 	}
 
-	res, err := pd.QueryDataService.QueryDataMultipleSources(ctx, anonymousUser, skipCache, metricReq, true)
+	res, err := pd.QueryDataService.QueryData(ctx, anonymousUser, skipCache, metricReq)
 
 	reqDatasources := metricReq.GetUniqueDatasourceTypes()
 	if err != nil {

pkg/services/query/errors.go

@@ -7,6 +7,5 @@ import (
 var (
 	ErrNoQueriesFound        = errutil.NewBase(errutil.StatusBadRequest, "query.noQueries", errutil.WithPublicMessage("No queries found")).Errorf("no queries found")
 	ErrInvalidDatasourceID   = errutil.NewBase(errutil.StatusBadRequest, "query.invalidDatasourceId", errutil.WithPublicMessage("Query does not contain a valid data source identifier")).Errorf("invalid data source identifier")
-	ErrMultipleDatasources   = errutil.NewBase(errutil.StatusBadRequest, "query.differentDatasources", errutil.WithPublicMessage("All queries must use the same datasource")).Errorf("all queries must use the same datasource")
 	ErrMissingDataSourceInfo = errutil.NewBase(errutil.StatusBadRequest, "query.missingDataSourceInfo").MustTemplate("query missing datasource info: {{ .Public.RefId }}", errutil.WithPublic("Query {{ .Public.RefId }} is missing datasource information"))
 )

pkg/services/query/query.go

@@ -16,7 +16,6 @@ import (
 	"github.com/grafana/grafana/pkg/plugins/adapters"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/oauthtoken"
-	publicDashboards "github.com/grafana/grafana/pkg/services/publicdashboards/queries"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/grafanads"
@@ -69,37 +68,37 @@ func (s *Service) Run(ctx context.Context) error {
 	return ctx.Err()
 }
 
-// QueryData can process queries and return query responses.
+// QueryData processes queries and returns query responses. It handles queries to single or mixed datasources, as well as expressions.
-func (s *Service) QueryData(ctx context.Context, user *user.SignedInUser, skipCache bool, reqDTO dtos.MetricRequest, handleExpressions bool) (*backend.QueryDataResponse, error) {
+func (s *Service) QueryData(ctx context.Context, user *user.SignedInUser, skipCache bool, reqDTO dtos.MetricRequest) (*backend.QueryDataResponse, error) {
+	// Parse the request into parsed queries grouped by datasource uid
 	parsedReq, err := s.parseMetricRequest(ctx, user, skipCache, reqDTO)
 	if err != nil {
 		return nil, err
 	}
-	if handleExpressions && parsedReq.hasExpression {
+	// If there are expressions, handle them and return
+	if parsedReq.hasExpression {
 		return s.handleExpressions(ctx, user, parsedReq)
 	}
-	return s.handleQueryData(ctx, user, parsedReq)
+	// If there is only one datasource, query it and return
-}
+	if len(parsedReq.parsedQueries) == 1 {
-
+		return s.handleQuerySingleDatasource(ctx, user, parsedReq)
-// QueryData can process queries and return query responses.
+	}
-func (s *Service) QueryDataMultipleSources(ctx context.Context, user *user.SignedInUser, skipCache bool, reqDTO dtos.MetricRequest, handleExpressions bool) (*backend.QueryDataResponse, error) {
+	// If there are multiple datasources, handle their queries concurrently and return the aggregate result
-	byDataSource := publicDashboards.GroupQueriesByDataSource(reqDTO.Queries)
+	byDataSource := parsedReq.parsedQueries
-
-	// The expression service will handle mixed datasources, so we don't need to group them when an expression is present.
-	if publicDashboards.HasExpressionQuery(reqDTO.Queries) || len(byDataSource) == 1 {
-		return s.QueryData(ctx, user, skipCache, reqDTO, handleExpressions)
-	} else {
 	resp := backend.NewQueryDataResponse()
 
 	g, ctx := errgroup.WithContext(ctx)
 	results := make([]backend.Responses, len(byDataSource))
 
 	for _, queries := range byDataSource {
-			dataSourceQueries := queries
+		rawQueries := make([]*simplejson.Json, len(queries))
+		for i := 0; i < len(queries); i++ {
+			rawQueries[i] = queries[i].rawQuery
+		}
 		g.Go(func() error {
-				subDTO := reqDTO.CloneWithQueries(dataSourceQueries)
+			subDTO := reqDTO.CloneWithQueries(rawQueries)
 
-				subResp, err := s.QueryData(ctx, user, skipCache, subDTO, handleExpressions)
+			subResp, err := s.QueryData(ctx, user, skipCache, subDTO)
 
 			if err == nil {
 				results = append(results, subResp.Responses)
@@ -120,7 +119,6 @@ func (s *Service) QueryDataMultipleSources(ctx context.Context, user *user.Signe
 	}
 
 	return resp, nil
-	}
 }
 
 // handleExpressions handles POST /api/ds/query when there is an expression.
@@ -130,7 +128,7 @@ func (s *Service) handleExpressions(ctx context.Context, user *user.SignedInUser
 		Queries: []expr.Query{},
 	}
 
-	for _, pq := range parsedReq.parsedQueries {
+	for _, pq := range parsedReq.getFlattenedQueries() {
 		if pq.datasource == nil {
 			return nil, ErrMissingDataSourceInfo.Build(errutil.TemplateData{
 				Public: map[string]interface{}{
@@ -160,12 +158,21 @@ func (s *Service) handleExpressions(ctx context.Context, user *user.SignedInUser
 	return qdr, nil
 }
 
-func (s *Service) handleQueryData(ctx context.Context, user *user.SignedInUser, parsedReq *parsedRequest) (*backend.QueryDataResponse, error) {
+// handleQuerySingleDatasource handles one or more queries to a single datasource
-	ds := parsedReq.parsedQueries[0].datasource
+func (s *Service) handleQuerySingleDatasource(ctx context.Context, user *user.SignedInUser, parsedReq *parsedRequest) (*backend.QueryDataResponse, error) {
+	queries := parsedReq.getFlattenedQueries()
+	ds := queries[0].datasource
 	if err := s.pluginRequestValidator.Validate(ds.Url, nil); err != nil {
 		return nil, datasources.ErrDataSourceAccessDenied
 	}
 
+	// ensure that each query passed to this function has the same datasource
+	for _, pq := range queries {
+		if ds.Uid != pq.datasource.Uid {
+			return nil, fmt.Errorf("all queries must have the same datasource - found %s and %s", ds.Uid, pq.datasource.Uid)
+		}
+	}
+
 	instanceSettings, err := adapters.ModelToInstanceSettings(ds, s.decryptSecureJsonDataFn(ctx))
 	if err != nil {
 		return nil, err
@@ -208,7 +215,7 @@ func (s *Service) handleQueryData(ctx context.Context, user *user.SignedInUser,
 		}
 	}
 
-	for _, q := range parsedReq.parsedQueries {
+	for _, q := range queries {
 		req.Queries = append(req.Queries, q.query)
 	}
 
@@ -220,14 +227,24 @@ func (s *Service) handleQueryData(ctx context.Context, user *user.SignedInUser,
 type parsedQuery struct {
 	datasource *datasources.DataSource
 	query      backend.DataQuery
+	rawQuery   *simplejson.Json
 }
 
 type parsedRequest struct {
 	hasExpression bool
-	parsedQueries []parsedQuery
+	parsedQueries map[string][]parsedQuery
 	httpRequest   *http.Request
 }
 
+func (pr parsedRequest) getFlattenedQueries() []parsedQuery {
+	queries := make([]parsedQuery, 0)
+	for _, pq := range pr.parsedQueries {
+		queries = append(queries, pq...)
+	}
+	return queries
+}
+
+// parseRequest parses a request into parsed queries grouped by datasource uid
 func (s *Service) parseMetricRequest(ctx context.Context, user *user.SignedInUser, skipCache bool, reqDTO dtos.MetricRequest) (*parsedRequest, error) {
 	if len(reqDTO.Queries) == 0 {
 		return nil, ErrNoQueriesFound
@@ -236,10 +253,10 @@ func (s *Service) parseMetricRequest(ctx context.Context, user *user.SignedInUse
 	timeRange := legacydata.NewDataTimeRange(reqDTO.From, reqDTO.To)
 	req := &parsedRequest{
 		hasExpression: false,
-		parsedQueries: []parsedQuery{},
+		parsedQueries: make(map[string][]parsedQuery),
 	}
 
-	// Parse the queries
+	// Parse the queries and store them by datasource
 	datasourcesByUid := map[string]*datasources.DataSource{}
 	for _, query := range reqDTO.Queries {
 		ds, err := s.getDataSourceFromQuery(ctx, user, skipCache, query, datasourcesByUid)
@@ -255,6 +272,10 @@ func (s *Service) parseMetricRequest(ctx context.Context, user *user.SignedInUse
 			req.hasExpression = true
 		}
 
+		if _, ok := req.parsedQueries[ds.Uid]; !ok {
+			req.parsedQueries[ds.Uid] = []parsedQuery{}
+		}
+
 		s.log.Debug("Processing metrics query", "query", query)
 
 		modelJSON, err := query.MarshalJSON()
@@ -262,7 +283,7 @@ func (s *Service) parseMetricRequest(ctx context.Context, user *user.SignedInUse
 			return nil, err
 		}
 
-		req.parsedQueries = append(req.parsedQueries, parsedQuery{
+		req.parsedQueries[ds.Uid] = append(req.parsedQueries[ds.Uid], parsedQuery{
 			datasource: ds,
 			query: backend.DataQuery{
 				TimeRange: backend.TimeRange{
@@ -275,16 +296,10 @@ func (s *Service) parseMetricRequest(ctx context.Context, user *user.SignedInUse
 				QueryType:     query.Get("queryType").MustString(""),
 				JSON:          modelJSON,
 			},
+			rawQuery: query,
 		})
 	}
 
-	if !req.hasExpression {
-		if len(datasourcesByUid) > 1 {
-			// We do not (yet) support mixed query type
-			return nil, ErrMultipleDatasources
-		}
-	}
-
 	if reqDTO.HTTPRequest != nil {
 		req.httpRequest = reqDTO.HTTPRequest
 	}

pkg/services/query/query_test.go

@@ -1,4 +1,4 @@
-package query_test
+package query
 
 import (
 	"context"
@@ -8,6 +8,8 @@ import (
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/expr"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
 
@@ -20,7 +22,6 @@ import (
 	fakeDatasources "github.com/grafana/grafana/pkg/services/datasources/fakes"
 	dsSvc "github.com/grafana/grafana/pkg/services/datasources/service"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
-	"github.com/grafana/grafana/pkg/services/query"
 	"github.com/grafana/grafana/pkg/services/secrets/fakes"
 	secretskvs "github.com/grafana/grafana/pkg/services/secrets/kvstore"
 	secretsmng "github.com/grafana/grafana/pkg/services/secrets/manager"
@@ -28,6 +29,146 @@ import (
 	"github.com/grafana/grafana/pkg/services/user"
 )
 
+func TestParseMetricRequest(t *testing.T) {
+	tc := setup(t)
+
+	t.Run("Test a simple single datasource query", func(t *testing.T) {
+		mr := metricRequestWithQueries(t, `{
+			"refId": "A",
+			"datasource": {
+				"uid": "gIEkMvIVz",
+				"type": "postgres"
+			}
+		}`, `{
+			"refId": "B",
+			"datasource": {
+				"uid": "gIEkMvIVz",
+				"type": "postgres"
+			}
+		}`)
+		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
+		require.NoError(t, err)
+		require.NotNil(t, parsedReq)
+		assert.False(t, parsedReq.hasExpression)
+		assert.Len(t, parsedReq.parsedQueries, 1)
+		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
+		assert.Len(t, parsedReq.getFlattenedQueries(), 2)
+	})
+
+	t.Run("Test a single datasource query with expressions", func(t *testing.T) {
+		mr := metricRequestWithQueries(t, `{
+			"refId": "A",
+			"datasource": {
+				"uid": "gIEkMvIVz",
+				"type": "postgres"
+			}
+		}`, `{
+			"refId": "B",
+			"datasource": {
+				"type": "__expr__",
+				"uid": "__expr__",
+				"name": "Expression"
+			},
+			"type": "math",
+			"expression": "$A - 50"
+		}`)
+		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
+		require.NoError(t, err)
+		require.NotNil(t, parsedReq)
+		assert.True(t, parsedReq.hasExpression)
+		assert.Len(t, parsedReq.parsedQueries, 2)
+		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
+		assert.Len(t, parsedReq.getFlattenedQueries(), 2)
+		// Make sure we end up with something valid
+		_, err = tc.queryService.handleExpressions(context.Background(), tc.signedInUser, parsedReq)
+		assert.NoError(t, err)
+	})
+
+	t.Run("Test a simple mixed datasource query", func(t *testing.T) {
+		mr := metricRequestWithQueries(t, `{
+			"refId": "A",
+			"datasource": {
+				"uid": "gIEkMvIVz",
+				"type": "postgres"
+			}
+		}`, `{
+			"refId": "B",
+			"datasource": {
+				"uid": "sEx6ZvSVk",
+				"type": "testdata"
+			}
+		}`)
+		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
+		require.NoError(t, err)
+		require.NotNil(t, parsedReq)
+		assert.False(t, parsedReq.hasExpression)
+		assert.Len(t, parsedReq.parsedQueries, 2)
+		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
+		assert.Contains(t, parsedReq.parsedQueries, "sEx6ZvSVk")
+		assert.Len(t, parsedReq.getFlattenedQueries(), 2)
+	})
+
+	t.Run("Test a mixed datasource query with expressions", func(t *testing.T) {
+		mr := metricRequestWithQueries(t, `{
+			"refId": "A",
+			"datasource": {
+				"uid": "gIEkMvIVz",
+				"type": "postgres"
+			}
+		}`, `{
+			"refId": "B",
+			"datasource": {
+				"uid": "sEx6ZvSVk",
+				"type": "testdata"
+			}
+		}`, `{
+			"refId": "A_resample",
+			"datasource": {
+				"type": "__expr__",
+				"uid": "__expr__",
+				"name": "Expression"
+			},
+			"expression": "A",
+			"type": "resample",
+			"downsampler": "mean",
+			"upsampler": "fillna",
+			"window": "10s"
+		}`, `{
+			"refId": "B_resample",
+			"datasource": {
+				"type": "__expr__",
+				"uid": "__expr__",
+				"name": "Expression"
+			},
+			"expression": "B",
+			"type": "resample",
+			"downsampler": "mean",
+			"upsampler": "fillna",
+			"window": "10s"
+		}`, `{
+			"refId": "C",
+			"datasource": {
+				"type": "__expr__",
+				"uid": "__expr__",
+				"name": "Expression"
+			},
+			"type": "math",
+			"expression": "$A_resample + $B_resample"
+		}`)
+		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
+		require.NoError(t, err)
+		require.NotNil(t, parsedReq)
+		assert.True(t, parsedReq.hasExpression)
+		assert.Len(t, parsedReq.parsedQueries, 3)
+		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
+		assert.Contains(t, parsedReq.parsedQueries, "sEx6ZvSVk")
+		assert.Len(t, parsedReq.getFlattenedQueries(), 5)
+		// Make sure we end up with something valid
+		_, err = tc.queryService.handleExpressions(context.Background(), tc.signedInUser, parsedReq)
+		assert.NoError(t, err)
+	})
+}
+
 func TestQueryDataMultipleSources(t *testing.T) {
 	t.Run("can query multiple datasources", func(t *testing.T) {
 		tc := setup(t)
@@ -59,19 +200,21 @@ func TestQueryDataMultipleSources(t *testing.T) {
 			HTTPRequest:                nil,
 		}
 
-		_, err = tc.queryService.QueryDataMultipleSources(context.Background(), nil, true, reqDTO, false)
+		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
 
 		require.NoError(t, err)
 	})
 
 	t.Run("can query multiple datasources with an expression present", func(t *testing.T) {
 		tc := setup(t)
+		// refId does get set if not included, but better to include it explicitly here
 		query1, err := simplejson.NewJson([]byte(`
 			{
 				"datasource": {
 					"type": "mysql",
 					"uid": "ds1"
-				}
+				},
+				"refId": "A"
 			}
 		`))
 		require.NoError(t, err)
@@ -108,7 +251,7 @@ func TestQueryDataMultipleSources(t *testing.T) {
 			HTTPRequest:                nil,
 		}
 
-		_, err = tc.queryService.QueryDataMultipleSources(context.Background(), nil, true, reqDTO, false)
+		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
 
 		require.NoError(t, err)
 	})
@@ -145,7 +288,7 @@ func TestQueryDataMultipleSources(t *testing.T) {
 			HTTPRequest:                nil,
 		}
 
-		_, err := tc.queryService.QueryDataMultipleSources(context.Background(), nil, true, reqDTO, false)
+		_, err := tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
 
 		require.Error(t, err)
 	})
@@ -163,7 +306,7 @@ func TestQueryData(t *testing.T) {
 		tc.oauthTokenService.passThruEnabled = true
 		tc.oauthTokenService.token = token
 
-		_, err := tc.queryService.QueryData(context.Background(), nil, true, metricRequest(), false)
+		_, err := tc.queryService.QueryData(context.Background(), nil, true, metricRequest())
 		require.Nil(t, err)
 
 		expected := map[string]string{
@@ -183,7 +326,7 @@ func TestQueryData(t *testing.T) {
 		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
 		require.NoError(t, err)
 		metricReq.HTTPRequest = httpReq
-		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricReq, false)
+		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, metricReq)
 		require.NoError(t, err)
 
 		require.Empty(t, tc.pluginContext.req.Headers)
@@ -204,7 +347,7 @@ func TestQueryData(t *testing.T) {
 		httpReq.AddCookie(&http.Cookie{Name: "foo", Value: "oof"})
 		httpReq.AddCookie(&http.Cookie{Name: "c"})
 		metricReq.HTTPRequest = httpReq
-		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricReq, false)
+		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, metricReq)
 		require.NoError(t, err)
 
 		require.Equal(t, map[string]string{"Cookie": "bar=rab; foo=oof"}, tc.pluginContext.req.Headers)
@@ -212,6 +355,7 @@ func TestQueryData(t *testing.T) {
 }
 
 func setup(t *testing.T) *testContext {
+	t.Helper()
 	pc := &fakePluginClient{}
 	dc := &fakeDataSourceCache{ds: &datasources.DataSource{}}
 	tc := &fakeOAuthTokenService{}
@@ -226,15 +370,16 @@ func setup(t *testing.T) *testContext {
 		DataSources:           nil,
 		SimulatePluginFailure: false,
 	}
-	exprService := expr.ProvideService(nil, pc, fakeDatasourceService)
+	exprService := expr.ProvideService(&setting.Cfg{ExpressionsEnabled: true}, pc, fakeDatasourceService)
-
+	queryService := ProvideService(nil, dc, exprService, rv, ds, pc, tc) // provider belonging to this package
 	return &testContext{
 		pluginContext:          pc,
 		secretStore:            ss,
 		dataSourceCache:        dc,
 		oauthTokenService:      tc,
 		pluginRequestValidator: rv,
-		queryService:           query.ProvideService(nil, dc, exprService, rv, ds, pc, tc),
+		queryService:           queryService,
+		signedInUser:           &user.SignedInUser{OrgID: 1},
 	}
 }
 
@@ -244,7 +389,8 @@ type testContext struct {
 	dataSourceCache        *fakeDataSourceCache
 	oauthTokenService      *fakeOAuthTokenService
 	pluginRequestValidator *fakePluginRequestValidator
-	queryService           *query.Service
+	queryService           *Service // implementation belonging to this package
+	signedInUser           *user.SignedInUser
 }
 
 func metricRequest() dtos.MetricRequest {
@@ -257,6 +403,22 @@ func metricRequest() dtos.MetricRequest {
 	}
 }
 
+func metricRequestWithQueries(t *testing.T, rawQueries ...string) dtos.MetricRequest {
+	t.Helper()
+	queries := make([]*simplejson.Json, 0)
+	for _, q := range rawQueries {
+		json, err := simplejson.NewJson([]byte(q))
+		require.NoError(t, err)
+		queries = append(queries, json)
+	}
+	return dtos.MetricRequest{
+		From:    "now-1h",
+		To:      "now",
+		Queries: queries,
+		Debug:   false,
+	}
+}
+
 type fakePluginRequestValidator struct {
 	err error
 }
@@ -287,7 +449,9 @@ func (c *fakeDataSourceCache) GetDatasource(ctx context.Context, datasourceID in
 }
 
 func (c *fakeDataSourceCache) GetDatasourceByUID(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
-	return c.ds, nil
+	return &datasources.DataSource{
+		Uid: datasourceUID,
+	}, nil
 }
 
 type fakePluginClient struct {