UpdateQueries: Fixes issues setting datasource on queries after changing type (#41702)

public/app/core/utils/query.ts

@@ -1,5 +1,4 @@
-import { DataQuery, DataSourceInstanceSettings, DataSourceRef, getDataSourceRef } from '@grafana/data';
-import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
+import { DataQuery, DataSourceRef } from '@grafana/data';
 
 export const getNextRefIdChar = (queries: DataQuery[]): string => {
   for (let num = 0; ; num++) {
@@ -22,28 +21,6 @@ export function addQuery(queries: DataQuery[], query?: Partial<DataQuery>, datas
   return [...queries, q as DataQuery];
 }
 
-export function updateQueries(
-  newSettings: DataSourceInstanceSettings,
-  queries: DataQuery[],
-  dsSettings?: DataSourceInstanceSettings
-): DataQuery[] {
-  const datasource = getDataSourceRef(newSettings);
-
-  if (!newSettings.meta.mixed && dsSettings?.meta.mixed) {
-    return queries.map((q) => {
-      if (!isExpressionReference(q.datasource)) {
-        q.datasource = datasource;
-      }
-      return q;
-    });
-  } else if (!newSettings.meta.mixed && dsSettings?.meta.id !== newSettings.meta.id) {
-    // we are changing data source type, clear queries
-    return [{ refId: 'A', datasource }];
-  }
-
-  return queries;
-}
-
 export function isDataQuery(url: string): boolean {
   if (
     url.indexOf('api/datasources/proxy') !== -1 ||

public/app/features/query/components/QueryGroup.tsx

@@ -27,7 +27,7 @@ import {
   PanelData,
 } from '@grafana/data';
 import { PluginHelp } from 'app/core/components/PluginHelp/PluginHelp';
-import { addQuery, updateQueries } from 'app/core/utils/query';
+import { addQuery } from 'app/core/utils/query';
 import { Unsubscribable } from 'rxjs';
 import { dataSource as expressionDatasource } from 'app/features/expressions/ExpressionDatasource';
 import { selectors } from '@grafana/e2e-selectors';
@@ -37,6 +37,7 @@ import { DashboardQueryEditor, isSharedDashboardQuery } from 'app/plugins/dataso
 import { css } from '@emotion/css';
 import { QueryGroupOptions } from 'app/types';
 import { GroupActionComponents } from './QueryActionComponent';
+import { updateQueries } from '../state/updateQueries';
 
 interface Props {
   queryRunner: PanelQueryRunner;

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

@@ -0,0 +1,102 @@
+import { ExpressionDatasourceRef } from '@grafana/runtime/src/utils/DataSourceWithBackend';
+import { updateQueries } from './updateQueries';
+
+describe('updateQueries', () => {
+  it('Should update all queries except expression query when changing data source with same type', () => {
+    const updated = updateQueries(
+      {
+        uid: 'new-uid',
+        type: 'same-type',
+        meta: {},
+      } as any,
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'same-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: ExpressionDatasourceRef,
+        },
+      ],
+      {
+        uid: 'old-uid',
+        type: 'same-type',
+      } as any
+    );
+
+    expect(updated[0].datasource).toEqual({ type: 'same-type', uid: 'new-uid' });
+    expect(updated[1].datasource).toEqual(ExpressionDatasourceRef);
+  });
+
+  it('Should clear queries when changing type', () => {
+    const updated = updateQueries(
+      {
+        uid: 'new-uid',
+        type: 'new-type',
+        meta: {},
+      } as any,
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+      ],
+      {
+        uid: 'old-uid',
+        type: 'old-type',
+      } as any
+    );
+
+    expect(updated.length).toEqual(1);
+    expect(updated[0].datasource).toEqual({ type: 'new-type', uid: 'new-uid' });
+  });
+
+  it('Should preserve query data source when changing to mixed', () => {
+    const updated = updateQueries(
+      {
+        uid: 'mixed',
+        type: 'mixed',
+        meta: {
+          mixed: true,
+        },
+      } as any,
+      [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ],
+      {
+        uid: 'old-uid',
+        type: 'old-type',
+      } as any
+    );
+
+    expect(updated[0].datasource).toEqual({ type: 'old-type', uid: 'old-uid' });
+    expect(updated[1].datasource).toEqual({ type: 'other-type', uid: 'other-uid' });
+  });
+});

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

@@ -0,0 +1,29 @@
+import { DataQuery, DataSourceInstanceSettings, getDataSourceRef } from '@grafana/data';
+import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
+
+export function updateQueries(
+  newSettings: DataSourceInstanceSettings,
+  queries: DataQuery[],
+  dsSettings?: DataSourceInstanceSettings
+): DataQuery[] {
+  const datasource = getDataSourceRef(newSettings);
+
+  // we are changing data source type
+  if (dsSettings?.type !== newSettings.type) {
+    // If changing to mixed do nothing
+    if (newSettings.meta.mixed) {
+      return queries;
+    } else {
+      // Changing to another datasource type clear queries
+      return [{ refId: 'A', datasource }];
+    }
+  }
+
+  // Set data source on all queries except expression queries
+  return queries.map((query) => {
+    if (!isExpressionReference(query.datasource)) {
+      query.datasource = datasource;
+    }
+    return query;
+  });
+}