QueryBuilder: Preserve queries when switching from Mixed (#71224)

* When switching from mixed to another datasource preserve relevant queries

- Update query refs
- Update tests

* Ensure templateVars are handled correctly

- Add test

* Comment

* Unused var

* For ds template variable only preserve if going from&to template

- Add test

public/app/features/query/state/updateQueries.test.ts

@@ -6,6 +6,7 @@ import {
 } from '@grafana/data';
 import { ExpressionDatasourceRef } from '@grafana/runtime/src/utils/DataSourceWithBackend';
 import { TestQuery } from 'app/core/utils/query.test';
+import { TemplateSrv } from 'app/features/templating/template_srv';
 
 import { updateQueries } from './updateQueries';
 
@@ -41,6 +42,13 @@ const newUidSameTypeDS = {
   },
 } as DataSourceApi;
 
+const templateSrv = new TemplateSrv();
+
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  getTemplateSrv: () => templateSrv,
+}));
+
 describe('updateQueries', () => {
   afterEach(() => {
     jest.clearAllMocks();
@@ -215,6 +223,144 @@ describe('updateQueries', () => {
     expect(updated[0].datasource).toEqual({ type: 'old-type', uid: 'old-uid' });
     expect(updated[1].datasource).toEqual({ type: 'other-type', uid: 'other-uid' });
   });
+
+  it('should preserve query when switching from mixed to a datasource where a query exists for the new datasource', async () => {
+    const updated = await updateQueries(
+      newUidDS,
+      'new-uid',
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'new-uid',
+            type: 'new-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ],
+      mixedDS
+    );
+
+    expect(updated[0].datasource).toEqual({ type: 'new-type', uid: 'new-uid' });
+    expect(updated.length).toEqual(1);
+  });
+
+  it('should preserve query when switching from mixed to a datasource where a query exists for the new datasource - when using datasource template variable', async () => {
+    templateSrv.init([
+      {
+        current: {
+          text: 'Azure Monitor',
+          value: 'ds-uid',
+        },
+        name: 'ds',
+        type: 'datasource',
+        id: 'ds',
+      },
+    ]);
+    const updated = await updateQueries(
+      newUidDS,
+      '$ds',
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: '$ds',
+            type: 'new-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ],
+      mixedDS
+    );
+
+    expect(updated[0].datasource).toEqual({ type: 'new-type', uid: '$ds' });
+    expect(updated.length).toEqual(1);
+  });
+
+  it('will not preserve query when switch from mixed with a ds variable query to the same datasource (non-variable)', async () => {
+    templateSrv.init([
+      {
+        current: {
+          text: 'Azure Monitor',
+          value: 'ds-uid',
+        },
+        name: 'ds',
+        type: 'datasource',
+        id: 'ds',
+      },
+    ]);
+    const updated = await updateQueries(
+      newUidDS,
+      'new-uid',
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: '$ds',
+            type: 'new-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ],
+      mixedDS
+    );
+
+    expect(updated[0].datasource).toEqual({ type: 'new-type', uid: 'new-uid' });
+    expect(updated.length).toEqual(1);
+  });
+
+  it('should update query refs when switching from mixed to a datasource where queries exist for new datasource', async () => {
+    const updated = await updateQueries(
+      newUidDS,
+      'new-uid',
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'new-uid',
+            type: 'new-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+        {
+          refId: 'C',
+          datasource: {
+            uid: 'new-uid',
+            type: 'new-type',
+          },
+        },
+      ],
+      mixedDS
+    );
+
+    expect(updated.length).toEqual(2);
+    expect(updated[0].refId).toEqual('A');
+    expect(updated[1].refId).toEqual('B');
+  });
 });
 
 describe('updateQueries with import', () => {

public/app/features/query/state/updateQueries.ts

@@ -1,5 +1,8 @@
-import { CoreApp, DataQuery, DataSourceApi, hasQueryExportSupport, hasQueryImportSupport } from '@grafana/data';
+import { CoreApp, DataSourceApi, hasQueryExportSupport, hasQueryImportSupport } from '@grafana/data';
+import { getTemplateSrv } from '@grafana/runtime';
 import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
+import { DataQuery } from '@grafana/schema';
+import { getNextRefIdChar } from 'app/core/utils/query';
 
 export async function updateQueries(
   nextDS: DataSourceApi,
@@ -26,8 +29,38 @@ export async function updateQueries(
     else if (currentDS && nextDS.importQueries) {
       nextQueries = await nextDS.importQueries(queries, currentDS);
     }
-    // Otherwise clear queries
+    // Otherwise clear queries that do not match the next datasource UID
     else {
+      if (currentDS) {
+        const templateSrv = getTemplateSrv();
+        const reducedQueries: DataQuery[] = [];
+        let nextUid = nextDS.uid;
+        const nextIsTemplate = templateSrv.containsTemplate(nextDSUidOrVariableExpression);
+        if (nextIsTemplate) {
+          nextUid = templateSrv.replace(nextDS.uid);
+        }
+        // Queries will only be preserved if the datasource UID of the query matches the UID
+        // of the next chosen datasource
+        const nextDsQueries = queries.reduce((reduced, currentQuery) => {
+          if (currentQuery.datasource) {
+            let currUid = currentQuery.datasource.uid;
+            const currIsTemplate = templateSrv.containsTemplate(currUid);
+            if (currIsTemplate) {
+              currUid = templateSrv.replace(currentQuery.datasource.uid);
+            }
+            if (currUid === nextUid && currIsTemplate === nextIsTemplate) {
+              currentQuery.refId = getNextRefIdChar(reduced);
+              return reduced.concat([currentQuery]);
+            }
+          }
+          return reduced;
+        }, reducedQueries);
+
+        if (nextDsQueries.length > 0) {
+          return nextDsQueries;
+        }
+      }
+
       return [DEFAULT_QUERY];
     }
   }