[MDS] TSVB Add import logic (opensearch-project#6543) (opensearch-pro…

…ject#6644)

* Add datasource references to the saved object



* Add import logic for TSVB visualizations



* Refactor repeated logic into a util function



* Address comments



* Address more comments



* Fix types



* Remove ts ignore



---------


(cherry picked from commit 22e2a70)

Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

src/core/server/saved_objects/import/check_conflict_for_data_source.test.ts

@@ -11,6 +11,7 @@ import {
   checkConflictsForDataSource,
   ConflictsForDataSourceParams,
 } from './check_conflict_for_data_source';
+import { VisualizationObject } from './types';
 
 type SavedObjectType = SavedObject<{ title?: string }>;
 
@@ -40,6 +41,23 @@ const createVegaVisualizationObject = (id: string): SavedObjectType => {
   } as SavedObjectType;
 };
 
+const createTSVBVisualizationObject = (id: string): VisualizationObject => {
+  const idParse = id.split('_');
+  const params = idParse.length > 1 ? { data_source_id: idParse[1] } : {};
+  const visState = {
+    type: 'metrics',
+    params,
+  };
+
+  return {
+    type: 'visualization',
+    id,
+    attributes: { title: 'some-title', visState: JSON.stringify(visState) },
+    references:
+      idParse.length > 1 ? [{ id: idParse[1], type: 'data-source', name: 'dataSource' }] : [],
+  } as VisualizationObject;
+};
+
 const getSavedObjectClient = (): SavedObjectsClientContract => {
   const savedObject = {} as SavedObjectsClientContract;
   savedObject.get = jest.fn().mockImplementation((type, id) => {
@@ -299,4 +317,57 @@ describe('#checkConflictsForDataSource', () => {
       })
     );
   });
+
+  /**
+   * TSVB test cases
+   */
+  it.each([
+    {
+      id: 'some-object-id',
+    },
+    {
+      id: 'old-datasource-id_some-object-id',
+    },
+  ])('will update datasource reference + visState of TSVB visualization', async ({ id }) => {
+    const tsvbSavedObject = createTSVBVisualizationObject(id);
+    const expectedVisState = JSON.parse(tsvbSavedObject.attributes.visState);
+    expectedVisState.params.data_source_id = 'some-datasource-id';
+    const newVisState = JSON.stringify(expectedVisState);
+    const params = setupParams({
+      objects: [tsvbSavedObject],
+      ignoreRegularConflicts: true,
+      dataSourceId: 'some-datasource-id',
+      savedObjectsClient: getSavedObjectClient(),
+    });
+    const checkConflictsForDataSourceResult = await checkConflictsForDataSource(params);
+
+    expect(checkConflictsForDataSourceResult).toEqual(
+      expect.objectContaining({
+        filteredObjects: [
+          {
+            ...tsvbSavedObject,
+            attributes: {
+              title: 'some-title',
+              visState: newVisState,
+            },
+            id: 'some-datasource-id_some-object-id',
+            references: [
+              {
+                id: 'some-datasource-id',
+                name: 'dataSource',
+                type: 'data-source',
+              },
+            ],
+          },
+        ],
+        errors: [],
+        importIdMap: new Map([
+          [
+            `visualization:some-object-id`,
+            { id: 'some-datasource-id_some-object-id', omitOriginId: true },
+          ],
+        ]),
+      })
+    );
+  });
 });

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

@@ -9,9 +9,11 @@ import {
   SavedObjectsImportError,
   SavedObjectsImportRetry,
 } from '../types';
+import { VisualizationObject } from './types';
 import {
   extractVegaSpecFromSavedObject,
   getDataSourceTitleFromId,
+  getUpdatedTSVBVisState,
   updateDataSourceNameInVegaSpec,
 } from './utils';
 
@@ -117,6 +119,16 @@ export async function checkConflictsForDataSource({
               });
             }
           }
+
+          if (!!dataSourceId) {
+            const visualizationObject = object as VisualizationObject;
+            const { visState, references } = getUpdatedTSVBVisState(
+              visualizationObject,
+              dataSourceId
+            );
+            visualizationObject.attributes.visState = visState;
+            object.references = references;
+          }
         }
 
         const omitOriginId = ignoreRegularConflicts;

src/core/server/saved_objects/import/create_saved_objects.test.ts

@@ -76,6 +76,7 @@ const dataSourceObj1 = createObject(DATA_SOURCE, 'ds-id1'); // -> success
 const dataSourceObj2 = createObject(DATA_SOURCE, 'ds-id2'); // -> conflict
 const dashboardObjWithDataSource = createObject('dashboard', 'ds_dashboard-id1'); // -> success
 const visualizationObjWithDataSource = createObject('visualization', 'ds_visualization-id1'); // -> success
+visualizationObjWithDataSource.attributes = { visState: '{}' };
 const searchObjWithDataSource = createObject('search', 'ds_search-id1'); // -> success
 
 // objs without data source id, used to test can get saved object with data source id
@@ -100,6 +101,7 @@ const visualizationObj = {
   type: 'visualization',
   attributes: {
     title: 'visualization-title',
+    visState: '{}',
   },
   references: [],
   source: {
@@ -130,6 +132,27 @@ const getVegaVisualizationObj = (id: string) => ({
   updated_at: 'some-date',
 });
 
+const getTSVBVisualizationObj = (id: string, dataSourceId?: string) => {
+  const params = dataSourceId ? { data_source_id: dataSourceId } : {};
+  const references = dataSourceId
+    ? [{ id: dataSourceId, name: 'dataSource', type: 'data-source' }]
+    : [];
+  return {
+    type: 'visualization',
+    id,
+    attributes: {
+      title: 'some-title',
+      visState: JSON.stringify({
+        type: 'metrics',
+        params,
+      }),
+    },
+    references,
+    namespaces: ['default'],
+    updated_at: 'some-date',
+  };
+};
+
 const getVegaMDSVisualizationObj = (id: string, dataSourceId: string) => ({
   type: 'visualization',
   id: dataSourceId ? `${dataSourceId}_${id}` : id,
@@ -500,6 +523,30 @@ describe('#createSavedObjects', () => {
     expect(results).toEqual(expectedResults);
   };
 
+  const testTSVBVisualizationsWithDataSources = async (params: {
+    objects: SavedObject[];
+    expectedFilteredObjects: SavedObject[];
+    dataSourceId?: string;
+    dataSourceTitle?: string;
+  }) => {
+    const savedObjectsCustomClient = savedObjectsClientMock.create();
+
+    const options = setupParams({
+      ...params,
+      savedObjectsCustomClient,
+    });
+
+    savedObjectsCustomClient.bulkCreate = jest.fn().mockImplementation((objectsToCreate, _) => {
+      return Promise.resolve({
+        saved_objects: objectsToCreate,
+      });
+    });
+
+    const results = await createSavedObjects(options);
+
+    expect(results.createdObjects).toMatchObject(params.expectedFilteredObjects);
+  };
+
   const testReturnValueWithDataSource = async (
     namespace?: string,
     dataSourceId?: string,
@@ -661,6 +708,46 @@ describe('#createSavedObjects', () => {
     });
   });
 
+  describe('with a data source for TSVB saved objects', () => {
+    test('can attach a TSVB datasource reference to a non-MDS ', async () => {
+      const objects = [getTSVBVisualizationObj('some-tsvb-id')];
+      const expectedObject = getTSVBVisualizationObj('some-tsvb-id', 'some-datasource-id');
+      const expectedFilteredObjects = [
+        {
+          ...expectedObject,
+          attributes: {
+            title: 'some-title_dataSourceName',
+          },
+        },
+      ];
+      await testTSVBVisualizationsWithDataSources({
+        objects,
+        expectedFilteredObjects,
+        dataSourceId: 'some-datasource-id',
+        dataSourceTitle: 'dataSourceName',
+      });
+    });
+
+    test('can update a TSVB datasource reference', async () => {
+      const objects = [getTSVBVisualizationObj('some-tsvb-id', 'old-datasource-id')];
+      const expectedObject = getTSVBVisualizationObj('some-tsvb-id', 'some-datasource-id');
+      const expectedFilteredObjects = [
+        {
+          ...expectedObject,
+          attributes: {
+            title: 'some-title_dataSourceName',
+          },
+        },
+      ];
+      await testTSVBVisualizationsWithDataSources({
+        objects,
+        expectedFilteredObjects,
+        dataSourceId: 'some-datasource-id',
+        dataSourceTitle: 'dataSourceName',
+      });
+    });
+  });
+
   describe('with a undefined workspaces', () => {
     test('calls bulkCreate once with input objects', async () => {
       const options = setupParams({ objects: objs });

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

@@ -35,8 +35,12 @@ import {
   SavedObjectsImportError,
 } from '../types';
 import { extractErrors } from './extract_errors';
-import { CreatedObject } from './types';
-import { extractVegaSpecFromSavedObject, updateDataSourceNameInVegaSpec } from './utils';
+import { CreatedObject, VisualizationObject } from './types';
+import {
+  extractVegaSpecFromSavedObject,
+  getUpdatedTSVBVisState,
+  updateDataSourceNameInVegaSpec,
+} from './utils';
 
 interface CreateSavedObjectsParams<T> {
   objects: Array<SavedObject<T>>;
@@ -125,6 +129,15 @@ export const createSavedObjects = async <T>({
               name: 'dataSource',
             });
           }
+
+          const visualizationObject = object as VisualizationObject;
+          const { visState, references } = getUpdatedTSVBVisState(
+            visualizationObject,
+            dataSourceId
+          );
+
+          visualizationObject.attributes.visState = visState;
+          object.references = references;
         }
 
         if (object.type === 'index-pattern') {

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

@@ -219,3 +219,5 @@ export interface SavedObjectsResolveImportErrorsOptions {
 }
 
 export type CreatedObject<T> = SavedObject<T> & { destinationId?: string };
+
+export type VisualizationObject<T = any> = SavedObject<T> & { attributes: { visState: string } };

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

@@ -7,6 +7,7 @@ import { readFileSync } from 'fs';
 import {
   extractVegaSpecFromSavedObject,
   getDataSourceTitleFromId,
+  getUpdatedTSVBVisState,
   updateDataSourceNameInVegaSpec,
 } from './utils';
 import { parse } from 'hjson';
@@ -245,3 +246,66 @@ describe('getDataSourceTitleFromId()', () => {
     expect(await getDataSourceTitleFromId('nonexistent-id', savedObjectsClient)).toBe(undefined);
   });
 });
+
+describe('getUpdatedTSVBVisState', () => {
+  const getTSVBSavedObject = (dataSourceId?: string) => {
+    const params = dataSourceId ? { data_source_id: dataSourceId } : {};
+    const references = dataSourceId
+      ? [{ id: dataSourceId, type: 'data-source', name: 'dataSource' }]
+      : [];
+
+    return {
+      type: 'visualization',
+      id: 'some-id',
+      attributes: {
+        title: 'Some Title',
+        visState: JSON.stringify({
+          type: 'metrics',
+          params,
+        }),
+      },
+      references,
+    };
+  };
+
+  test('non-TSVB object should return the old references and visState', () => {
+    const visState = {
+      type: 'area',
+      params: {},
+    };
+
+    const object = {
+      type: 'visualization',
+      id: 'some-id',
+      attributes: {
+        title: 'Some title',
+        visState: JSON.stringify(visState),
+      },
+      references: [],
+    };
+
+    expect(getUpdatedTSVBVisState(object, 'some-datasource-id')).toMatchObject({
+      visState: JSON.stringify(visState),
+      references: [],
+    });
+  });
+
+  test.each(['old-datasource-id', undefined])(
+    `non-MDS TSVB object should update the datasource when the old datasource is "%s"`,
+    (oldDataSourceId) => {
+      const object = getTSVBSavedObject(oldDataSourceId);
+      const dataSourceId = 'some-datasource-id';
+      const expectedVisState = JSON.stringify({
+        type: 'metrics',
+        params: {
+          data_source_id: dataSourceId,
+        },
+      });
+
+      expect(getUpdatedTSVBVisState(object, dataSourceId)).toMatchObject({
+        visState: expectedVisState,
+        references: [{ id: dataSourceId, name: 'dataSource', type: 'data-source' }],
+      });
+    }
+  );
+});