Add Support for Temporary Credentials (#284)

Add Support for Temporary Credentials, publish to 2.12.0

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 2.12.0
+
+- Add support for Temporary Credentials [#284] https://github.com/grafana/athena-datasource/pull/284
+
 ## 2.11.1
 
 - Update @grafana/aws-sdk to 0.1.2 to fix bug with temporary credentials

package.json

@@ -1,6 +1,6 @@
 {
   "name": "grafana-athena-datasource",
-  "version": "2.11.1",
+  "version": "2.12.0",
   "description": "Use Amazon Athena with Grafana",
   "scripts": {
     "build": "webpack -c ./.config/webpack/webpack.config.ts --env production",
@@ -29,7 +29,7 @@
   "devDependencies": {
     "@babel/core": "^7.16.7",
     "@emotion/css": "^11.1.3",
-    "@grafana/aws-sdk": "0.1.2",
+    "@grafana/aws-sdk": "0.2.0",
     "@grafana/data": "9.3.2",
     "@grafana/e2e": "9.3.2",
     "@grafana/e2e-selectors": "9.3.2",
@@ -42,6 +42,7 @@
     "@swc/jest": "^0.2.20",
     "@testing-library/jest-dom": "^5.16.2",
     "@testing-library/react": "^12.1.3",
+    "@testing-library/user-event": "^14.5.1",
     "@types/glob": "^8.0.0",
     "@types/jest": "^27.4.1",
     "@types/lodash": "latest",

pkg/athena/routes/routes.go

@@ -2,8 +2,10 @@ package routes
 
 import (
 	"net/http"
+	"os"
 
 	"github.com/grafana/athena-datasource/pkg/athena"
+	"github.com/grafana/grafana-aws-sdk/pkg/awsds"
 	"github.com/grafana/grafana-aws-sdk/pkg/sql/routes"
 )
 
@@ -49,10 +51,22 @@ func (r *(AthenaResourceHandler)) workgroupEngineVersion(rw http.ResponseWriter,
 	routes.SendResources(rw, res, err)
 }
 
+type ExternalIdResponse struct {
+	ExternalId string `json:"externalId"`
+}
+
+func (r *AthenaResourceHandler) externalId(rw http.ResponseWriter, req *http.Request) {
+	res := ExternalIdResponse{
+		ExternalId: os.Getenv(awsds.GrafanaAssumeRoleExternalIdKeyName),
+	}
+	routes.SendResources(rw, res, nil)
+}
+
 func (r *AthenaResourceHandler) Routes() map[string]func(http.ResponseWriter, *http.Request) {
 	routes := r.DefaultRoutes()
 	routes["/catalogs"] = r.catalogs
 	routes["/workgroups"] = r.workgroups
 	routes["/workgroupEngineVersion"] = r.workgroupEngineVersion
+	routes["/externalId"] = r.externalId
 	return routes
 }

pkg/athena/routes/routes_test.go

@@ -9,6 +9,8 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/grafana/athena-datasource/pkg/athena/fake"
+	"github.com/grafana/grafana-aws-sdk/pkg/awsds"
+	"github.com/stretchr/testify/require"
 )
 
 var ds = &fake.AthenaFakeDatasource{
@@ -126,6 +128,13 @@ func TestRoutes(t *testing.T) {
 			reqBody:      []byte{},
 			expectedCode: http.StatusBadRequest,
 		},
+		{
+			description:    "externalId",
+			route:          "/externalId",
+			reqBody:        []byte{},
+			expectedCode:   http.StatusOK,
+			expectedResult: `{"externalId":""}`,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.description, func(t *testing.T) {
@@ -149,3 +158,40 @@ func TestRoutes(t *testing.T) {
 		})
 	}
 }
+
+func setupHandler() AthenaResourceHandler {
+	rh := AthenaResourceHandler{athena: ds}
+	rh.API = ds
+	return rh
+}
+
+func hitRoute(rh AthenaResourceHandler, route string, reqBody []byte) (*http.Response, []byte, error) {
+	req := httptest.NewRequest("GET", route, bytes.NewReader(reqBody))
+	rw := httptest.NewRecorder()
+	rh.Routes()[route](rw, req)
+	resp := rw.Result()
+	body, err := io.ReadAll(resp.Body)
+	return resp, body, err
+}
+func TestRoutes_ExternalId(t *testing.T) {
+
+	t.Run("it returns an externalId if one is set in the env", func(t *testing.T) {
+		t.Setenv(awsds.GrafanaAssumeRoleExternalIdKeyName, "a fake external id")
+
+		rh := setupHandler()
+		resp, body, err := hitRoute(rh, "/externalId", []byte{})
+
+		require.NoError(t, err)
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		require.Equal(t, `{"externalId":"a fake external id"}`, string(body))
+	})
+	t.Run("it returns an empty string if there is no external id set in the env", func(t *testing.T) {
+		rh := setupHandler()
+		resp, body, err := hitRoute(rh, "/externalId", []byte{})
+
+		require.NoError(t, err)
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		require.Equal(t, `{"externalId":""}`, string(body))
+	})
+
+}

src/ConfigEditor.test.tsx

@@ -1,40 +1,34 @@
 import React from 'react';
-import { render, screen, fireEvent } from '@testing-library/react';
+import { render, screen, fireEvent, waitFor } from '@testing-library/react';
 import { ConfigEditor } from './ConfigEditor';
 import { mockDatasourceOptions } from './__mocks__/datasource';
 import { select } from 'react-select-event';
 import { selectors } from 'tests/selectors';
+import { AwsAuthType } from '@grafana/aws-sdk';
+import * as runtime from '@grafana/runtime';
+import userEvent from '@testing-library/user-event';
 
 const resourceName = 'foo';
-
-jest.mock('@grafana/aws-sdk', () => {
-  return {
-    ...(jest.requireActual('@grafana/aws-sdk') as any),
-    ConnectionConfig: function ConnectionConfig() {
-      return <></>;
-    },
-  };
-});
-jest.mock('@grafana/runtime', () => {
-  return {
-    ...(jest.requireActual('@grafana/runtime') as any),
-    getBackendSrv: () => ({
-      put: jest.fn().mockResolvedValue({ datasource: {} }),
-      post: jest.fn().mockResolvedValue([resourceName]),
-    }),
-  };
-});
 const props = mockDatasourceOptions;
 
+const setUpMockBackendServer = (mockBackendSrv: { put: () => void; post: () => void }) => {
+  jest.spyOn(runtime, 'getBackendSrv').mockImplementation(() => mockBackendSrv as unknown as runtime.BackendSrv);
+};
+
 describe('ConfigEditor', () => {
   it('should save and request catalogs', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue([resourceName]),
+    });
+
     const onChange = jest.fn();
     render(<ConfigEditor {...props} onOptionsChange={onChange} />);
 
     const selectEl = screen.getByLabelText(selectors.components.ConfigEditor.catalog.input);
     expect(selectEl).toBeInTheDocument();
 
-    await select(selectEl, resourceName, { container: document.body });
+    await waitFor(() => select(selectEl, resourceName, { container: document.body }));
 
     expect(onChange).toHaveBeenCalledWith({
       ...props.options,
@@ -43,6 +37,11 @@ describe('ConfigEditor', () => {
   });
 
   it('should save and request databases', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue([resourceName]),
+    });
+
     const onChange = jest.fn();
     render(<ConfigEditor {...props} onOptionsChange={onChange} />);
 
@@ -53,7 +52,7 @@ describe('ConfigEditor', () => {
     const selectEl = screen.getByLabelText(selectors.components.ConfigEditor.database.input);
     expect(selectEl).toBeInTheDocument();
 
-    await select(selectEl, resourceName, { container: document.body });
+    await waitFor(() => select(selectEl, resourceName, { container: document.body }));
 
     expect(onChange).toHaveBeenCalledWith({
       ...props.options,
@@ -62,6 +61,11 @@ describe('ConfigEditor', () => {
   });
 
   it('should save and request workgroups', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue([resourceName]),
+    });
+
     const onChange = jest.fn();
     render(<ConfigEditor {...props} onOptionsChange={onChange} />);
 
@@ -72,7 +76,7 @@ describe('ConfigEditor', () => {
     const selectEl = screen.getByLabelText(selectors.components.ConfigEditor.workgroup.input);
     expect(selectEl).toBeInTheDocument();
 
-    await select(selectEl, resourceName, { container: document.body });
+    await waitFor(() => select(selectEl, resourceName, { container: document.body }));
 
     expect(onChange).toHaveBeenCalledWith({
       ...props.options,
@@ -81,6 +85,11 @@ describe('ConfigEditor', () => {
   });
 
   it('should use an output location', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue([resourceName]),
+    });
+
     const onChange = jest.fn();
     render(<ConfigEditor {...props} onOptionsChange={onChange} />);
     const input = screen.getByTestId(selectors.components.ConfigEditor.OutputLocation.wrapper);
@@ -91,4 +100,85 @@ describe('ConfigEditor', () => {
       jsonData: { ...props.options.jsonData, outputLocation: bucket },
     });
   });
+
+  it('should fetch and display externalId when the auth type is grafana_assume_role', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue({ externalId: 'fake-external-id' }),
+    });
+
+    render(
+      <ConfigEditor
+        {...props}
+        options={{
+          ...props.options,
+          jsonData: {
+            ...props.options.jsonData,
+            authType: AwsAuthType.GrafanaAssumeRole,
+          },
+        }}
+      />
+    );
+
+    expect(screen.queryByText('Grafana Assume Role')).toBeInTheDocument();
+    const instructionsButton = await screen.findByRole('button', {
+      name: /How to create an IAM role for grafana to assume/i,
+    });
+    await userEvent.click(instructionsButton);
+    expect(screen.queryByText('fake-external-id')).toBeInTheDocument();
+  });
+
+  it('gracefully handles when the fetch for external id throws an error', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockRejectedValue('the server exploded for some reason'),
+    });
+
+    render(
+      <ConfigEditor
+        {...props}
+        options={{
+          ...props.options,
+          jsonData: {
+            ...props.options.jsonData,
+            authType: AwsAuthType.GrafanaAssumeRole,
+          },
+        }}
+      />
+    );
+
+    expect(screen.queryByText('Grafana Assume Role')).toBeInTheDocument();
+    const instructionsButton = await screen.findByRole('button', {
+      name: /How to create an IAM role for grafana to assume/i,
+    });
+    await userEvent.click(instructionsButton);
+    expect(screen.queryByText('External Id is currently unavailable')).toBeInTheDocument();
+  });
+
+  it('gracefully handles when the fetch for external id return an empty string', async () => {
+    setUpMockBackendServer({
+      put: jest.fn().mockResolvedValue({ datasource: {} }),
+      post: jest.fn().mockResolvedValue({ externalId: '' }),
+    });
+
+    render(
+      <ConfigEditor
+        {...props}
+        options={{
+          ...props.options,
+          jsonData: {
+            ...props.options.jsonData,
+            authType: AwsAuthType.GrafanaAssumeRole,
+          },
+        }}
+      />
+    );
+
+    expect(screen.queryByText('Grafana Assume Role')).toBeInTheDocument();
+    const instructionsButton = await screen.findByRole('button', {
+      name: /How to create an IAM role for grafana to assume/i,
+    });
+    await userEvent.click(instructionsButton);
+    expect(screen.queryByText('External Id is currently unavailable')).toBeInTheDocument();
+  });
 });

src/ConfigEditor.tsx

@@ -1,8 +1,8 @@
-import React, { useState, FormEvent } from 'react';
+import React, { FormEvent, useCallback, useEffect, useState } from 'react';
 import { DataSourcePluginOptionsEditorProps, DataSourceSettings, SelectableValue } from '@grafana/data';
 import { AthenaDataSourceOptions, AthenaDataSourceSecureJsonData, AthenaDataSourceSettings, defaultKey } from './types';
 import { getBackendSrv } from '@grafana/runtime';
-import { InlineInput, ConfigSelect, ConnectionConfig } from '@grafana/aws-sdk';
+import { AwsAuthType, ConfigSelect, ConnectionConfig, InlineInput } from '@grafana/aws-sdk';
 import { selectors } from 'tests/selectors';
 
 type Props = DataSourcePluginOptionsEditorProps<AthenaDataSourceOptions, AthenaDataSourceSecureJsonData>;
@@ -13,6 +13,7 @@ export function ConfigEditor(props: Props) {
   const baseURL = `/api/datasources/${props.options.id}`;
   const resourcesURL = `${baseURL}/resources`;
   const [saved, setSaved] = useState(!!props.options.jsonData.defaultRegion);
+  const [externalId, setExternalId] = useState('');
   const saveOptions = async () => {
     if (saved) {
       return;
@@ -51,6 +52,21 @@ export function ConfigEditor(props: Props) {
     return loadedWorkgroups;
   };
 
+  const fetchExternalId = useCallback(async () => {
+    try {
+      const response = await getBackendSrv().post(resourcesURL + '/externalId', {});
+      setExternalId(response.externalId);
+    } catch {
+      setExternalId('');
+    }
+  }, [resourcesURL]);
+
+  useEffect(() => {
+    if (props.options.jsonData.authType === AwsAuthType.GrafanaAssumeRole) {
+      fetchExternalId();
+    }
+  }, [props.options.jsonData.authType, fetchExternalId]);
+
   const onOptionsChange = (options: DataSourceSettings<AthenaDataSourceOptions, AthenaDataSourceSecureJsonData>) => {
     setSaved(false);
     props.onOptionsChange(options);
@@ -80,7 +96,7 @@ export function ConfigEditor(props: Props) {
 
   return (
     <div className="gf-form-group">
-      <ConnectionConfig {...props} onOptionsChange={onOptionsChange} />
+      <ConnectionConfig {...props} onOptionsChange={onOptionsChange} externalId={externalId} />
       <h3>Athena Details</h3>
       <ConfigSelect
         {...props}