[MDS] Add Vega support for importing saved objects

CHANGELOG.md

@@ -46,7 +46,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Workspace] Add workspace id in basePath ([#6060](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6060))
 - Implement new home page ([#6065](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6065))
 - Add sidecar service ([#5920](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5920))
-
+- [Multiple Datasource] Add import support for Vega when specifying a datasource ([#6123](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6123))
 
 ### 🐛 Bug Fixes
 

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,82 @@ 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(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(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 },
+          ],
+        ]),
+      })
+    );
+  });
 });

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

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

@@ -115,6 +115,37 @@ const visualizationObj = {
     },
   },
 };
+
+const getVegaVisualizationObj = (id: string) => ({
+  type: 'visualization',
+  id,
+  attributes: {
+    title: 'some-title',
+    visState:
+      '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n    }\\n  }\\n}"}}',
+  },
+  references: [],
+  namespaces: ['default'],
+  version: 'some-version',
+  updated_at: 'some-date',
+});
+
+const getVegaMDSVisualizationObj = (id: string, dataSourceId: string) => ({
+  type: 'visualization',
+  id: dataSourceId ? `${dataSourceId}_${id}` : id,
+  attributes: {
+    title: 'some-other-title',
+    visState:
+      '{"title":"some-other-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: old-datasource-title\\n    }\\n  }\\n}"}}',
+  },
+  references: [
+    {
+      id: dataSourceId,
+      name: 'dataSource',
+      type: 'data-source',
+    },
+  ],
+});
 // non-multi-namespace types shouldn't have origin IDs, but we include test cases to ensure it's handled gracefully
 // non-multi-namespace types by definition cannot result in an unresolvable conflict, so we don't include test cases for those
 const importId3 = 'id-foo';
@@ -142,8 +173,11 @@ describe('#createSavedObjects', () => {
     overwrite?: boolean;
     dataSourceId?: string;
     dataSourceTitle?: string;
+    savedObjectsCustomClient?: jest.Mocked<SavedObjectsClientContract>;
   }): CreateSavedObjectsParams => {
-    savedObjectsClient = savedObjectsClientMock.create();
+    savedObjectsClient = !!partial.savedObjectsCustomClient
+      ? partial.savedObjectsCustomClient
+      : savedObjectsClientMock.create();
     bulkCreate = savedObjectsClient.bulkCreate;
     return { accumulatedErrors: [], ...partial, savedObjectsClient, importIdMap };
   };
@@ -490,6 +524,29 @@ describe('#createSavedObjects', () => {
     expect(results).toEqual(expectedResultsWithDataSource);
   };
 
+  const testVegaVisualizationsWithDataSources = async (params: {
+    objects: SavedObject[];
+    expectedFilteredObjects: Array<Record<string, unknown>>;
+    dataSourceId?: string;
+    dataSourceTitle?: string;
+  }) => {
+    const savedObjectsCustomClient = savedObjectsClientMock.create();
+
+    const options = setupParams({
+      ...params,
+      savedObjectsCustomClient,
+    });
+    savedObjectsCustomClient.bulkCreate = jest.fn().mockResolvedValue({
+      saved_objects: params.objects.map((obj) => {
+        return getResultMock.success(obj, options);
+      }),
+    });
+
+    const results = await createSavedObjects(options);
+
+    expect(results.createdObjects).toMatchObject(params.expectedFilteredObjects);
+  };
+
   describe('with an undefined namespace', () => {
     test('calls bulkCreate once with input objects', async () => {
       await testBulkCreateObjects();
@@ -546,4 +603,61 @@ describe('#createSavedObjects', () => {
       );
     });
   });
+
+  describe('with a data source for Vega saved objects', () => {
+    test('can attach a data source name to the Vega spec if there is a local query', async () => {
+      const objects = [getVegaVisualizationObj('some-vega-id')];
+      const expectedObject = getVegaVisualizationObj('some-vega-id');
+      const expectedFilteredObjects = [
+        {
+          ...expectedObject,
+          attributes: {
+            title: 'some-title_dataSourceName',
+            visState:
+              '{"title":"some-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: dataSourceName\\n    }\\n  }\\n}"}}',
+          },
+          id: 'some-vega-id',
+          references: [
+            {
+              id: 'some-datasource-id',
+              type: 'data-source',
+              name: 'dataSource',
+            },
+          ],
+        },
+      ];
+      await testVegaVisualizationsWithDataSources({
+        objects,
+        expectedFilteredObjects,
+        dataSourceId: 'some-datasource-id',
+        dataSourceTitle: 'dataSourceName',
+      });
+    });
+
+    test('will not update the data source name in the Vega spec if no local cluster queries', async () => {
+      const objects = [getVegaMDSVisualizationObj('some-vega-id', 'old-datasource-id')];
+      const expectedObject = getVegaMDSVisualizationObj('some-vega-id', 'old-datasource-id');
+      expectedObject.references.push({
+        id: 'some-datasource-id',
+        name: 'dataSource',
+        type: 'data-source',
+      });
+      const expectedFilteredObjects = [
+        {
+          ...expectedObject,
+          attributes: {
+            title: 'some-other-title_dataSourceName',
+            visState:
+              '{"title":"some-other-title","type":"vega","aggs":[],"params":{"spec":"{\\n  data: {\\n    url: {\\n      index: example_index\\n      data_source_name: old-datasource-title\\n    }\\n  }\\n}"}}',
+          },
+        },
+      ];
+      await testVegaVisualizationsWithDataSources({
+        objects,
+        expectedFilteredObjects,
+        dataSourceId: 'some-datasource-id',
+        dataSourceTitle: 'dataSourceName',
+      });
+    });
+  });
 });