MixedDataSource: Support multi value data source variable that issues a query to each data source (#87586)

* Revert "Revert "MixedDataSource: Support multi value data source variable tha…"

This reverts commit 979d87f46f.

* fix issue

public/app/plugins/datasource/mixed/MixedDataSource.test.ts

@@ -2,7 +2,11 @@ import { lastValueFrom } from 'rxjs';
 import { getQueryOptions } from 'test/helpers/getQueryOptions';
 import { DatasourceSrvMock, MockObservableDataSourceApi } from 'test/mocks/datasource_srv';
 
-import { DataQueryRequest, DataSourceInstanceSettings, LoadingState } from '@grafana/data';
+import { DataQueryRequest, DataSourceInstanceSettings, DataSourceRef, LoadingState } from '@grafana/data';
+import { DataSourceSrv, setDataSourceSrv, setTemplateSrv } from '@grafana/runtime';
+import { CustomVariable, SceneFlexLayout, SceneVariableSet } from '@grafana/scenes';
+
+import { TemplateSrv } from '../../../features/templating/template_srv';
 
 import { MIXED_DATASOURCE_NAME } from './MixedDataSource';
 import { MixedDatasource } from './module';
@@ -21,13 +25,19 @@ const datasourceSrv = new DatasourceSrvMock(defaultDS, {
   ]),
 });
 
-const getDataSourceSrvMock = jest.fn().mockReturnValue(datasourceSrv);
-jest.mock('@grafana/runtime', () => ({
-  ...jest.requireActual('@grafana/runtime'),
-  getDataSourceSrv: () => getDataSourceSrvMock(),
-}));
-
 describe('MixedDatasource', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    setDataSourceSrv({
+      ...datasourceSrv,
+      get: (uid: DataSourceRef) => datasourceSrv.get(uid),
+      getInstanceSettings: jest.fn().mockReturnValue({ meta: {} }),
+      getList: jest.fn(),
+      reload: jest.fn(),
+    });
+    setTemplateSrv(new TemplateSrv());
+  });
+
   describe('with no errors', () => {
     it('direct query should return results', async () => {
       const ds = new MixedDatasource({} as DataSourceInstanceSettings);
@@ -83,6 +93,100 @@ describe('MixedDatasource', () => {
     });
   });
 
+  describe('with multi template variable', () => {
+    beforeAll(() => {
+      setDataSourceSrv({
+        getInstanceSettings() {
+          return {};
+        },
+      } as DataSourceSrv);
+    });
+
+    const scene = new SceneFlexLayout({
+      children: [],
+      $variables: new SceneVariableSet({
+        variables: [new CustomVariable({ name: 'ds', value: ['B', 'C'] })],
+      }),
+    });
+
+    it('should run query for each datasource when there is a multi value template variable', async () => {
+      const ds = new MixedDatasource({} as DataSourceInstanceSettings);
+
+      const request = {
+        targets: [{ refId: 'AA', datasource: { uid: '$ds' } }],
+        scopedVars: {
+          __sceneObject: { value: scene },
+        },
+      } as unknown as DataQueryRequest;
+
+      await expect(ds.query(request)).toEmitValuesWith((results) => {
+        expect(results).toHaveLength(2);
+        expect(results[0].key).toBe('mixed-0-');
+        expect(results[0].state).toBe(LoadingState.Loading);
+        expect(results[1].key).toBe('mixed-1-');
+        expect(results[1].state).toBe(LoadingState.Done);
+      });
+    });
+
+    it('should run query for picked datasource and template variable datasource', async () => {
+      const ds = new MixedDatasource({} as DataSourceInstanceSettings);
+      const request = {
+        targets: [
+          { refId: 'AA', datasource: { uid: '$ds' } },
+          { refId: 'BB', datasource: { uid: 'Loki' } },
+        ],
+        scopedVars: {
+          __sceneObject: { value: scene },
+        },
+      } as unknown as DataQueryRequest;
+
+      await expect(ds.query(request)).toEmitValuesWith((results) => {
+        expect(results).toHaveLength(4);
+        expect(results[0].key).toBe('mixed-0-');
+        expect(results[0].state).toBe(LoadingState.Loading);
+        expect(results[1].key).toBe('mixed-1-');
+        expect(results[1].state).toBe(LoadingState.Loading);
+        expect(results[2].key).toBe('mixed-2-A');
+        expect(results[2].state).toBe(LoadingState.Loading);
+        expect(results[3].key).toBe('mixed-2-B');
+        expect(results[3].state).toBe(LoadingState.Done);
+      });
+    });
+  });
+
+  describe('with single value template variable', () => {
+    beforeAll(() => {
+      setDataSourceSrv({
+        getInstanceSettings() {
+          return {};
+        },
+      } as DataSourceSrv);
+    });
+
+    const scene = new SceneFlexLayout({
+      children: [],
+      $variables: new SceneVariableSet({
+        variables: [new CustomVariable({ name: 'ds', value: 'B' })],
+      }),
+    });
+
+    it('should run query for correct datasource', async () => {
+      const ds = new MixedDatasource({} as DataSourceInstanceSettings);
+
+      const request = {
+        targets: [{ refId: 'AA', datasource: { uid: '$ds' } }],
+        scopedVars: {
+          __sceneObject: { value: scene },
+        },
+      } as unknown as DataQueryRequest;
+
+      await expect(ds.query(request)).toEmitValuesWith((results) => {
+        expect(results).toHaveLength(1);
+        expect(results[0].data).toEqual(['BBBB']);
+      });
+    });
+  });
+
   it('should return both query results from the same data source', async () => {
     const ds = new MixedDatasource({} as DataSourceInstanceSettings);
     const request = {

public/app/plugins/datasource/mixed/MixedDataSource.ts

@@ -10,8 +10,10 @@ import {
   DataSourceApi,
   DataSourceInstanceSettings,
   LoadingState,
+  ScopedVars,
 } from '@grafana/data';
-import { getDataSourceSrv, toDataQueryError } from '@grafana/runtime';
+import { getDataSourceSrv, getTemplateSrv, toDataQueryError } from '@grafana/runtime';
+import { CustomFormatterVariable } from '@grafana/scenes';
 
 export const MIXED_DATASOURCE_NAME = '-- Mixed --';
 
@@ -19,7 +21,8 @@ export const mixedRequestId = (queryIdx: number, requestId?: string) => `mixed-$
 
 export interface BatchedQueries {
   datasource: Promise<DataSourceApi>;
-  targets: DataQuery[];
+  queries: DataQuery[];
+  scopedVars: ScopedVars;
 }
 
 export class MixedDatasource extends DataSourceApi<DataQuery> {
@@ -39,23 +42,71 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
 
     // Build groups of queries to run in parallel
     const sets: { [key: string]: DataQuery[] } = groupBy(queries, 'datasource.uid');
-    const mixed: BatchedQueries[] = [];
+    const batches: BatchedQueries[] = [];
 
     for (const key in sets) {
-      const targets = sets[key];
-
-      mixed.push({
-        datasource: getDataSourceSrv().get(targets[0].datasource, request.scopedVars),
-        targets,
-      });
+      batches.push(...this.getBatchesForQueries(sets[key], request));
     }
 
     // Missing UIDs?
-    if (!mixed.length) {
+    if (!batches.length) {
       return of({ data: [] }); // nothing
     }
 
-    return this.batchQueries(mixed, request);
+    return this.batchQueries(batches, request);
+  }
+
+  /**
+   * Almost always returns a single batch for each set of queries.
+   * Unless the query is using a multi value variable.
+   */
+  private getBatchesForQueries(queries: DataQuery[], request: DataQueryRequest<DataQuery>) {
+    const dsRef = queries[0].datasource;
+    const batches: BatchedQueries[] = [];
+
+    // Using the templateSrv.replace function here with a custom formatter as that is the cleanest way
+    // to access the raw value or value array of a variable.
+    const datasourceUid = getTemplateSrv().replace(
+      dsRef?.uid,
+      request.scopedVars,
+      (value: string | string[], variable: CustomFormatterVariable) => {
+        // If it's not a data source variable, or single value
+        if (!Array.isArray(value)) {
+          return value;
+        }
+
+        for (const uid of value) {
+          if (uid === 'default') {
+            continue;
+          }
+
+          const dsSettings = getDataSourceSrv().getInstanceSettings(uid);
+
+          batches.push({
+            datasource: getDataSourceSrv().get(uid),
+            queries: cloneDeep(queries),
+            scopedVars: {
+              ...request.scopedVars,
+              [variable.name]: { value: uid, text: dsSettings?.name },
+            },
+          });
+        }
+
+        return '';
+      }
+    );
+
+    if (datasourceUid !== '') {
+      batches.push({
+        datasource: getDataSourceSrv().get(datasourceUid),
+        queries: cloneDeep(queries),
+        scopedVars: {
+          ...request.scopedVars,
+        },
+      });
+    }
+
+    return batches;
   }
 
   batchQueries(mixed: BatchedQueries[], request: DataQueryRequest<DataQuery>): Observable<DataQueryResponse> {
@@ -64,7 +115,8 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
         mergeMap((api: DataSourceApi) => {
           const dsRequest = cloneDeep(request);
           dsRequest.requestId = mixedRequestId(i, dsRequest.requestId);
-          dsRequest.targets = query.targets;
+          dsRequest.targets = query.queries;
+          dsRequest.scopedVars = query.scopedVars;
 
           return from(api.query(dsRequest)).pipe(
             map((response) => {
@@ -102,7 +154,7 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
   }
 
   private isQueryable(query: BatchedQueries): boolean {
-    return query && Array.isArray(query.targets) && query.targets.length > 0;
+    return query && Array.isArray(query.queries) && query.queries.length > 0;
   }
 
   private finalizeResponses(responses: DataQueryResponse[]): DataQueryResponse[] {