From 9b1c0455c455c73807f863bd57f8c9999e470f79 Mon Sep 17 00:00:00 2001
From: Andrej Ocenas <mr.ocenas@gmail.com>
Date: Tue, 5 May 2020 14:21:23 +0200
Subject: [PATCH] Jaeger: Fix root span label in cascader (#24164)

---
 .../datasource/jaeger/QueryField.test.tsx     | 90 ++++++++++++++++---
 .../plugins/datasource/jaeger/QueryField.tsx  | 14 +--
 2 files changed, 88 insertions(+), 16 deletions(-)

diff --git a/public/app/plugins/datasource/jaeger/QueryField.test.tsx b/public/app/plugins/datasource/jaeger/QueryField.test.tsx
index 77f93924e21..944a13864ca 100644
--- a/public/app/plugins/datasource/jaeger/QueryField.test.tsx
+++ b/public/app/plugins/datasource/jaeger/QueryField.test.tsx
@@ -1,24 +1,15 @@
 import React from 'react';
 import { JaegerQueryField } from './QueryField';
-import { shallow } from 'enzyme';
+import { shallow, mount } from 'enzyme';
 import { JaegerDatasource, JaegerQuery } from './datasource';
 import { ButtonCascader } from '@grafana/ui';
 
 describe('JaegerQueryField', function() {
   it('shows empty value if no services returned', function() {
-    const dsMock: JaegerDatasource = {
-      metadataRequest(url: string) {
-        if (url.indexOf('/services') > 0) {
-          return Promise.resolve([]);
-        }
-        throw new Error(`Unexpected url: ${url}`);
-      },
-    } as any;
-
     const wrapper = shallow(
       <JaegerQueryField
         history={[]}
-        datasource={dsMock}
+        datasource={makeDatasourceMock({})}
         query={{ query: '1234' } as JaegerQuery}
         onRunQuery={() => {}}
         onChange={() => {}}
@@ -26,4 +17,81 @@ describe('JaegerQueryField', function() {
     );
     expect(wrapper.find(ButtonCascader).props().options[0].label).toBe('No traces found');
   });
+
+  it('shows root span as 3rd level in cascader', async function() {
+    const wrapper = mount(
+      <JaegerQueryField
+        history={[]}
+        datasource={makeDatasourceMock({
+          service1: {
+            op1: [
+              {
+                traceID: '12345',
+                spans: [
+                  {
+                    spanID: 's2',
+                    operationName: 'nonRootOp',
+                    references: [{ refType: 'CHILD_OF', traceID: '12345', spanID: 's1' }],
+                    duration: 10,
+                  },
+                  {
+                    operationName: 'rootOp',
+                    spanID: 's1',
+                    references: [],
+                    duration: 99,
+                  },
+                ],
+              },
+            ],
+          },
+        })}
+        query={{ query: '1234' } as JaegerQuery}
+        onRunQuery={() => {}}
+        onChange={() => {}}
+      />
+    );
+
+    // Simulating selection options. We need this as the function depends on the intermediate state of the component
+    await wrapper
+      .find(ButtonCascader)
+      .props()
+      .loadData([{ value: 'service1', label: 'service1' }]);
+
+    await wrapper
+      .find(ButtonCascader)
+      .props()
+      .loadData([
+        { value: 'service1', label: 'service1' },
+        { value: 'op1', label: 'op1' },
+      ]);
+
+    wrapper.update();
+    expect(wrapper.find(ButtonCascader).props().options[0].children[1].children[0]).toEqual({
+      label: 'rootOp [99 ms]',
+      value: '12345',
+    });
+  });
 });
+
+function makeDatasourceMock(data: { [service: string]: { [operation: string]: any } }): JaegerDatasource {
+  return {
+    metadataRequest(url: string, params: Record<string, any>) {
+      if (url.match(/\/services$/)) {
+        return Promise.resolve(Object.keys(data));
+      }
+      let match = url.match(/\/services\/(\w+)\/operations/);
+      if (match) {
+        return Promise.resolve(Object.keys(data[match[1]]));
+      }
+
+      if (url.match(/\/traces?/)) {
+        return Promise.resolve(data[params.service][params.operation]);
+      }
+      throw new Error(`Unexpected url: ${url}`);
+    },
+
+    getTimeRange(): { start: number; end: number } {
+      return { start: 1, end: 100 };
+    },
+  } as any;
+}
diff --git a/public/app/plugins/datasource/jaeger/QueryField.tsx b/public/app/plugins/datasource/jaeger/QueryField.tsx
index b505c01c14e..75afee9ce29 100644
--- a/public/app/plugins/datasource/jaeger/QueryField.tsx
+++ b/public/app/plugins/datasource/jaeger/QueryField.tsx
@@ -4,6 +4,7 @@ import { ButtonCascader, CascaderOption } from '@grafana/ui';
 
 import { AppEvents, ExploreQueryFieldProps } from '@grafana/data';
 import { appEvents } from '../../../core/core';
+import { Span, TraceData } from '@jaegertracing/jaeger-ui-components';
 
 const ALL_OPERATIONS_KEY = '__ALL__';
 const NO_TRACES_KEY = '__NO_TRACES__';
@@ -13,11 +14,14 @@ interface State {
   serviceOptions: CascaderOption[];
 }
 
-function getLabelFromTrace(trace: any): string {
-  // TODO: seems like the spans are not ordered so this may not be actually a root span
-  const firstSpan = trace.spans && trace.spans[0];
-  if (firstSpan) {
-    return `${firstSpan.operationName} [${firstSpan.duration} ms]`;
+function findRootSpan(spans: Span[]): Span | undefined {
+  return spans.find(s => !s.references?.length);
+}
+
+function getLabelFromTrace(trace: TraceData & { spans: Span[] }): string {
+  const rootSpan = findRootSpan(trace.spans);
+  if (rootSpan) {
+    return `${rootSpan.operationName} [${rootSpan.duration} ms]`;
   }
   return trace.traceID;
 }