Cloudwatch: Add support for template variables in new log group picker (#61243)

pkg/tsdb/cloudwatch/services/log_groups.go

@@ -33,20 +33,28 @@ func (s *LogGroupsService) GetLogGroups(req resources.LogGroupsRequest) ([]resou
 			input.AccountIdentifiers = []*string{req.AccountId}
 		}
 	}
-	response, err := s.logGroupsAPI.DescribeLogGroups(input)
-	if err != nil || response == nil {
-		return nil, err
-	}
+	result := []resources.ResourceResponse[resources.LogGroup]{}
 
-	var result []resources.ResourceResponse[resources.LogGroup]
-	for _, logGroup := range response.LogGroups {
-		result = append(result, resources.ResourceResponse[resources.LogGroup]{
-			Value: resources.LogGroup{
-				Arn:  *logGroup.Arn,
-				Name: *logGroup.LogGroupName,
-			},
-			AccountId: utils.Pointer(getAccountId(*logGroup.Arn)),
-		})
+	for {
+		response, err := s.logGroupsAPI.DescribeLogGroups(input)
+		if err != nil || response == nil {
+			return nil, err
+		}
+
+		for _, logGroup := range response.LogGroups {
+			result = append(result, resources.ResourceResponse[resources.LogGroup]{
+				Value: resources.LogGroup{
+					Arn:  *logGroup.Arn,
+					Name: *logGroup.LogGroupName,
+				},
+				AccountId: utils.Pointer(getAccountId(*logGroup.Arn)),
+			})
+		}
+
+		if !req.ListAllLogGroups || response.NextToken == nil {
+			break
+		}
+		input.NextToken = response.NextToken
 	}
 
 	return result, nil

pkg/tsdb/cloudwatch/services/log_groups_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
 	"github.com/grafana/grafana/pkg/tsdb/cloudwatch/mocks"
 	"github.com/grafana/grafana/pkg/tsdb/cloudwatch/models/resources"
@@ -44,6 +45,17 @@ func TestGetLogGroups(t *testing.T) {
 		}, resp)
 	})
 
+	t.Run("Should return an empty error if api doesn't return any data", func(t *testing.T) {
+		mockLogsAPI := &mocks.LogsAPI{}
+		mockLogsAPI.On("DescribeLogGroups", mock.Anything).Return(&cloudwatchlogs.DescribeLogGroupsOutput{}, nil)
+		service := NewLogGroupsService(mockLogsAPI, false)
+
+		resp, err := service.GetLogGroups(resources.LogGroupsRequest{})
+
+		assert.NoError(t, err)
+		assert.Equal(t, []resources.ResourceResponse[resources.LogGroup]{}, resp)
+	})
+
 	t.Run("Should only use LogGroupNamePrefix even if LogGroupNamePattern passed in resource call", func(t *testing.T) {
 		// TODO: use LogGroupNamePattern when we have accounted for its behavior, still a little unexpected at the moment
 		mockLogsAPI := &mocks.LogsAPI{}
@@ -86,6 +98,82 @@ func TestGetLogGroups(t *testing.T) {
 		assert.Error(t, err)
 		assert.Equal(t, "some error", err.Error())
 	})
+
+	t.Run("Should only call the api once in case ListAllLogGroups is set to false", func(t *testing.T) {
+		mockLogsAPI := &mocks.LogsAPI{}
+		req := resources.LogGroupsRequest{
+			Limit:              2,
+			LogGroupNamePrefix: utils.Pointer("test"),
+			ListAllLogGroups:   false,
+		}
+
+		mockLogsAPI.On("DescribeLogGroups", &cloudwatchlogs.DescribeLogGroupsInput{
+			Limit:              aws.Int64(req.Limit),
+			LogGroupNamePrefix: req.LogGroupNamePrefix,
+		}).Return(&cloudwatchlogs.DescribeLogGroupsOutput{
+			LogGroups: []*cloudwatchlogs.LogGroup{
+				{Arn: utils.Pointer("arn:aws:logs:us-east-1:111:log-group:group_a"), LogGroupName: utils.Pointer("group_a")},
+			},
+			NextToken: aws.String("next_token"),
+		}, nil)
+
+		service := NewLogGroupsService(mockLogsAPI, false)
+		resp, err := service.GetLogGroups(req)
+
+		assert.NoError(t, err)
+		mockLogsAPI.AssertNumberOfCalls(t, "DescribeLogGroups", 1)
+		assert.Equal(t, []resources.ResourceResponse[resources.LogGroup]{
+			{
+				AccountId: utils.Pointer("111"),
+				Value:     resources.LogGroup{Arn: "arn:aws:logs:us-east-1:111:log-group:group_a", Name: "group_a"},
+			},
+		}, resp)
+	})
+
+	t.Run("Should keep on calling the api until NextToken is empty in case ListAllLogGroups is set to true", func(t *testing.T) {
+		mockLogsAPI := &mocks.LogsAPI{}
+		req := resources.LogGroupsRequest{
+			Limit:              2,
+			LogGroupNamePrefix: utils.Pointer("test"),
+			ListAllLogGroups:   true,
+		}
+
+		// first call
+		mockLogsAPI.On("DescribeLogGroups", &cloudwatchlogs.DescribeLogGroupsInput{
+			Limit:              aws.Int64(req.Limit),
+			LogGroupNamePrefix: req.LogGroupNamePrefix,
+		}).Return(&cloudwatchlogs.DescribeLogGroupsOutput{
+			LogGroups: []*cloudwatchlogs.LogGroup{
+				{Arn: utils.Pointer("arn:aws:logs:us-east-1:111:log-group:group_a"), LogGroupName: utils.Pointer("group_a")},
+			},
+			NextToken: utils.Pointer("token"),
+		}, nil)
+
+		// second call
+		mockLogsAPI.On("DescribeLogGroups", &cloudwatchlogs.DescribeLogGroupsInput{
+			Limit:              aws.Int64(req.Limit),
+			LogGroupNamePrefix: req.LogGroupNamePrefix,
+			NextToken:          utils.Pointer("token"),
+		}).Return(&cloudwatchlogs.DescribeLogGroupsOutput{
+			LogGroups: []*cloudwatchlogs.LogGroup{
+				{Arn: utils.Pointer("arn:aws:logs:us-east-1:222:log-group:group_b"), LogGroupName: utils.Pointer("group_b")},
+			},
+		}, nil)
+		service := NewLogGroupsService(mockLogsAPI, false)
+		resp, err := service.GetLogGroups(req)
+		assert.NoError(t, err)
+		mockLogsAPI.AssertNumberOfCalls(t, "DescribeLogGroups", 2)
+		assert.Equal(t, []resources.ResourceResponse[resources.LogGroup]{
+			{
+				AccountId: utils.Pointer("111"),
+				Value:     resources.LogGroup{Arn: "arn:aws:logs:us-east-1:111:log-group:group_a", Name: "group_a"},
+			},
+			{
+				AccountId: utils.Pointer("222"),
+				Value:     resources.LogGroup{Arn: "arn:aws:logs:us-east-1:222:log-group:group_b", Name: "group_b"},
+			},
+		}, resp)
+	})
 }
 
 func TestGetLogGroupsCrossAccountQuerying(t *testing.T) {