[Multiple Datasource] Enhanced data source selector with default data…

…source shows as first choice (#6293)

* edit datasource selector to show dafault datasource and return default to customer

Signed-off-by: Yuanqi(Ella) Zhu <zhyuanqi@amazon.com>

* Update src/plugins/data_source_management/public/components/data_source_selector/data_source_selector.tsx

Co-authored-by: Lu Yu <nluyu@amazon.com>
Signed-off-by: Yuanqi(Ella) Zhu <53279298+zhyuanqi@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Lu Yu <nluyu@amazon.com>
Signed-off-by: Yuanqi(Ella) Zhu <53279298+zhyuanqi@users.noreply.github.com>

* fix syntax

Signed-off-by: Yuanqi(Ella) Zhu <zhyuanqi@amazon.com>

---------

Signed-off-by: Yuanqi(Ella) Zhu <zhyuanqi@amazon.com>
Signed-off-by: Yuanqi(Ella) Zhu <53279298+zhyuanqi@users.noreply.github.com>
Co-authored-by: Lu Yu <nluyu@amazon.com>
(cherry picked from commit 557fbcf)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

src/plugins/data_source_management/public/components/data_source_selector/create_data_source_selector.test.tsx

@@ -7,9 +7,11 @@ import { SavedObjectsClientContract } from '../../../../../core/public';
 import { notificationServiceMock } from '../../../../../core/public/mocks';
 import React from 'react';
 import { render } from '@testing-library/react';
+import { coreMock } from '../../../../../core/public/mocks';
 
 describe('create data source selector', () => {
   let client: SavedObjectsClientContract;
+  const { uiSettings } = coreMock.createSetup();
   const { toasts } = notificationServiceMock.createStartContract();
 
   beforeEach(() => {
@@ -27,7 +29,7 @@ describe('create data source selector', () => {
       hideLocalCluster: false,
       fullWidth: false,
     };
-    const TestComponent = createDataSourceSelector();
+    const TestComponent = createDataSourceSelector(uiSettings);
     const component = render(<TestComponent {...props} />);
     expect(component).toMatchSnapshot();
     expect(client.find).toBeCalledWith({

src/plugins/data_source_management/public/components/data_source_selector/create_data_source_selector.tsx

@@ -4,8 +4,11 @@
  */
 
 import React from 'react';
+import { IUiSettingsClient } from 'src/core/public';
 import { DataSourceSelector, DataSourceSelectorProps } from './data_source_selector';
 
-export function createDataSourceSelector() {
-  return (props: DataSourceSelectorProps) => <DataSourceSelector {...props} />;
+export function createDataSourceSelector(uiSettings: IUiSettingsClient) {
+  return (props: DataSourceSelectorProps) => (
+    <DataSourceSelector {...props} uiSettings={uiSettings} />
+  );
 }

src/plugins/data_source_management/public/components/data_source_selector/data_source_selector.test.tsx

@@ -8,8 +8,13 @@ import { DataSourceSelector } from './data_source_selector';
 import { SavedObjectsClientContract } from '../../../../../core/public';
 import { notificationServiceMock } from '../../../../../core/public/mocks';
 import React from 'react';
-import { getDataSourcesWithFieldsResponse, mockResponseForSavedObjectsCalls } from '../../mocks';
+import {
+  getDataSourcesWithFieldsResponse,
+  mockManagementPlugin,
+  mockResponseForSavedObjectsCalls,
+} from '../../mocks';
 import { AuthType } from 'src/plugins/data_source/common/data_sources';
+import * as utils from '../utils';
 
 describe('DataSourceSelector', () => {
   let component: ShallowWrapper<any, Readonly<{}>, React.Component<{}, {}, any>>;
@@ -69,8 +74,11 @@ describe('DataSourceSelector: check dataSource options', () => {
   let client: SavedObjectsClientContract;
   const { toasts } = notificationServiceMock.createStartContract();
   const nextTick = () => new Promise((res) => process.nextTick(res));
+  const mockedContext = mockManagementPlugin.createDataSourceManagementContext();
+  const uiSettings = mockedContext.uiSettings;
 
   beforeEach(async () => {
+    jest.clearAllMocks();
     client = {
       find: jest.fn().mockResolvedValue([]),
     } as any;
@@ -168,6 +176,47 @@ describe('DataSourceSelector: check dataSource options', () => {
     component.instance().componentDidMount!();
     await nextTick();
     expect(component).toMatchSnapshot();
-    expect(toasts.addWarning).toBeCalledTimes(0);
+    expect(toasts.addWarning).toHaveBeenCalled();
+  });
+
+  it('should get default datasource if uiSettings exists', async () => {
+    spyOn(uiSettings, 'get').and.returnValue('test1');
+    spyOn(utils, 'getFilteredDataSources').and.returnValue([]);
+    spyOn(utils, 'getDefaultDataSource').and.returnValue([]);
+    component = shallow(
+      <DataSourceSelector
+        savedObjectsClient={client}
+        notifications={toasts}
+        onSelectedDataSource={jest.fn()}
+        disabled={false}
+        hideLocalCluster={false}
+        fullWidth={false}
+        uiSettings={uiSettings}
+      />
+    );
+
+    component.instance().componentDidMount!();
+    await nextTick();
+    expect(component).toMatchSnapshot();
+    expect(uiSettings.get).toBeCalledWith('defaultDataSource', null);
+    expect(utils.getFilteredDataSources).toHaveBeenCalled();
+    expect(utils.getDefaultDataSource).toHaveBeenCalled();
+    expect(toasts.addWarning).toHaveBeenCalled();
+  });
+
+  it('should not render options with default badge when id does not matches defaultDataSource', () => {
+    component = shallow(
+      <DataSourceSelector
+        savedObjectsClient={client}
+        notifications={toasts}
+        onSelectedDataSource={jest.fn()}
+        disabled={false}
+        hideLocalCluster={false}
+        fullWidth={false}
+        uiSettings={uiSettings}
+      />
+    );
+    expect(component).toMatchSnapshot();
+    expect(component.find('EuiComboBox').exists()).toBe(true);
   });
 });

src/plugins/data_source_management/public/components/data_source_selector/data_source_selector.tsx

@@ -5,9 +5,10 @@
 
 import React from 'react';
 import { i18n } from '@osd/i18n';
-import { EuiComboBox } from '@elastic/eui';
+import { EuiComboBox, EuiBadge, EuiFlexItem, EuiFlexGroup } from '@elastic/eui';
 import { SavedObjectsClientContract, ToastsStart, SavedObject } from 'opensearch-dashboards/public';
-import { getDataSourcesWithFields } from '../utils';
+import { IUiSettingsClient } from 'src/core/public';
+import { getDataSourcesWithFields, getDefaultDataSource, getFilteredDataSources } from '../utils';
 import { DataSourceAttributes } from '../../types';
 
 export const LocalCluster: DataSourceOption = {
@@ -29,11 +30,13 @@ export interface DataSourceSelectorProps {
   removePrepend?: boolean;
   dataSourceFilter?: (dataSource: SavedObject<DataSourceAttributes>) => boolean;
   compressed?: boolean;
+  uiSettings?: IUiSettingsClient;
 }
 
 interface DataSourceSelectorState {
   selectedOption: DataSourceOption[];
   allDataSources: Array<SavedObject<DataSourceAttributes>>;
+  defaultDataSource: string | null;
 }
 
 export interface DataSourceOption {
@@ -53,6 +56,7 @@ export class DataSourceSelector extends React.Component<
 
     this.state = {
       allDataSources: [],
+      defaultDataSource: '',
       selectedOption: this.props.defaultOption
         ? this.props.defaultOption
         : this.props.hideLocalCluster
@@ -67,6 +71,13 @@ export class DataSourceSelector extends React.Component<
 
   async componentDidMount() {
     this._isMounted = true;
+
+    const currentDefaultDataSource = this.props.uiSettings?.get('defaultDataSource', null) ?? null;
+    this.setState({
+      ...this.state,
+      defaultDataSource: currentDefaultDataSource,
+    });
+
     getDataSourcesWithFields(this.props.savedObjectsClient, ['id', 'title', 'auth.type'])
       .then((fetchedDataSources) => {
         if (fetchedDataSources?.length) {
@@ -76,6 +87,25 @@ export class DataSourceSelector extends React.Component<
             allDataSources: fetchedDataSources,
           });
         }
+        const dataSources = getFilteredDataSources(
+          this.state.allDataSources,
+          this.props.dataSourceFilter
+        );
+        const selectedDataSource = getDefaultDataSource(
+          dataSources,
+          LocalCluster,
+          this.props.uiSettings,
+          this.props.hideLocalCluster,
+          this.props.defaultOption
+        );
+        if (selectedDataSource.length === 0) {
+          this.props.notifications.addWarning('No connected data source available.');
+        } else {
+          this.props.onSelectedDataSource(selectedDataSource);
+          this.setState({
+            selectedOption: selectedDataSource,
+          });
+        }
       })
       .catch(() => {
         this.props.notifications.addWarning(
@@ -100,9 +130,11 @@ export class DataSourceSelector extends React.Component<
         ? 'Select a data source'
         : this.props.placeholderText;
 
-    const dataSources = this.props.dataSourceFilter
-      ? this.state.allDataSources.filter((ds) => this.props.dataSourceFilter!(ds))
-      : this.state.allDataSources;
+    // The filter condition can be changed, thus we filter again here to make sure each time we will get the filtered data sources before rendering
+    const dataSources = getFilteredDataSources(
+      this.state.allDataSources,
+      this.props.dataSourceFilter
+    );
 
     const options = dataSources.map((ds) => ({ id: ds.id, label: ds.attributes?.title || '' }));
     if (!this.props.hideLocalCluster) {
@@ -140,6 +172,16 @@ export class DataSourceSelector extends React.Component<
         isDisabled={this.props.disabled}
         fullWidth={this.props.fullWidth || false}
         data-test-subj={'dataSourceSelectorComboBox'}
+        renderOption={(option) => (
+          <EuiFlexGroup alignItems="center">
+            <EuiFlexItem grow={1}>{option.label}</EuiFlexItem>
+            {option.id === this.state.defaultDataSource && (
+              <EuiFlexItem grow={false}>
+                <EuiBadge iconSide="left">Default</EuiBadge>
+              </EuiFlexItem>
+            )}
+          </EuiFlexGroup>
+        )}
       />
     );
   }