From eb3efa8b7ea5756268feed83336b838450bd1fc9 Mon Sep 17 00:00:00 2001
From: Lin Wang <wonglam@amazon.com>
Date: Thu, 30 May 2024 00:31:46 +0800
Subject: [PATCH] Fix index pattern data source reference not updated in sample
 data (#6851)

* Fix index pattern data source reference not updated in sample data

Signed-off-by: Lin Wang <wonglam@amazon.com>

* Changeset file for PR #6851 created/updated

---------

Signed-off-by: Lin Wang <wonglam@amazon.com>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
---
 changelogs/fragments/6851.yml                 |  2 ++
 .../sample_data/data_sets/util.test.ts        | 35 ++++++++++++++++++-
 .../services/sample_data/data_sets/util.ts    | 11 ++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 changelogs/fragments/6851.yml

diff --git a/changelogs/fragments/6851.yml b/changelogs/fragments/6851.yml
new file mode 100644
index 000000000000..c4e131e73c9d
--- /dev/null
+++ b/changelogs/fragments/6851.yml
@@ -0,0 +1,2 @@
+fix:
+- Update index pattern references with data source when import sample data ([#6851](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6851))
\ No newline at end of file
diff --git a/src/plugins/home/server/services/sample_data/data_sets/util.test.ts b/src/plugins/home/server/services/sample_data/data_sets/util.test.ts
index 0cefac395b9c..6a7e7814761f 100644
--- a/src/plugins/home/server/services/sample_data/data_sets/util.test.ts
+++ b/src/plugins/home/server/services/sample_data/data_sets/util.test.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { getSavedObjectsWithDataSource, getFinalSavedObjects, appendDataSourceId } from './util';
+import { getSavedObjectsWithDataSource, getFinalSavedObjects } from './util';
 import { SavedObject, updateDataSourceNameInVegaSpec } from '../../../../../../core/server';
 import visualizationObjects from './test_utils/visualization_objects.json';
 
@@ -62,6 +62,39 @@ describe('getSavedObjectsWithDataSource()', () => {
 
     expect(updatedVegaVisualizationsFields).toEqual(expect.arrayContaining(expectedUpdatedFields));
   });
+
+  it('should update index-pattern id and references with given data source', () => {
+    const dataSourceId = 'some-datasource-id';
+    const dataSourceName = 'Data Source Name';
+
+    expect(
+      getSavedObjectsWithDataSource(
+        [
+          {
+            id: 'saved-object-1',
+            type: 'index-pattern',
+            attributes: {},
+            references: [],
+          },
+        ],
+        dataSourceId,
+        dataSourceName
+      )
+    ).toEqual([
+      {
+        id: 'some-datasource-id_saved-object-1',
+        type: 'index-pattern',
+        attributes: {},
+        references: [
+          {
+            id: `${dataSourceId}`,
+            type: 'data-source',
+            name: 'dataSource',
+          },
+        ],
+      },
+    ]);
+  });
 });
 
 describe('getFinalSavedObjects()', () => {
diff --git a/src/plugins/home/server/services/sample_data/data_sets/util.ts b/src/plugins/home/server/services/sample_data/data_sets/util.ts
index 4991b02272cc..04eff6036e88 100644
--- a/src/plugins/home/server/services/sample_data/data_sets/util.ts
+++ b/src/plugins/home/server/services/sample_data/data_sets/util.ts
@@ -71,6 +71,17 @@ export const getSavedObjectsWithDataSource = (
     return saveObjectList.map((saveObject) => {
       overrideSavedObjectId(saveObject, idGenerator);
 
+      // update reference
+      if (saveObject.type === 'index-pattern') {
+        saveObject.references = [
+          {
+            id: `${dataSourceId}`,
+            type: 'data-source',
+            name: 'dataSource',
+          },
+        ];
+      }
+
       if (dataSourceTitle) {
         if (
           saveObject.type === 'dashboard' ||