From df413784729179932dd3267f99254ace6eec89dc Mon Sep 17 00:00:00 2001 From: Shabeeb Khalid Date: Fri, 27 Oct 2023 11:49:37 +0300 Subject: [PATCH] CloudWatch: Use context in aws ListSinks and ListAttachedLinks (#77083) * Use context in aws ListSinks and ListAttachedLinks In the current way, ListSinks and ListAttachedLinks is used which doesn't allow cancelling the request if the context changes. Using ListSinksWithContext and ListAttachedLinksWithContext is the preferred way. Adding context for GetAccountsForCurrentUserOrRole is required to pass it to ListSinks method. --- pkg/tsdb/cloudwatch/mocks/accounts_service.go | 4 +- pkg/tsdb/cloudwatch/mocks/oam_client.go | 7 ++- pkg/tsdb/cloudwatch/models/api.go | 6 +-- pkg/tsdb/cloudwatch/routes/accounts.go | 2 +- pkg/tsdb/cloudwatch/services/accounts.go | 7 +-- pkg/tsdb/cloudwatch/services/accounts_test.go | 53 ++++++++++--------- 6 files changed, 43 insertions(+), 36 deletions(-) diff --git a/pkg/tsdb/cloudwatch/mocks/accounts_service.go b/pkg/tsdb/cloudwatch/mocks/accounts_service.go index 52f3aff63b5..d7b742d6848 100644 --- a/pkg/tsdb/cloudwatch/mocks/accounts_service.go +++ b/pkg/tsdb/cloudwatch/mocks/accounts_service.go @@ -1,6 +1,8 @@ package mocks import ( + "context" + "github.com/grafana/grafana/pkg/tsdb/cloudwatch/models/resources" "github.com/stretchr/testify/mock" ) @@ -9,7 +11,7 @@ type AccountsServiceMock struct { mock.Mock } -func (a *AccountsServiceMock) GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error) { +func (a *AccountsServiceMock) GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error) { args := a.Called() return args.Get(0).([]resources.ResourceResponse[resources.Account]), args.Error(1) diff --git a/pkg/tsdb/cloudwatch/mocks/oam_client.go b/pkg/tsdb/cloudwatch/mocks/oam_client.go index 2d617cc143a..932ca1579b4 100644 --- a/pkg/tsdb/cloudwatch/mocks/oam_client.go +++ b/pkg/tsdb/cloudwatch/mocks/oam_client.go @@ -1,6 +1,9 @@ package mocks import ( + "context" + + "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/service/oam" "github.com/stretchr/testify/mock" ) @@ -9,12 +12,12 @@ type FakeOAMClient struct { mock.Mock } -func (o *FakeOAMClient) ListSinks(input *oam.ListSinksInput) (*oam.ListSinksOutput, error) { +func (o *FakeOAMClient) ListSinksWithContext(ctx context.Context, input *oam.ListSinksInput, opts ...request.Option) (*oam.ListSinksOutput, error) { args := o.Called(input) return args.Get(0).(*oam.ListSinksOutput), args.Error(1) } -func (o *FakeOAMClient) ListAttachedLinks(input *oam.ListAttachedLinksInput) (*oam.ListAttachedLinksOutput, error) { +func (o *FakeOAMClient) ListAttachedLinksWithContext(ctx context.Context, input *oam.ListAttachedLinksInput, opts ...request.Option) (*oam.ListAttachedLinksOutput, error) { args := o.Called(input) return args.Get(0).(*oam.ListAttachedLinksOutput), args.Error(1) } diff --git a/pkg/tsdb/cloudwatch/models/api.go b/pkg/tsdb/cloudwatch/models/api.go index 42d54c84e81..dc540f54d6f 100644 --- a/pkg/tsdb/cloudwatch/models/api.go +++ b/pkg/tsdb/cloudwatch/models/api.go @@ -42,7 +42,7 @@ type LogGroupsProvider interface { } type AccountsProvider interface { - GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error) + GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error) } type RegionsAPIProvider interface { @@ -65,8 +65,8 @@ type CloudWatchLogsAPIProvider interface { } type OAMAPIProvider interface { - ListSinks(*oam.ListSinksInput) (*oam.ListSinksOutput, error) - ListAttachedLinks(*oam.ListAttachedLinksInput) (*oam.ListAttachedLinksOutput, error) + ListSinksWithContext(ctx context.Context, in *oam.ListSinksInput, opts ...request.Option) (*oam.ListSinksOutput, error) + ListAttachedLinksWithContext(ctx context.Context, in *oam.ListAttachedLinksInput, opts ...request.Option) (*oam.ListAttachedLinksOutput, error) } type EC2APIProvider interface { diff --git a/pkg/tsdb/cloudwatch/routes/accounts.go b/pkg/tsdb/cloudwatch/routes/accounts.go index 3c18591b55f..565a8c67ac4 100644 --- a/pkg/tsdb/cloudwatch/routes/accounts.go +++ b/pkg/tsdb/cloudwatch/routes/accounts.go @@ -24,7 +24,7 @@ func AccountsHandler(ctx context.Context, pluginCtx backend.PluginContext, reqCt return nil, models.NewHttpError("error in AccountsHandler", http.StatusInternalServerError, err) } - accounts, err := service.GetAccountsForCurrentUserOrRole() + accounts, err := service.GetAccountsForCurrentUserOrRole(ctx) if err != nil { msg := "error getting accounts for current user or role" switch { diff --git a/pkg/tsdb/cloudwatch/services/accounts.go b/pkg/tsdb/cloudwatch/services/accounts.go index 7f30f038050..f05b18335a3 100644 --- a/pkg/tsdb/cloudwatch/services/accounts.go +++ b/pkg/tsdb/cloudwatch/services/accounts.go @@ -1,6 +1,7 @@ package services import ( + "context" "errors" "fmt" @@ -20,11 +21,11 @@ func NewAccountsService(oamClient models.OAMAPIProvider) models.AccountsProvider return &AccountsService{oamClient} } -func (a *AccountsService) GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error) { +func (a *AccountsService) GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error) { var nextToken *string sinks := []*oam.ListSinksItem{} for { - response, err := a.ListSinks(&oam.ListSinksInput{NextToken: nextToken}) + response, err := a.ListSinksWithContext(ctx, &oam.ListSinksInput{NextToken: nextToken}) if err != nil { var aerr awserr.Error if errors.As(err, &aerr) { @@ -61,7 +62,7 @@ func (a *AccountsService) GetAccountsForCurrentUserOrRole() ([]resources.Resourc nextToken = nil for { - links, err := a.ListAttachedLinks(&oam.ListAttachedLinksInput{ + links, err := a.ListAttachedLinksWithContext(ctx, &oam.ListAttachedLinksInput{ SinkIdentifier: sinkIdentifier, NextToken: nextToken, }) diff --git a/pkg/tsdb/cloudwatch/services/accounts_test.go b/pkg/tsdb/cloudwatch/services/accounts_test.go index 7e74e6fbec9..dbbadaf660c 100644 --- a/pkg/tsdb/cloudwatch/services/accounts_test.go +++ b/pkg/tsdb/cloudwatch/services/accounts_test.go @@ -1,6 +1,7 @@ package services import ( + "context" "fmt" "testing" @@ -17,11 +18,11 @@ import ( func TestHandleGetAccounts(t *testing.T) { t.Run("Should return an error in case of insufficient permissions from ListSinks", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, awserr.New("AccessDeniedException", + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, awserr.New("AccessDeniedException", "AWS message", nil)) accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.Error(t, err) assert.Nil(t, resp) @@ -31,10 +32,10 @@ func TestHandleGetAccounts(t *testing.T) { t.Run("Should return an error in case of any error from ListSinks", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, fmt.Errorf("some error")) + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, fmt.Errorf("some error")) accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.Error(t, err) assert.Nil(t, resp) @@ -43,10 +44,10 @@ func TestHandleGetAccounts(t *testing.T) { t.Run("Should return empty array in case no monitoring account exists", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, nil) + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, nil) accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.NoError(t, err) assert.Empty(t, resp) @@ -54,26 +55,26 @@ func TestHandleGetAccounts(t *testing.T) { t.Run("Should return one monitoring account (the first) even though ListSinks returns multiple sinks", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")}, {Name: aws.String("Account 2"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group2")}, }, NextToken: new(string), }, nil).Once() - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")}, }, NextToken: nil, }, nil) - fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil) + fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil) accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.NoError(t, err) - fakeOAMClient.AssertNumberOfCalls(t, "ListSinks", 2) + fakeOAMClient.AssertNumberOfCalls(t, "ListSinksWithContext", 2) require.Len(t, resp, 1) assert.True(t, resp[0].Value.IsMonitoringAccount) assert.Equal(t, "Account 1", resp[0].Value.Label) @@ -82,27 +83,27 @@ func TestHandleGetAccounts(t *testing.T) { t.Run("Should merge the first sink with attached links", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")}, {Name: aws.String("Account 2"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group2")}, }, NextToken: new(string), }, nil).Once() - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")}, }, NextToken: nil, }, nil) - fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{ + fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{ Items: []*oam.ListAttachedLinksItem{ {Label: aws.String("Account 10"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789013:log-group:my-log-group10")}, {Label: aws.String("Account 11"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789014:log-group:my-log-group11")}, }, NextToken: new(string), }, nil).Once() - fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{ + fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{ Items: []*oam.ListAttachedLinksItem{ {Label: aws.String("Account 12"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group12")}, }, @@ -110,11 +111,11 @@ func TestHandleGetAccounts(t *testing.T) { }, nil) accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.NoError(t, err) - fakeOAMClient.AssertNumberOfCalls(t, "ListSinks", 2) - fakeOAMClient.AssertNumberOfCalls(t, "ListAttachedLinks", 2) + fakeOAMClient.AssertNumberOfCalls(t, "ListSinksWithContext", 2) + fakeOAMClient.AssertNumberOfCalls(t, "ListAttachedLinksWithContext", 2) expectedAccounts := []resources.ResourceResponse[resources.Account]{ {Value: resources.Account{Id: "123456789012", Label: "Account 1", Arn: "arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1", IsMonitoringAccount: true}}, {Value: resources.Account{Id: "123456789013", Label: "Account 10", Arn: "arn:aws:logs:us-east-1:123456789013:log-group:my-log-group10", IsMonitoringAccount: false}}, @@ -126,37 +127,37 @@ func TestHandleGetAccounts(t *testing.T) { t.Run("Should call ListAttachedLinks with arn of first sink", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")}, }, NextToken: new(string), }, nil).Once() - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{ {Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")}, }, NextToken: nil, }, nil).Once() - fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil) + fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil) accounts := NewAccountsService(fakeOAMClient) - _, _ = accounts.GetAccountsForCurrentUserOrRole() + _, _ = accounts.GetAccountsForCurrentUserOrRole(context.Background()) - fakeOAMClient.AssertCalled(t, "ListAttachedLinks", &oam.ListAttachedLinksInput{ + fakeOAMClient.AssertCalled(t, "ListAttachedLinksWithContext", &oam.ListAttachedLinksInput{ SinkIdentifier: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1"), }) }) t.Run("Should return an error in case of any error from ListAttachedLinks", func(t *testing.T) { fakeOAMClient := &mocks.FakeOAMClient{} - fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{ + fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{ Items: []*oam.ListSinksItem{{Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")}}, }, nil) - fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, fmt.Errorf("some error")).Once() + fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, fmt.Errorf("some error")).Once() accounts := NewAccountsService(fakeOAMClient) - resp, err := accounts.GetAccountsForCurrentUserOrRole() + resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background()) assert.Error(t, err) assert.Nil(t, resp)