From 61e29b57b33d41c31011df2ff52204552ab739dc Mon Sep 17 00:00:00 2001
From: Huy Nguyen <73027756+huyaboo@users.noreply.github.com>
Date: Wed, 5 Jun 2024 21:35:32 -0700
Subject: [PATCH] [MDS] Fix sample data to use datasources for TSVB
 visualizations (#6940)

* Fix sample data to use datasources for TSVB visualizations

Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com>

* Changeset file for PR #6940 created/updated

---------

Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
---
 changelogs/fragments/6940.yml                 |  2 +
 .../sample_data/data_sets/util.test.ts        | 42 +++++++++++++++++++
 .../services/sample_data/data_sets/util.ts    | 18 +++++++-
 3 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 changelogs/fragments/6940.yml

diff --git a/changelogs/fragments/6940.yml b/changelogs/fragments/6940.yml
new file mode 100644
index 00000000000..18c34077f38
--- /dev/null
+++ b/changelogs/fragments/6940.yml
@@ -0,0 +1,2 @@
+fix:
+- Add TSVB Support for adding sample data ([#6940](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6940))
\ 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 1b02fa778f9..c51e289231d 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
@@ -13,6 +13,19 @@ describe('getSavedObjectsWithDataSource()', () => {
     return visualizationObjects.saved_objects;
   };
 
+  const TSVBVisualizationSavedObject = {
+    type: 'visualization',
+    id: 'some-id',
+    attributes: {
+      title: 'some-title',
+      visState: JSON.stringify({
+        type: 'metrics',
+        params: {},
+      }),
+    },
+    references: [],
+  };
+
   test('when processing Vega Visualization saved objects, it should attach data_source_name to each OpenSearch query', () => {
     const dataSourceId = 'some-datasource-id';
     const dataSourceName = 'Data Source Name';
@@ -127,6 +140,35 @@ describe('getSavedObjectsWithDataSource()', () => {
       },
     ]);
   });
+
+  test('when processing TSVB Visualization saved objects, it should attach data_source_id to the visState and add datasource reference', () => {
+    const dataSourceId = 'some-datasource-id';
+    const dataSourceTitle = 'Data Source Name';
+    const expectedTSVBVisualizationSavedObject = {
+      ...TSVBVisualizationSavedObject,
+      id: `${dataSourceId}_some-id`,
+      attributes: {
+        title: `some-title_${dataSourceTitle}`,
+        visState: JSON.stringify({
+          type: 'metrics',
+          params: {
+            data_source_id: dataSourceId,
+          },
+        }),
+      },
+      references: [
+        {
+          id: dataSourceId,
+          type: 'data-source',
+          name: 'dataSource',
+        },
+      ],
+    };
+
+    expect(
+      getSavedObjectsWithDataSource([TSVBVisualizationSavedObject], dataSourceId, dataSourceTitle)
+    ).toMatchObject([expectedTSVBVisualizationSavedObject]);
+  });
 });
 
 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 5585720649d..871a9574b59 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
@@ -96,6 +96,11 @@ export const getSavedObjectsWithDataSource = (
         if (saveObject.type === 'visualization') {
           const vegaSpec = extractVegaSpecFromSavedObject(saveObject);
 
+          const visualizationSavedObject = saveObject as SavedObject & {
+            attributes: { visState: string };
+          };
+          const visStateObject = JSON.parse(visualizationSavedObject.attributes.visState);
+
           if (!!vegaSpec) {
             const updatedVegaSpec = updateDataSourceNameInVegaSpec({
               spec: vegaSpec,
@@ -104,8 +109,6 @@ export const getSavedObjectsWithDataSource = (
               spacing: 1,
             });
 
-            // @ts-expect-error
-            const visStateObject = JSON.parse(saveObject.attributes?.visState);
             visStateObject.params.spec = updatedVegaSpec;
 
             // @ts-expect-error
@@ -131,6 +134,17 @@ export const getSavedObjectsWithDataSource = (
             // @ts-expect-error
             saveObject.attributes.visState = JSON.stringify(timelineStateObject);
           }
+
+          if (visStateObject.type === 'metrics') {
+            visStateObject.params.data_source_id = dataSourceId;
+
+            visualizationSavedObject.attributes.visState = JSON.stringify(visStateObject);
+            visualizationSavedObject.references.push({
+              id: `${dataSourceId}`,
+              type: 'data-source',
+              name: 'dataSource',
+            });
+          }
         }
       }