Modify the adding sample data part for timeline

Signed-off-by: Yuanqi(Ella) Zhu <zhyuanqi@amazon.com>
opensearch-project · Jun 5, 2024 · d8fc98e · d8fc98e
1 parent 766a39a
commit d8fc98e
Show file tree

Hide file tree

Showing 5 changed files with 82 additions and 4 deletions.
diff --git a/src/core/server/index.ts b/src/core/server/index.ts
@@ -328,6 +328,8 @@ export {
   SavedObjectsDeleteByWorkspaceOptions,
   updateDataSourceNameInVegaSpec,
   extractVegaSpecFromSavedObject,
+  extractTimelineExpression,
+  updateDataSourceNameInTimeline,
 } from './saved_objects';
 
 export {

diff --git a/src/core/server/saved_objects/import/index.ts b/src/core/server/saved_objects/import/index.ts
@@ -43,4 +43,9 @@ export {
   SavedObjectsResolveImportErrorsOptions,
   SavedObjectsImportRetry,
 } from './types';
-export { updateDataSourceNameInVegaSpec, extractVegaSpecFromSavedObject } from './utils';
+export {
+  updateDataSourceNameInVegaSpec,
+  extractVegaSpecFromSavedObject,
+  extractTimelineExpression,
+  updateDataSourceNameInTimeline,
+} from './utils';
diff --git a/src/core/server/saved_objects/import/utils.test.ts b/src/core/server/saved_objects/import/utils.test.ts
@@ -9,6 +9,7 @@ import {
   getDataSourceTitleFromId,
   getUpdatedTSVBVisState,
   updateDataSourceNameInVegaSpec,
+  updateDataSourceNameInTimeline,
 } from './utils';
 import { parse } from 'hjson';
 import { isEqual } from 'lodash';
@@ -199,6 +200,30 @@ describe('updateDataSourceNameInVegaSpec()', () => {
   });
 });
 
+describe('updateDataSourceNameInTimeline()', () => {
+  test('When a timeline expression does not contain a data source name, modify the expression', () => {
+    const expression =
+      '.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    const expectedExpression =
+      '.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name=newDataSource).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    expect(updateDataSourceNameInTimeline(expression, 'newDataSource')).toBe(expectedExpression);
+  });
+
+  test('When a timeline expression contains a data source name, then do nothing', () => {
+    const expression =
+      '.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name=newDataSource).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    expect(updateDataSourceNameInTimeline(expression, 'newDataSource')).toBe(expression);
+  });
+
+  test('When a timeline expression contains multiple timeline expression, modify each of them', () => {
+    const expression =
+      '.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp,data_source_name=aos211).lines(show=true).points(show=true).yaxis(label="Average bytes"),.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    const expectedExpression =
+      '.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp,data_source_name=aos211).lines(show=true).points(show=true).yaxis(label="Average bytes"),.opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name=aos211).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    expect(updateDataSourceNameInTimeline(expression, 'aos211')).toBe(expectedExpression);
+  });
+});
+
 describe('extractVegaSpecFromSavedObject()', () => {
   test('For a Vega visualization saved object, return its spec', () => {
     const spec = 'some-vega-spec';

diff --git a/src/core/server/saved_objects/import/utils.ts b/src/core/server/saved_objects/import/utils.ts
@@ -91,6 +91,25 @@ export const updateDataSourceNameInVegaSpec = (
       });
 };
 
+export const updateDataSourceNameInTimeline = (
+  timelineExpression: string,
+  dataSourceTitle: string
+) => {
+  const expressionRegex = /\.(opensearch|es|elasticsearch)\(([^)]*)\)/g;
+
+  const replaceCallback = (match: string, funcName: string, args: string) => {
+    if (!args.includes('data_source_name')) {
+      args = args.trim();
+      args = `${args}, data_source_name=${dataSourceTitle}`;
+      return `.${funcName}(${args})`;
+    }
+    return match;
+  };
+
+  const modifiedExpression = timelineExpression.replace(expressionRegex, replaceCallback);
+  return modifiedExpression;
+};
+
 export const getDataSourceTitleFromId = async (
   dataSourceId: string,
   savedObjectsClient: SavedObjectsClientContract
@@ -102,7 +121,7 @@ export const getDataSourceTitleFromId = async (
 };
 
 export const extractVegaSpecFromSavedObject = (savedObject: SavedObject) => {
-  if (isVegaVisualization(savedObject)) {
+  if (checkVisualization(savedObject, 'vega')) {
     // @ts-expect-error
     const visStateObject = JSON.parse(savedObject.attributes?.visState);
     return visStateObject.params.spec;
@@ -111,12 +130,22 @@ export const extractVegaSpecFromSavedObject = (savedObject: SavedObject) => {
   return undefined;
 };
 
-const isVegaVisualization = (savedObject: SavedObject) => {
+export const extractTimelineExpression = (savedObject: SavedObject) => {
+  if (checkVisualization(savedObject, 'timelion')) {
+    // @ts-expect-error
+    const visStateObject = JSON.parse(savedObject.attributes?.visState);
+    return visStateObject.params.expression;
+  }
+
+  return undefined;
+};
+
+const checkVisualization = (savedObject: SavedObject, visualizationType: string) => {
   // @ts-expect-error
   const visState = savedObject.attributes?.visState;
   if (!!visState) {
     const visStateObject = JSON.parse(visState);
-    return !!visStateObject.type && visStateObject.type === 'vega';
+    return !!visStateObject.type && visStateObject.type === visualizationType;
   }
   return false;
 };

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
@@ -6,7 +6,9 @@
 import { SavedObject } from 'opensearch-dashboards/server';
 import {
   extractVegaSpecFromSavedObject,
+  extractTimelineExpression,
   updateDataSourceNameInVegaSpec,
+  updateDataSourceNameInTimeline,
 } from '../../../../../../core/server';
 import { SampleDatasetSchema } from '../lib/sample_dataset_registry_types';
 
@@ -114,6 +116,21 @@ export const getSavedObjectsWithDataSource = (
               name: 'dataSource',
             });
           }
+
+          const timelineExpression = extractTimelineExpression(saveObject);
+          if (!!timelineExpression) {
+            // Get the timeline expression with the updated data source name
+            const modifiedExpression = updateDataSourceNameInTimeline(
+              timelineExpression,
+              dataSourceTitle
+            );
+
+            // @ts-expect-error
+            const timelineStateObject = JSON.parse(saveObject.attributes?.visState);
+            timelineStateObject.params.expression = modifiedExpression;
+            // @ts-expect-error
+            saveObject.attributes.visState = JSON.stringify(timelineStateObject);
+          }
         }
       }