Do not allow to copy saved objects to the workspace which unassigned …

…data source

Signed-off-by: yubonluo <[email protected]>

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

@@ -35,6 +35,7 @@ import {
   SavedObjectsType,
   SavedObject,
   SavedObjectsImportError,
+  SavedObjectsBaseOptions,
 } from '../types';
 import { savedObjectsClientMock } from '../../mocks';
 import { SavedObjectsImportOptions, ISavedObjectTypeRegistry } from '..';
@@ -48,6 +49,7 @@ import { checkConflicts } from './check_conflicts';
 import { checkOriginConflicts } from './check_origin_conflicts';
 import { createSavedObjects } from './create_saved_objects';
 import { checkConflictsForDataSource } from './check_conflict_for_data_source';
+import { findDataSourceForObject } from './utils';
 
 jest.mock('./collect_saved_objects');
 jest.mock('./regenerate_ids');
@@ -56,6 +58,7 @@ jest.mock('./check_conflicts');
 jest.mock('./check_origin_conflicts');
 jest.mock('./create_saved_objects');
 jest.mock('./check_conflict_for_data_source');
+jest.mock('./utils');
 
 const getMockFn = <T extends (...args: any[]) => any, U>(fn: (...args: Parameters<T>) => U) =>
   fn as jest.MockedFunction<(...args: Parameters<T>) => U>;
@@ -101,7 +104,9 @@ describe('#importSavedObjectsFromStream', () => {
   const setupOptions = (
     createNewCopies: boolean = false,
     dataSourceId: string | undefined = undefined,
-    dataSourceEnabled: boolean | undefined = false
+    dataSourceEnabled: boolean | undefined = false,
+    workspaces: SavedObjectsBaseOptions['workspaces'] = undefined,
+    assignedDataSources: string[] = []
   ): SavedObjectsImportOptions => {
     readStream = new Readable();
     savedObjectsClient = savedObjectsClientMock.create();
@@ -122,6 +127,8 @@ describe('#importSavedObjectsFromStream', () => {
       namespace,
       createNewCopies,
       dataSourceId,
+      workspaces,
+      assignedDataSources,
     };
   };
   const createObject = (
@@ -397,6 +404,48 @@ describe('#importSavedObjectsFromStream', () => {
       expect(result).toEqual({ success: false, successCount: 0, errors: [expect.any(Object)] });
     });
 
+    test('validates workspace with assigned data source', async () => {
+      const options = setupOptions(false, undefined, false, ['workspace-1'], ['dataSource-1']);
+      const collectedObjects = [createObject()];
+      getMockFn(collectSavedObjects).mockResolvedValue({
+        errors: [],
+        collectedObjects,
+        importIdMap: new Map(),
+      });
+      getMockFn(findDataSourceForObject).mockResolvedValue('dataSource-1');
+
+      const result = await importSavedObjectsFromStream(options);
+      expect(result).toEqual({ success: true, successCount: 0 });
+    });
+
+    test('validates workspace with unassigned data source', async () => {
+      const options = setupOptions(false, undefined, false, ['workspace-1'], ['dataSource-1']);
+      const collectedObjects = [createObject()];
+      getMockFn(collectSavedObjects).mockResolvedValue({
+        errors: [],
+        collectedObjects,
+        importIdMap: new Map(),
+      });
+      getMockFn(findDataSourceForObject).mockResolvedValue('dataSource-2');
+
+      const result = await importSavedObjectsFromStream(options);
+      expect(result).toEqual({ success: false, successCount: 0, errors: [expect.any(Object)] });
+    });
+
+    test('validates workspace with catch error', async () => {
+      const options = setupOptions(false, undefined, false, ['workspace-1'], ['dataSource-1']);
+      const collectedObjects = [createObject()];
+      getMockFn(collectSavedObjects).mockResolvedValue({
+        errors: [],
+        collectedObjects,
+        importIdMap: new Map(),
+      });
+      getMockFn(findDataSourceForObject).mockRejectedValue(null);
+
+      const result = await importSavedObjectsFromStream(options);
+      expect(result).toEqual({ success: false, successCount: 0, errors: [expect.any(Object)] });
+    });
+
     describe('handles a mix of successes and errors and injects metadata', () => {
       const obj1 = createObject();
       const tmp = createObject();

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

@@ -42,6 +42,7 @@ import { checkConflicts } from './check_conflicts';
 import { regenerateIds } from './regenerate_ids';
 import { checkConflictsForDataSource } from './check_conflict_for_data_source';
 import { isSavedObjectWithDataSource } from './validate_object_id';
+import { findDataSourceForObject } from './utils';
 
 /**
  * Import saved objects from given stream. See the {@link SavedObjectsImportOptions | options} for more
@@ -61,6 +62,7 @@ export async function importSavedObjectsFromStream({
   dataSourceTitle,
   workspaces,
   dataSourceEnabled,
+  assignedDataSources,
 }: SavedObjectsImportOptions): Promise<SavedObjectsImportResponse> {
   let errorAccumulator: SavedObjectsImportError[] = [];
   const supportedTypes = typeRegistry.getImportableAndExportableTypes().map((type) => type.name);
@@ -95,6 +97,39 @@ export async function importSavedObjectsFromStream({
     }
   }
 
+  // If target workspace and assigned data source is exists, it will check whether
+  // the assigned data sources in the target workspace include the data sources of the imported saved objects.
+  if (workspaces && workspaces.length > 0 && assignedDataSources) {
+    const errors: SavedObjectsImportError[] = [];
+    const duplicateObjects: typeof collectSavedObjectsResult.collectedObjects = [];
+    for (const object of collectSavedObjectsResult.collectedObjects) {
+      try {
+        const referenceDSId = await findDataSourceForObject(object, savedObjectsClient);
+        if (referenceDSId && !assignedDataSources.some((ds) => ds === referenceDSId)) {
+          const error: SavedObjectsImportError = {
+            type: object.type,
+            id: object.id,
+            error: { type: 'unsupported_type' },
+            meta: { title: object.attributes?.title },
+          };
+          errors.push(error);
+        } else {
+          duplicateObjects.push(object);
+        }
+      } catch (err) {
+        const error: SavedObjectsImportError = {
+          type: object.type,
+          id: object.id,
+          error: { type: 'unsupported_type' },
+          meta: { title: object.attributes?.title },
+        };
+        errors.push(error);
+      }
+    }
+    errorAccumulator = [...errorAccumulator, ...errors];
+    collectSavedObjectsResult.collectedObjects = duplicateObjects;
+  }
+
   errorAccumulator = [...errorAccumulator, ...collectSavedObjectsResult.errors];
   /** Map of all IDs for objects that we are attempting to import; each value is empty by default */
   let importIdMap = collectSavedObjectsResult.importIdMap;

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

@@ -191,6 +191,7 @@ export interface SavedObjectsImportOptions {
   dataSourceTitle?: string;
   dataSourceEnabled?: boolean;
   workspaces?: SavedObjectsBaseOptions['workspaces'];
+  assignedDataSources?: string[];
 }
 
 /**

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

@@ -10,11 +10,13 @@ import {
   getUpdatedTSVBVisState,
   updateDataSourceNameInVegaSpec,
   updateDataSourceNameInTimeline,
+  findDataSourceForObject,
 } from './utils';
 import { parse } from 'hjson';
 import { isEqual } from 'lodash';
 import { join } from 'path';
 import { SavedObject, SavedObjectsClientContract } from '../types';
+import { savedObjectsClientMock } from '../../mocks';
 
 describe('updateDataSourceNameInVegaSpec()', () => {
   const loadHJSONStringFromFile = (filepath: string) => {
@@ -342,3 +344,107 @@ describe('getUpdatedTSVBVisState', () => {
     }
   );
 });
+
+describe('findDataSourceForObject', () => {
+  let savedObjectsClient: jest.Mocked<SavedObjectsClientContract>;
+  const indexPatternObject: SavedObject = {
+    id: 'indexPattern',
+    type: 'index-pattern',
+    references: [{ type: 'data-source', id: 'dataSource', name: 'dataSource' }],
+    attributes: {},
+  };
+
+  beforeEach(() => {
+    savedObjectsClient = savedObjectsClientMock.create();
+  });
+
+  it('should return the data source id for an index-pattern object', async () => {
+    const dataSourceId = await findDataSourceForObject(indexPatternObject, savedObjectsClient);
+    expect(dataSourceId).toBe('dataSource');
+  });
+
+  it('should use bulkGet to resolve multiple references and return data source', async () => {
+    const savedObject: SavedObject = {
+      type: 'dashboard',
+      id: 'dashboard',
+      references: [{ type: 'visualization', id: 'visualization', name: 'visualization' }],
+      attributes: {},
+    };
+
+    const visualizationObject: SavedObject = {
+      type: 'visualization',
+      id: 'visualization',
+      references: [{ type: 'index-pattern', id: 'indexPattern', name: 'indexPattern' }],
+      attributes: {},
+    };
+
+    savedObjectsClient.bulkGet.mockResolvedValueOnce({
+      saved_objects: [visualizationObject],
+    });
+    savedObjectsClient.bulkGet.mockResolvedValueOnce({
+      saved_objects: [indexPatternObject],
+    });
+
+    const dataSourceId = await findDataSourceForObject(savedObject, savedObjectsClient);
+    expect(dataSourceId).toBe('dataSource');
+  });
+
+  it('should use bulkGet to resolve multiple references and return the first found data source', async () => {
+    const savedObject: SavedObject = {
+      id: 'visualization',
+      type: 'visualization',
+      references: [
+        // { type: 'index-pattern', id: 'indexPattern', name: 'index-pattern' },
+        { type: 'dashboard', id: 'dashboard', name: 'dashboard' },
+      ],
+      attributes: {},
+    };
+
+    savedObjectsClient.bulkGet.mockResolvedValue({
+      saved_objects: [indexPatternObject],
+    });
+
+    const dataSourceId = await findDataSourceForObject(savedObject, savedObjectsClient);
+    expect(dataSourceId).toBe('dataSource');
+    expect(savedObjectsClient.bulkGet).toHaveBeenCalledWith([
+      { type: 'dashboard', id: 'dashboard' },
+    ]);
+  });
+
+  it('should return null if no data source is found', async () => {
+    const savedObject: SavedObject = {
+      id: 'visualization',
+      type: 'visualization',
+      references: [],
+      attributes: {},
+    };
+
+    const dataSourceId = await findDataSourceForObject(savedObject, savedObjectsClient);
+    expect(dataSourceId).toBeNull();
+  });
+
+  it('should return null if there is an error in bulkGet', async () => {
+    const savedObject: SavedObject = {
+      id: 'visualization',
+      type: 'visualization',
+      references: [{ type: 'index-pattern', id: 'indexPattern', name: 'indexPattern' }],
+      attributes: {},
+    };
+
+    savedObjectsClient.bulkGet.mockResolvedValue({
+      saved_objects: [
+        {
+          id: 'indexPattern',
+          type: 'index-pattern',
+          error: { error: '', statusCode: 400, message: '' },
+          attributes: undefined,
+          references: [],
+        },
+      ],
+    });
+
+    await expect(findDataSourceForObject(savedObject, savedObjectsClient)).rejects.toThrow(
+      'Bad Request'
+    );
+  });
+});

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

@@ -4,6 +4,7 @@
  */
 
 import { parse, stringify } from 'hjson';
+import Boom from '@hapi/boom';
 import { SavedObject, SavedObjectReference, SavedObjectsClientContract } from '../types';
 import { VisualizationObject } from './types';
 
@@ -183,3 +184,53 @@ const parseJSONSpec = (spec: string) => {
 
   return undefined;
 };
+
+export async function findDataSourceForObject(
+  savedObject: SavedObject,
+  savedObjectsClient: SavedObjectsClientContract,
+  visitedObjects: Set<string> = new Set()
+): Promise<string | null> {
+  const references = savedObject.references;
+  if (!references || references.length === 0) {
+    return null;
+  }
+
+  const objectKey = `${savedObject.type}:${savedObject.id}`;
+  if (visitedObjects.has(objectKey)) {
+    return null;
+  }
+  visitedObjects.add(objectKey);
+
+  const dataSourceReference = references.find((ref) => ref.type === 'data-source');
+  if (dataSourceReference) {
+    return dataSourceReference.id;
+  }
+
+  const bulkGetResponse = await savedObjectsClient.bulkGet(
+    references.map((reference) => ({
+      type: reference.type,
+      id: reference.id,
+    }))
+  );
+
+  const referencedObjects = bulkGetResponse.saved_objects;
+  const erroredObjects = referencedObjects.filter(
+    (obj) => obj.error && obj.error.statusCode !== 404
+  );
+
+  if (erroredObjects.length > 0) {
+    const err = Boom.badRequest();
+    err.output.payload.attributes = {
+      objects: erroredObjects,
+    };
+    throw err;
+  }
+
+  for (const referencedObject of referencedObjects) {
+    const dataSource = await findDataSourceForObject(referencedObject, savedObjectsClient);
+    if (dataSource) {
+      return dataSource;
+    }
+  }
+  return null;
+}

src/plugins/workspace/opensearch_dashboards.json

@@ -7,6 +7,6 @@
     "savedObjects",
     "opensearchDashboardsReact"
   ],
-  "optionalPlugins": ["savedObjectsManagement","management","dataSourceManagement","contentManagement"],
+  "optionalPlugins": ["savedObjectsManagement","management","dataSourceManagement","contentManagement","dataSource"],
   "requiredBundles": ["opensearchDashboardsReact", "home"]
 }

src/plugins/workspace/server/plugin.ts

@@ -41,6 +41,11 @@ import {
 import { getOSDAdminConfigFromYMLConfig, updateDashboardAdminStateForRequest } from './utils';
 import { WorkspaceIdConsumerWrapper } from './saved_objects/workspace_id_consumer_wrapper';
 import { WorkspaceUiSettingsClientWrapper } from './saved_objects/workspace_ui_settings_client_wrapper';
+import { DataSourcePluginSetup } from '../../data_source/server';
+
+export interface WorkspacePluginDependencies {
+  dataSource: DataSourcePluginSetup;
+}
 
 export class WorkspacePlugin implements Plugin<WorkspacePluginSetup, WorkspacePluginStart> {
   private readonly logger: Logger;
@@ -113,10 +118,11 @@ export class WorkspacePlugin implements Plugin<WorkspacePluginSetup, WorkspacePl
     this.globalConfig$ = initializerContext.config.legacy.globalConfig$;
   }
 
-  public async setup(core: CoreSetup) {
+  public async setup(core: CoreSetup, deps: WorkspacePluginDependencies) {
     this.logger.debug('Setting up Workspaces service');
     const globalConfig = await this.globalConfig$.pipe(first()).toPromise();
     const isPermissionControlEnabled = globalConfig.savedObjects.permission.enabled === true;
+    const isDataSourceEnabled = !!deps.dataSource;
 
     this.client = new WorkspaceClient(core, this.logger);
 
@@ -156,6 +162,7 @@ export class WorkspacePlugin implements Plugin<WorkspacePluginSetup, WorkspacePl
       maxImportExportSize,
       permissionControlClient: this.permissionControl,
       isPermissionControlEnabled,
+      isDataSourceEnabled,
     });
 
     core.capabilities.registerProvider(() => ({