[BUG]fix add sample notebooks

public/components/notebooks/components/helpers/modal_containers.tsx

@@ -3,23 +3,27 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useState } from 'react';
 import {
-  EuiOverlayMask,
-  EuiConfirmModal,
-  EuiSmallButton,
-  EuiSmallButtonEmpty,
   EuiCompressedFieldText,
-  EuiForm,
   EuiCompressedFormRow,
+  EuiConfirmModal,
+  EuiForm,
   EuiModal,
   EuiModalBody,
   EuiModalFooter,
   EuiModalHeader,
   EuiModalHeaderTitle,
-  EuiText,
+  EuiOverlayMask,
+  EuiSmallButton,
+  EuiSmallButtonEmpty,
   EuiSpacer,
+  EuiText,
+  EuiTitle,
 } from '@elastic/eui';
+import React, { useState } from 'react';
+import { CoreStart, SavedObjectsStart } from '../../../../../../../src/core/public';
+import { DataSourceManagementPluginSetup } from '../../../../../../../src/plugins/data_source_management/public/plugin';
+import { dataSourceFilterFn } from '../../../../../common/utils/shared';
 import { CustomInputModal } from './custom_modals/custom_input_modal';
 
 /* The file contains helper functions for modal layouts
@@ -80,8 +84,26 @@ export const getSampleNotebooksModal = (
   onCancel: (
     event?: React.KeyboardEvent<HTMLDivElement> | React.MouseEvent<HTMLButtonElement, MouseEvent>
   ) => void,
-  onConfirm: (event?: React.MouseEvent<HTMLButtonElement, MouseEvent>) => void
+  onConfirm: (event?: React.MouseEvent<HTMLButtonElement, MouseEvent>) => void,
+  dataSourceEnabled: boolean,
+  dataSourceManagement: DataSourceManagementPluginSetup,
+  savedObjectsMDSClient: SavedObjectsStart,
+  notifications: CoreStart['notifications'],
+  handleSelectedDataSourceChange: (
+    dataSourceMDSId: string | undefined,
+    dataSourceMDSLabel: string | undefined
+  ) => void
 ) => {
+  let DataSourceSelector;
+  const onSelectedDataSource = (e) => {
+    const dataConnectionId = e[0] ? e[0].id : undefined;
+    const dataConnectionLabel = e[0] ? e[0].label : undefined;
+    handleSelectedDataSourceChange(dataConnectionId, dataConnectionLabel);
+  };
+
+  if (dataSourceEnabled) {
+    DataSourceSelector = dataSourceManagement.ui.DataSourceSelector;
+  }
   return (
     <EuiOverlayMask>
       <EuiConfirmModal
@@ -92,12 +114,27 @@ export const getSampleNotebooksModal = (
         confirmButtonText="Yes"
         defaultFocusedButton="confirm"
       >
-        <EuiText size="s">
-          <p>
-            Do you want to add sample notebooks? This will also add Dashboards sample flights and
-            logs data if they have not been added.
-          </p>
-        </EuiText>
+        {dataSourceEnabled && (
+          <>
+            <EuiTitle size="s">
+              <h4>Select a Data source</h4>
+            </EuiTitle>
+            <DataSourceSelector
+              savedObjectsClient={savedObjectsMDSClient.client}
+              notifications={notifications}
+              onSelectedDataSource={onSelectedDataSource}
+              disabled={false}
+              fullWidth={false}
+              removePrepend={false}
+              dataSourceFilter={dataSourceFilterFn}
+            />
+          </>
+        )}
+        <EuiSpacer />
+        <p>
+          Do you want to add sample notebooks? This will also add Dashboards sample flights and logs
+          data if they have not been added.
+        </p>
       </EuiConfirmModal>
     </EuiOverlayMask>
   );

public/components/notebooks/components/main.tsx

@@ -95,6 +95,18 @@
 
   // Fetches path and id for all stored notebooks
   fetchNotebooks = () => {
+    if (this.props.dataSourceEnabled) {
+      // If `MDS` is enabled, only fetch from the first endpoint.
+      return this.props.http
+        .get(`${NOTEBOOKS_API_PREFIX}/savedNotebook/`)
+        .then((savedNotebooksResponse) => {
+          this.setState({ data: savedNotebooksResponse.data });
+        })
+        .catch((err) => {
+          console.error('Issue in fetching the notebooks', err.body.message);
+        });
+    }
+    // If `MDS` is not enabled /savedNotebook/ API returns notebooks stored as saved objects, and the other one returns notebooks stored as observability objects.
     return Promise.all([
       this.props.http.get(`${NOTEBOOKS_API_PREFIX}/savedNotebook/`),
       this.props.http.get(`${NOTEBOOKS_API_PREFIX}/`),
@@ -142,7 +154,7 @@
   };
 
   // Renames an existing notebook
   renameNotebook = async (editedNoteName: string, editedNoteID: string): Promise<any> => {
     if (editedNoteName.length >= 50 || editedNoteName.length === 0) {
       this.setToast('Invalid notebook name', 'danger');
       return;
@@ -285,7 +297,7 @@
         console.error(err.body.message);
       });
   };
-  addSampleNotebooks = async () => {
+  addSampleNotebooks = async (dataSourceMDSId?: string, dataSourceMDSLabel?: string) => {
     try {
       this.setState({ loading: true });
       const flights = await this.props.http
@@ -296,7 +308,17 @@
             search: 'opensearch_dashboards_sample_data_flights',
           },
         })
-        .then((resp) => resp.total === 0);
+        .then((resp) => {
+          if (resp.total === 0) {
+            return true;
+          }
+          const hasDataSourceMDSId = resp.saved_objects.some((obj) =>
+            obj.references.some((ref) => ref.type === 'data-source' && ref.id === dataSourceMDSId)
+          );
+
+          // Return true if dataSourceMDSId is not found in any references
+          return !hasDataSourceMDSId;
+        });
       const logs = await this.props.http
         .get('../api/saved_objects/_find', {
           query: {
@@ -305,46 +327,105 @@
             search: 'opensearch_dashboards_sample_data_logs',
           },
         })
-        .then((resp) => resp.total === 0);
-      if (flights || logs) this.setToast('Adding sample data. This can take some time.');
-      await Promise.all([
-        flights ? this.props.http.post('../api/sample_data/flights') : Promise.resolve(),
-        logs ? this.props.http.post('../api/sample_data/logs') : Promise.resolve(),
-      ]);
+        .then((resp) => {
+          if (resp.total === 0) {
+            return true;
+          }
+          const hasDataSourceMDSId = resp.saved_objects.some((obj) =>
+            obj.references.some((ref) => ref.type === 'data-source' && ref.id === dataSourceMDSId)
+          );
+
+          // Return true if dataSourceMDSId is not found in any references
+          return !hasDataSourceMDSId;
+        });
+      if (flights) {
+        this.setToast('Adding sample data for flights. This can take some time.');
+        await this.props.http.post('../api/sample_data/flights', {
+          query: { data_source_id: dataSourceMDSId },
+        });
+      }
+      if (logs) {
+        this.setToast('Adding sample data for logs. This can take some time.');
+        await this.props.http.post('../api/sample_data/logs', {
+          query: { data_source_id: dataSourceMDSId },
+        });
+      }
       const visIds: string[] = [];
       await this.props.http
         .get('../api/saved_objects/_find', {
           query: {
             type: 'visualization',
             search_fields: 'title',
-            search: '[Logs] Response Codes Over Time + Annotations',
+            search:
+              `[Logs] Response Codes Over Time + Annotations` +
+              (dataSourceMDSLabel ? `_${dataSourceMDSLabel}` : ''),
           },
         })
-        .then((resp) => visIds.push(resp.saved_objects[0].id));
+        .then((resp) => {
+          if (this.props.dataSourceEnabled) {
+            const searchTitle = `[Logs] Response Codes Over Time + Annotations_${dataSourceMDSLabel}`;
+            const savedObjects = resp.saved_objects;
+
+            const foundObject = savedObjects.find((obj) => obj.attributes.title === searchTitle);
+            if (foundObject) {
+              visIds.push(foundObject.id);
+            }
+          } else {
+            visIds.push(resp.saved_objects[0].id);
+          }
+        });
       await this.props.http
         .get('../api/saved_objects/_find', {
           query: {
             type: 'visualization',
             search_fields: 'title',
-            search: '[Logs] Unique Visitors vs. Average Bytes',
+            search:
+              `[Logs] Unique Visitors vs. Average Bytes` +
+              (dataSourceMDSLabel ? `_${dataSourceMDSLabel}` : ''),
           },
         })
-        .then((resp) => visIds.push(resp.saved_objects[0].id));
+        .then((resp) => {
+          if (this.props.dataSourceEnabled) {
+            const searchTitle = `[Logs] Unique Visitors vs. Average Bytes_${dataSourceMDSLabel}`;
+            const savedObjects = resp.saved_objects;
+
+            const foundObject = savedObjects.find((obj) => obj.attributes.title === searchTitle);
+            if (foundObject) {
+              visIds.push(foundObject.id);
+            }
+          } else {
+            visIds.push(resp.saved_objects[0].id);
+          }
+        });
       await this.props.http
         .get('../api/saved_objects/_find', {
           query: {
             type: 'visualization',
             search_fields: 'title',
-            search: '[Flights] Flight Count and Average Ticket Price',
+            search:
+              `[Flights] Flight Count and Average Ticket Price` +
+              (dataSourceMDSLabel ? `_${dataSourceMDSLabel}` : ''),
           },
         })
-        .then((resp) => visIds.push(resp.saved_objects[0].id));
+        .then((resp) => {
+          if (this.props.dataSourceEnabled) {
+            const searchTitle = `[Flights] Flight Count and Average Ticket Price_${dataSourceMDSLabel}`;
+            const savedObjects = resp.saved_objects;
+
+            const foundObject = savedObjects.find((obj) => obj.attributes.title === searchTitle);
+            if (foundObject) {
+              visIds.push(foundObject.id);
+            }
+          } else {
+            visIds.push(resp.saved_objects[0].id);
+          }
+        });
       await this.props.http
         .post(`${NOTEBOOKS_API_PREFIX}/note/savedNotebook/addSampleNotebooks`, {
           body: JSON.stringify({ visIds }),
         })
         .then((res) => {
           const newData = res.body.map((notebook: any) => ({
             path: notebook.name,
             id: notebook.id,
             dateCreated: notebook.dateCreated,
@@ -355,7 +436,7 @@
           }));
         });
       this.setToast(`Sample notebooks successfully added.`);
     } catch (err: any) {
       this.setToast('Error adding sample notebooks.', 'danger');
       console.error(err.body.message);
     } finally {
@@ -393,6 +474,10 @@
                   parentBreadcrumb={this.props.parentBreadcrumb}
                   setBreadcrumbs={this.props.setBreadcrumbs}
                   setToast={this.setToast}
+                  dataSourceManagement={this.props.dataSourceManagement}
+                  notifications={this.props.notifications}
+                  dataSourceEnabled={this.props.dataSourceEnabled}
+                  savedObjectsMDSClient={this.props.savedObjectsMDSClient}
                 />
               )}
             />

public/components/notebooks/components/note_table.tsx

@@ -4,7 +4,6 @@
  */
 
 import {
-  EuiSmallButton,
   EuiCompressedFieldSearch,
   EuiFlexGroup,
   EuiFlexItem,
@@ -19,6 +18,7 @@
   EuiPageContentHeaderSection,
   EuiPageHeader,
   EuiPageHeaderSection,
+  EuiSmallButton,
   EuiSpacer,
   EuiTableFieldDataColumnType,
   EuiText,
@@ -28,35 +28,49 @@
 import moment from 'moment';
 import React, { useEffect, useState } from 'react';
 import { useHistory, useLocation } from 'react-router-dom';
-import { ChromeBreadcrumb } from '../../../../../../src/core/public';
+import {
+  ChromeBreadcrumb,
+  CoreStart,
+  MountPoint,
+  SavedObjectsStart,
+} from '../../../../../../src/core/public';
+import { DataSourceManagementPluginSetup } from '../../../../../../src/plugins/data_source_management/public';
 import {
   CREATE_NOTE_MESSAGE,
   NOTEBOOKS_DOCUMENTATION_URL,
 } from '../../../../common/constants/notebooks';
 import { UI_DATE_FORMAT } from '../../../../common/constants/shared';
+import { setNavBreadCrumbs } from '../../../../common/utils/set_nav_bread_crumbs';
+import { HeaderControlledComponentsWrapper } from '../../../../public/plugin_helpers/plugin_headerControl';
+import { coreRefs } from '../../../framework/core_refs';
 import {
   DeleteNotebookModal,
   getCustomModal,
   getSampleNotebooksModal,
 } from './helpers/modal_containers';
 import { NotebookType } from './main';
-import { setNavBreadCrumbs } from '../../../../common/utils/set_nav_bread_crumbs';
-import { HeaderControlledComponentsWrapper } from '../../../../public/plugin_helpers/plugin_headerControl';
-import { coreRefs } from '../../../framework/core_refs';
 
 const newNavigation = coreRefs.chrome?.navGroup.getNavGroupEnabled();
 
 interface NoteTableProps {
   loading: boolean;
   fetchNotebooks: () => void;
-  addSampleNotebooks: () => void;
+  addSampleNotebooks: (
+    dataSourceMDSId: string | undefined,
+    dataSourceLabel: string | undefined
+  ) => void;
   notebooks: NotebookType[];
   createNotebook: (newNoteName: string) => void;
   renameNotebook: (newNoteName: string, noteId: string) => void;
   cloneNotebook: (newNoteName: string, noteId: string) => void;
   deleteNotebook: (noteList: string[], toastMessage?: string) => void;
   parentBreadcrumb: ChromeBreadcrumb;
   setBreadcrumbs: (newBreadcrumbs: ChromeBreadcrumb[]) => void;
+  dataSourceEnabled: boolean;
+  dataSourceManagement: DataSourceManagementPluginSetup;
+  setActionMenu: (menuMount: MountPoint | undefined) => void;
+  savedObjectsMDSClient: SavedObjectsStart;
+  notifications: CoreStart['notifications'];
   // setToast: (title: string, color?: string, text?: string) => void;
 }
 
@@ -69,6 +83,10 @@
   deleteNotebook,
   parentBreadcrumb,
   setBreadcrumbs,
+  dataSourceEnabled,
+  dataSourceManagement,
+  savedObjectsMDSClient,
+  notifications,
 }: NoteTableProps) {
   const [isModalVisible, setIsModalVisible] = useState(false); // Modal Toggle
   const [modalLayout, setModalLayout] = useState(<EuiOverlayMask />); // Modal Layout
@@ -96,7 +114,7 @@
     if (url[url.length - 1] === 'create') {
       createNote();
     }
   }, [location]);
 
   const closeModal = () => {
     setIsModalVisible(false);
@@ -154,11 +172,25 @@
   };
 
   const addSampleNotebooksModal = async () => {
+    let selectedDataSourceId: string | undefined;
+    let selectedDataSourceLabel: string | undefined;
+    const handleSelectedDataSourceChange = (id?: string, label?: string) => {
+      selectedDataSourceId = id;
+      selectedDataSourceLabel = label;
+    };
     setModalLayout(
-      getSampleNotebooksModal(closeModal, async () => {
-        closeModal();
-        await addSampleNotebooks();
-      })
+      getSampleNotebooksModal(
+        closeModal,
+        async () => {
+          closeModal();
+          await addSampleNotebooks(selectedDataSourceId, selectedDataSourceLabel);
+        },
+        dataSourceEnabled,
+        dataSourceManagement,
+        savedObjectsMDSClient,
+        notifications,
+        handleSelectedDataSourceChange
+      )
     );
     showModal();
   };