Chore: Fix for query variables using DataSourceRef (#41255)

public/app/features/variables/pickers/OptionsPicker/actions.test.ts

@@ -459,7 +459,7 @@ function createMultiVariable(extend?: Partial<QueryVariableModel>): QueryVariabl
     options: [],
     query: 'options-query',
     name: 'Constant',
-    datasource: 'datasource',
+    datasource: { uid: 'datasource' },
     definition: '',
     sort: VariableSort.alphabeticalAsc,
     refresh: VariableRefresh.never,

public/app/features/variables/query/QueryVariableEditor.tsx

@@ -4,7 +4,7 @@ import { css } from '@emotion/css';
 import { InlineField, InlineFieldRow, VerticalGroup } from '@grafana/ui';
 import { selectors } from '@grafana/e2e-selectors';
 import { DataSourcePicker, getTemplateSrv } from '@grafana/runtime';
-import { DataSourceInstanceSettings, LoadingState, SelectableValue } from '@grafana/data';
+import { DataSourceInstanceSettings, getDataSourceRef, LoadingState, SelectableValue } from '@grafana/data';
 
 import { SelectionOptionsEditor } from '../editor/SelectionOptionsEditor';
 import { QueryVariableModel, VariableRefresh, VariableSort, VariableWithMultiSupport } from '../types';
@@ -68,7 +68,7 @@ export class QueryVariableEditorUnConnected extends PureComponent<Props, State>
   onDataSourceChange = (dsSettings: DataSourceInstanceSettings) => {
     this.props.onPropChange({
       propName: 'datasource',
-      propValue: dsSettings.isDefault ? null : dsSettings.name,
+      propValue: dsSettings.isDefault ? null : getDataSourceRef(dsSettings),
     });
   };
 

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

@@ -1,4 +1,4 @@
-import { getDefaultTimeRange, LoadingState } from '@grafana/data';
+import { DataSourceRef, getDefaultTimeRange, LoadingState } from '@grafana/data';
 
 import { variableAdapters } from '../adapters';
 import { createQueryVariableAdapter } from './adapter';
@@ -43,7 +43,7 @@ const mocks: Record<string, any> = {
     metricFindQuery: jest.fn().mockResolvedValue([]),
   },
   dataSourceSrv: {
-    get: (name: string) => Promise.resolve(mocks[name]),
+    get: (ref: DataSourceRef) => Promise.resolve(mocks[ref.uid!]),
     getList: jest.fn().mockReturnValue([]),
   },
   pluginLoader: {
@@ -186,7 +186,7 @@ describe('query actions', () => {
       const variable = createVariable({ includeAll: true });
       const error = { message: 'failed to fetch metrics' };
 
-      mocks[variable.datasource!].metricFindQuery = jest.fn(() => Promise.reject(error));
+      mocks[variable.datasource!.uid!].metricFindQuery = jest.fn(() => Promise.reject(error));
 
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
@@ -320,7 +320,7 @@ describe('query actions', () => {
 
   describe('when changeQueryVariableDataSource is dispatched', () => {
     it('then correct actions are dispatched', async () => {
-      const variable = createVariable({ datasource: 'other' });
+      const variable = createVariable({ datasource: { uid: 'other' } });
       const editor = {};
 
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -330,7 +330,10 @@ describe('query actions', () => {
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(addVariable(toVariablePayload(variable, { global: false, index: 0, model: variable })))
-        .whenAsyncActionIsDispatched(changeQueryVariableDataSource(toVariablePayload(variable), 'datasource'), true);
+        .whenAsyncActionIsDispatched(
+          changeQueryVariableDataSource(toVariablePayload(variable), { uid: 'datasource' }),
+          true
+        );
 
       tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
         const [updateDatasource, updateEditor] = actions;
@@ -349,7 +352,7 @@ describe('query actions', () => {
 
     describe('and data source type changed', () => {
       it('then correct actions are dispatched', async () => {
-        const variable = createVariable({ datasource: 'other' });
+        const variable = createVariable({ datasource: { uid: 'other' } });
         const editor = {};
         const preloadedState: any = { templating: { editor: { extended: { dataSource: { type: 'previous' } } } } };
 
@@ -362,7 +365,10 @@ describe('query actions', () => {
           .whenActionIsDispatched(
             addVariable(toVariablePayload(variable, { global: false, index: 0, model: variable }))
           )
-          .whenAsyncActionIsDispatched(changeQueryVariableDataSource(toVariablePayload(variable), 'datasource'), true);
+          .whenAsyncActionIsDispatched(
+            changeQueryVariableDataSource(toVariablePayload(variable), { uid: 'datasource' }),
+            true
+          );
 
         tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
           const [changeVariable, updateDatasource, updateEditor] = actions;
@@ -386,7 +392,7 @@ describe('query actions', () => {
 
   describe('when changeQueryVariableDataSource is dispatched and editor is not configured', () => {
     it('then correct actions are dispatched', async () => {
-      const variable = createVariable({ datasource: 'other' });
+      const variable = createVariable({ datasource: { uid: 'other' } });
       const editor = LegacyVariableQueryEditor;
 
       mocks.pluginLoader.importDataSourcePlugin = jest.fn().mockResolvedValue({
@@ -396,7 +402,10 @@ describe('query actions', () => {
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(addVariable(toVariablePayload(variable, { global: false, index: 0, model: variable })))
-        .whenAsyncActionIsDispatched(changeQueryVariableDataSource(toVariablePayload(variable), 'datasource'), true);
+        .whenAsyncActionIsDispatched(
+          changeQueryVariableDataSource(toVariablePayload(variable), { uid: 'datasource' }),
+          true
+        );
 
       tester.thenDispatchedActionsPredicateShouldEqual((actions) => {
         const [updateDatasource, updateEditor] = actions;
@@ -417,7 +426,7 @@ describe('query actions', () => {
   describe('when changeQueryVariableQuery is dispatched', () => {
     it('then correct actions are dispatched', async () => {
       const optionsMetrics = [createMetric('A'), createMetric('B')];
-      const variable = createVariable({ datasource: 'datasource', includeAll: true });
+      const variable = createVariable({ datasource: { uid: 'datasource' }, includeAll: true });
 
       const query = '$datasource';
       const definition = 'depends on datasource variable';
@@ -447,7 +456,7 @@ describe('query actions', () => {
   describe('when changeQueryVariableQuery is dispatched for variable without tags', () => {
     it('then correct actions are dispatched', async () => {
       const optionsMetrics = [createMetric('A'), createMetric('B')];
-      const variable = createVariable({ datasource: 'datasource', includeAll: true });
+      const variable = createVariable({ datasource: { uid: 'datasource' }, includeAll: true });
 
       const query = '$datasource';
       const definition = 'depends on datasource variable';
@@ -477,7 +486,7 @@ describe('query actions', () => {
   describe('when changeQueryVariableQuery is dispatched for variable without tags and all', () => {
     it('then correct actions are dispatched', async () => {
       const optionsMetrics = [createMetric('A'), createMetric('B')];
-      const variable = createVariable({ datasource: 'datasource', includeAll: false });
+      const variable = createVariable({ datasource: { uid: 'datasource' }, includeAll: false });
       const query = '$datasource';
       const definition = 'depends on datasource variable';
 
@@ -505,7 +514,7 @@ describe('query actions', () => {
 
   describe('when changeQueryVariableQuery is dispatched with invalid query', () => {
     it('then correct actions are dispatched', async () => {
-      const variable = createVariable({ datasource: 'datasource', includeAll: false });
+      const variable = createVariable({ datasource: { uid: 'datasource' }, includeAll: false });
       const query = `$${variable.name}`;
       const definition = 'depends on datasource variable';
 
@@ -688,7 +697,7 @@ function mockDatasourceMetrics(variable: QueryVariableModel, optionsMetrics: any
     [variable.query]: optionsMetrics,
   };
 
-  const { metricFindQuery } = mocks[variable.datasource!];
+  const { metricFindQuery } = mocks[variable.datasource?.uid!];
 
   metricFindQuery.mockReset();
   metricFindQuery.mockImplementation((query: string) => Promise.resolve(metrics[query] ?? []));
@@ -707,7 +716,7 @@ function createVariable(extend?: Partial<QueryVariableModel>): QueryVariableMode
     hide: VariableHide.dontHide,
     skipUrlSync: false,
     index: 0,
-    datasource: 'datasource',
+    datasource: { uid: 'datasource' },
     definition: '',
     sort: VariableSort.alphabeticalAsc,
     refresh: VariableRefresh.onDashboardLoad,

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

@@ -1,4 +1,7 @@
+import { Subscription } from 'rxjs';
 import { getDataSourceSrv, toDataQueryError } from '@grafana/runtime';
+import { DataSourceRef } from '@grafana/data';
+
 import { updateOptions } from '../state/actions';
 import { QueryVariableModel } from '../types';
 import { ThunkResult } from '../../../types';
@@ -12,7 +15,6 @@ import {
 import { changeVariableProp } from '../state/sharedReducer';
 import { toVariableIdentifier, toVariablePayload, VariableIdentifier } from '../state/types';
 import { getVariableQueryEditor } from '../editor/getVariableQueryEditor';
-import { Subscription } from 'rxjs';
 import { getVariableQueryRunner } from './VariableQueryRunner';
 import { variableQueryObserver } from './variableQueryObserver';
 import { QueryVariableEditorState } from './reducer';
@@ -60,7 +62,7 @@ export const initQueryVariableEditor = (identifier: VariableIdentifier): ThunkRe
 
 export const changeQueryVariableDataSource = (
   identifier: VariableIdentifier,
-  name: string | null
+  name: DataSourceRef | null
 ): ThunkResult<void> => {
   return async (dispatch, getState) => {
     try {

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

@@ -21,7 +21,7 @@ export const createQueryVariableAdapter = (): VariableAdapter<QueryVariableModel
     picker: optionPickerFactory<QueryVariableModel>(),
     editor: QueryVariableEditor,
     dependsOn: (variable, variableToTest) => {
-      return containsVariable(variable.query, variable.datasource, variable.regex, variableToTest.name);
+      return containsVariable(variable.query, variable.datasource?.uid, variable.regex, variableToTest.name);
     },
     setValue: async (variable, option, emitChanges = false) => {
       await dispatch(setOptionAsCurrent(toVariableIdentifier(variable), option, emitChanges));

public/app/features/variables/shared/testing/queryVariableBuilder.ts

@@ -1,8 +1,10 @@
+import { DataSourceRef } from '@grafana/data';
+
 import { QueryVariableModel } from 'app/features/variables/types';
 import { DatasourceVariableBuilder } from './datasourceVariableBuilder';
 
 export class QueryVariableBuilder<T extends QueryVariableModel> extends DatasourceVariableBuilder<T> {
-  withDatasource(datasource: string) {
+  withDatasource(datasource: DataSourceRef) {
     this.variable.datasource = datasource;
     return this;
   }

public/app/features/variables/state/upgradeLegacyQueries.test.ts

@@ -13,7 +13,13 @@ interface Args {
 }
 function getTestContext({ query = '', variable, datasource }: Args = {}) {
   variable =
-    variable ?? queryBuilder().withId('query').withName('query').withQuery(query).withDatasource('test-data').build();
+    variable ??
+    queryBuilder()
+      .withId('query')
+      .withName('query')
+      .withQuery(query)
+      .withDatasource({ uid: 'test-data', type: 'test-data' })
+      .build();
   const state = {
     templating: {
       variables: {
@@ -56,7 +62,7 @@ describe('upgradeLegacyQueries', () => {
         }),
       ]);
       expect(get).toHaveBeenCalledTimes(1);
-      expect(get).toHaveBeenCalledWith('test-data');
+      expect(get).toHaveBeenCalledWith({ uid: 'test-data', type: 'test-data' });
     });
   });
 
@@ -70,7 +76,7 @@ describe('upgradeLegacyQueries', () => {
 
       expect(dispatchedActions).toEqual([]);
       expect(get).toHaveBeenCalledTimes(1);
-      expect(get).toHaveBeenCalledWith('test-data');
+      expect(get).toHaveBeenCalledWith({ uid: 'test-data', type: 'test-data' });
     });
   });
 
@@ -88,7 +94,7 @@ describe('upgradeLegacyQueries', () => {
 
       expect(dispatchedActions).toEqual([]);
       expect(get).toHaveBeenCalledTimes(1);
-      expect(get).toHaveBeenCalledWith('test-data');
+      expect(get).toHaveBeenCalledWith({ uid: 'test-data', type: 'test-data' });
     });
   });
 
@@ -107,7 +113,7 @@ describe('upgradeLegacyQueries', () => {
 
       expect(dispatchedActions).toEqual([]);
       expect(get).toHaveBeenCalledTimes(1);
-      expect(get).toHaveBeenCalledWith('test-data');
+      expect(get).toHaveBeenCalledWith({ uid: 'test-data', type: 'test-data' });
     });
   });
 
@@ -126,7 +132,7 @@ describe('upgradeLegacyQueries', () => {
 
       expect(dispatchedActions).toEqual([]);
       expect(get).toHaveBeenCalledTimes(1);
-      expect(get).toHaveBeenCalledWith('test-data');
+      expect(get).toHaveBeenCalledWith({ uid: 'test-data', type: 'test-data' });
     });
   });
 

public/app/features/variables/types.ts

@@ -68,7 +68,7 @@ export interface DataSourceVariableModel extends VariableWithMultiSupport {
 }
 
 export interface QueryVariableModel extends DataSourceVariableModel {
-  datasource: string | null;
+  datasource: DataSourceRef | null;
   definition: string;
   sort: VariableSort;
   queryValue?: string;