[Feature] Complied saved_objects create/find (opensearch-project#18)

* temp: save

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: make create/find support workspaces

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: extract management code

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* fix: type check

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* fix: build error

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: enable workspaces on saved client server side

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: some optimization

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: extract management code

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: merge fix

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: optimize code

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: reuse common function

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: optimize code when create

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

* feat: remove useless test code

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

---------

Signed-off-by: SuZhoue-Joe <suzhou@amazon.com>

Author: SuZhou-Joe

src/core/public/saved_objects/saved_objects_client.ts

@@ -42,6 +42,7 @@ import {
 
 import { SimpleSavedObject } from './simple_saved_object';
 import { HttpFetchOptions, HttpSetup } from '../http';
+import { WorkspacesStart } from '../workspace';
 
 type SavedObjectsFindOptions = Omit<
   SavedObjectFindOptionsServer,
@@ -61,6 +62,7 @@ export interface SavedObjectsCreateOptions {
   /** {@inheritDoc SavedObjectsMigrationVersion} */
   migrationVersion?: SavedObjectsMigrationVersion;
   references?: SavedObjectReference[];
+  workspaces?: string[];
 }
 
 /**
@@ -183,6 +185,7 @@ const getObjectsToFetch = (queue: BatchQueueEntry[]): ObjectTypeAndId[] => {
 export class SavedObjectsClient {
   private http: HttpSetup;
   private batchQueue: BatchQueueEntry[];
+  private currentWorkspaceId?: string;
 
   /**
    * Throttled processing of get requests into bulk requests at 100ms interval
@@ -227,6 +230,15 @@ export class SavedObjectsClient {
     this.batchQueue = [];
   }
 
+  private async _getCurrentWorkspace(): Promise<string | null> {
+    return this.currentWorkspaceId || null;
+  }
+
+  public async setCurrentWorkspace(workspaceId: string): Promise<boolean> {
+    this.currentWorkspaceId = workspaceId;
+    return true;
+  }
+
   /**
    * Persists an object
    *
@@ -235,7 +247,7 @@ export class SavedObjectsClient {
    * @param options
    * @returns
    */
-  public create = <T = unknown>(
+  public create = async <T = unknown>(
     type: string,
     attributes: T,
     options: SavedObjectsCreateOptions = {}
@@ -248,6 +260,7 @@ export class SavedObjectsClient {
     const query = {
       overwrite: options.overwrite,
     };
+    const currentWorkspaceId = await this._getCurrentWorkspace();
 
     const createRequest: Promise<SavedObject<T>> = this.savedObjectsFetch(path, {
       method: 'POST',
@@ -256,6 +269,11 @@ export class SavedObjectsClient {
         attributes,
         migrationVersion: options.migrationVersion,
         references: options.references,
+        ...(options.workspaces || currentWorkspaceId
+          ? {
+              workspaces: options.workspaces || [currentWorkspaceId],
+            }
+          : {}),
       }),
     });
 
@@ -328,7 +346,7 @@ export class SavedObjectsClient {
    * @property {object} [options.hasReference] - { type, id }
    * @returns A find result with objects matching the specified search.
    */
-  public find = <T = unknown>(
+  public find = async <T = unknown>(
     options: SavedObjectsFindOptions
   ): Promise<SavedObjectsFindResponsePublic<T>> => {
     const path = this.getPath(['_find']);
@@ -345,9 +363,18 @@ export class SavedObjectsClient {
       filter: 'filter',
       namespaces: 'namespaces',
       preference: 'preference',
+      workspaces: 'workspaces',
     };
 
-    const renamedQuery = renameKeys<SavedObjectsFindOptions, any>(renameMap, options);
+    const workspaces = [
+      ...(options.workspaces || [await this._getCurrentWorkspace()]),
+      'public',
+    ].filter((item) => item);
+
+    const renamedQuery = renameKeys<SavedObjectsFindOptions, any>(renameMap, {
+      ...options,
+      workspaces,
+    });
     const query = pick.apply(null, [renamedQuery, ...Object.values<string>(renameMap)]) as Partial<
       Record<string, any>
     >;

src/core/server/saved_objects/export/get_sorted_objects_for_export.ts

@@ -60,6 +60,8 @@ export interface SavedObjectsExportOptions {
   excludeExportDetails?: boolean;
   /** optional namespace to override the namespace used by the savedObjectsClient. */
   namespace?: string;
+  /** optional workspaces to override the workspaces used by the savedObjectsClient. */
+  workspaces?: string[];
 }
 
 /**
@@ -87,13 +89,15 @@ async function fetchObjectsToExport({
   exportSizeLimit,
   savedObjectsClient,
   namespace,
+  workspaces,
 }: {
   objects?: SavedObjectsExportOptions['objects'];
   types?: string[];
   search?: string;
   exportSizeLimit: number;
   savedObjectsClient: SavedObjectsClientContract;
   namespace?: string;
+  workspaces?: string[];
 }) {
   if ((types?.length ?? 0) > 0 && (objects?.length ?? 0) > 0) {
     throw Boom.badRequest(`Can't specify both "types" and "objects" properties when exporting`);
@@ -105,7 +109,7 @@ async function fetchObjectsToExport({
     if (typeof search === 'string') {
       throw Boom.badRequest(`Can't specify both "search" and "objects" properties when exporting`);
     }
-    const bulkGetResult = await savedObjectsClient.bulkGet(objects, { namespace });
+    const bulkGetResult = await savedObjectsClient.bulkGet(objects, { namespace, workspaces });
     const erroredObjects = bulkGetResult.saved_objects.filter((obj) => !!obj.error);
     if (erroredObjects.length) {
       const err = Boom.badRequest();
@@ -121,6 +125,7 @@ async function fetchObjectsToExport({
       search,
       perPage: exportSizeLimit,
       namespaces: namespace ? [namespace] : undefined,
+      workspaces,
     });
     if (findResponse.total > exportSizeLimit) {
       throw Boom.badRequest(`Can't export more than ${exportSizeLimit} objects`);
@@ -153,6 +158,7 @@ export async function exportSavedObjectsToStream({
   includeReferencesDeep = false,
   excludeExportDetails = false,
   namespace,
+  workspaces,
 }: SavedObjectsExportOptions) {
   const rootObjects = await fetchObjectsToExport({
     types,
@@ -161,6 +167,7 @@ export async function exportSavedObjectsToStream({
     savedObjectsClient,
     exportSizeLimit,
     namespace,
+    workspaces,
   });
   let exportedObjects: Array<SavedObject<unknown>> = [];
   let missingReferences: SavedObjectsExportResultDetails['missingReferences'] = [];

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

@@ -39,6 +39,7 @@ interface CreateSavedObjectsParams<T> {
   importIdMap: Map<string, { id?: string; omitOriginId?: boolean }>;
   namespace?: string;
   overwrite?: boolean;
+  workspaces?: string[];
 }
 interface CreateSavedObjectsResult<T> {
   createdObjects: Array<CreatedObject<T>>;
@@ -56,6 +57,7 @@ export const createSavedObjects = async <T>({
   importIdMap,
   namespace,
   overwrite,
+  workspaces,
 }: CreateSavedObjectsParams<T>): Promise<CreateSavedObjectsResult<T>> => {
   // filter out any objects that resulted in errors
   const errorSet = accumulatedErrors.reduce(
@@ -103,14 +105,15 @@ export const createSavedObjects = async <T>({
     const bulkCreateResponse = await savedObjectsClient.bulkCreate(objectsToCreate, {
       namespace,
       overwrite,
+      workspaces,
     });
     expectedResults = bulkCreateResponse.saved_objects;
   }
 
   // remap results to reflect the object IDs that were submitted for import
   // this ensures that consumers understand the results
   const remappedResults = expectedResults.map<CreatedObject<T>>((result) => {
-    const { id } = objectIdMap.get(`${result.type}:${result.id}`)!;
+    const { id } = objectIdMap.get(`${result.type}:${result.id}`) || ({} as SavedObject<T>);
     // also, include a `destinationId` field if the object create attempt was made with a different ID
     return { ...result, id, ...(id !== result.id && { destinationId: result.id }) };
   });

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

@@ -54,6 +54,7 @@ export async function importSavedObjectsFromStream({
   savedObjectsClient,
   typeRegistry,
   namespace,
+  workspaces,
 }: SavedObjectsImportOptions): Promise<SavedObjectsImportResponse> {
   let errorAccumulator: SavedObjectsImportError[] = [];
   const supportedTypes = typeRegistry.getImportableAndExportableTypes().map((type) => type.name);
@@ -118,6 +119,7 @@ export async function importSavedObjectsFromStream({
     importIdMap,
     overwrite,
     namespace,
+    workspaces,
   };
   const createSavedObjectsResult = await createSavedObjects(createSavedObjectsParams);
   errorAccumulator = [...errorAccumulator, ...createSavedObjectsResult.errors];

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

@@ -187,6 +187,8 @@ export interface SavedObjectsImportOptions {
   namespace?: string;
   /** If true, will create new copies of import objects, each with a random `id` and undefined `originId`. */
   createNewCopies: boolean;
+  /** if specified, will import in given workspaces, else will import as global object */
+  workspaces?: string[];
 }
 
 /**

src/core/server/saved_objects/migrations/core/build_active_mappings.ts

@@ -175,6 +175,9 @@ function defaultMapping(): IndexMapping {
           },
         },
       },
+      workspaces: {
+        type: 'keyword',
+      },
     },
   };
 }

src/core/server/saved_objects/routes/bulk_create.ts

@@ -30,6 +30,7 @@
 
 import { schema } from '@osd/config-schema';
 import { IRouter } from '../../http';
+import { formatWorkspaces, workspacesValidator } from './utils';
 
 export const registerBulkCreateRoute = (router: IRouter) => {
   router.post(
@@ -38,6 +39,7 @@ export const registerBulkCreateRoute = (router: IRouter) => {
       validate: {
         query: schema.object({
           overwrite: schema.boolean({ defaultValue: false }),
+          workspaces: workspacesValidator,
         }),
         body: schema.arrayOf(
           schema.object({
@@ -62,7 +64,11 @@ export const registerBulkCreateRoute = (router: IRouter) => {
     },
     router.handleLegacyErrors(async (context, req, res) => {
       const { overwrite } = req.query;
-      const result = await context.core.savedObjects.client.bulkCreate(req.body, { overwrite });
+      const workspaces = formatWorkspaces(req.query.workspaces);
+      const result = await context.core.savedObjects.client.bulkCreate(req.body, {
+        overwrite,
+        workspaces,
+      });
       return res.ok({ body: result });
     })
   );

src/core/server/saved_objects/routes/create.ts

@@ -56,15 +56,23 @@ export const registerCreateRoute = (router: IRouter) => {
             )
           ),
           initialNamespaces: schema.maybe(schema.arrayOf(schema.string(), { minSize: 1 })),
+          workspaces: schema.maybe(schema.arrayOf(schema.string(), { minSize: 1 })),
         }),
       },
     },
     router.handleLegacyErrors(async (context, req, res) => {
       const { type, id } = req.params;
       const { overwrite } = req.query;
-      const { attributes, migrationVersion, references, initialNamespaces } = req.body;
+      const { attributes, migrationVersion, references, initialNamespaces, workspaces } = req.body;
 
-      const options = { id, overwrite, migrationVersion, references, initialNamespaces };
+      const options = {
+        id,
+        overwrite,
+        migrationVersion,
+        references,
+        initialNamespaces,
+        workspaces,
+      };
       const result = await context.core.savedObjects.client.create(type, attributes, options);
       return res.ok({ body: result });
     })

src/core/server/saved_objects/routes/export.ts

@@ -57,12 +57,20 @@ export const registerExportRoute = (router: IRouter, config: SavedObjectConfig)
           search: schema.maybe(schema.string()),
           includeReferencesDeep: schema.boolean({ defaultValue: false }),
           excludeExportDetails: schema.boolean({ defaultValue: false }),
+          workspaces: schema.maybe(schema.arrayOf(schema.string())),
         }),
       },
     },
     router.handleLegacyErrors(async (context, req, res) => {
       const savedObjectsClient = context.core.savedObjects.client;
-      const { type, objects, search, excludeExportDetails, includeReferencesDeep } = req.body;
+      const {
+        type,
+        objects,
+        search,
+        excludeExportDetails,
+        includeReferencesDeep,
+        workspaces,
+      } = req.body;
       const types = typeof type === 'string' ? [type] : type;
 
       // need to access the registry for type validation, can't use the schema for this
@@ -98,6 +106,7 @@ export const registerExportRoute = (router: IRouter, config: SavedObjectConfig)
         exportSizeLimit: maxImportExportSize,
         includeReferencesDeep,
         excludeExportDetails,
+        workspaces,
       });
 
       const docsToExport: string[] = await createPromiseFromStreams([

src/core/server/saved_objects/routes/find.ts

@@ -30,6 +30,7 @@
 
 import { schema } from '@osd/config-schema';
 import { IRouter } from '../../http';
+import { formatWorkspaces, workspacesValidator } from './utils';
 
 export const registerFindRoute = (router: IRouter) => {
   router.get(
@@ -59,6 +60,7 @@ export const registerFindRoute = (router: IRouter) => {
           namespaces: schema.maybe(
             schema.oneOf([schema.string(), schema.arrayOf(schema.string())])
           ),
+          workspaces: workspacesValidator,
         }),
       },
     },
@@ -67,6 +69,7 @@ export const registerFindRoute = (router: IRouter) => {
 
       const namespaces =
         typeof req.query.namespaces === 'string' ? [req.query.namespaces] : req.query.namespaces;
+      const workspaces = formatWorkspaces(query.workspaces);
 
       const result = await context.core.savedObjects.client.find({
         perPage: query.per_page,
@@ -81,6 +84,7 @@ export const registerFindRoute = (router: IRouter) => {
         fields: typeof query.fields === 'string' ? [query.fields] : query.fields,
         filter: query.filter,
         namespaces,
+        workspaces,
       });
 
       return res.ok({ body: result });

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

@@ -34,7 +34,7 @@ import { schema } from '@osd/config-schema';
 import { IRouter } from '../../http';
 import { importSavedObjectsFromStream } from '../import';
 import { SavedObjectConfig } from '../saved_objects_config';
-import { createSavedObjectsStreamFromNdJson } from './utils';
+import { createSavedObjectsStreamFromNdJson, formatWorkspaces, workspacesValidator } from './utils';
 
 interface FileStream extends Readable {
   hapi: {
@@ -60,6 +60,7 @@ export const registerImportRoute = (router: IRouter, config: SavedObjectConfig)
           {
             overwrite: schema.boolean({ defaultValue: false }),
             createNewCopies: schema.boolean({ defaultValue: false }),
+            workspaces: workspacesValidator,
           },
           {
             validate: (object) => {
@@ -91,13 +92,16 @@ export const registerImportRoute = (router: IRouter, config: SavedObjectConfig)
         });
       }
 
+      const workspaces = formatWorkspaces(req.query.workspaces);
+
       const result = await importSavedObjectsFromStream({
         savedObjectsClient: context.core.savedObjects.client,
         typeRegistry: context.core.savedObjects.typeRegistry,
         readStream,
         objectLimit: maxImportExportSize,
         overwrite,
         createNewCopies,
+        workspaces,
       });
 
       return res.ok({ body: result });