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

…t#6123) (opensearch-project#6214) (opensearch-project#6222)

* Add MDS support for Vega



* Refactor field to data_source_id



* Add to CHANGELOG.md



* Added test cases and renamed field to use data_source_name



* Add prefix datasource name test case and add example in default hjson



* Move CHANGELOG to appropriate section



* Increased test coverage of search() method



* Add test cases for util function



* Add util function to modify Vega Spec



* Add method to verify Vega saved object type



* Add import saved object support for Vega



* Add unit tests for Vega objects in create and conflict modes



* Refactored utils test file



* Add to CHANGELOG



* Use bulkget instead of single get



* Add datasource references to the specs



* Fix bootstrap errors



* Add edge case where title is potentially undefined



* Address PR comments



* Add more test coverage for checking conflict



* Fix unit test



---------


(cherry picked from commit de978d4)


# Conflicts:
#	CHANGELOG.md



(cherry picked from commit d144637)

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>
Co-authored-by: Ashwin P Chandran <ashwinpc@amazon.com>

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

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