From 65fa94b16b6ba4735fb77717975093e075e17008 Mon Sep 17 00:00:00 2001
From: harshabaddam <baddamharsha33@gmail.com>
Date: Wed, 4 Oct 2023 13:59:46 +0100
Subject: [PATCH] Loki: Support X-ray as internal link in derived fields
 (#75756)

Update Loki datasource to support X-ray as internal link in derived fields
---
 .../datasource/loki/getDerivedFields.test.ts  | 50 ++++++++++++++++++-
 .../datasource/loki/getDerivedFields.ts       | 12 ++++-
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/public/app/plugins/datasource/loki/getDerivedFields.test.ts b/public/app/plugins/datasource/loki/getDerivedFields.test.ts
index 554975630e4..78e14d002e7 100644
--- a/public/app/plugins/datasource/loki/getDerivedFields.test.ts
+++ b/public/app/plugins/datasource/loki/getDerivedFields.test.ts
@@ -7,7 +7,13 @@ jest.mock('@grafana/runtime', () => ({
   ...jest.requireActual('@grafana/runtime'),
   getDataSourceSrv: () => {
     return {
-      getInstanceSettings: () => {
+      getInstanceSettings: (datasourceUid?: string) => {
+        switch (datasourceUid) {
+          case 'tempo-datasource-uid':
+            return { name: 'Tempo', type: 'tempo' };
+          case 'xray-datasource-uid':
+            return { name: 'X-ray', type: 'grafana-x-ray-datasource' };
+        }
         return { name: 'Loki1' };
       },
     };
@@ -36,8 +42,22 @@ describe('getDerivedFields', () => {
         datasourceUid: 'uid2',
         urlDisplayLabel: 'Custom Label',
       },
+      {
+        matcherRegex: 'trace=(\\w+)',
+        name: 'tempoTraceId',
+        url: 'test',
+        datasourceUid: 'tempo-datasource-uid',
+        urlDisplayLabel: 'Tempo',
+      },
+      {
+        matcherRegex: 'trace=(\\w+)',
+        name: 'xrayTraceId',
+        url: 'test',
+        datasourceUid: 'xray-datasource-uid',
+        urlDisplayLabel: 'AWS X-ray',
+      },
     ]);
-    expect(newFields.length).toBe(2);
+    expect(newFields.length).toBe(4);
     const trace1 = newFields.find((f) => f.name === 'trace1');
     expect(trace1!.values).toEqual([null, '1234', null]);
     expect(trace1!.config.links![0]).toEqual({
@@ -58,5 +78,31 @@ describe('getDerivedFields', () => {
       internal: { datasourceName: 'Loki1', datasourceUid: 'uid2', query: { query: 'test' } },
       url: '',
     });
+
+    const tempoTraceId = newFields.find((f) => f.name === 'tempoTraceId');
+    expect(tempoTraceId!.values).toEqual([null, null, null]);
+    expect(tempoTraceId!.config.links!.length).toBe(1);
+    expect(tempoTraceId!.config.links![0]).toEqual({
+      title: 'Tempo',
+      internal: {
+        datasourceName: 'Tempo',
+        datasourceUid: 'tempo-datasource-uid',
+        query: { query: 'test', queryType: 'traceql' },
+      },
+      url: '',
+    });
+
+    const xrayTraceId = newFields.find((f) => f.name === 'xrayTraceId');
+    expect(xrayTraceId!.values).toEqual([null, null, null]);
+    expect(xrayTraceId!.config.links!.length).toBe(1);
+    expect(xrayTraceId!.config.links![0]).toEqual({
+      title: 'AWS X-ray',
+      internal: {
+        datasourceName: 'X-ray',
+        datasourceUid: 'xray-datasource-uid',
+        query: { query: 'test', queryType: 'getTrace' },
+      },
+      url: '',
+    });
   });
 });
diff --git a/public/app/plugins/datasource/loki/getDerivedFields.ts b/public/app/plugins/datasource/loki/getDerivedFields.ts
index bb17b33cd0f..e2434c23fd8 100644
--- a/public/app/plugins/datasource/loki/getDerivedFields.ts
+++ b/public/app/plugins/datasource/loki/getDerivedFields.ts
@@ -42,6 +42,16 @@ function fieldFromDerivedFieldConfig(derivedFieldConfigs: DerivedFieldConfig[]):
     // Having field.datasourceUid means it is an internal link.
     if (derivedFieldConfig.datasourceUid) {
       const dsSettings = dataSourceSrv.getInstanceSettings(derivedFieldConfig.datasourceUid);
+      const queryType = (type: string | undefined): string | undefined => {
+        switch (type) {
+          case 'tempo':
+            return 'traceql';
+          case 'grafana-x-ray-datasource':
+            return 'getTrace';
+          default:
+            return undefined;
+        }
+      };
 
       acc.push({
         // Will be filled out later
@@ -49,7 +59,7 @@ function fieldFromDerivedFieldConfig(derivedFieldConfigs: DerivedFieldConfig[]):
         url: '',
         // This is hardcoded for Jaeger or Zipkin not way right now to specify datasource specific query object
         internal: {
-          query: { query: derivedFieldConfig.url, queryType: dsSettings?.type === 'tempo' ? 'traceql' : undefined },
+          query: { query: derivedFieldConfig.url, queryType: queryType(dsSettings?.type) },
           datasourceUid: derivedFieldConfig.datasourceUid,
           datasourceName: dsSettings?.name ?? 'Data source not found',
         },