Filter on account id (#157)

* Filter service graph by account id

* Refactor to create a separate resources route

* Use a map rather than looping everytime

* Fixes after CR:
- refactor /accountIds to /accounts to be more future proof
- make a for loop more readable
- put work behind feature flag
- move call to fetch account ids within QueryEditorForm
- fix bug to reset selection, if account id is no longer a valid option
- fix bug to handle edges with account ids with targets that do not have account ids
- remove accidentally left comments.

* WIP

* More pr fixes and adding in autocomplete

* Fix test

* Fix backend linter

* Fix backend test

* Comment out frontend tests

* Modify plugin.json id

* Fix bug about group name and fix broken test

* Fix go test

* prepare release

* Merge pull request #3 from grafana/lattice-docs

Add some documentation for cross account queries

* Merge main, fix conflicts

* Remove duplicate test from merge conflict

* revert plugin id

Co-authored-by: Sarah Zinger <sarah.zinger@grafana.com>
Co-authored-by: Shirley Leu <4163034+fridgepoet@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -2,8 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
-## 2.2.0
+## 2.3.0
+- Feature: Make it possible to filter on account id in https://github.com/grafana/x-ray-datasource/pull/157
 
+## 2.2.0
 - Make properties of `SummaryStatistics` optional
 
 ## 2.1.2

README.md

@@ -146,7 +146,7 @@ In Insights you can see the summary table for Insights. Clicking the InsightId w
 
 Service map in Grafana enables customers to view the map of their applications built using microservices architecture. Each node on the map represents a service such as an AWS Lambda function or API running on API Gateway or a DynamoDB table. With this map, customers can easily detect performance issues, or increase in error, fault or throttle rates in any of their services and dive deep into corresponding traces and root cause.
 
-![Service map](https://grafana.com/static/img/docs/node-graph/node-graph-7-4.png "Service map")
+![Service map](https://grafana.com/static/img/docs/node-graph/node-graph-7-4.png 'Service map')
 
 Service Map query type shows the same data as a service map inside X-ray console.
 
@@ -158,8 +158,7 @@ To display the service map:
 
 You can pan and zoom the view with buttons or you mouse. For details about the visualization, refer to [Node graph panel](https://grafana.com/docs/grafana/latest/panels/visualizations/node-graph/).
 
-![Service map navigation](https://storage.googleapis.com/integration-artifacts/grafana-x-ray-datasource/screenshots/x-ray-service-map-nav.gif "Service map navigation")
-
+![Service map navigation](https://storage.googleapis.com/integration-artifacts/grafana-x-ray-datasource/screenshots/x-ray-service-map-nav.gif 'Service map navigation')
 
 Similar to X-ray root nodes, nodes in the service map representing the client application are on the left side of the map.
 
@@ -221,3 +220,22 @@ datasources:
       accessKey: '<your access key>'
       secretKey: '<your secret key>'
 ```
+
+## Cross-Account Observability
+
+The X-Ray plugin allows you to monitor traces across multiple AWS accounts within a region with a feature called Cross-Account Observability. Using cross-account observability, you can seamlessly search, visualize and analyze AWS traces without worrying about account boundaries.
+
+### Getting started
+
+To enable cross-account observability, first enable the feature in AWS using the official [CloudWatch docs](http://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-Unified-Cross-Account.html), then add [two new API actions]({{< relref "../Metrics_and_Logs" >}}) to the IAM policy attached to the role/user running the plugin.
+
+This feature is currently behind the `cloudWatchCrossAccountQuerying` feature toggle.
+
+> You can enable feature toggles through configuration file or environment variables. See configuration [docs]({{< relref "../setup-grafana/configure-grafana/#feature_toggles" >}}) for details.
+> Grafana Cloud users can access this feature by [opening a support ticket in the Cloud Portal](https://grafana.com/profile/org#support).
+
+## Filtering Traces by Account Id
+
+Once the feature is enabled, you will be able to both display traces across multiple accounts and filter those traces by account ID. When you select the Service Map query type in Grafana, an account dropdown displays and populates with the account IDs shown in the traces returned in the selected time frame.
+
+You can also add account ID as part of a query filter expression in the Trace List query type.

pkg/datasource/datasource.go

@@ -103,6 +103,8 @@ func NewDatasource(xrayClientFactory XrayClientFactory) *Datasource {
 	resMux := http.NewServeMux()
 	resMux.HandleFunc("/groups", ds.getGroups)
 	resMux.HandleFunc("/regions", ds.GetRegions)
+	resMux.HandleFunc("/accounts", ds.GetAccounts)
+
 	ds.ResourceMux = httpadapter.New(resMux)
 	return ds
 }

pkg/datasource/datasource_test.go

@@ -27,7 +27,14 @@ func (client *XrayClientMock) GetServiceGraphPagesWithContext(ctx aws.Context, i
 	output := &xray.GetServiceGraphOutput{
 		NextToken: nil,
 		Services: []*xray.Service{
-			{Name: aws.String(serviceName)},
+			{
+				Name:      aws.String(serviceName),
+				AccountId: aws.String("testAccount1"),
+			},
+			{
+				Name:      aws.String(serviceName + "2"),
+				AccountId: aws.String("testAccount2"),
+			},
 		},
 	}
 	fn(output, false)
@@ -527,7 +534,7 @@ func TestDatasource(t *testing.T) {
 
 		// Bit simplistic test but right now we just send each service as a json to frontend and do transform there.
 		frame := response.Responses["A"].Frames[0]
-		require.Equal(t, 1, frame.Fields[0].Len())
+		require.Equal(t, 2, frame.Fields[0].Len()) // 2 because of the 2 services added to the mock
 	})
 
 	t.Run("getServiceMap query with region", func(t *testing.T) {
@@ -537,7 +544,7 @@ func TestDatasource(t *testing.T) {
 
 		// Bit simplistic test but right now we just send each service as a json to frontend and do transform there.
 		frame := response.Responses["A"].Frames[0]
-		require.Equal(t, 1, frame.Fields[0].Len())
+		require.Equal(t, 2, frame.Fields[0].Len())
 		require.True(t, strings.Contains(frame.Fields[0].At(0).(string), "mockServiceName-us-east-1"))
 	})
 

pkg/datasource/getAccounts.go

@@ -0,0 +1,90 @@
+package datasource
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/url"
+	"time"
+
+	"github.com/aws/aws-sdk-go/service/xray"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/resource/httpadapter"
+)
+
+type Account struct {
+	Id string
+}
+
+func (ds *Datasource) GetAccounts(rw http.ResponseWriter, req *http.Request) {
+	if req.Method != "GET" {
+		rw.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+	urlQuery, err := url.ParseQuery(req.URL.RawQuery)
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+	region := urlQuery.Get("region")
+
+	pluginConfig := httpadapter.PluginConfigFromContext(req.Context())
+	xrayClient, err := ds.xrayClientFactory(&pluginConfig, RequestSettings{Region: region})
+
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+
+	group := urlQuery.Get("group")
+
+	layout := "2006-01-02T15:04:05.000Z"
+	startTime, err := time.Parse(layout, urlQuery.Get("startTime"))
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+
+	endTime, err := time.Parse(layout, urlQuery.Get("endTime"))
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+
+	input := &xray.GetServiceGraphInput{
+		StartTime: &startTime,
+		EndTime:   &endTime,
+		GroupName: &group,
+	}
+
+	accounts := []Account{}
+
+	err = xrayClient.GetServiceGraphPagesWithContext(req.Context(), input, func(page *xray.GetServiceGraphOutput, lastPage bool) bool {
+		for _, service := range page.Services {
+			if service.AccountId != nil {
+				account := Account{
+					Id: *service.AccountId,
+				}
+				accounts = append(accounts, account)
+			}
+		}
+
+		// Not sure how many pages there can possibly be but for now try to iterate over all the pages.
+		return true
+	})
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+
+	body, err := json.Marshal(accounts)
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+
+	rw.Header().Set("content-type", "application/json")
+	_, err = rw.Write(body)
+	if err != nil {
+		sendError(rw, err)
+		return
+	}
+}

pkg/datasource/getAccounts_test.go

@@ -0,0 +1,30 @@
+package datasource_test
+
+import (
+	"encoding/json"
+	"io"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/x-ray-datasource/pkg/datasource"
+	"github.com/stretchr/testify/require"
+)
+
+type Account struct {
+  Id string
+}
+
+func TestAccounts(t *testing.T) {
+  t.Run("when passed a get request it returns a list of all accountIds in the traces in the selected time frame", func(t *testing.T) {
+    ds := datasource.NewDatasource(xrayClientFactory)
+    req := httptest.NewRequest("GET", "http://example.com/accounts?startTime=2022-09-23T00:15:14.365Z&endTime=2022-09-23T01:15:14.365Z&group=somegroup", nil)
+    w := httptest.NewRecorder()
+    ds.GetAccounts(w, req)
+    resp := w.Result()
+    body, _ := io.ReadAll(resp.Body)
+    accounts := []Account{}
+    require.NoError(t, json.Unmarshal(body, &accounts))
+    require.Contains(t, accounts[0].Id, "testAccount1")
+    require.Contains(t, accounts[1].Id, "testAccount2")
+  })
+}

pkg/datasource/getServiceMap.go

@@ -11,8 +11,9 @@ import (
 )
 
 type GetServiceMapQueryData struct {
-	Region string      `json:"region"`
-	Group  *xray.Group `json:"group"`
+	Region     string      `json:"region"`
+	Group      *xray.Group `json:"group"`
+	AccountIds []string    `json:"accountIds,omitempty"`
 }
 
 func (ds *Datasource) getServiceMap(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
@@ -57,8 +58,25 @@ func (ds *Datasource) getSingleServiceMap(ctx context.Context, query backend.Dat
 		EndTime:   &query.TimeRange.To,
 		GroupName: queryData.Group.GroupName,
 	}
+
+	accountIdsToFilterBy := make(map[string]bool)
+	for _, value := range queryData.AccountIds {
+		accountIdsToFilterBy[value] = true
+	}
+
 	err = xrayClient.GetServiceGraphPagesWithContext(ctx, input, func(page *xray.GetServiceGraphOutput, lastPage bool) bool {
 		for _, service := range page.Services {
+			// filter out non matching account ids, if user has selected them
+			if len(queryData.AccountIds) > 0 {
+				// sometimes traces don't have accountId data, without knowing where it came from we have to filter it out
+				if service.AccountId == nil {
+					continue
+				}
+
+				if !accountIdsToFilterBy[*service.AccountId] {
+					continue
+				}
+			}
 			bytes, err := json.Marshal(service)
 			if err != nil {
 				// TODO: probably does not make sense to fail just because of one service but I assume the layout will fail

src/DataSource.ts

@@ -10,7 +10,7 @@ import {
   TimeRange,
   toDuration,
 } from '@grafana/data';
-import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
+import { DataSourceWithBackend, getTemplateSrv, TemplateSrv, config } from '@grafana/runtime';
 import { Observable } from 'rxjs';
 import { map } from 'rxjs/operators';
 
@@ -61,11 +61,11 @@ export class XrayDataSource extends DataSourceWithBackend<XrayQuery, XrayJsonDat
       const params = new URLSearchParams({ region });
       searchString = '?' + params.toString();
     }
-    return this.getResource(`/groups${searchString}`);
+    return this.getResource(`groups${searchString}`);
   }
 
   async getRegions(): Promise<Region[]> {
-    const response = await this.getResource('/regions');
+    const response = await this.getResource('regions');
     return [
       ...sortBy(
         response.map((name: string) => ({
@@ -78,6 +78,22 @@ export class XrayDataSource extends DataSourceWithBackend<XrayQuery, XrayJsonDat
     ];
   }
 
+  async getAccountIdsForServiceMap(range?: TimeRange, group?: Group): Promise<string[]> {
+    if (!config.featureToggles.cloudWatchCrossAccountQuerying) {
+      return [];
+    }
+    const params = new URLSearchParams({
+      startTime: range ? range.from.toISOString() : '',
+      endTime: range ? range.to.toISOString() : '',
+      group: group?.GroupName || 'Default',
+    });
+
+    const searchString = '?' + params.toString();
+
+    const response = await this.getResource(`accounts${searchString}`);
+    return response.map((account: { Id: string }) => account.Id);
+  }
+
   getServiceMapUrl(region?: string): string {
     return `${this.getXrayUrl(region)}#/service-map/`;
   }

src/components/QueryEditor/QueryEditor.test.tsx

@@ -39,6 +39,9 @@ const defaultProps = {
     async getRegions(): Promise<Region[]> {
       return [{ label: 'region1', text: 'region1', value: 'region1' }];
     },
+    async getAccountIdsForServiceMap(): Promise<string[]> {
+      return ['account1', 'account2'];
+    },
     getServiceMapUrl() {
       return 'service-map';
     },
@@ -153,9 +156,7 @@ describe('QueryEditor', () => {
 
     fireEvent.change(field, { target: { value: '1-5f160a8b-83190adad07f429219c0e259' } });
 
-    // First call would be the query init call. We do not update the query based on that so when doing this second one
-    // it's done without default region or group.
-    expect(onChange.mock.calls[1][0]).toEqual({
+    expect(onChange.mock.calls[3][0]).toEqual({
       refId: 'A',
       query: '1-5f160a8b-83190adad07f429219c0e259',
       queryType: XrayQueryType.getTrace,
@@ -216,6 +217,26 @@ describe('QueryEditor', () => {
       serviceMap: 'https://region2.console.aws.amazon.com/xray/home?region=region2#/service-map/',
     });
   });
+
+  it('shows the accountIds in a dropdown on service map selection', async () => {
+    await act(async () => {
+      render(
+        <QueryEditor
+          {...{
+            ...defaultProps,
+            query: {
+              refId: 'A',
+              queryType: 'getServiceMap',
+              accountIds: ['account1'],
+            } as any,
+          }}
+          onChange={() => {}}
+        />
+      );
+      expect(screen.getByText('', { selector: '.fa-spinner' })).toBeDefined();
+      await waitFor(() => expect(screen.getByText('account1')).toBeDefined());
+    });
+  });
 });
 
 function makeDataSource(settings: DataSourceInstanceSettings<XrayJsonData>) {

src/components/QueryEditor/QueryEditor.tsx

@@ -3,7 +3,6 @@ import { Spinner } from '@grafana/ui';
 import { useGroups } from './useGroups';
 import { QueryEditorForm, XrayQueryEditorFormProps } from './QueryEditorForm';
 import { useRegions } from './useRegions';
-
 /**
  * Simple wrapper that is only responsible to load groups and delay actual render of the QueryEditorForm. Main reason
  * for that is that there is queryInit code that requires groups to be already loaded and is separate hook and it