From 870ef7b058429748bbe19af34134d907dea4909b Mon Sep 17 00:00:00 2001
From: Isabella Siu <Isabella.siu@grafana.com>
Date: Tue, 15 Aug 2023 17:31:55 -0400
Subject: [PATCH] CloudWatch Logs: Wrap sync error from executeGetQueryResults
 (#73252)

---
 pkg/tsdb/cloudwatch/log_actions.go         |  9 +++++-
 pkg/tsdb/cloudwatch/log_sync_query.go      |  2 +-
 pkg/tsdb/cloudwatch/log_sync_query_test.go | 32 ++++++++++++++++++++++
 pkg/tsdb/cloudwatch/test_utils.go          | 21 ++++++++++++++
 4 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/pkg/tsdb/cloudwatch/log_actions.go b/pkg/tsdb/cloudwatch/log_actions.go
index 207d61f5be5..fb7bf922524 100644
--- a/pkg/tsdb/cloudwatch/log_actions.go
+++ b/pkg/tsdb/cloudwatch/log_actions.go
@@ -299,7 +299,14 @@ func (e *cloudWatchExecutor) executeGetQueryResults(ctx context.Context, logsCli
 		QueryId: aws.String(logsQuery.QueryId),
 	}
 
-	return logsClient.GetQueryResultsWithContext(ctx, queryInput)
+	getQueryResultsResponse, err := logsClient.GetQueryResultsWithContext(ctx, queryInput)
+	if err != nil {
+		var awsErr awserr.Error
+		if errors.As(err, &awsErr) {
+			return getQueryResultsResponse, &AWSError{Code: awsErr.Code(), Message: err.Error()}
+		}
+	}
+	return getQueryResultsResponse, err
 }
 
 func (e *cloudWatchExecutor) handleGetQueryResults(ctx context.Context, logsClient cloudwatchlogsiface.CloudWatchLogsAPI,
diff --git a/pkg/tsdb/cloudwatch/log_sync_query.go b/pkg/tsdb/cloudwatch/log_sync_query.go
index 5d6fbd08774..b64a2db8d92 100644
--- a/pkg/tsdb/cloudwatch/log_sync_query.go
+++ b/pkg/tsdb/cloudwatch/log_sync_query.go
@@ -108,7 +108,7 @@ func (e *cloudWatchExecutor) syncQuery(ctx context.Context, logsClient cloudwatc
 	for range ticker.C {
 		res, err := e.executeGetQueryResults(ctx, logsClient, requestParams)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("CloudWatch Error: %w", err)
 		}
 		if isTerminated(*res.Status) {
 			return res, err
diff --git a/pkg/tsdb/cloudwatch/log_sync_query_test.go b/pkg/tsdb/cloudwatch/log_sync_query_test.go
index e43b616f386..96b27cd848f 100644
--- a/pkg/tsdb/cloudwatch/log_sync_query_test.go
+++ b/pkg/tsdb/cloudwatch/log_sync_query_test.go
@@ -336,4 +336,36 @@ func Test_executeSyncLogQuery_handles_RefId_from_input_queries(t *testing.T) {
 		assert.Error(t, err)
 		cli.AssertNumberOfCalls(t, "GetQueryResultsWithContext", 1)
 	})
+
+	t.Run("when getQueryResults returns aws error is returned, it keeps the context", func(t *testing.T) {
+		cli = &mockLogsSyncClient{}
+		cli.On("StartQueryWithContext", mock.Anything, mock.Anything, mock.Anything).Return(&cloudwatchlogs.StartQueryOutput{
+			QueryId: aws.String("abcd-efgh-ijkl-mnop"),
+		}, nil)
+		cli.On("GetQueryResultsWithContext", mock.Anything, mock.Anything, mock.Anything).Return(
+			&cloudwatchlogs.GetQueryResultsOutput{Status: aws.String("Complete")},
+			&fakeAWSError{code: "foo", message: "bar"},
+		)
+		im := datasource.NewInstanceManager(func(s backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+			return DataSource{Settings: models.CloudWatchSettings{}}, nil
+		})
+		executor := newExecutor(im, newTestConfig(), &fakeSessionCache{}, featuremgmt.WithFeatures())
+
+		res, err := executor.QueryData(context.Background(), &backend.QueryDataRequest{
+			Headers:       map[string]string{ngalertmodels.FromAlertHeaderName: "some value"},
+			PluginContext: backend.PluginContext{DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{}},
+			Queries: []backend.DataQuery{
+				{
+					TimeRange: backend.TimeRange{From: time.Unix(0, 0), To: time.Unix(1, 0)},
+					JSON: json.RawMessage(`{
+						"queryMode":    "Logs"
+					}`),
+				},
+			},
+		})
+
+		require.Nil(t, res)
+		require.Error(t, err)
+		require.Equal(t, "CloudWatch Error: foo: bar", err.Error())
+	})
 }
diff --git a/pkg/tsdb/cloudwatch/test_utils.go b/pkg/tsdb/cloudwatch/test_utils.go
index dea49df53fa..967d8e624e2 100644
--- a/pkg/tsdb/cloudwatch/test_utils.go
+++ b/pkg/tsdb/cloudwatch/test_utils.go
@@ -254,3 +254,24 @@ func (s *mockedCallResourceResponseSenderForOauth) Send(resp *backend.CallResour
 	s.Response = resp
 	return nil
 }
+
+type fakeAWSError struct {
+	code    string
+	message string
+}
+
+func (e fakeAWSError) OrigErr() error {
+	return nil
+}
+
+func (e fakeAWSError) Error() string {
+	return e.message
+}
+
+func (e fakeAWSError) Code() string {
+	return e.code
+}
+
+func (e fakeAWSError) Message() string {
+	return e.message
+}