Fix sample data to use datasources for TSVB visualizations

Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com>
opensearch-project · Jun 6, 2024 · a6507ab · a6507ab
1 parent cb84bfb
commit a6507ab
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
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
@@ -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
@@ -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',
+            });
+          }
         }
       }