From 69c49f6ba2b3e7219001b1cab0f9eebd2be6d8b9 Mon Sep 17 00:00:00 2001
From: juanicabanas <juan.cabanas@grafana.com>
Date: Wed, 5 Oct 2022 18:45:17 -0300
Subject: [PATCH] PublicDashboards: Make mixed datasource calls concurrently
 (#56421)

PublicDashboards: Make mixed datasource calls concurrently
---
 pkg/services/query/query.go      | 30 +++++++++++++++++------
 pkg/services/query/query_test.go | 42 +++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/pkg/services/query/query.go b/pkg/services/query/query.go
index 2cffca07861..c2474b44269 100644
--- a/pkg/services/query/query.go
+++ b/pkg/services/query/query.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"golang.org/x/sync/errgroup"
 )
 
 func ProvideService(
@@ -90,18 +91,31 @@ func (s *Service) QueryDataMultipleSources(ctx context.Context, user *user.Signe
 	} else {
 		resp := backend.NewQueryDataResponse()
 
-		// create new reqDTO with only the queries for that datasource
+		g, ctx := errgroup.WithContext(ctx)
+		results := make([]backend.Responses, len(byDataSource))
+
 		for _, queries := range byDataSource {
-			subDTO := reqDTO.CloneWithQueries(queries)
+			dataSourceQueries := queries
+			g.Go(func() error {
+				subDTO := reqDTO.CloneWithQueries(dataSourceQueries)
 
-			subResp, err := s.QueryData(ctx, user, skipCache, subDTO, handleExpressions)
+				subResp, err := s.QueryData(ctx, user, skipCache, subDTO, handleExpressions)
 
-			if err != nil {
-				return nil, err
-			}
+				if err == nil {
+					results = append(results, subResp.Responses)
+				}
 
-			for refId, queryResponse := range subResp.Responses {
-				resp.Responses[refId] = queryResponse
+				return err
+			})
+		}
+
+		if err := g.Wait(); err != nil {
+			return nil, err
+		}
+
+		for _, result := range results {
+			for refId, dataResponse := range result {
+				resp.Responses[refId] = dataResponse
 			}
 		}
 
diff --git a/pkg/services/query/query_test.go b/pkg/services/query/query_test.go
index 33c8fc1f6bf..cd6af3250af 100644
--- a/pkg/services/query/query_test.go
+++ b/pkg/services/query/query_test.go
@@ -112,6 +112,43 @@ func TestQueryDataMultipleSources(t *testing.T) {
 
 		require.NoError(t, err)
 	})
+
+	t.Run("error is returned when one of the queries fails", func(t *testing.T) {
+		tc := setup(t)
+
+		query1, _ := simplejson.NewJson([]byte(`
+			{
+				"datasource": {
+					"type": "mysql",
+					"uid": "ds1"
+				}
+			}
+		`))
+		query2, _ := simplejson.NewJson([]byte(`
+			{
+				"datasource": {
+					"type": "prometheus",
+					"uid": "ds2"
+				},
+				"queryType": "FAIL"
+			}
+		`))
+
+		queries := []*simplejson.Json{query1, query2}
+
+		reqDTO := dtos.MetricRequest{
+			From:                       "2022-01-01",
+			To:                         "2022-01-02",
+			Queries:                    queries,
+			Debug:                      false,
+			PublicDashboardAccessToken: "abc123",
+			HTTPRequest:                nil,
+		}
+
+		_, err := tc.queryService.QueryDataMultipleSources(context.Background(), nil, true, reqDTO, false)
+
+		require.Error(t, err)
+	})
 }
 
 func TestQueryData(t *testing.T) {
@@ -255,7 +292,6 @@ func (c *fakeDataSourceCache) GetDatasourceByUID(ctx context.Context, datasource
 
 type fakePluginClient struct {
 	plugins.Client
-
 	req *backend.QueryDataRequest
 }
 
@@ -267,5 +303,9 @@ func (c *fakePluginClient) QueryData(ctx context.Context, req *backend.QueryData
 		return nil, errors.New("cant query an expression datasource")
 	}
 
+	if req.Queries[0].QueryType == "FAIL" {
+		return nil, errors.New("plugin client failed")
+	}
+
 	return &backend.QueryDataResponse{Responses: make(backend.Responses)}, nil
 }