AdHoc variable: Correctly preselect datasource when provisioning (#54088)

* Adhoc variable: Correctly preselect datasource when provisioning

* Fix test

* Remove data sources from ad hoc variable state in favor of DataSourcePicker

.betterer.results

@@ -7970,8 +7970,7 @@ exports[`better eslint`] = {
       [0, 0, 0, "Unexpected any. Specify a different type.", "17"]
     ],
     "public/app/plugins/datasource/tempo/language_provider.ts:5381": [
-      [0, 0, 0, "Unexpected any. Specify a different type.", "0"],
-      [0, 0, 0, "Unexpected any. Specify a different type.", "1"]
+      [0, 0, 0, "Unexpected any. Specify a different type.", "0"]
     ],
     "public/app/plugins/datasource/tempo/resultTransformer.test.ts:5381": [
       [0, 0, 0, "Unexpected any. Specify a different type.", "0"]
@@ -7996,6 +7995,28 @@ exports[`better eslint`] = {
       [0, 0, 0, "Do not use any type assertions.", "2"],
       [0, 0, 0, "Unexpected any. Specify a different type.", "3"]
     ],
+    "public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx:5381": [
+      [0, 0, 0, "Do not use any type assertions.", "0"],
+      [0, 0, 0, "Do not use any type assertions.", "1"]
+    ],
+    "public/app/plugins/datasource/tempo/traceql/autocomplete.test.ts:5381": [
+      [0, 0, 0, "Unexpected any. Specify a different type.", "0"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "1"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "2"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "3"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "4"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "5"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "6"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "7"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "8"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "9"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "10"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "11"],
+      [0, 0, 0, "Unexpected any. Specify a different type.", "12"]
+    ],
+    "public/app/plugins/datasource/tempo/traceql/autocomplete.ts:5381": [
+      [0, 0, 0, "Do not use any type assertions.", "0"]
+    ],
     "public/app/plugins/datasource/testdata/ConfigEditor.tsx:5381": [
       [0, 0, 0, "Unexpected any. Specify a different type.", "0"]
     ],

public/app/features/variables/adhoc/AdHocVariableEditor.test.tsx

@@ -1,11 +1,42 @@
 import { render, screen } from '@testing-library/react';
 import React from 'react';
-import { selectOptionInTest, getSelectParent } from 'test/helpers/selectOptionInTest';
+import { selectOptionInTest } from 'test/helpers/selectOptionInTest';
+
+import { selectors } from '@grafana/e2e-selectors';
+import { mockDataSource } from 'app/features/alerting/unified/mocks';
+import { DataSourceType } from 'app/features/alerting/unified/utils/datasource';
 
 import { adHocBuilder } from '../shared/testing/builders';
 
 import { AdHocVariableEditorUnConnected as AdHocVariableEditor } from './AdHocVariableEditor';
 
+const promDsMock = mockDataSource({
+  name: 'Prometheus',
+  type: DataSourceType.Prometheus,
+});
+
+const lokiDsMock = mockDataSource({
+  name: 'Loki',
+  type: DataSourceType.Loki,
+});
+
+jest.mock('@grafana/runtime/src/services/dataSourceSrv', () => {
+  return {
+    getDataSourceSrv: () => ({
+      get: () => {
+        return Promise.resolve(promDsMock);
+      },
+      getList: () => [promDsMock, lokiDsMock],
+      getInstanceSettings: (v: string) => {
+        if (v === 'Prometheus') {
+          return promDsMock;
+        }
+        return lokiDsMock;
+      },
+    }),
+  };
+});
+
 const props = {
   extended: {
     dataSources: [
@@ -29,17 +60,17 @@ describe('AdHocVariableEditor', () => {
   it('has a datasource select menu', async () => {
     render(<AdHocVariableEditor {...props} />);
 
-    const selectContainer = getSelectParent(screen.getByLabelText('Data source'));
-    expect(selectContainer).toHaveTextContent('Prometheus');
+    expect(await screen.findByLabelText(selectors.components.DataSourcePicker.inputV2)).toBeInTheDocument();
   });
 
   it('calls the callback when changing the datasource', async () => {
     render(<AdHocVariableEditor {...props} />);
-    await selectOptionInTest(screen.getByLabelText('Data source'), 'Loki');
+    const selectEl = screen.getByLabelText(selectors.components.DataSourcePicker.inputV2);
+    await selectOptionInTest(selectEl, 'Loki');
 
     expect(props.changeVariableDatasource).toBeCalledWith(
       { type: 'adhoc', id: 'adhoc', rootStateKey: 'key' },
-      { type: 'loki-ds', uid: 'abc' }
+      { type: 'loki', uid: 'mock-ds-3' }
     );
   });
 

public/app/features/variables/adhoc/AdHocVariableEditor.tsx

@@ -1,12 +1,12 @@
 import React, { PureComponent } from 'react';
 import { connect, ConnectedProps } from 'react-redux';
 
-import { DataSourceRef, SelectableValue } from '@grafana/data';
-import { Alert, InlineFieldRow, VerticalGroup } from '@grafana/ui';
+import { DataSourceInstanceSettings, getDataSourceRef } from '@grafana/data';
+import { DataSourcePicker } from '@grafana/runtime';
+import { Alert, InlineField, InlineFieldRow, VerticalGroup } from '@grafana/ui';
 import { StoreState } from 'app/types';
 
 import { VariableSectionHeader } from '../editor/VariableSectionHeader';
-import { VariableSelectField } from '../editor/VariableSelectField';
 import { initialVariableEditorState } from '../editor/reducer';
 import { getAdhocVariableEditorState } from '../editor/selectors';
 import { VariableEditorProps } from '../editor/types';
@@ -14,7 +14,7 @@ import { getVariablesState } from '../state/selectors';
 import { AdHocVariableModel } from '../types';
 import { toKeyedVariableIdentifier } from '../utils';
 
-import { changeVariableDatasource, initAdHocVariableEditor } from './actions';
+import { changeVariableDatasource } from './actions';
 
 const mapStateToProps = (state: StoreState, ownProps: OwnProps) => {
   const { rootStateKey } = ownProps.variable;
@@ -34,7 +34,6 @@ const mapStateToProps = (state: StoreState, ownProps: OwnProps) => {
 };
 
 const mapDispatchToProps = {
-  initAdHocVariableEditor,
   changeVariableDatasource,
 };
 
@@ -51,33 +50,24 @@ export class AdHocVariableEditorUnConnected extends PureComponent<Props> {
       console.error('AdHocVariableEditor: variable has no rootStateKey');
       return;
     }
-
-    this.props.initAdHocVariableEditor(rootStateKey);
   }
 
-  onDatasourceChanged = (option: SelectableValue<DataSourceRef>) => {
-    this.props.changeVariableDatasource(toKeyedVariableIdentifier(this.props.variable), option.value);
+  onDatasourceChanged = (ds: DataSourceInstanceSettings) => {
+    this.props.changeVariableDatasource(toKeyedVariableIdentifier(this.props.variable), getDataSourceRef(ds));
   };
 
   render() {
     const { variable, extended } = this.props;
-    const dataSources = extended?.dataSources ?? [];
     const infoText = extended?.infoText ?? null;
-    const options = dataSources.map((ds) => ({ label: ds.text, value: ds.value }));
-    const value = options.find((o) => o.value?.uid === variable.datasource?.uid) ?? options[0];
 
     return (
       <VerticalGroup spacing="xs">
         <VariableSectionHeader name="Options" />
         <VerticalGroup spacing="sm">
           <InlineFieldRow>
-            <VariableSelectField
-              name="Data source"
-              value={value}
-              options={options}
-              onChange={this.onDatasourceChanged}
-              labelWidth={10}
-            />
+            <InlineField label="Data source" labelWidth={20} htmlFor="data-source-picker">
+              <DataSourcePicker current={variable.datasource} onChange={this.onDatasourceChanged} noDefault />
+            </InlineField>
           </InlineFieldRow>
 
           {infoText ? <Alert title={infoText} severity="info" /> : null}

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

@@ -16,7 +16,6 @@ import {
   applyFilterFromTable,
   changeFilter,
   changeVariableDatasource,
-  initAdHocVariableEditor,
   removeFilter,
   setFiltersFromUrl,
 } from './actions';
@@ -45,14 +44,6 @@ const datasources = [
   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 () => {
@@ -405,23 +396,6 @@ describe('adhoc actions', () => {
     });
   });
 
-  describe('when initAdHocVariableEditor is dispatched', () => {
-    it('then correct actions are dispatched', async () => {
-      const key = 'key';
-
-      getList.mockRestore();
-      getList.mockReturnValue(datasources);
-
-      const tester = reduxTester<RootReducerType>()
-        .givenRootReducer(getRootReducer())
-        .whenActionIsDispatched(initAdHocVariableEditor(key));
-
-      tester.thenDispatchedActionsShouldEqual(
-        toKeyedAction(key, changeVariableEditorExtended({ dataSources: expectedDatasources }))
-      );
-    });
-  });
-
   describe('when changeVariableDatasource is dispatched with unsupported datasource', () => {
     it('then correct actions are dispatched', async () => {
       const key = 'key';
@@ -441,7 +415,6 @@ describe('adhoc actions', () => {
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(createAddVariableAction(variable))
-        .whenActionIsDispatched(initAdHocVariableEditor(key))
         .whenAsyncActionIsDispatched(changeVariableDatasource(toKeyedVariableIdentifier(variable), datasource), true);
 
       tester.thenDispatchedActionsShouldEqual(
@@ -453,7 +426,6 @@ describe('adhoc actions', () => {
           key,
           changeVariableEditorExtended({
             infoText: 'This data source does not support ad hoc filters yet.',
-            dataSources: expectedDatasources,
           })
         )
       );
@@ -482,7 +454,6 @@ describe('adhoc actions', () => {
       const tester = await reduxTester<RootReducerType>()
         .givenRootReducer(getRootReducer())
         .whenActionIsDispatched(createAddVariableAction(variable))
-        .whenActionIsDispatched(initAdHocVariableEditor(key))
         .whenAsyncActionIsDispatched(changeVariableDatasource(toKeyedVariableIdentifier(variable), datasource), true);
 
       tester.thenDispatchedActionsShouldEqual(
@@ -490,7 +461,7 @@ describe('adhoc actions', () => {
           key,
           changeVariableProp(toVariablePayload(variable, { propName: 'datasource', propValue: datasource }))
         ),
-        toKeyedAction(key, changeVariableEditorExtended({ infoText: loadingText, dataSources: expectedDatasources }))
+        toKeyedAction(key, changeVariableEditorExtended({ infoText: loadingText }))
       );
     });
   });

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

@@ -1,12 +1,11 @@
 import { cloneDeep } from 'lodash';
 
-import { DataSourceRef, getDataSourceRef } from '@grafana/data';
+import { DataSourceRef } from '@grafana/data';
 import { getDatasourceSrv } from 'app/features/plugins/datasource_srv';
 import { AdHocVariableFilter, AdHocVariableModel } from 'app/features/variables/types';
 import { StoreState, ThunkResult } from 'app/types';
 
 import { changeVariableEditorExtended } from '../editor/reducer';
-import { getAdhocVariableEditorState } from '../editor/selectors';
 import { isAdHoc } from '../guard';
 import { variableUpdated } from '../state/actions';
 import { toKeyedAction } from '../state/keyedVariablesReducer';
@@ -101,8 +100,6 @@ export const changeVariableDatasource = (
   datasource?: DataSourceRef
 ): ThunkResult<void> => {
   return async (dispatch, getState) => {
-    const { editor } = getVariablesState(identifier.rootStateKey, getState());
-    const extended = getAdhocVariableEditorState(editor);
     const variable = getVariable(identifier, getState());
     dispatch(
       toKeyedAction(
@@ -123,42 +120,12 @@ export const changeVariableDatasource = (
         identifier.rootStateKey,
         changeVariableEditorExtended({
           infoText: message,
-          dataSources: extended?.dataSources ?? [],
         })
       )
     );
   };
 };
 
-export const initAdHocVariableEditor =
-  (key: string): ThunkResult<void> =>
-  (dispatch) => {
-    const dataSources = getDatasourceSrv().getList({ metrics: true, variables: true });
-    const selectable = dataSources.reduce(
-      (all: Array<{ text: string; value: DataSourceRef | null }>, ds) => {
-        if (ds.meta.mixed) {
-          return all;
-        }
-
-        const text = ds.isDefault ? `${ds.name} (default)` : ds.name;
-        const value = getDataSourceRef(ds);
-        all.push({ text, value });
-
-        return all;
-      },
-      [{ text: '', value: {} }]
-    );
-
-    dispatch(
-      toKeyedAction(
-        key,
-        changeVariableEditorExtended({
-          dataSources: selectable,
-        })
-      )
-    );
-  };
-
 const createAdHocVariable = (options: AdHocTableOptions): ThunkResult<void> => {
   return (dispatch, getState) => {
     const key = getLastKey(getState());

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

@@ -1,13 +1,12 @@
 import { createSlice, PayloadAction } from '@reduxjs/toolkit';
 
-import { DataSourceApi, DataSourceRef } from '@grafana/data';
+import { DataSourceApi } from '@grafana/data';
 
 import { VariablePayload } from '../state/types';
 import { VariableQueryEditorType } from '../types';
 
 export interface AdHocVariableEditorState {
   infoText?: string;
-  dataSources: Array<{ text: string; value: DataSourceRef | null }>;
 }
 
 export interface DataSourceVariableEditorState {

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

@@ -14,10 +14,7 @@ import {
 } from './selectors';
 
 const adhocExtended: AdHocVariableEditorState = {
-  dataSources: [
-    { text: 'Prometheus', value: null }, // default datasource
-    { text: 'Loki', value: { type: 'loki-ds', uid: 'abc' } },
-  ],
+  infoText: 'infoText',
 };
 
 const datasourceExtended: DataSourceVariableEditorState = {

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

@@ -9,7 +9,7 @@ import {
  * Narrows generic variable editor state down to specific Adhoc variable extended editor state
  */
 export function getAdhocVariableEditorState(editorState: VariableEditorState): AdHocVariableEditorState | null {
-  if (editorState.extended && 'dataSources' in editorState.extended) {
+  if (editorState.extended && 'infoText' in editorState.extended) {
     return editorState.extended;
   }