Jaeger: Run traceID queries through backend when node graph is enabled (#104029)

* Refactor current solution

* Run traceid queries with enabled node graph through backend

* Update test

* Fix lint

* Uodate based on feedback, don't spread

public/app/plugins/datasource/jaeger/datasource.test.ts

@@ -9,7 +9,7 @@ import {
   PluginType,
   ScopedVars,
 } from '@grafana/data';
-import { BackendSrv } from '@grafana/runtime';
+import { BackendSrv, config, DataSourceWithBackend } from '@grafana/runtime';
 
 import { ALL_OPERATIONS_KEY } from './components/SearchForm';
 import { JaegerDatasource, JaegerJsonData } from './datasource';
@@ -307,6 +307,102 @@ describe('JaegerDatasource', () => {
       url: `${defaultSettings.url}/api/traces?service=interpolationText&operation=interpolationText&minDuration=interpolationText&maxDuration=interpolationText&${defaultSearchRangeParams}&lookback=custom`,
     });
   });
+
+  describe('when jaegerBackendMigration feature toggle is enabled', () => {
+    let originalFeatureToggleValue: boolean | undefined;
+
+    beforeEach(() => {
+      originalFeatureToggleValue = config.featureToggles.jaegerBackendMigration;
+      config.featureToggles.jaegerBackendMigration = true;
+    });
+
+    afterEach(() => {
+      config.featureToggles.jaegerBackendMigration = originalFeatureToggleValue;
+    });
+
+    it('should add node graph frames to response when nodeGraph is enabled and query is a trace ID query', async () => {
+      // Create a datasource with nodeGraph enabled
+      const settings = {
+        ...defaultSettings,
+        jsonData: {
+          ...defaultSettings.jsonData,
+          nodeGraph: { enabled: true },
+        },
+      };
+
+      const ds = new JaegerDatasource(settings);
+
+      // Mock the super.query method to return our mock response
+      jest.spyOn(DataSourceWithBackend.prototype, 'query').mockImplementation(() => {
+        return of({
+          data: [
+            {
+              fields: testResponseDataFrameFields,
+              values: testResponseDataFrameFields.values,
+            },
+          ],
+        });
+      });
+
+      // Create a query without queryType (trace ID query)
+      const query = {
+        ...defaultQuery,
+        targets: [
+          {
+            query: '12345',
+            refId: '1',
+          },
+        ],
+      };
+
+      // Execute the query
+      const response = await lastValueFrom(ds.query(query));
+      // Verify that the response contains the original data plus node graph frames
+      expect(response.data.length).toBe(3);
+    });
+
+    it('should not add node graph frames when nodeGraph is disabled', async () => {
+      // Create a datasource with nodeGraph disabled
+      const settings = {
+        ...defaultSettings,
+        jsonData: {
+          ...defaultSettings.jsonData,
+          nodeGraph: { enabled: false },
+        },
+      };
+
+      const ds = new JaegerDatasource(settings);
+
+      // Mock the super.query method to return our mock response
+      jest.spyOn(DataSourceWithBackend.prototype, 'query').mockImplementation(() => {
+        return of({
+          data: [
+            {
+              fields: testResponseDataFrameFields,
+              values: testResponseDataFrameFields.values,
+            },
+          ],
+        });
+      });
+
+      // Create a query without queryType (trace ID query)
+      const query = {
+        ...defaultQuery,
+        targets: [
+          {
+            query: '12345',
+            refId: '1',
+          },
+        ],
+      };
+
+      // Execute the query
+      const response = await lastValueFrom(ds.query(query));
+      // Verify that the response contains only the original data
+      expect(response.data.length).toBe(1);
+      expect(response.data[0].fields).toMatchObject(testResponseDataFrameFields);
+    });
+  });
 });
 
 describe('when performing testDataSource', () => {

public/app/plugins/datasource/jaeger/datasource.ts

@@ -13,9 +13,10 @@ import {
   getDefaultTimeRange,
   MutableDataFrame,
   ScopedVars,
+  toDataFrame,
   urlUtil,
 } from '@grafana/data';
-import { NodeGraphOptions, SpanBarOptions } from '@grafana/o11y-ds-frontend';
+import { createNodeGraphFrames, NodeGraphOptions, SpanBarOptions } from '@grafana/o11y-ds-frontend';
 import {
   BackendSrvRequest,
   config,
@@ -69,29 +70,29 @@ export class JaegerDatasource extends DataSourceWithBackend<JaegerQuery, JaegerJ
   }
 
   query(options: DataQueryRequest<JaegerQuery>): Observable<DataQueryResponse> {
-    // No query type means that the query is a trace ID query
-    // If all targets are trace ID queries, we can use the backend querying
-    const allTargetsTraceIdQuery = options.targets.every((target) => !target.queryType);
-    const allTargetsDependencyGraph = options.targets.every((target) => target.queryType === 'dependencyGraph');
-    // We have not migrated the node graph to the backend
-    // If the node graph is disabled, we can use the backend migration
-    const nodeGraphDisabled = !this.nodeGraph?.enabled;
-    if (
-      config.featureToggles.jaegerBackendMigration &&
-      (allTargetsTraceIdQuery || allTargetsDependencyGraph) &&
-      nodeGraphDisabled
-    ) {
-      return super.query(options);
-    }
-
     // At this moment we expect only one target. In case we somehow change the UI to be able to show multiple
     // traces at one we need to change this.
     const target: JaegerQuery = options.targets[0];
-
     if (!target) {
       return of({ data: [emptyTraceDataFrame] });
     }
 
+    if (
+      config.featureToggles.jaegerBackendMigration &&
+      // No query type means that the query is a trace ID query
+      (!target.queryType || target.queryType === 'dependencyGraph')
+    ) {
+      return super.query({ ...options, targets: [target] }).pipe(
+        map((response) => {
+          // If the node graph is enabled and the query is a trace ID query, add the node graph frames to the response
+          if (this.nodeGraph?.enabled && !target.queryType) {
+            return addNodeGraphFramesToResponse(response);
+          }
+          return response;
+        })
+      );
+    }
+
     // Use the internal Jaeger /dependencies API for rendering the dependency graph.
     if (target.queryType === 'dependencyGraph') {
       const timeRange = options.range ?? getDefaultTimeRange();
@@ -309,3 +310,18 @@ const emptyTraceDataFrame = new MutableDataFrame({
     },
   },
 });
+
+export function addNodeGraphFramesToResponse(response: DataQueryResponse): DataQueryResponse {
+  if (!response.data || response.data.length === 0) {
+    return response;
+  }
+
+  // Convert the first frame to a DataFrame for node graph processing
+  const frame = toDataFrame(response.data[0]);
+  // Add the node graph frames to the response
+  const data = response.data.concat(createNodeGraphFrames(frame));
+  return {
+    ...response,
+    data,
+  };
+}