Variables: Fix crash when changing query variable datasource (#44957)

* Variables: refactor variable editor state to dispatch typed & atomic extended props

* Mock variable editor in tests

* Make dataSources mandatory in AdHocVariableEditorState

* nit

public/app/features/variables/adhoc/actions.test.ts

@@ -35,6 +35,23 @@ jest.mock('app/features/plugins/datasource_srv', () => ({
 
 variableAdapters.setInit(() => [createAdHocVariableAdapter()]);
 
+const datasources = [
+  { ...createDatasource('default', true, true), value: null },
+  createDatasource('elasticsearch-v1'),
+  createDatasource('loki', false),
+  createDatasource('influx'),
+  createDatasource('google-sheets', false),
+  createDatasource('elasticsearch-v7'),
+];
+
+const expectedDatasources = [
+  { text: '', value: {} },
+  { text: 'default (default)', value: { uid: 'default', type: 'default' } },
+  { text: 'elasticsearch-v1', value: { uid: 'elasticsearch-v1', type: 'elasticsearch-v1' } },
+  { text: 'influx', value: { uid: 'influx', type: 'influx' } },
+  { text: 'elasticsearch-v7', value: { uid: 'elasticsearch-v7', type: 'elasticsearch-v7' } },
+];
+
 describe('adhoc actions', () => {
   describe('when applyFilterFromTable is dispatched and filter already exist', () => {
     it('then correct actions are dispatched', async () => {
@@ -350,15 +367,6 @@ describe('adhoc actions', () => {
 
   describe('when initAdHocVariableEditor is dispatched', () => {
     it('then correct actions are dispatched', async () => {
-      const datasources = [
-        { ...createDatasource('default', true, true), value: null },
-        createDatasource('elasticsearch-v1'),
-        createDatasource('loki', false),
-        createDatasource('influx'),
-        createDatasource('google-sheets', false),
-        createDatasource('elasticsearch-v7'),
-      ];
-
       getList.mockRestore();
       getList.mockReturnValue(datasources);
 
@@ -366,41 +374,32 @@ describe('adhoc actions', () => {
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(initAdHocVariableEditor());
 
-      const expectedDatasources = [
+      tester.thenDispatchedActionsShouldEqual(changeVariableEditorExtended({ dataSources: expectedDatasources }));
-        { text: '', value: {} },
-        { text: 'default (default)', value: { uid: 'default', type: 'default' } },
-        { text: 'elasticsearch-v1', value: { uid: 'elasticsearch-v1', type: 'elasticsearch-v1' } },
-        { text: 'influx', value: { uid: 'influx', type: 'influx' } },
-        { text: 'elasticsearch-v7', value: { uid: 'elasticsearch-v7', type: 'elasticsearch-v7' } },
-      ];
-
-      tester.thenDispatchedActionsShouldEqual(
-        changeVariableEditorExtended({ propName: 'dataSources', propValue: expectedDatasources })
-      );
     });
   });
 
   describe('when changeVariableDatasource is dispatched with unsupported datasource', () => {
     it('then correct actions are dispatched', async () => {
       const datasource = { uid: 'mysql' };
-      const loadingText = 'Ad hoc filters are applied automatically to all queries that target this data source';
       const variable = adHocBuilder().withId('Filters').withName('Filters').withDatasource({ uid: 'influxdb' }).build();
 
       getDatasource.mockRestore();
       getDatasource.mockResolvedValue(null);
+      getList.mockRestore();
+      getList.mockReturnValue(datasources);
 
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(createAddVariableAction(variable))
         .whenActionIsDispatched(setIdInEditor({ id: variable.id }))
+        .whenActionIsDispatched(initAdHocVariableEditor())
         .whenAsyncActionIsDispatched(changeVariableDatasource(datasource), true);
 
       tester.thenDispatchedActionsShouldEqual(
-        changeVariableEditorExtended({ propName: 'infoText', propValue: loadingText }),
         changeVariableProp(toVariablePayload(variable, { propName: 'datasource', propValue: datasource })),
         changeVariableEditorExtended({
-          propName: 'infoText',
+          infoText: 'This data source does not support ad hoc filters yet.',
-          propValue: 'This data source does not support ad hoc filters yet.',
+          dataSources: expectedDatasources,
         })
       );
     });
@@ -416,16 +415,19 @@ describe('adhoc actions', () => {
       getDatasource.mockResolvedValue({
         getTagKeys: () => {},
       });
+      getList.mockRestore();
+      getList.mockReturnValue(datasources);
 
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(createAddVariableAction(variable))
         .whenActionIsDispatched(setIdInEditor({ id: variable.id }))
+        .whenActionIsDispatched(initAdHocVariableEditor())
         .whenAsyncActionIsDispatched(changeVariableDatasource(datasource), true);
 
       tester.thenDispatchedActionsShouldEqual(
-        changeVariableEditorExtended({ propName: 'infoText', propValue: loadingText }),
+        changeVariableProp(toVariablePayload(variable, { propName: 'datasource', propValue: datasource })),
-        changeVariableProp(toVariablePayload(variable, { propName: 'datasource', propValue: datasource }))
+        changeVariableEditorExtended({ infoText: loadingText, dataSources: expectedDatasources })
       );
     });
   });

public/app/features/variables/adhoc/actions.ts

@@ -17,6 +17,7 @@ import { AdHocVariableFilter, AdHocVariableModel } from 'app/features/variables/
 import { variableUpdated } from '../state/actions';
 import { isAdHoc } from '../guard';
 import { DataSourceRef, getDataSourceRef } from '@grafana/data';
+import { getAdhocVariableEditorState } from '../editor/selectors';
 
 export interface AdHocTableOptions {
   datasource: DataSourceRef;
@@ -84,28 +85,23 @@ export const setFiltersFromUrl = (id: string, filters: AdHocVariableFilter[]): T
 export const changeVariableDatasource = (datasource?: DataSourceRef): ThunkResult<void> => {
   return async (dispatch, getState) => {
     const { editor } = getState().templating;
+    const extended = getAdhocVariableEditorState(editor);
     const variable = getVariable(editor.id, getState());
-
-    const loadingText = 'Ad hoc filters are applied automatically to all queries that target this data source';
-
-    dispatch(
-      changeVariableEditorExtended({
-        propName: 'infoText',
-        propValue: loadingText,
-      })
-    );
     dispatch(changeVariableProp(toVariablePayload(variable, { propName: 'datasource', propValue: datasource })));
 
     const ds = await getDatasourceSrv().get(datasource);
 
-    if (!ds || !ds.getTagKeys) {
+    // TS TODO: ds is not typed to be optional - is this check unnecessary or is the type incorrect?
+    const message = ds?.getTagKeys
+      ? 'Ad hoc filters are applied automatically to all queries that target this data source'
+      : 'This data source does not support ad hoc filters yet.';
+
     dispatch(
       changeVariableEditorExtended({
-          propName: 'infoText',
+        infoText: message,
-          propValue: 'This data source does not support ad hoc filters yet.',
+        dataSources: extended?.dataSources ?? [],
       })
     );
-    }
   };
 };
 
@@ -128,8 +124,7 @@ export const initAdHocVariableEditor = (): ThunkResult<void> => (dispatch) => {
 
   dispatch(
     changeVariableEditorExtended({
-      propName: 'dataSources',
+      dataSources: selectable,
-      propValue: selectable,
     })
   );
 };

public/app/features/variables/datasource/actions.test.ts

@@ -58,7 +58,7 @@ describe('data source actions', () => {
             true
           );
 
-        await tester.thenDispatchedActionsShouldEqual(
+        tester.thenDispatchedActionsShouldEqual(
           createDataSourceOptions(
             toVariablePayload(
               { type: 'datasource', id: '0' },
@@ -106,7 +106,7 @@ describe('data source actions', () => {
             true
           );
 
-        await tester.thenDispatchedActionsShouldEqual(
+        tester.thenDispatchedActionsShouldEqual(
           createDataSourceOptions(
             toVariablePayload(
               { type: 'datasource', id: '0' },
@@ -132,7 +132,7 @@ describe('data source actions', () => {
   });
 
   describe('when initDataSourceVariableEditor is dispatched', () => {
-    it('then the correct actions are dispatched', async () => {
+    it('then the correct actions are dispatched', () => {
       const meta = getMockPlugin({ name: 'mock-data-name', id: 'mock-data-id' });
       const sources: DataSourceInstanceSettings[] = [
         getDataSourceInstanceSetting('first-name', meta),
@@ -141,13 +141,12 @@ describe('data source actions', () => {
 
       const { dependencies, getListMock, getDatasourceSrvMock } = getTestContext({ sources });
 
-      await reduxTester<RootReducerType>()
+      reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(initDataSourceVariableEditor(dependencies))
         .thenDispatchedActionsShouldEqual(
           changeVariableEditorExtended({
-            propName: 'dataSourceTypes',
+            dataSourceTypes: [
-            propValue: [
               { text: '', value: '' },
               { text: 'mock-data-name', value: 'mock-data-id' },
             ],

public/app/features/variables/datasource/actions.ts

@@ -47,10 +47,5 @@ export const initDataSourceVariableEditor =
 
     dataSourceTypes.unshift({ text: '', value: '' });
 
-    dispatch(
+    dispatch(changeVariableEditorExtended({ dataSourceTypes }));
-      changeVariableEditorExtended({
-        propName: 'dataSourceTypes',
-        propValue: dataSourceTypes,
-      })
-    );
   };

public/app/features/variables/editor/reducer.test.ts

@@ -181,15 +181,15 @@ describe('variableEditorReducer', () => {
 
   describe('when changeVariableEditorExtended is dispatched', () => {
     it('then state should be correct', () => {
-      const payload = { propName: 'someProp', propValue: [{}] };
+      const payload = { dataSourceTypes: [] };
+
       reducerTester<VariableEditorState>()
         .givenReducer(variableEditorReducer, { ...initialVariableEditorState })
         .whenActionIsDispatched(changeVariableEditorExtended(payload))
         .thenStateShouldEqual({
           ...initialVariableEditorState,
           extended: {
-            // @ts-ignore - temp ignoring this, we'll fix it soon
+            dataSourceTypes: [],
-            someProp: [{}],
           },
         });
     });

public/app/features/variables/editor/reducer.ts

@@ -78,14 +78,10 @@ const variableEditorReducerSlice = createSlice({
       delete state.errors[action.payload.errorProp];
       state.isValid = Object.keys(state.errors).length === 0;
     },
-    changeVariableEditorExtended: (
+    changeVariableEditorExtended: (state: VariableEditorState, action: PayloadAction<VariableEditorExtension>) => {
-      state: VariableEditorState,
-      action: PayloadAction<{ propName: string; propValue: any }>
-    ) => {
-      // @ts-ignore - temp ignoring the errors now the state type is more strict
       state.extended = {
         ...state.extended,
-        [action.payload.propName]: action.payload.propValue,
+        ...action.payload,
       };
     },
     cleanEditorState: () => initialVariableEditorState,

public/app/features/variables/query/actions.test.tsx

@@ -1,10 +1,11 @@
+import React from 'react';
 import { DataSourceRef, getDefaultTimeRange, LoadingState } from '@grafana/data';
 
 import { variableAdapters } from '../adapters';
 import { createQueryVariableAdapter } from './adapter';
 import { reduxTester } from '../../../../test/core/redux/reduxTester';
 import { getRootReducer, RootReducerType } from '../state/helpers';
-import { QueryVariableModel, VariableHide, VariableRefresh, VariableSort } from '../types';
+import { QueryVariableModel, VariableHide, VariableQueryEditorProps, VariableRefresh, VariableSort } from '../types';
 import { toVariablePayload } from '../state/types';
 import {
   addVariable,
@@ -51,6 +52,9 @@ const mocks: Record<string, any> = {
   pluginLoader: {
     importDataSourcePlugin: jest.fn().mockResolvedValue({ components: {} }),
   },
+  VariableQueryEditor(props: VariableQueryEditorProps) {
+    return <div>this is a variable query editor</div>;
+  },
 };
 
 setDataSourceSrv(mocks.dataSourceSrv as any);
@@ -240,7 +244,7 @@ describe('query actions', () => {
     it('then correct actions are dispatched', async () => {
       const variable = createVariable({ includeAll: true });
       const testMetricSource = { name: 'test', value: 'test', meta: {} };
-      const editor = {};
+      const editor = mocks.VariableQueryEditor;
 
       mocks.dataSourceSrv.getList = jest.fn().mockReturnValue([testMetricSource]);
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -253,16 +257,9 @@ describe('query actions', () => {
         .whenActionIsDispatched(variablesInitTransaction({ uid: 'a uid' }))
         .whenAsyncActionIsDispatched(initQueryVariableEditor(toVariablePayload(variable)), true);
 
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+      tester.thenDispatchedActionsShouldEqual(
-        const [setDatasource, setEditor] = actions;
+        changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-        const expectedNumberOfActions = 2;
-
-        expect(setDatasource).toEqual(
-          changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks['datasource'] })
       );
-        expect(setEditor).toEqual(changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor }));
-        return actions.length === expectedNumberOfActions;
-      });
     });
   });
 
@@ -270,7 +267,7 @@ describe('query actions', () => {
     it('then correct actions are dispatched', async () => {
       const variable = createVariable({ includeAll: true });
       const testMetricSource = { name: 'test', value: null as unknown as string, meta: {} };
-      const editor = {};
+      const editor = mocks.VariableQueryEditor;
 
       mocks.dataSourceSrv.getList = jest.fn().mockReturnValue([testMetricSource]);
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -283,23 +280,16 @@ describe('query actions', () => {
         .whenActionIsDispatched(variablesInitTransaction({ uid: 'a uid' }))
         .whenAsyncActionIsDispatched(initQueryVariableEditor(toVariablePayload(variable)), true);
 
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+      tester.thenDispatchedActionsShouldEqual(
-        const [setDatasource, setEditor] = actions;
+        changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-        const expectedNumberOfActions = 2;
-
-        expect(setDatasource).toEqual(
-          changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks['datasource'] })
       );
-        expect(setEditor).toEqual(changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor }));
-        return actions.length === expectedNumberOfActions;
-      });
     });
   });
 
   describe('when initQueryVariableEditor is dispatched and no metric sources was found', () => {
     it('then correct actions are dispatched', async () => {
       const variable = createVariable({ includeAll: true });
-      const editor = {};
+      const editor = mocks.VariableQueryEditor;
 
       mocks.dataSourceSrv.getList = jest.fn().mockReturnValue([]);
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -312,43 +302,16 @@ describe('query actions', () => {
         .whenActionIsDispatched(variablesInitTransaction({ uid: 'a uid' }))
         .whenAsyncActionIsDispatched(initQueryVariableEditor(toVariablePayload(variable)), true);
 
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+      tester.thenDispatchedActionsShouldEqual(
-        const [setDatasource, setEditor] = actions;
+        changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-        const expectedNumberOfActions = 2;
-
-        expect(setDatasource).toEqual(
-          changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks['datasource'] })
       );
-        expect(setEditor).toEqual(changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor }));
-        return actions.length === expectedNumberOfActions;
-      });
-    });
-  });
-
-  describe('when initQueryVariableEditor is dispatched and variable dont have datasource', () => {
-    it('then correct actions are dispatched', async () => {
-      const variable = createVariable({ datasource: undefined });
-
-      const tester = await reduxTester<RootReducerType>()
-        .givenRootReducer(getRootReducer())
-        .whenActionIsDispatched(addVariable(toVariablePayload(variable, { global: false, index: 0, model: variable })))
-        .whenActionIsDispatched(variablesInitTransaction({ uid: 'a uid' }))
-        .whenAsyncActionIsDispatched(initQueryVariableEditor(toVariablePayload(variable)), true);
-
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
-        const [setDatasource] = actions;
-        const expectedNumberOfActions = 1;
-
-        expect(setDatasource).toEqual(changeVariableEditorExtended({ propName: 'dataSource', propValue: undefined }));
-        return actions.length === expectedNumberOfActions;
-      });
     });
   });
 
   describe('when changeQueryVariableDataSource is dispatched', () => {
     it('then correct actions are dispatched', async () => {
       const variable = createVariable({ datasource: { uid: 'other' } });
-      const editor = {};
+      const editor = mocks.VariableQueryEditor;
 
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
         components: { VariableQueryEditor: editor },
@@ -363,25 +326,15 @@ describe('query actions', () => {
           true
         );
 
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+      tester.thenDispatchedActionsShouldEqual(
-        const [updateDatasource, updateEditor] = actions;
+        changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-        const expectedNumberOfActions = 2;
-
-        expect(updateDatasource).toEqual(
-          changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks.datasource })
       );
-        expect(updateEditor).toEqual(
-          changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor })
-        );
-
-        return actions.length === expectedNumberOfActions;
-      });
     });
 
     describe('and data source type changed', () => {
       it('then correct actions are dispatched', async () => {
         const variable = createVariable({ datasource: { uid: 'other' } });
-        const editor = {};
+        const editor = mocks.VariableQueryEditor;
         const preloadedState: any = { templating: { editor: { extended: { dataSource: { type: 'previous' } } } } };
 
         mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -399,22 +352,10 @@ describe('query actions', () => {
             true
           );
 
-        tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+        tester.thenDispatchedActionsShouldEqual(
-          const [changeVariable, updateDatasource, updateEditor] = actions;
+          changeVariableProp(toVariablePayload(variable, { propName: 'query', propValue: '' })),
-          const expectedNumberOfActions = 3;
+          changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-
-          expect(changeVariable).toEqual(
-            changeVariableProp(toVariablePayload(variable, { propName: 'query', propValue: '' }))
         );
-          expect(updateDatasource).toEqual(
-            changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks.datasource })
-          );
-          expect(updateEditor).toEqual(
-            changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor })
-          );
-
-          return actions.length === expectedNumberOfActions;
-        });
       });
     });
   });
@@ -437,19 +378,9 @@ describe('query actions', () => {
           true
         );
 
-      tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
+      tester.thenDispatchedActionsShouldEqual(
-        const [updateDatasource, updateEditor] = actions;
+        changeVariableEditorExtended({ dataSource: mocks.datasource, VariableQueryEditor: editor })
-        const expectedNumberOfActions = 2;
-
-        expect(updateDatasource).toEqual(
-          changeVariableEditorExtended({ propName: 'dataSource', propValue: mocks.datasource })
       );
-        expect(updateEditor).toEqual(
-          changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: editor })
-        );
-
-        return actions.length === expectedNumberOfActions;
-      });
     });
   });
 

public/app/features/variables/query/actions.ts

@@ -69,13 +69,19 @@ export const changeQueryVariableDataSource = (
       const extendedEditorState = getQueryVariableEditorState(getState().templating.editor);
       const previousDatasource = extendedEditorState?.dataSource;
       const dataSource = await getDataSourceSrv().get(name ?? '');
+
       if (previousDatasource && previousDatasource.type !== dataSource?.type) {
         dispatch(changeVariableProp(toVariablePayload(identifier, { propName: 'query', propValue: '' })));
       }
-      dispatch(changeVariableEditorExtended({ propName: 'dataSource', propValue: dataSource }));
 
       const VariableQueryEditor = await getVariableQueryEditor(dataSource);
-      dispatch(changeVariableEditorExtended({ propName: 'VariableQueryEditor', propValue: VariableQueryEditor }));
+
+      dispatch(
+        changeVariableEditorExtended({
+          dataSource: dataSource,
+          VariableQueryEditor: VariableQueryEditor,
+        })
+      );
     } catch (err) {
       console.error(err);
     }