Add service map support

pkg/datasource/datasource.go

@@ -63,6 +63,7 @@ type Datasource struct {
 	ec2ClientFactory  Ec2ClientFactory
 }
 
+// Needs to match XrayQueryType in frontend code
 const (
 	QueryGetTrace                                 = "getTrace"
 	QueryGetTraceSummaries                        = "getTraceSummaries"
@@ -79,6 +80,7 @@ const (
 	QueryGetAnalyticsUser                         = "getAnalyticsUser"
 	QueryGetAnalyticsStatusCode                   = "getAnalyticsStatusCode"
 	QueryGetInsights                              = "getInsights"
+  QueryGetServiceMap                            = "getServiceMap"
 )
 
 func NewDatasource(
@@ -106,6 +108,7 @@ func NewDatasource(
 	mux.HandleFunc(QueryGetAnalyticsUrl, ds.getAnalytics)
 	mux.HandleFunc(QueryGetAnalyticsStatusCode, ds.getAnalytics)
 	mux.HandleFunc(QueryGetInsights, ds.getInsights)
+  mux.HandleFunc(QueryGetServiceMap, ds.getServiceMap)
 
 	ds.QueryMux = mux
 
@@ -154,4 +157,6 @@ type XrayClient interface {
 	) error
 	GetInsightSummaries(input *xray.GetInsightSummariesInput) (*xray.GetInsightSummariesOutput, error)
 	GetGroupsPages(input *xray.GetGroupsInput, fn func(*xray.GetGroupsOutput, bool) bool) error
+	GetServiceGraphPagesWithContext(ctx aws.Context, input *xray.GetServiceGraphInput, fn func(*xray.GetServiceGraphOutput, bool) bool, opts ...request.Option) error
+  GetTraceGraphPages(input *xray.GetTraceGraphInput, fn func(*xray.GetTraceGraphOutput, bool) bool) error
 }

pkg/datasource/datasource_test.go

@@ -17,6 +17,19 @@ import (
 
 type XrayClientMock struct{}
 
+func (client *XrayClientMock) GetServiceGraphPagesWithContext(ctx aws.Context, input *xray.GetServiceGraphInput, fn func(*xray.GetServiceGraphOutput, bool) bool, opts ...request.Option) error {
+  panic("implement me")
+}
+
+func (client *XrayClientMock) GetTraceGraphPages(input *xray.GetTraceGraphInput, fn func(*xray.GetTraceGraphOutput, bool) bool) error {
+  output := &xray.GetTraceGraphOutput{
+    NextToken: nil,
+    Services: []*xray.Service{} ,
+  }
+  fn(output, false)
+  return nil
+}
+
 func makeSummary() *xray.TraceSummary {
 	http := &xray.Http{
 		ClientIp:   aws.String("127.0.0.1"),

pkg/datasource/getAnalytics_test.go

@@ -19,6 +19,14 @@ type XrayClientMock struct {
   traces []*xray.TraceSummary
 }
 
+func (client *XrayClientMock) GetServiceGraphPagesWithContext(ctx aws.Context, input *xray.GetServiceGraphInput, fn func(*xray.GetServiceGraphOutput, bool) bool, opts ...request.Option) error {
+  panic("implement me")
+}
+
+func (client *XrayClientMock) GetTraceGraphPages(input *xray.GetTraceGraphInput, fn func(*xray.GetTraceGraphOutput, bool) bool) error {
+  panic("implement me")
+}
+
 func NewXrayClientMock(traces ...[]*xray.TraceSummary) *XrayClientMock {
   var allTraces []*xray.TraceSummary
   for _, l := range traces {

pkg/datasource/getServiceMap.go

@@ -0,0 +1,84 @@
+package datasource
+
+import (
+  "context"
+  "encoding/json"
+  "github.com/grafana/grafana-plugin-sdk-go/backend"
+  "github.com/grafana/grafana-plugin-sdk-go/backend/log"
+  "github.com/grafana/grafana-plugin-sdk-go/data"
+  xray "github.com/grafana/x-ray-datasource/pkg/xray"
+)
+
+type GetServiceMapQueryData struct {
+  Query  string `json:"query"`
+  Region string `json:"region"`
+  Group  *xray.Group `json:"group"`
+}
+
+func (ds *Datasource) getServiceMap(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+  response := &backend.QueryDataResponse{
+    Responses: make(map[string]backend.DataResponse),
+  }
+
+  // TODO not used in the app but this could actually be done in one call for multiple traceIDs
+  for _, query := range req.Queries {
+    response.Responses[query.RefID] = ds.getSingleServiceMap(ctx, query, &req.PluginContext)
+  }
+
+  return response, nil
+}
+
+// getSingleTrace returns single trace from BatchGetTraces API and unmarshals it.
+func (ds *Datasource) getSingleServiceMap(ctx context.Context, query backend.DataQuery, pluginContext *backend.PluginContext) backend.DataResponse {
+  queryData := &GetServiceMapQueryData{}
+  err := json.Unmarshal(query.JSON, queryData)
+
+  if err != nil {
+    return backend.DataResponse{
+      Error: err,
+    }
+  }
+
+  xrayClient, err := ds.xrayClientFactory(pluginContext, queryData.Region)
+  if err != nil {
+    return backend.DataResponse{
+      Error: err,
+    }
+  }
+
+  var frame = data.NewFrame(
+    "ServiceMap",
+    data.NewField("Service", nil, []string{}),
+  )
+
+
+  log.DefaultLogger.Debug("getSingleServiceMap", "RefID", query.RefID, "query", queryData.Query)
+  input := &xray.GetServiceGraphInput{
+    StartTime:        &query.TimeRange.From,
+    EndTime:          &query.TimeRange.To,
+    GroupName:        queryData.Group.GroupName,
+  }
+  err = xrayClient.GetServiceGraphPagesWithContext(ctx, input, func(page *xray.GetServiceGraphOutput, lastPage bool) bool {
+    for _, service := range page.Services {
+      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
+        //  because of some edge not connected to anything.
+        log.DefaultLogger.Error("getSingleServiceMap failed to marshal service", "Name", service.Name, "ReferenceId", service.ReferenceId)
+      }
+      frame.AppendRow(string(bytes))
+    }
+    // Not sure how many pages there can possibly be but for now try to iterate over all the pages.
+    return true
+  })
+
+  if err != nil {
+    return backend.DataResponse{
+      Error: err,
+    }
+  }
+
+  return backend.DataResponse{
+    Frames: []*data.Frame{frame},
+  }
+}

pkg/datasource/getTrace.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+  "sync"
 
-	"github.com/grafana/grafana-plugin-sdk-go/backend"
+  "github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	xray "github.com/grafana/x-ray-datasource/pkg/xray"
@@ -50,10 +51,53 @@ func (ds *Datasource) getSingleTrace(query backend.DataQuery, pluginContext *bac
 
 	log.DefaultLogger.Debug("getSingleTrace", "RefID", query.RefID, "query", queryData.Query)
 
-	tracesResponse, err := xrayClient.BatchGetTraces(&xray.BatchGetTracesInput{TraceIds: []*string{&queryData.Query}})
-	if err != nil {
+  var wg sync.WaitGroup
+  var tracesResponse *xray.BatchGetTracesOutput
+  var tracesError error
+
+  var traceGraphFrame = data.NewFrame(
+    "TraceGraph",
+    data.NewField("Service", nil, []string{}),
+  )
+  var traceGraphError error
+
+
+  wg.Add(1)
+  go func() {
+    defer wg.Done()
+    tracesResponse, tracesError = xrayClient.BatchGetTraces(&xray.BatchGetTracesInput{TraceIds: []*string{&queryData.Query}})
+  }()
+
+  // We get the trace graph in parallel but if this fails we still return the trace
+  wg.Add(1)
+  go func() {
+    defer wg.Done()
+    traceGraphError = xrayClient.GetTraceGraphPages(
+      &xray.GetTraceGraphInput{TraceIds: []*string{&queryData.Query}},
+      func(page *xray.GetTraceGraphOutput, hasMore bool) bool {
+        for _, service := range page.Services {
+          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
+            //  because of some edge not connected to anything.
+            log.DefaultLogger.Error(
+              "getSingleTrace failed to marshal service from trace graph",
+              "Name", service.Name,
+              "ReferenceId", service.ReferenceId,
+            )
+          }
+          traceGraphFrame.AppendRow(string(bytes))
+        }
+        return true
+      },
+    )
+  }()
+
+  wg.Wait()
+
+	if tracesError != nil {
 		return backend.DataResponse{
-			Error: err,
+			Error: tracesError,
 		}
 	}
 
@@ -72,11 +116,19 @@ func (ds *Datasource) getSingleTrace(query backend.DataQuery, pluginContext *bac
 		}
 	}
 
+	frames := []*data.Frame{
+		data.NewFrame(
+		  "Traces", data.NewField("Trace", nil, []string{string(traceBytes)}),
+		),
+	}
+
+	if traceGraphError == nil {
+	  frames = append(frames, traceGraphFrame)
+  }
+
 	return backend.DataResponse{
-		Frames: []*data.Frame{
-			data.NewFrame(
-				"Traces", data.NewField("Trace", nil, []string{string(traceBytes)}),
-			),
-		},
+		Frames: frames,
+		// TODO not sure what will this show if we have both data and error
+		Error: traceGraphError,
 	}
 }