From 9cbfa64f07d3478c4daa04a960d909dc41648c9b Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Tue, 16 Apr 2024 10:43:21 -0700
Subject: [PATCH] Refactor read-only component to cover more edge cases (#6416)
 (#6466)

* Refactor read-only component to cover more edge case

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

* fix test

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

* fix snapshot

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

---------

Signed-off-by: Yuanqi(Ella) Zhu <zhyuanqi@amazon.com>
(cherry picked from commit b9acf57def95e1a77e76d19f6df34f408ae8dd5a)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../components/data_source_view_example.tsx   |  37 +++--
 .../public/components/constants.tsx           |  14 ++
 .../data_source_menu.test.tsx.snap            |   4 +
 .../data_source_menu/data_source_menu.tsx     |  13 +-
 .../components/data_source_menu/types.ts      |   2 +
 .../data_source_view.test.tsx.snap            | 150 +++++++++++++++++-
 .../data_source_view.test.tsx                 | 104 +++++++++---
 .../data_source_view/data_source_view.tsx     |  77 ++++++---
 .../public/components/utils.test.ts           |  31 +++-
 .../public/components/utils.ts                |  42 +++--
 .../data_source_management/public/mocks.ts    |  11 ++
 11 files changed, 414 insertions(+), 71 deletions(-)
 create mode 100644 src/plugins/data_source_management/public/components/constants.tsx

diff --git a/examples/multiple_data_source_examples/public/components/data_source_view_example.tsx b/examples/multiple_data_source_examples/public/components/data_source_view_example.tsx
index 11ebe9ee343a..b80a65bb1b1a 100644
--- a/examples/multiple_data_source_examples/public/components/data_source_view_example.tsx
+++ b/examples/multiple_data_source_examples/public/components/data_source_view_example.tsx
@@ -2,7 +2,7 @@
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
  */
-import React from 'react';
+import React, { useState, useMemo } from 'react';
 import {
   EuiBasicTable,
   EuiPageBody,
@@ -34,8 +34,11 @@ export const DataSourceViewExample = ({
   dataSourceEnabled,
   setActionMenu,
   dataSourceManagement,
+  notifications,
+  savedObjects,
 }: DataSourceViewExampleProps) => {
   const DataSourceMenu = dataSourceManagement.ui.getDataSourceMenu<DataSourceViewConfig>();
+  const [selectedDataSources, setSelectedDataSources] = useState<string[]>([]);
   const data: ComponentProp[] = [
     {
       name: 'savedObjects',
@@ -68,19 +71,31 @@ export const DataSourceViewExample = ({
     },
   ];
 
+  const renderDataSourceComponent = useMemo(() => {
+    return (
+      <DataSourceMenu
+        setMenuMountPoint={setActionMenu}
+        componentType={'DataSourceView'}
+        componentConfig={{
+          notifications,
+          savedObjects: savedObjects.client,
+          fullWidth: false,
+          activeOption: [{ id: '' }],
+          dataSourceFilter: (ds) => {
+            return true;
+          },
+          onSelectedDataSources: (ds) => {
+            setSelectedDataSources(ds);
+          },
+        }}
+      />
+    );
+  }, [setActionMenu, notifications, savedObjects]);
+
   return (
     <EuiPageBody component="main">
       <EuiPageHeader>
-        {dataSourceEnabled && (
-          <DataSourceMenu
-            setMenuMountPoint={setActionMenu}
-            componentType={'DataSourceView'}
-            componentConfig={{
-              fullWidth: false,
-              activeOption: [{ id: 'example id', label: 'example data source' }],
-            }}
-          />
-        )}
+        {dataSourceEnabled && renderDataSourceComponent}
         <EuiPageHeaderSection>
           <EuiTitle size="l">
             <h1>Data Source View Example</h1>
diff --git a/src/plugins/data_source_management/public/components/constants.tsx b/src/plugins/data_source_management/public/components/constants.tsx
new file mode 100644
index 000000000000..0d22aed50179
--- /dev/null
+++ b/src/plugins/data_source_management/public/components/constants.tsx
@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { i18n } from '@osd/i18n';
+import { DataSourceOption } from './data_source_menu/types';
+
+export const LocalCluster: DataSourceOption = {
+  label: i18n.translate('dataSource.localCluster', {
+    defaultMessage: 'Local cluster',
+  }),
+  id: '',
+};
diff --git a/src/plugins/data_source_management/public/components/data_source_menu/__snapshots__/data_source_menu.test.tsx.snap b/src/plugins/data_source_management/public/components/data_source_menu/__snapshots__/data_source_menu.test.tsx.snap
index 8f8ea1a0f55d..517f8de733ba 100644
--- a/src/plugins/data_source_management/public/components/data_source_menu/__snapshots__/data_source_menu.test.tsx.snap
+++ b/src/plugins/data_source_management/public/components/data_source_menu/__snapshots__/data_source_menu.test.tsx.snap
@@ -3,6 +3,7 @@
 exports[`DataSourceMenu can render data source view when only pass id in the activeOption 1`] = `
 <DataSourceView
   fullWidth={true}
+  hideLocalCluster={false}
   notifications={
     Object {
       "add": [MockFunction],
@@ -33,6 +34,7 @@ exports[`DataSourceMenu can render data source view when only pass id in the act
 exports[`DataSourceMenu can render data source view when provide activeOption 1`] = `
 <DataSourceView
   fullWidth={true}
+  hideLocalCluster={false}
   notifications={
     Object {
       "add": [MockFunction],
@@ -408,6 +410,7 @@ exports[`DataSourceMenu should render data source selectable only with local clu
 exports[`DataSourceMenu should render data source view if not pass saved object or notification 1`] = `
 <DataSourceView
   fullWidth={true}
+  hideLocalCluster={false}
   notifications={
     Object {
       "add": [MockFunction],
@@ -426,6 +429,7 @@ exports[`DataSourceMenu should render data source view if not pass saved object
 exports[`DataSourceMenu should render data source view only 1`] = `
 <DataSourceView
   fullWidth={true}
+  hideLocalCluster={false}
   notifications={
     Object {
       "add": [MockFunction],
diff --git a/src/plugins/data_source_management/public/components/data_source_menu/data_source_menu.tsx b/src/plugins/data_source_management/public/components/data_source_menu/data_source_menu.tsx
index fdc5d2c7508b..9511090e6c5a 100644
--- a/src/plugins/data_source_management/public/components/data_source_menu/data_source_menu.tsx
+++ b/src/plugins/data_source_management/public/components/data_source_menu/data_source_menu.tsx
@@ -22,13 +22,24 @@ export function DataSourceMenu<T>(props: DataSourceMenuProps<T>): ReactElement |
   const { componentType, componentConfig, uiSettings, hideLocalCluster } = props;
 
   function renderDataSourceView(config: DataSourceViewConfig): ReactElement | null {
-    const { activeOption, fullWidth, savedObjects, notifications } = config;
+    const {
+      activeOption,
+      fullWidth,
+      savedObjects,
+      notifications,
+      dataSourceFilter,
+      onSelectedDataSources,
+    } = config;
     return (
       <DataSourceView
         fullWidth={fullWidth}
         selectedOption={activeOption}
         savedObjectsClient={savedObjects}
         notifications={notifications?.toasts}
+        hideLocalCluster={hideLocalCluster || false}
+        dataSourceFilter={dataSourceFilter}
+        onSelectedDataSources={onSelectedDataSources}
+        uiSettings={uiSettings}
       />
     );
   }
diff --git a/src/plugins/data_source_management/public/components/data_source_menu/types.ts b/src/plugins/data_source_management/public/components/data_source_menu/types.ts
index 7c32cd3cd3df..a37d7d0dd0c3 100644
--- a/src/plugins/data_source_management/public/components/data_source_menu/types.ts
+++ b/src/plugins/data_source_management/public/components/data_source_menu/types.ts
@@ -48,6 +48,8 @@ export interface DataSourceViewConfig extends DataSourceBaseConfig {
   activeOption: DataSourceOption[];
   savedObjects?: SavedObjectsClientContract;
   notifications?: NotificationsStart;
+  dataSourceFilter?: (dataSource: SavedObject<DataSourceAttributes>) => boolean;
+  onSelectedDataSources?: (dataSources: DataSourceOption[]) => void;
 }
 
 export interface DataSourceAggregatedViewConfig extends DataSourceBaseConfig {
diff --git a/src/plugins/data_source_management/public/components/data_source_view/__snapshots__/data_source_view.test.tsx.snap b/src/plugins/data_source_management/public/components/data_source_view/__snapshots__/data_source_view.test.tsx.snap
index 183ebc563501..9e0f04cea122 100644
--- a/src/plugins/data_source_management/public/components/data_source_view/__snapshots__/data_source_view.test.tsx.snap
+++ b/src/plugins/data_source_management/public/components/data_source_view/__snapshots__/data_source_view.test.tsx.snap
@@ -1,6 +1,6 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`DataSourceView should call getDataSourceById when only pass id no label 1`] = `
+exports[`DataSourceView Should render successfully when provided datasource has not been filtered out 1`] = `
 <Fragment>
   <EuiButtonEmpty
     aria-label="dataSourceViewMenuButton"
@@ -51,7 +51,99 @@ exports[`DataSourceView should call getDataSourceById when only pass id no label
 </Fragment>
 `;
 
-exports[`DataSourceView should call notification warning when there is data source fetch error 1`] = `
+exports[`DataSourceView Should return error when provided datasource has been filtered out 1`] = `
+<Fragment>
+  <EuiButtonEmpty
+    aria-label="dataSourceViewMenuButton"
+    className="euiHeaderLink"
+    data-test-subj="dataSourceViewContextMenuHeaderLink"
+    disabled={true}
+    iconSide="left"
+    iconType="database"
+    size="s"
+  />
+  <EuiPopover
+    anchorPosition="downLeft"
+    button={
+      <EuiButtonIcon
+        aria-label="Next"
+        display="empty"
+        iconType="iInCircle"
+        onClick={[Function]}
+      />
+    }
+    closePopover={[Function]}
+    display="inlineBlock"
+    hasArrow={true}
+    id="dataSourceViewContextMenuPopover"
+    isOpen={false}
+    ownFocus={true}
+    panelPaddingSize="none"
+  >
+    <EuiContextMenu
+      initialPanelId={0}
+      panels={
+        Array [
+          Object {
+            "id": 0,
+            "items": Array [],
+            "title": "Selected data source",
+          },
+        ]
+      }
+      size="m"
+    />
+  </EuiPopover>
+</Fragment>
+`;
+
+exports[`DataSourceView When selected option is local cluster and hide local Cluster is true, should return error 1`] = `
+<Fragment>
+  <EuiButtonEmpty
+    aria-label="dataSourceViewMenuButton"
+    className="euiHeaderLink"
+    data-test-subj="dataSourceViewContextMenuHeaderLink"
+    disabled={true}
+    iconSide="left"
+    iconType="database"
+    size="s"
+  />
+  <EuiPopover
+    anchorPosition="downLeft"
+    button={
+      <EuiButtonIcon
+        aria-label="Next"
+        display="empty"
+        iconType="iInCircle"
+        onClick={[Function]}
+      />
+    }
+    closePopover={[Function]}
+    display="inlineBlock"
+    hasArrow={true}
+    id="dataSourceViewContextMenuPopover"
+    isOpen={false}
+    ownFocus={true}
+    panelPaddingSize="none"
+  >
+    <EuiContextMenu
+      initialPanelId={0}
+      panels={
+        Array [
+          Object {
+            "id": 0,
+            "items": Array [],
+            "title": "Selected data source",
+          },
+        ]
+      }
+      size="m"
+    />
+  </EuiPopover>
+</Fragment>
+`;
+
+exports[`DataSourceView should call getDataSourceById when only pass id with no label 1`] = `
 <Fragment>
   <EuiButtonEmpty
     aria-label="dataSourceViewMenuButton"
@@ -102,6 +194,52 @@ exports[`DataSourceView should call notification warning when there is data sour
 </Fragment>
 `;
 
+exports[`DataSourceView should call notification warning when there is data source fetch error 1`] = `
+<Fragment>
+  <EuiButtonEmpty
+    aria-label="dataSourceViewMenuButton"
+    className="euiHeaderLink"
+    data-test-subj="dataSourceViewContextMenuHeaderLink"
+    disabled={true}
+    iconSide="left"
+    iconType="database"
+    size="s"
+  />
+  <EuiPopover
+    anchorPosition="downLeft"
+    button={
+      <EuiButtonIcon
+        aria-label="Next"
+        display="empty"
+        iconType="iInCircle"
+        onClick={[Function]}
+      />
+    }
+    closePopover={[Function]}
+    display="inlineBlock"
+    hasArrow={true}
+    id="dataSourceViewContextMenuPopover"
+    isOpen={false}
+    ownFocus={true}
+    panelPaddingSize="none"
+  >
+    <EuiContextMenu
+      initialPanelId={0}
+      panels={
+        Array [
+          Object {
+            "id": 0,
+            "items": Array [],
+            "title": "Selected data source",
+          },
+        ]
+      }
+      size="m"
+    />
+  </EuiPopover>
+</Fragment>
+`;
+
 exports[`DataSourceView should render normally with local cluster not hidden 1`] = `
 <Fragment>
   <EuiButtonEmpty
@@ -177,9 +315,7 @@ Object {
           />
           <span
             class="euiButtonEmpty__text"
-          >
-            test1
-          </span>
+          />
         </span>
       </button>
       <div
@@ -223,9 +359,7 @@ Object {
         />
         <span
           class="euiButtonEmpty__text"
-        >
-          test1
-        </span>
+        />
       </span>
     </button>
     <div
diff --git a/src/plugins/data_source_management/public/components/data_source_view/data_source_view.test.tsx b/src/plugins/data_source_management/public/components/data_source_view/data_source_view.test.tsx
index a3cfd5223f66..a518fee29b78 100644
--- a/src/plugins/data_source_management/public/components/data_source_view/data_source_view.test.tsx
+++ b/src/plugins/data_source_management/public/components/data_source_view/data_source_view.test.tsx
@@ -3,17 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ShallowWrapper, shallow } from 'enzyme';
+import { ShallowWrapper, shallow, mount } from 'enzyme';
 import React from 'react';
 import { DataSourceView } from './data_source_view';
 import { SavedObjectsClientContract } from 'opensearch-dashboards/public';
 import { notificationServiceMock } from '../../../../../core/public/mocks';
-import {
-  getSingleDataSourceResponse,
-  mockErrorResponseForSavedObjectsCalls,
-  mockResponseForSavedObjectsCalls,
-} from '../../mocks';
+import { getSingleDataSourceResponse, mockResponseForSavedObjectsCalls } from '../../mocks';
 import { render } from '@testing-library/react';
+import * as utils from '../utils';
 
 describe('DataSourceView', () => {
   let component: ShallowWrapper<any, Readonly<{}>, React.Component<{}, {}, any>>;
@@ -21,6 +18,7 @@ describe('DataSourceView', () => {
   const { toasts } = notificationServiceMock.createStartContract();
 
   beforeEach(() => {
+    jest.clearAllMocks();
     client = {
       get: jest.fn().mockResolvedValue([]),
     } as any;
@@ -28,48 +26,114 @@ describe('DataSourceView', () => {
   });
 
   it('should render normally with local cluster not hidden', () => {
+    spyOn(utils, 'getDataSourceById').and.returnValue([{ id: 'test1', label: 'test1' }]);
     component = shallow(
-      <DataSourceView fullWidth={false} selectedOption={[{ id: 'test1', label: 'test1' }]} />
+      <DataSourceView
+        fullWidth={false}
+        selectedOption={[{ id: 'test1', label: 'test1' }]}
+        hideLocalCluster={false}
+        onSelectedDataSources={jest.fn()}
+      />
     );
     expect(component).toMatchSnapshot();
     expect(toasts.addWarning).toBeCalledTimes(0);
   });
-  it('should show popover when click on button', async () => {
-    const container = render(
-      <DataSourceView fullWidth={false} selectedOption={[{ id: 'test1', label: 'test1' }]} />
+  it('When selected option is local cluster and hide local Cluster is true, should return error', () => {
+    component = shallow(
+      <DataSourceView
+        fullWidth={false}
+        selectedOption={[{ id: '' }]}
+        hideLocalCluster={true}
+        notifications={toasts}
+        onSelectedDataSources={jest.fn()}
+      />
     );
-    const button = await container.findByTestId('dataSourceViewContextMenuHeaderLink');
-    button.click();
-    expect(container).toMatchSnapshot();
+    expect(component).toMatchSnapshot();
+    expect(toasts.addWarning).toBeCalledTimes(1);
+  });
+  it('Should return error when provided datasource has been filtered out', async () => {
+    component = shallow(
+      <DataSourceView
+        fullWidth={false}
+        selectedOption={[{ id: 'test1', label: 'test1' }]}
+        hideLocalCluster={false}
+        notifications={toasts}
+        onSelectedDataSources={jest.fn()}
+        dataSourceFilter={(ds) => {
+          return false;
+        }}
+      />
+    );
+    expect(component).toMatchSnapshot();
+    expect(toasts.addWarning).toBeCalledTimes(1);
   });
-  it('should call getDataSourceById when only pass id no label', async () => {
+  it('Should render successfully when provided datasource has not been filtered out', async () => {
+    spyOn(utils, 'getDataSourceById').and.returnValue([{ id: 'test1', label: 'test1' }]);
+    component = shallow(
+      <DataSourceView
+        fullWidth={false}
+        selectedOption={[{ id: 'test1' }]}
+        hideLocalCluster={false}
+        notifications={toasts}
+        onSelectedDataSources={jest.fn()}
+        dataSourceFilter={(ds) => {
+          return true;
+        }}
+      />
+    );
+    expect(component).toMatchSnapshot();
+    expect(toasts.addWarning).toBeCalledTimes(0);
+    expect(utils.getDataSourceById).toBeCalledTimes(1);
+  });
+  it('should call getDataSourceById when only pass id with no label', async () => {
+    spyOn(utils, 'getDataSourceById').and.returnValue([{ id: 'test1', label: 'test1' }]);
     component = shallow(
       <DataSourceView
         fullWidth={false}
         selectedOption={[{ id: 'test1' }]}
         savedObjectsClient={client}
         notifications={toasts}
+        hideLocalCluster={false}
+        onSelectedDataSources={jest.fn()}
       />
     );
     expect(component).toMatchSnapshot();
-    expect(client.get).toBeCalledWith('data-source', 'test1');
+    expect(utils.getDataSourceById).toBeCalledTimes(1);
     expect(toasts.addWarning).toBeCalledTimes(0);
   });
   it('should call notification warning when there is data source fetch error', async () => {
-    jest.mock('../utils', () => ({
-      getDataSourceById: jest.fn(),
-    }));
+    spyOn(utils, 'getDataSourceById').and.throwError('Data source is not available');
     component = shallow(
       <DataSourceView
         fullWidth={false}
         selectedOption={[{ id: 'test1' }]}
         savedObjectsClient={client}
         notifications={toasts}
+        hideLocalCluster={false}
+        onSelectedDataSources={jest.fn()}
       />
     );
     expect(component).toMatchSnapshot();
-    mockErrorResponseForSavedObjectsCalls(client, 'get');
     expect(toasts.addWarning).toBeCalledTimes(1);
-    expect(toasts.addWarning).toBeCalledWith(`Data source with id test1 is not available`);
+    expect(utils.getDataSourceById).toBeCalledTimes(1);
+  });
+
+  it('should show popover when click on button', async () => {
+    const onSelectedDataSource = jest.fn();
+    spyOn(utils, 'getDataSourceById').and.returnValue([{ id: 'test1', label: 'test1' }]);
+    spyOn(utils, 'handleDataSourceFetchError').and.returnValue('');
+    const container = render(
+      <DataSourceView
+        savedObjectsClient={client}
+        notifications={toasts}
+        selectedOption={[{ id: 'test1' }]}
+        onSelectedDataSources={onSelectedDataSource}
+        hideLocalCluster={false}
+        fullWidth={false}
+      />
+    );
+    const button = await container.findByTestId('dataSourceViewContextMenuHeaderLink');
+    button.click();
+    expect(container).toMatchSnapshot();
   });
 });
diff --git a/src/plugins/data_source_management/public/components/data_source_view/data_source_view.tsx b/src/plugins/data_source_management/public/components/data_source_view/data_source_view.tsx
index 005d54e216c6..ab48e925d18f 100644
--- a/src/plugins/data_source_management/public/components/data_source_view/data_source_view.tsx
+++ b/src/plugins/data_source_management/public/components/data_source_view/data_source_view.tsx
@@ -7,15 +7,25 @@ import React from 'react';
 import { i18n } from '@osd/i18n';
 import { EuiPopover, EuiButtonEmpty, EuiButtonIcon, EuiContextMenu } from '@elastic/eui';
 import { SavedObjectsClientContract, ToastsStart } from 'opensearch-dashboards/public';
+import { IUiSettingsClient } from 'src/core/public';
 import { DataSourceOption } from '../data_source_menu/types';
-import { getDataSourceById } from '../utils';
+import {
+  getDataSourceById,
+  handleDataSourceFetchError,
+  handleNoAvailableDataSourceError,
+} from '../utils';
 import { MenuPanelItem } from '../../types';
+import { LocalCluster } from '../constants';
 
 interface DataSourceViewProps {
   fullWidth: boolean;
   selectedOption: DataSourceOption[];
+  hideLocalCluster: boolean;
   savedObjectsClient?: SavedObjectsClientContract;
   notifications?: ToastsStart;
+  uiSettings?: IUiSettingsClient;
+  dataSourceFilter?: (dataSource: any) => boolean;
+  onSelectedDataSources?: (dataSources: DataSourceOption[]) => void;
 }
 
 interface DataSourceViewState {
@@ -40,33 +50,60 @@ export class DataSourceView extends React.Component<DataSourceViewProps, DataSou
   }
   async componentDidMount() {
     this._isMounted = true;
-    const selectedOption = this.props.selectedOption;
-    // early return if not possible to fetch data source
 
+    const selectedOption = this.props.selectedOption;
     const option = selectedOption[0];
     const optionId = option.id;
+
+    if (optionId === '' && !this.props.hideLocalCluster) {
+      this.setState({
+        selectedOption: [LocalCluster],
+      });
+      if (this.props.onSelectedDataSources) {
+        this.props.onSelectedDataSources([LocalCluster]);
+      }
+      return;
+    }
+
+    if (
+      (optionId === '' && this.props.hideLocalCluster) ||
+      (this.props.dataSourceFilter &&
+        this.props.selectedOption.filter(this.props.dataSourceFilter).length === 0)
+    ) {
+      this.setState({
+        selectedOption: [],
+      });
+      if (this.props.onSelectedDataSources) {
+        this.props.onSelectedDataSources([]);
+      }
+      handleNoAvailableDataSourceError(this.props.notifications!);
+      return;
+    }
+
     if (!option.label) {
       try {
-        const title = (await getDataSourceById(optionId, this.props.savedObjectsClient)).title;
-        if (!title) {
-          this.props.notifications.addWarning(
-            i18n.translate('dataSource.fetchDataSourceError', {
-              defaultMessage: `Data source with id ${optionId} is not available`,
-            })
-          );
-        } else {
-          if (!this._isMounted) return;
-          this.setState({
-            selectedOption: [{ id: optionId, label: title }],
-          });
+        const selectedDataSource = await getDataSourceById(
+          optionId,
+          this.props.savedObjectsClient!
+        );
+        if (!this._isMounted) return;
+        this.setState({
+          selectedOption: [{ id: optionId, label: selectedDataSource.title }],
+        });
+        if (this.props.onSelectedDataSources) {
+          this.props.onSelectedDataSources([{ id: optionId, label: selectedDataSource.title }]);
         }
       } catch (error) {
-        this.props.notifications.addWarning(
-          i18n.translate('dataSource.fetchDataSourceError', {
-            defaultMessage: `Failed to fetch data source due to ${error}`,
-          })
-        );
+        this.setState({
+          selectedOption: [],
+        });
+        if (this.props.onSelectedDataSources) {
+          this.props.onSelectedDataSources([]);
+        }
+        handleDataSourceFetchError(this.props.notifications!);
       }
+    } else if (this.props.onSelectedDataSources) {
+      this.props.onSelectedDataSources([option]);
     }
   }
 
diff --git a/src/plugins/data_source_management/public/components/utils.test.ts b/src/plugins/data_source_management/public/components/utils.test.ts
index 63fac96be57e..397f55333ed2 100644
--- a/src/plugins/data_source_management/public/components/utils.test.ts
+++ b/src/plugins/data_source_management/public/components/utils.test.ts
@@ -18,8 +18,10 @@ import {
   setFirstDataSourceAsDefault,
   getFilteredDataSources,
   getDefaultDataSource,
+  handleDataSourceFetchError,
+  handleNoAvailableDataSourceError,
 } from './utils';
-import { coreMock } from '../../../../core/public/mocks';
+import { coreMock, notificationServiceMock } from '../../../../core/public/mocks';
 import {
   getDataSourceByIdWithCredential,
   getDataSourceByIdWithoutCredential,
@@ -32,6 +34,7 @@ import {
   getSingleDataSourceResponse,
   getDataSource,
   getDataSourceOptions,
+  getDataSourceByIdWithError,
 } from '../mocks';
 import {
   AuthType,
@@ -69,6 +72,24 @@ describe('DataSourceManagement: Utils.ts', () => {
     });
   });
 
+  describe('Handle fetch data source error', () => {
+    const { toasts } = notificationServiceMock.createStartContract();
+
+    test('should send warning when data source fetch failed', () => {
+      handleDataSourceFetchError(toasts);
+      expect(toasts.addWarning).toHaveBeenCalledWith(`Failed to fetch data source`);
+    });
+  });
+
+  describe('Handle no available data source error', () => {
+    const { toasts } = notificationServiceMock.createStartContract();
+
+    test('should  send warning when data source is not available', () => {
+      handleNoAvailableDataSourceError(toasts);
+      expect(toasts.addWarning).toHaveBeenCalledWith(`Data source is not available`);
+    });
+  });
+
   describe('Get data source by ID', () => {
     test('Success: getting data source by ID with credential', async () => {
       mockResponseForSavedObjectsCalls(savedObjects.client, 'get', getDataSourceByIdWithCredential);
@@ -99,6 +120,14 @@ describe('DataSourceManagement: Utils.ts', () => {
         expect(e).toBeTruthy();
       }
     });
+    test('failure: gets error when response contains error', async () => {
+      try {
+        mockResponseForSavedObjectsCalls(savedObjects.client, 'get', getDataSourceByIdWithError);
+        await getDataSourceById('alpha-test', savedObjects.client);
+      } catch (e) {
+        expect(e).toBeTruthy();
+      }
+    });
   });
 
   describe('Create data source', () => {
diff --git a/src/plugins/data_source_management/public/components/utils.ts b/src/plugins/data_source_management/public/components/utils.ts
index 81c0cd5c707a..e433cdff52c5 100644
--- a/src/plugins/data_source_management/public/components/utils.ts
+++ b/src/plugins/data_source_management/public/components/utils.ts
@@ -8,7 +8,9 @@ import {
   SavedObjectsClientContract,
   SavedObject,
   IUiSettingsClient,
+  ToastsStart,
 } from 'src/core/public';
+import { i18n } from '@osd/i18n';
 import {
   DataSourceAttributes,
   DataSourceTableItem,
@@ -79,6 +81,22 @@ export async function setFirstDataSourceAsDefault(
   }
 }
 
+export function handleDataSourceFetchError(notifications: ToastsStart) {
+  notifications.addWarning(
+    i18n.translate('dataSource.fetchDataSourceError', {
+      defaultMessage: `Failed to fetch data source`,
+    })
+  );
+}
+
+export function handleNoAvailableDataSourceError(notifications: ToastsStart) {
+  notifications.addWarning(
+    i18n.translate('dataSource.noAvailableDataSourceError', {
+      defaultMessage: `Data source is not available`,
+    })
+  );
+}
+
 export function getFilteredDataSources(
   dataSources: Array<SavedObject<DataSourceAttributes>>,
   filter = (ds: SavedObject<DataSourceAttributes>) => true
@@ -129,16 +147,20 @@ export async function getDataSourceById(
   id: string,
   savedObjectsClient: SavedObjectsClientContract
 ) {
-  return savedObjectsClient.get('data-source', id).then((response) => {
-    const attributes: any = response?.attributes || {};
-    return {
-      id: response.id,
-      title: attributes.title,
-      endpoint: attributes.endpoint,
-      description: attributes.description || '',
-      auth: attributes.auth,
-    };
-  });
+  const response = await savedObjectsClient.get('data-source', id);
+
+  if (!response || response.error) {
+    throw new Error('Unable to find data source');
+  }
+
+  const attributes: any = response?.attributes || {};
+  return {
+    id: response.id,
+    title: attributes.title,
+    endpoint: attributes.endpoint,
+    description: attributes.description || '',
+    auth: attributes.auth,
+  };
 }
 
 export async function createSingleDataSource(
diff --git a/src/plugins/data_source_management/public/mocks.ts b/src/plugins/data_source_management/public/mocks.ts
index d6c89dc5c7dd..3d991bd3e10a 100644
--- a/src/plugins/data_source_management/public/mocks.ts
+++ b/src/plugins/data_source_management/public/mocks.ts
@@ -315,6 +315,17 @@ export const getDataSourceByIdWithoutCredential = {
   references: [],
 };
 
+export const getDataSourceByIdWithError = {
+  attributes: {
+    ...getDataSourceByIdWithCredential.attributes,
+    Error: {
+      statusCode: 404,
+      errorMessage: 'Unable to find data source',
+    },
+  },
+  references: [],
+};
+
 export const mockResponseForSavedObjectsCalls = (
   savedObjectsClient: SavedObjectsClientContract,
   savedObjectsMethodName: 'get' | 'find' | 'create' | 'delete' | 'update',