Allow queries import when changing data source type (#47435)

* Enable queries import when changing datasource

* Supporting empty imports

* Review applied

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

@@ -111,7 +111,10 @@ export class QueryGroup extends PureComponent<Props, State> {
 
   onChangeDataSource = async (newSettings: DataSourceInstanceSettings) => {
     const { dsSettings } = this.state;
-    const queries = updateQueries(newSettings, this.state.queries, dsSettings);
+    const currentDS = dsSettings ? await getDataSourceSrv().get(dsSettings.uid) : undefined;
+    const nextDS = await getDataSourceSrv().get(newSettings.uid);
+
+    const queries = await updateQueries(nextDS, this.state.queries, currentDS);
 
     const dataSource = await this.dataSourceSrv.get(newSettings.name);
     this.onChange({

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

@@ -1,20 +1,54 @@
+import {
+  DataQuery,
+  DataSourceApi,
+  DataSourceWithQueryExportSupport,
+  DataSourceWithQueryImportSupport,
+} from '@grafana/data';
 import { ExpressionDatasourceRef } from '@grafana/runtime/src/utils/DataSourceWithBackend';
 import { updateQueries } from './updateQueries';
 
+const oldUidDS = {
+  uid: 'old-uid',
+  type: 'old-type',
+  meta: {
+    id: 'old-type',
+  },
+} as DataSourceApi;
+
+const mixedDS = {
+  uid: 'mixed',
+  meta: {
+    id: 'mixed',
+    mixed: true,
+  },
+} as DataSourceApi;
+
+const newUidDS = {
+  uid: 'new-uid',
+  type: 'new-type',
+  meta: {
+    id: 'new-type',
+  },
+} as DataSourceApi;
+
+const newUidSameTypeDS = {
+  uid: 'new-uid-same-type',
+  type: 'old-type',
+  meta: {
+    id: 'old-type',
+  },
+} as DataSourceApi;
+
 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,
+  it('Should update all queries except expression query when changing data source with same type', async () => {
+    const updated = await updateQueries(
+      newUidSameTypeDS,
       [
         {
           refId: 'A',
           datasource: {
             uid: 'old-uid',
-            type: 'same-type',
+            type: 'old-type',
           },
         },
         {
@@ -22,23 +56,16 @@ describe('updateQueries', () => {
           datasource: ExpressionDatasourceRef,
         },
       ],
-      {
-        uid: 'old-uid',
-        type: 'same-type',
-      } as any
+      oldUidDS
     );
 
-    expect(updated[0].datasource).toEqual({ type: 'same-type', uid: 'new-uid' });
+    expect(updated[0].datasource).toEqual({ type: 'old-type', uid: 'new-uid-same-type' });
     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,
+  it('Should clear queries when changing type', async () => {
+    const updated = await updateQueries(
+      newUidDS,
       [
         {
           refId: 'A',
@@ -55,25 +82,16 @@ describe('updateQueries', () => {
           },
         },
       ],
-      {
-        uid: 'old-uid',
-        type: 'old-type',
-      } as any
+      oldUidDS
     );
 
     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,
+  it('Should preserve query data source when changing to mixed', async () => {
+    const updated = await updateQueries(
+      mixedDS,
       [
         {
           refId: 'A',
@@ -90,25 +108,16 @@ describe('updateQueries', () => {
           },
         },
       ],
-      {
-        uid: 'old-uid',
-        type: 'old-type',
-      } as any
+      oldUidDS
     );
 
     expect(updated[0].datasource).toEqual({ type: 'old-type', uid: 'old-uid' });
     expect(updated[1].datasource).toEqual({ type: 'other-type', uid: 'other-uid' });
   });
 
-  it('should change nothing mixed updated to mixed', () => {
-    const updated = updateQueries(
-      {
-        uid: 'mixed',
-        type: 'mixed',
-        meta: {
-          mixed: true,
-        },
-      } as any,
+  it('should change nothing mixed updated to mixed', async () => {
+    const updated = await updateQueries(
+      mixedDS,
       [
         {
           refId: 'A',
@@ -125,16 +134,249 @@ describe('updateQueries', () => {
           },
         },
       ],
-      {
-        uid: 'mixed',
-        type: 'mixed',
-        meta: {
-          mixed: true,
-        },
-      } as any
+      mixedDS
     );
 
     expect(updated[0].datasource).toEqual({ type: 'old-type', uid: 'old-uid' });
     expect(updated[1].datasource).toEqual({ type: 'other-type', uid: 'other-uid' });
   });
 });
+
+describe('updateQueries with import', () => {
+  describe('abstract queries support', () => {
+    it('should migrate abstract queries', async () => {
+      const exportSpy = jest.fn();
+      const importSpy = jest.fn();
+
+      const newUidDSWithAbstract = {
+        uid: 'new-uid',
+        type: 'new-type',
+        meta: {
+          id: 'new-type',
+        },
+        importFromAbstractQueries: (queries) => {
+          importSpy(queries);
+          const importedQueries = queries.map((q) => ({ ...q, imported: true }));
+          return Promise.resolve(importedQueries);
+        },
+      } as DataSourceWithQueryImportSupport<any>;
+
+      const oldUidDSWithAbstract = {
+        uid: 'old-uid',
+        type: 'old-type',
+        meta: {
+          id: 'old-type',
+        },
+        exportToAbstractQueries: (queries) => {
+          exportSpy(queries);
+          const exportedQueries = queries.map((q) => ({ ...q, exported: true }));
+          return Promise.resolve(exportedQueries);
+        },
+      } as DataSourceWithQueryExportSupport<any>;
+
+      const queries = [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ];
+
+      const updated = await updateQueries(newUidDSWithAbstract as any, queries, oldUidDSWithAbstract as any);
+
+      expect(exportSpy).toBeCalledWith(queries);
+      expect(importSpy).toBeCalledWith(queries.map((q) => ({ ...q, exported: true })));
+
+      expect(updated).toMatchInlineSnapshot(`
+        Array [
+          Object {
+            "datasource": Object {
+              "type": "new-type",
+              "uid": "new-uid",
+            },
+            "exported": true,
+            "imported": true,
+            "refId": "A",
+          },
+          Object {
+            "datasource": Object {
+              "type": "new-type",
+              "uid": "new-uid",
+            },
+            "exported": true,
+            "imported": true,
+            "refId": "B",
+          },
+        ]
+      `);
+    });
+
+    it('should clear queries when no queries were imported', async () => {
+      const newUidDSWithAbstract = {
+        uid: 'new-uid',
+        type: 'new-type',
+        meta: {
+          id: 'new-type',
+        },
+        importFromAbstractQueries: () => {
+          return Promise.resolve([]);
+        },
+      } as DataSourceWithQueryImportSupport<any>;
+
+      const oldUidDSWithAbstract = {
+        uid: 'old-uid',
+        type: 'old-type',
+        meta: {
+          id: 'old-type',
+        },
+        exportToAbstractQueries: (queries) => {
+          const exportedQueries = queries.map((q) => ({ ...q, exported: true }));
+          return Promise.resolve(exportedQueries);
+        },
+      } as DataSourceWithQueryExportSupport<any>;
+
+      const queries = [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ];
+
+      const updated = await updateQueries(newUidDSWithAbstract as any, queries, oldUidDSWithAbstract as any);
+
+      expect(updated.length).toEqual(1);
+      expect(updated[0].datasource).toEqual({ type: 'new-type', uid: 'new-uid' });
+    });
+  });
+
+  describe('importQueries support', () => {
+    it('should import queries when abstract queries are not supported by datasources', async () => {
+      const importSpy = jest.fn();
+
+      const newUidDSWithImport = {
+        uid: 'new-uid',
+        type: 'new-type',
+        meta: {
+          id: 'new-type',
+        },
+        importQueries: (queries, origin) => {
+          importSpy(queries, origin);
+          const importedQueries = queries.map((q) => ({ ...q, imported: true }));
+          return Promise.resolve(importedQueries);
+        },
+      } as DataSourceApi<any>;
+
+      const oldUidDS = {
+        uid: 'old-uid',
+        type: 'old-type',
+        meta: {
+          id: 'old-type',
+        },
+      } as DataSourceApi;
+
+      const queries = [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ];
+
+      const updated = await updateQueries(newUidDSWithImport, queries, oldUidDS);
+
+      expect(importSpy).toBeCalledWith(queries, { uid: 'old-uid', type: 'old-type', meta: { id: 'old-type' } });
+
+      expect(updated).toMatchInlineSnapshot(`
+        Array [
+          Object {
+            "datasource": Object {
+              "type": "new-type",
+              "uid": "new-uid",
+            },
+            "imported": true,
+            "refId": "A",
+          },
+          Object {
+            "datasource": Object {
+              "type": "new-type",
+              "uid": "new-uid",
+            },
+            "imported": true,
+            "refId": "B",
+          },
+        ]
+      `);
+    });
+
+    it('should clear queries when no queries were imported', async () => {
+      const newUidDSWithImport = {
+        uid: 'new-uid',
+        type: 'new-type',
+        meta: {
+          id: 'new-type',
+        },
+        importQueries: (queries, origin) => {
+          return Promise.resolve([] as DataQuery[]);
+        },
+      } as DataSourceApi<any>;
+
+      const oldUidDS = {
+        uid: 'old-uid',
+        type: 'old-type',
+        meta: {
+          id: 'old-type',
+        },
+      } as DataSourceApi;
+
+      const queries = [
+        {
+          refId: 'A',
+          datasource: {
+            uid: 'old-uid',
+            type: 'old-type',
+          },
+        },
+        {
+          refId: 'B',
+          datasource: {
+            uid: 'other-uid',
+            type: 'other-type',
+          },
+        },
+      ];
+
+      const updated = await updateQueries(newUidDSWithImport, queries, oldUidDS);
+
+      expect(updated.length).toEqual(1);
+      expect(updated[0].datasource).toEqual({ type: 'new-type', uid: 'new-uid' });
+    });
+  });
+});

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

@@ -1,27 +1,42 @@
-import { DataQuery, DataSourceInstanceSettings, getDataSourceRef } from '@grafana/data';
+import { DataQuery, DataSourceApi, hasQueryExportSupport, hasQueryImportSupport } from '@grafana/data';
 import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
 
-export function updateQueries(
-  newSettings: DataSourceInstanceSettings,
+export async function updateQueries(
+  nextDS: DataSourceApi,
   queries: DataQuery[],
-  dsSettings?: DataSourceInstanceSettings
-): DataQuery[] {
-  const datasource = getDataSourceRef(newSettings);
+  currentDS?: DataSourceApi
+): Promise<DataQuery[]> {
+  let nextQueries = queries;
+  const datasource = { type: nextDS.type, uid: nextDS.uid };
 
   // we are changing data source type
-  if (dsSettings?.type !== newSettings.type) {
+  if (currentDS?.meta.id !== nextDS.meta.id) {
     // If changing to mixed do nothing
-    if (newSettings.meta.mixed) {
+    if (nextDS.meta.mixed) {
       return queries;
-    } else {
-      // Changing to another datasource type clear queries
+    }
+    // when both data sources support abstract queries
+    else if (hasQueryExportSupport(currentDS) && hasQueryImportSupport(nextDS)) {
+      const abstractQueries = await currentDS.exportToAbstractQueries(queries);
+      nextQueries = await nextDS.importFromAbstractQueries(abstractQueries);
+    }
+    // when datasource supports query import
+    else if (currentDS && nextDS.importQueries) {
+      nextQueries = await nextDS.importQueries(queries, currentDS);
+    }
+    // Otherwise clear queries
+    else {
       return [{ refId: 'A', datasource }];
     }
   }
 
+  if (nextQueries.length === 0) {
+    return [{ refId: 'A', datasource }];
+  }
+
   // Set data source on all queries except expression queries
-  return queries.map((query) => {
-    if (!isExpressionReference(query.datasource) && !newSettings.meta.mixed) {
+  return nextQueries.map((query) => {
+    if (!isExpressionReference(query.datasource) && !nextDS.meta.mixed) {
       query.datasource = datasource;
     }
     return query;