From eaba50283ef25eb447dfca2e0a47c3dce5feb104 Mon Sep 17 00:00:00 2001
From: Ryan McKinley <ryantxu@gmail.com>
Date: Wed, 6 Nov 2019 22:18:27 -0800
Subject: [PATCH] MixedQuery: refactor so other components could also batch
 queries (#20219)

* query mixed

* query mixed

* batch queries

* check that the values only have 3 rows at the *end* as well
---
 .../datasource/mixed/MixedDataSource.test.ts  |  1 +
 .../datasource/mixed/MixedDataSource.ts       | 38 +++++++++++++------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/public/app/plugins/datasource/mixed/MixedDataSource.test.ts b/public/app/plugins/datasource/mixed/MixedDataSource.test.ts
index fb74bffb030..4c18ce82f6e 100644
--- a/public/app/plugins/datasource/mixed/MixedDataSource.test.ts
+++ b/public/app/plugins/datasource/mixed/MixedDataSource.test.ts
@@ -47,5 +47,6 @@ describe('MixedDatasource', () => {
     expect(results[1].data).toEqual(['BBBB']);
     expect(results[2].data).toEqual(['CCCC']);
     expect(results[2].state).toEqual(LoadingState.Done);
+    expect(results.length).toBe(3);
   });
 });
diff --git a/public/app/plugins/datasource/mixed/MixedDataSource.ts b/public/app/plugins/datasource/mixed/MixedDataSource.ts
index ff96ab20a7b..a95bfad833b 100644
--- a/public/app/plugins/datasource/mixed/MixedDataSource.ts
+++ b/public/app/plugins/datasource/mixed/MixedDataSource.ts
@@ -16,6 +16,11 @@ import { mergeMap, map } from 'rxjs/operators';
 
 export const MIXED_DATASOURCE_NAME = '-- Mixed --';
 
+export interface BatchedQueries {
+  datasource: Promise<DataSourceApi>;
+  targets: DataQuery[];
+}
+
 export class MixedDatasource extends DataSourceApi<DataQuery> {
   constructor(instanceSettings: DataSourceInstanceSettings) {
     super(instanceSettings);
@@ -31,24 +36,35 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
       return of({ data: [] } as DataQueryResponse); // nothing
     }
 
+    // Build groups of queries to run in parallel
     const sets: { [key: string]: DataQuery[] } = groupBy(queries, 'datasource');
-    const observables: Array<Observable<DataQueryResponse>> = [];
-    let runningSubRequests = 0;
-
+    const mixed: BatchedQueries[] = [];
     for (const key in sets) {
       const targets = sets[key];
       const dsName = targets[0].datasource;
+      mixed.push({
+        datasource: getDataSourceSrv().get(dsName),
+        targets,
+      });
+    }
+    return this.batchQueries(mixed, request);
+  }
 
-      const observable = from(getDataSourceSrv().get(dsName)).pipe(
+  batchQueries(mixed: BatchedQueries[], request: DataQueryRequest<DataQuery>): Observable<DataQueryResponse> {
+    const observables: Array<Observable<DataQueryResponse>> = [];
+    let runningSubRequests = 0;
+
+    for (let i = 0; i < mixed.length; i++) {
+      const query = mixed[i];
+      if (!query.targets || !query.targets.length) {
+        continue;
+      }
+      const observable = from(query.datasource).pipe(
         mergeMap((dataSourceApi: DataSourceApi) => {
           const datasourceRequest = cloneDeep(request);
 
-          datasourceRequest.requestId = `${dsName}${datasourceRequest.requestId || ''}`;
-
-          // all queries hidden return empty result for for this requestId
-          if (datasourceRequest.targets.length === 0) {
-            return of({ data: [], key: datasourceRequest.requestId });
-          }
+          datasourceRequest.requestId = `mixed-${i}-${datasourceRequest.requestId || ''}`;
+          datasourceRequest.targets = query.targets;
 
           runningSubRequests++;
           let hasCountedAsDone = false;
@@ -79,7 +95,7 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
                 ...response,
                 data: response.data || [],
                 state: runningSubRequests === 0 ? LoadingState.Done : LoadingState.Loading,
-                key: `${dsName}${response.key || ''}`,
+                key: `mixed-${i}-${response.key || ''}`,
               } as DataQueryResponse;
             })
           );