Modify the adding sample data part for timeline

src/core/server/index.ts

@@ -328,6 +328,8 @@ export {
   SavedObjectsDeleteByWorkspaceOptions,
   updateDataSourceNameInVegaSpec,
   extractVegaSpecFromSavedObject,
+  extractTimelineExpression,
+  updateDataSourceNameInTimeline,
 } from './saved_objects';
 
 export {

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';

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,38 @@ 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);
+  });
+
+  test('When a timeline expression contains multiple timeline expression and the datasource name contains space, we modify each of them', () => {
+    const expression =
+      '.es(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp).lines(show=true).points(show=true).yaxis(label="Average bytes"),.elasticsearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp).lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    const expectedExpression =
+      '.es(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name="aos 211").lines(show=true).points(show=true).yaxis(label="Average bytes"),.elasticsearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name="aos 211").lines(show=true).points(show=true).yaxis(label="Average bytes")';
+    expect(updateDataSourceNameInTimeline(expression, 'aos 211')).toBe(expectedExpression);
+  });
+});
+
 describe('extractVegaSpecFromSavedObject()', () => {
   test('For a Vega visualization saved object, return its spec', () => {
     const spec = 'some-vega-spec';

src/core/server/saved_objects/import/utils.ts

@@ -91,6 +91,25 @@
       });
 };
 
+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')) {
+      let expressionArgs = args.trim();
+      expressionArgs = `${expressionArgs}, data_source_name="${dataSourceTitle}"`;
+      return `.${funcName}(${expressionArgs})`;
+    }
+    return match;
+  };
+
+  const modifiedExpression = timelineExpression.replace(expressionRegex, replaceCallback);
+  return modifiedExpression;
+};
+
 export const getDataSourceTitleFromId = async (
   dataSourceId: string,
   savedObjectsClient: SavedObjectsClientContract
@@ -102,7 +121,7 @@
 };
 
 export const extractVegaSpecFromSavedObject = (savedObject: SavedObject) => {
-  if (isVegaVisualization(savedObject)) {
+  if (confirmVisualizationType(savedObject, 'vega')) {
     // @ts-expect-error
     const visStateObject = JSON.parse(savedObject.attributes?.visState);
     return visStateObject.params.spec;
@@ -111,12 +130,26 @@
   return undefined;
 };
 
-const isVegaVisualization = (savedObject: SavedObject) => {
+export const extractTimelineExpression = (savedObject: SavedObject) => {
+  if (!confirmVisualizationType(savedObject, 'timelion')) {
+    return undefined;
+  }
+  // @ts-expect-error
+  const visStateString = savedObject.attributes?.visState;
+  if (!visStateString) {
+    return undefined;
+  }
+
+  const visStateObject = JSON.parse(visStateString);
+  return visStateObject.params.expression;
+};
+
+const confirmVisualizationType = (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;
 };

src/plugins/home/server/services/sample_data/data_sets/util.test.ts

@@ -63,6 +63,38 @@ describe('getSavedObjectsWithDataSource()', () => {
     expect(updatedVegaVisualizationsFields).toEqual(expect.arrayContaining(expectedUpdatedFields));
   });
 
+  it('should processing timeline saved object and add datasource name in the end', () => {
+    const dataSourceId = 'some-datasource-id';
+    const dataSourceName = 'dataSourceName';
+    const savedObjects = [
+      {
+        id: 'saved-object-1',
+        type: 'visualization',
+        title: 'example',
+        attributes: {
+          title: 'example',
+          visState:
+            '{"title":"(Timeline) Avg bytes over time","type":"timelion","aggs":[],"params":{"expression":".opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp).lines(show=true).points(show=true).yaxis(label=\\"Average bytes\\")","interval":"auto"}}',
+        },
+        references: [],
+      },
+    ];
+
+    expect(getSavedObjectsWithDataSource(savedObjects, dataSourceId, dataSourceName)).toEqual([
+      {
+        id: 'some-datasource-id_saved-object-1',
+        type: 'visualization',
+        title: 'example',
+        attributes: {
+          title: 'example_dataSourceName',
+          visState:
+            '{"title":"(Timeline) Avg bytes over time","type":"timelion","aggs":[],"params":{"expression":".opensearch(opensearch_dashboards_sample_data_logs, metric=avg:bytes, timefield=@timestamp, data_source_name=\\"dataSourceName\\").lines(show=true).points(show=true).yaxis(label=\\"Average bytes\\")","interval":"auto"}}',
+        },
+        references: [],
+      },
+    ]);
+  });
+
   it('should update index-pattern id and references with given data source', () => {
     const dataSourceId = 'some-datasource-id';
     const dataSourceName = 'Data Source Name';

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);
+          }
         }
       }