[MDS] Add Vega support for importing saved objects (opensearch-projec…

…t#6123) * Add MDS support for Vega Signed-off-by: Huy Nguyen <[email protected]> * Refactor field to data_source_id Signed-off-by: Huy Nguyen <[email protected]> * Add to CHANGELOG.md Signed-off-by: Huy Nguyen <[email protected]> * Added test cases and renamed field to use data_source_name Signed-off-by: Huy Nguyen <[email protected]> * Add prefix datasource name test case and add example in default hjson Signed-off-by: Huy Nguyen <[email protected]> * Move CHANGELOG to appropriate section Signed-off-by: Huy Nguyen <[email protected]> * Increased test coverage of search() method Signed-off-by: Huy Nguyen <[email protected]> * Add test cases for util function Signed-off-by: Huy Nguyen <[email protected]> * Add util function to modify Vega Spec Signed-off-by: Huy Nguyen <[email protected]> * Add method to verify Vega saved object type Signed-off-by: Huy Nguyen <[email protected]> * Add import saved object support for Vega Signed-off-by: Huy Nguyen <[email protected]> * Add unit tests for Vega objects in create and conflict modes Signed-off-by: Huy Nguyen <[email protected]> * Refactored utils test file Signed-off-by: Huy Nguyen <[email protected]> * Add to CHANGELOG Signed-off-by: Huy Nguyen <[email protected]> * Use bulkget instead of single get Signed-off-by: Huy Nguyen <[email protected]> * Add datasource references to the specs Signed-off-by: Huy Nguyen <[email protected]> * Fix bootstrap errors Signed-off-by: Huy Nguyen <[email protected]> * Add edge case where title is potentially undefined Signed-off-by: Huy Nguyen <[email protected]> * Address PR comments Signed-off-by: Huy Nguyen <[email protected]> * Add more test coverage for checking conflict Signed-off-by: Huy Nguyen <[email protected]> * Fix unit test Signed-off-by: Huy Nguyen <[email protected]> --------- Signed-off-by: Huy Nguyen <[email protected]>
huyaboo · Mar 19, 2024 · de978d4 · de978d4
1 parent 1f74ab3
commit de978d4
Show file tree

Hide file tree

Showing 36 changed files with 3,568 additions and 72 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -58,7 +58,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Dynamic Configurations] Pass request headers when making application config calls ([#6164](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6164))
 - [Discover] Options button to configure legacy mode and remove the top navigation option ([#6170](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6170))
 - [Multiple Datasource] Add default functionality for customer to choose default datasource ([#6058](https://github.com/opensearch-project/OpenSearch-Dashboards/issues/6058))
-
+- [Multiple Datasource] Add import support for Vega when specifying a datasource ([#6123](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6123))
 
 ### 🐛 Bug Fixes
 

diff --git a/src/core/server/saved_objects/import/check_conflict_for_data_source.test.ts b/src/core/server/saved_objects/import/check_conflict_for_data_source.test.ts
@@ -5,7 +5,7 @@
 
 import { mockUuidv4 } from './__mocks__';
 import { SavedObjectReference, SavedObjectsImportRetry } from 'opensearch-dashboards/public';
-import { SavedObject } from '../types';
+import { SavedObject, SavedObjectsClientContract } from '../types';
 import { SavedObjectsErrorHelpers } from '..';
 import {
   checkConflictsForDataSource,
@@ -24,6 +24,45 @@ const createObject = (type: string, id: string): SavedObjectType => ({
   references: (Symbol() as unknown) as SavedObjectReference[],
 });
 
+const createVegaVisualizationObject = (id: string): SavedObjectType => {
+  const visState =
+    id.split('_').length > 1
+      ? '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: old-datasource-title\\n    }\\n  }\\n}"}}'
+      : '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n    }\\n  }\\n}"}}';
+  return {
+    type: 'visualization',
+    id,
+    attributes: { title: 'some-title', visState },
+    references:
+      id.split('_').length > 1
+        ? [{ id: id.split('_')[0], type: 'data-source', name: 'dataSource' }]
+        : [],
+  } as SavedObjectType;
+};
+
+const getSavedObjectClient = (): SavedObjectsClientContract => {
+  const savedObject = {} as SavedObjectsClientContract;
+  savedObject.get = jest.fn().mockImplementation((type, id) => {
+    if (type === 'data-source' && id === 'old-datasource-id') {
+      return Promise.resolve({
+        attributes: {
+          title: 'old-datasource-title',
+        },
+      });
+    } else if (type === 'data-source') {
+      return Promise.resolve({
+        attributes: {
+          title: 'some-datasource-title',
+        },
+      });
+    }
+
+    return Promise.resolve(undefined);
+  });
+
+  return savedObject;
+};
+
 const getResultMock = {
   conflict: (type: string, id: string) => {
     const error = SavedObjectsErrorHelpers.createConflictError(type, id).output.payload;
@@ -56,6 +95,7 @@ describe('#checkConflictsForDataSource', () => {
     retries?: SavedObjectsImportRetry[];
     createNewCopies?: boolean;
     dataSourceId?: string;
+    savedObjectsClient?: SavedObjectsClientContract;
   }): ConflictsForDataSourceParams => {
     return { ...partial };
   };
@@ -140,4 +180,123 @@ describe('#checkConflictsForDataSource', () => {
       importIdMap: new Map(),
     });
   });
+
+  /*
+  Vega test cases
+  */
+  it('will attach datasource name to Vega spec when importing from local to datasource', async () => {
+    const vegaSavedObject = createVegaVisualizationObject('some-object-id');
+    const params = setupParams({
+      objects: [vegaSavedObject],
+      ignoreRegularConflicts: true,
+      dataSourceId: 'some-datasource-id',
+      savedObjectsClient: getSavedObjectClient(),
+    });
+    const checkConflictsForDataSourceResult = await checkConflictsForDataSource(params);
+
+    expect(params.savedObjectsClient?.get).toHaveBeenCalledWith(
+      'data-source',
+      'some-datasource-id'
+    );
+    expect(checkConflictsForDataSourceResult).toEqual(
+      expect.objectContaining({
+        filteredObjects: [
+          {
+            ...vegaSavedObject,
+            attributes: {
+              title: 'some-title',
+              visState:
+                '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: some-datasource-title\\n    }\\n  }\\n}"}}',
+            },
+            id: 'some-datasource-id_some-object-id',
+            references: [
+              {
+                id: 'some-datasource-id',
+                type: 'data-source',
+                name: 'dataSource',
+              },
+            ],
+          },
+        ],
+        errors: [],
+        importIdMap: new Map([
+          [
+            `visualization:some-object-id`,
+            { id: 'some-datasource-id_some-object-id', omitOriginId: true },
+          ],
+        ]),
+      })
+    );
+  });
+
+  it('will not change Vega spec when importing from datasource to different datasource', async () => {
+    const vegaSavedObject = createVegaVisualizationObject('old-datasource-id_some-object-id');
+    const params = setupParams({
+      objects: [vegaSavedObject],
+      ignoreRegularConflicts: true,
+      dataSourceId: 'some-datasource-id',
+      savedObjectsClient: getSavedObjectClient(),
+    });
+    const checkConflictsForDataSourceResult = await checkConflictsForDataSource(params);
+
+    expect(params.savedObjectsClient?.get).toHaveBeenCalledWith(
+      'data-source',
+      'some-datasource-id'
+    );
+    expect(checkConflictsForDataSourceResult).toEqual(
+      expect.objectContaining({
+        filteredObjects: [
+          {
+            ...vegaSavedObject,
+            attributes: {
+              title: 'some-title',
+              visState:
+                '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: old-datasource-title\\n    }\\n  }\\n}"}}',
+            },
+            id: 'some-datasource-id_some-object-id',
+          },
+        ],
+        errors: [],
+        importIdMap: new Map([
+          [
+            `visualization:some-object-id`,
+            { id: 'some-datasource-id_some-object-id', omitOriginId: true },
+          ],
+        ]),
+      })
+    );
+  });
+
+  it('will not change Vega spec when dataSourceTitle is undefined', async () => {
+    const vegaSavedObject = createVegaVisualizationObject('old-datasource-id_some-object-id');
+    const params = setupParams({
+      objects: [vegaSavedObject],
+      ignoreRegularConflicts: true,
+      dataSourceId: 'nonexistent-datasource-title-id',
+      savedObjectsClient: getSavedObjectClient(),
+    });
+    const checkConflictsForDataSourceResult = await checkConflictsForDataSource(params);
+
+    expect(params.savedObjectsClient?.get).toHaveBeenCalledWith(
+      'data-source',
+      'nonexistent-datasource-title-id'
+    );
+    expect(checkConflictsForDataSourceResult).toEqual(
+      expect.objectContaining({
+        filteredObjects: [
+          {
+            ...vegaSavedObject,
+            id: 'nonexistent-datasource-title-id_some-object-id',
+          },
+        ],
+        errors: [],
+        importIdMap: new Map([
+          [
+            `visualization:some-object-id`,
+            { id: 'nonexistent-datasource-title-id_some-object-id', omitOriginId: true },
+          ],
+        ]),
+      })
+    );
+  });
 });
diff --git a/src/core/server/saved_objects/import/check_conflict_for_data_source.ts b/src/core/server/saved_objects/import/check_conflict_for_data_source.ts
@@ -3,13 +3,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { SavedObject, SavedObjectsImportError, SavedObjectsImportRetry } from '../types';
+import {
+  SavedObject,
+  SavedObjectsClientContract,
+  SavedObjectsImportError,
+  SavedObjectsImportRetry,
+} from '../types';
+import {
+  extractVegaSpecFromSavedObject,
+  getDataSourceTitleFromId,
+  updateDataSourceNameInVegaSpec,
+} from './utils';
 
 export interface ConflictsForDataSourceParams {
   objects: Array<SavedObject<{ title?: string }>>;
   ignoreRegularConflicts?: boolean;
   retries?: SavedObjectsImportRetry[];
   dataSourceId?: string;
+  savedObjectsClient?: SavedObjectsClientContract;
 }
 
 interface ImportIdMapEntry {
@@ -31,6 +42,7 @@ export async function checkConflictsForDataSource({
   ignoreRegularConflicts,
   retries = [],
   dataSourceId,
+  savedObjectsClient,
 }: ConflictsForDataSourceParams) {
   const filteredObjects: Array<SavedObject<{ title?: string }>> = [];
   const errors: SavedObjectsImportError[] = [];
@@ -43,6 +55,12 @@ export async function checkConflictsForDataSource({
     (acc, cur) => acc.set(`${cur.type}:${cur.id}`, cur),
     new Map<string, SavedObjectsImportRetry>()
   );
+
+  const dataSourceTitle =
+    !!dataSourceId && !!savedObjectsClient
+      ? await getDataSourceTitleFromId(dataSourceId, savedObjectsClient)
+      : undefined;
+
   objects.forEach((object) => {
     const {
       type,
@@ -74,6 +92,33 @@ export async function checkConflictsForDataSource({
         /**
          * Only update importIdMap and filtered objects
          */
+
+        // Some visualization types will need special modifications, like Vega visualizations
+        if (object.type === 'visualization') {
+          const vegaSpec = extractVegaSpecFromSavedObject(object);
+
+          if (!!vegaSpec && !!dataSourceTitle) {
+            const updatedVegaSpec = updateDataSourceNameInVegaSpec({
+              spec: vegaSpec,
+              newDataSourceName: dataSourceTitle,
+            });
+
+            // @ts-expect-error
+            const visStateObject = JSON.parse(object.attributes?.visState);
+            visStateObject.params.spec = updatedVegaSpec;
+
+            // @ts-expect-error
+            object.attributes.visState = JSON.stringify(visStateObject);
+            if (!!dataSourceId) {
+              object.references.push({
+                id: dataSourceId,
+                name: 'dataSource',
+                type: 'data-source',
+              });
+            }
+          }
+        }
+
         const omitOriginId = ignoreRegularConflicts;
         importIdMap.set(`${type}:${id}`, { id: `${dataSourceId}_${rawId}`, omitOriginId });
         filteredObjects.push({ ...object, id: `${dataSourceId}_${rawId}` });