CloudWatch: Refactor around handleGetRegions (#65713)

* Create minimal EC2 interface

* convert regions variable to a function returning a copy

* Add test for multiple calls to regions to check regionsCache

* Add returns to handler after error

* Refactor handleGetRegions

pkg/tsdb/cloudwatch/client_factory.go

@@ -8,7 +8,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
 	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/aws/aws-sdk-go/service/oam"
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi"
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi/resourcegroupstaggingapiiface"
@@ -53,7 +52,7 @@ var NewCWLogsClient = func(sess *session.Session) cloudwatchlogsiface.CloudWatch
 // EC2 client factory.
 //
 // Stubbable by tests.
-var newEC2Client = func(provider client.ConfigProvider) ec2iface.EC2API {
+var newEC2Client = func(provider client.ConfigProvider) models.EC2APIProvider {
 	return ec2.New(provider)
 }
 

pkg/tsdb/cloudwatch/cloudwatch.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"sync"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
@@ -13,7 +14,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
-	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi/resourcegroupstaggingapiiface"
 	"github.com/grafana/grafana-aws-sdk/pkg/awsds"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -119,10 +119,11 @@ func NewInstanceSettings(httpClientProvider httpclient.Provider) datasource.Inst
 
 // cloudWatchExecutor executes CloudWatch requests.
 type cloudWatchExecutor struct {
-	im       instancemgmt.InstanceManager
-	cfg      *setting.Cfg
-	sessions SessionCache
-	features featuremgmt.FeatureToggles
+	im          instancemgmt.InstanceManager
+	cfg         *setting.Cfg
+	sessions    SessionCache
+	features    featuremgmt.FeatureToggles
+	regionCache sync.Map
 
 	resourceHandler backend.CallResourceHandler
 }
@@ -310,7 +311,7 @@ func (e *cloudWatchExecutor) getCWLogsClient(pluginCtx backend.PluginContext, re
 	return logsClient, nil
 }
 
-func (e *cloudWatchExecutor) getEC2Client(pluginCtx backend.PluginContext, region string) (ec2iface.EC2API, error) {
+func (e *cloudWatchExecutor) getEC2Client(pluginCtx backend.PluginContext, region string) (models.EC2APIProvider, error) {
 	sess, err := e.newSession(pluginCtx, region)
 	if err != nil {
 		return nil, err

pkg/tsdb/cloudwatch/constants/metrics.go

@@ -523,9 +523,37 @@ var NamespaceDimensionKeysMap = map[string][]string{
 	"CloudWatchSynthetics":        {"CanaryName", "StepName"},
 }
 
-var Regions = []string{
-	"af-south-1", "ap-east-1", "ap-northeast-1", "ap-northeast-2", "ap-northeast-3", "ap-south-1", "ap-southeast-1",
-	"ap-southeast-2", "ap-southeast-3", "ca-central-1", "cn-north-1", "cn-northwest-1", "eu-central-1", "eu-north-1", "eu-south-1", "eu-west-1",
-	"eu-west-2", "eu-west-3", "me-south-1", "sa-east-1", "us-east-1", "us-east-2", "us-gov-east-1", "us-gov-west-1",
-	"us-iso-east-1", "us-isob-east-1", "us-west-1", "us-west-2",
+type RegionsSet map[string]struct{}
+
+func Regions() RegionsSet {
+	return RegionsSet{
+		"af-south-1":     {},
+		"ap-east-1":      {},
+		"ap-northeast-1": {},
+		"ap-northeast-2": {},
+		"ap-northeast-3": {},
+		"ap-south-1":     {},
+		"ap-southeast-1": {},
+		"ap-southeast-2": {},
+		"ap-southeast-3": {},
+		"ca-central-1":   {},
+		"cn-north-1":     {},
+		"cn-northwest-1": {},
+		"eu-central-1":   {},
+		"eu-north-1":     {},
+		"eu-south-1":     {},
+		"eu-west-1":      {},
+		"eu-west-2":      {},
+		"eu-west-3":      {},
+		"me-south-1":     {},
+		"sa-east-1":      {},
+		"us-east-1":      {},
+		"us-east-2":      {},
+		"us-gov-east-1":  {},
+		"us-gov-west-1":  {},
+		"us-iso-east-1":  {},
+		"us-isob-east-1": {},
+		"us-west-1":      {},
+		"us-west-2":      {},
+	}
 }

pkg/tsdb/cloudwatch/metric_find_query.go

@@ -8,7 +8,6 @@ import (
 	"reflect"
 	"sort"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -24,8 +23,6 @@ type suggestData struct {
 	Label string `json:"label,omitempty"`
 }
 
-var regionCache sync.Map
-
 func parseMultiSelectValue(input string) []string {
 	trimmedInput := strings.TrimSpace(input)
 	if strings.HasPrefix(trimmedInput, "{") {
@@ -49,7 +46,7 @@ func (e *cloudWatchExecutor) handleGetRegions(pluginCtx backend.PluginContext, p
 	}
 
 	profile := instance.Settings.Profile
-	if cache, ok := regionCache.Load(profile); ok {
+	if cache, ok := e.regionCache.Load(profile); ok {
 		if cache2, ok2 := cache.([]suggestData); ok2 {
 			return cache2, nil
 		}
@@ -59,38 +56,35 @@ func (e *cloudWatchExecutor) handleGetRegions(pluginCtx backend.PluginContext, p
 	if err != nil {
 		return nil, err
 	}
-	regions := constants.Regions
-	r, err := client.DescribeRegions(&ec2.DescribeRegionsInput{})
+	regions := constants.Regions()
+	ec2Regions, err := client.DescribeRegions(&ec2.DescribeRegionsInput{})
 	if err != nil {
 		// ignore error for backward compatibility
 		logger.Error("Failed to get regions", "error", err)
 	} else {
-		for _, region := range r.Regions {
-			exists := false
-
-			for _, existingRegion := range regions {
-				if existingRegion == *region.RegionName {
-					exists = true
-					break
-				}
-			}
-
-			if !exists {
-				regions = append(regions, *region.RegionName)
-			}
-		}
+		mergeEC2RegionsAndConstantRegions(regions, ec2Regions.Regions)
 	}
-	sort.Strings(regions)
 
 	result := make([]suggestData, 0)
-	for _, region := range regions {
+	for region := range regions {
 		result = append(result, suggestData{Text: region, Value: region, Label: region})
 	}
-	regionCache.Store(profile, result)
+	sort.Slice(result, func(i, j int) bool {
+		return result[i].Text < result[j].Text
+	})
+	e.regionCache.Store(profile, result)
 
 	return result, nil
 }
 
+func mergeEC2RegionsAndConstantRegions(regions map[string]struct{}, ec2Regions []*ec2.Region) {
+	for _, region := range ec2Regions {
+		if _, ok := regions[*region.RegionName]; !ok {
+			regions[*region.RegionName] = struct{}{}
+		}
+	}
+}
+
 func (e *cloudWatchExecutor) handleGetEbsVolumeIds(pluginCtx backend.PluginContext, parameters url.Values) ([]suggestData, error) {
 	region := parameters.Get("region")
 	instanceId := parameters.Get("instanceId")

pkg/tsdb/cloudwatch/metric_find_query_test.go

@@ -3,12 +3,12 @@ package cloudwatch
 import (
 	"encoding/json"
 	"net/url"
+	"sort"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/client"
 	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi"
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi/resourcegroupstaggingapiiface"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -19,6 +19,7 @@ import (
 	"github.com/grafana/grafana/pkg/tsdb/cloudwatch/constants"
 	"github.com/grafana/grafana/pkg/tsdb/cloudwatch/models"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -30,7 +31,7 @@ func TestQuery_Regions(t *testing.T) {
 
 	var cli fakeEC2Client
 
-	newEC2Client = func(client.ConfigProvider) ec2iface.EC2API {
+	newEC2Client = func(client.ConfigProvider) models.EC2APIProvider {
 		return cli
 	}
 
@@ -55,7 +56,7 @@ func TestQuery_Regions(t *testing.T) {
 		)
 		require.NoError(t, err)
 
-		expRegions := append(constants.Regions, regionName)
+		expRegions := buildSortedSliceOfDefaultAndExtraRegions(t, regionName)
 		expFrame := data.NewFrame(
 			"",
 			data.NewField("text", nil, expRegions),
@@ -63,7 +64,7 @@ func TestQuery_Regions(t *testing.T) {
 		)
 		expFrame.Meta = &data.FrameMeta{
 			Custom: map[string]interface{}{
-				"rowCount": len(constants.Regions) + 1,
+				"rowCount": len(constants.Regions()) + 1,
 			},
 		}
 
@@ -75,6 +76,46 @@ func TestQuery_Regions(t *testing.T) {
 	})
 }
 
+func buildSortedSliceOfDefaultAndExtraRegions(t *testing.T, regionName string) []string {
+	t.Helper()
+	regions := constants.Regions()
+	regions[regionName] = struct{}{}
+	var expRegions []string
+	for region := range regions {
+		expRegions = append(expRegions, region)
+	}
+	sort.Strings(expRegions)
+	return expRegions
+}
+
+func Test_handleGetRegions_regionCache(t *testing.T) {
+	origNewEC2Client := newEC2Client
+	t.Cleanup(func() {
+		newEC2Client = origNewEC2Client
+	})
+	cli := mockEC2Client{}
+	newEC2Client = func(client.ConfigProvider) models.EC2APIProvider {
+		return &cli
+	}
+	im := datasource.NewInstanceManager(func(s backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+		return DataSource{Settings: models.CloudWatchSettings{}}, nil
+	})
+
+	t.Run("AWS only called once for multiple calls to handleGetRegions", func(t *testing.T) {
+		cli.On("DescribeRegions", mock.Anything, mock.Anything).Return(&ec2.DescribeRegionsOutput{}, nil)
+		executor := newExecutor(im, newTestConfig(), &fakeSessionCache{}, featuremgmt.WithFeatures())
+		_, err := executor.handleGetRegions(
+			backend.PluginContext{DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{}}, nil)
+		require.NoError(t, err)
+
+		_, err = executor.handleGetRegions(
+			backend.PluginContext{DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{}}, nil)
+		require.NoError(t, err)
+
+		cli.AssertNumberOfCalls(t, "DescribeRegions", 1)
+	})
+}
+
 func TestQuery_InstanceAttributes(t *testing.T) {
 	origNewEC2Client := newEC2Client
 	t.Cleanup(func() {
@@ -83,7 +124,7 @@ func TestQuery_InstanceAttributes(t *testing.T) {
 
 	var cli fakeEC2Client
 
-	newEC2Client = func(client.ConfigProvider) ec2iface.EC2API {
+	newEC2Client = func(client.ConfigProvider) models.EC2APIProvider {
 		return cli
 	}
 
@@ -144,7 +185,7 @@ func TestQuery_EBSVolumeIDs(t *testing.T) {
 
 	var cli fakeEC2Client
 
-	newEC2Client = func(client.ConfigProvider) ec2iface.EC2API {
+	newEC2Client = func(client.ConfigProvider) models.EC2APIProvider {
 		return cli
 	}
 

pkg/tsdb/cloudwatch/models/api.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/service/cloudwatch"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
+	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/oam"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
@@ -58,3 +59,8 @@ type OAMAPIProvider interface {
 	ListSinks(*oam.ListSinksInput) (*oam.ListSinksOutput, error)
 	ListAttachedLinks(*oam.ListAttachedLinksInput) (*oam.ListAttachedLinksOutput, error)
 }
+
+type EC2APIProvider interface {
+	DescribeRegions(in *ec2.DescribeRegionsInput) (*ec2.DescribeRegionsOutput, error)
+	DescribeInstancesPages(in *ec2.DescribeInstancesInput, fn func(*ec2.DescribeInstancesOutput, bool) bool) error
+}

pkg/tsdb/cloudwatch/resource_handler.go

@@ -37,19 +37,23 @@ func handleResourceReq(handleFunc handleFn) func(rw http.ResponseWriter, req *ht
 		err := req.ParseForm()
 		if err != nil {
 			writeResponse(rw, http.StatusBadRequest, fmt.Sprintf("unexpected error %v", err))
+			return
 		}
 		data, err := handleFunc(pluginContext, req.URL.Query())
 		if err != nil {
 			writeResponse(rw, http.StatusBadRequest, fmt.Sprintf("unexpected error %v", err))
+			return
 		}
 		body, err := json.Marshal(data)
 		if err != nil {
 			writeResponse(rw, http.StatusBadRequest, fmt.Sprintf("unexpected error %v", err))
+			return
 		}
 		rw.WriteHeader(http.StatusOK)
 		_, err = rw.Write(body)
 		if err != nil {
 			logger.Error("Unable to write HTTP response", "error", err)
+			return
 		}
 	}
 }

pkg/tsdb/cloudwatch/test_utils.go

@@ -106,6 +106,20 @@ func (c *fakeCWAnnotationsClient) DescribeAlarms(params *cloudwatch.DescribeAlar
 	return c.describeAlarmsOutput, nil
 }
 
+type mockEC2Client struct {
+	mock.Mock
+}
+
+func (c *mockEC2Client) DescribeRegions(in *ec2.DescribeRegionsInput) (*ec2.DescribeRegionsOutput, error) {
+	args := c.Called(in)
+	return args.Get(0).(*ec2.DescribeRegionsOutput), args.Error(1)
+}
+
+func (c *mockEC2Client) DescribeInstancesPages(in *ec2.DescribeInstancesInput, fn func(*ec2.DescribeInstancesOutput, bool) bool) error {
+	args := c.Called(in, fn)
+	return args.Error(0)
+}
+
 type fakeEC2Client struct {
 	ec2iface.EC2API