From 91605cc22f5aeb70673ba958f2dec67739e29faa Mon Sep 17 00:00:00 2001
From: SuZhou-Joe <suzhou@amazon.com>
Date: Sun, 27 Aug 2023 15:19:56 +0800
Subject: [PATCH] feat: update

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>
---
 src/plugins/workspace/server/plugin.ts        | 13 +++++
 .../workspace_saved_objects_client_wrapper.ts | 49 ++++++-------------
 2 files changed, 29 insertions(+), 33 deletions(-)

diff --git a/src/plugins/workspace/server/plugin.ts b/src/plugins/workspace/server/plugin.ts
index 5d6fe4ad6987..77231db2eb29 100644
--- a/src/plugins/workspace/server/plugin.ts
+++ b/src/plugins/workspace/server/plugin.ts
@@ -85,6 +85,19 @@ export class WorkspacePlugin implements Plugin<{}, {}> {
       this.workspaceSavedObjectsClientWrapper.wrapperFactory
     );
 
+    core.savedObjects.setClientFactoryProvider(
+      (repositoryFactory) => ({ request, includedHiddenTypes }) => {
+        const enabled = this.isEnabled;
+        if (enabled) {
+          return new SavedObjectsClient(repositoryFactory.createInternalRepository());
+        }
+
+        return new SavedObjectsClient(
+          repositoryFactory.createScopedRepository(request, includedHiddenTypes)
+        );
+      }
+    );
+
     this.proxyWorkspaceTrafficToRealHandler(core);
 
     registerRoutes({
diff --git a/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts b/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts
index d88ec4d63039..c14ffa7343af 100644
--- a/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts
+++ b/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts
@@ -163,23 +163,6 @@ export class WorkspaceSavedObjectsClientWrapper {
   }
 
   public wrapperFactory: SavedObjectsClientWrapperFactory = (wrapperOptions) => {
-    /**
-     * The client here is scopedSavedObjectsClient by default
-     */
-    let client = wrapperOptions.client;
-    const featureFlagEnabled = this.options.enabled$.getValue();
-
-    if (!featureFlagEnabled) {
-      return client;
-    }
-
-    /**
-     * If featureFlag is open and we have internalRepositoryFactory
-     * Use internal repository as access control will be provided by ACL.
-     */
-    if (featureFlagEnabled && this.internalRepositoryFactory) {
-      client = new SavedObjectsClient(this.internalRepositoryFactory());
-    }
     const deleteWithWorkspacePermissionControl = async (
       type: string,
       id: string,
@@ -191,13 +174,13 @@ export class WorkspaceSavedObjectsClientWrapper {
         ]);
       }
 
-      const objectToDeleted = await client.get(type, id, options);
+      const objectToDeleted = await wrapperOptions.client.get(type, id, options);
       await this.validateMultiWorkspacesPermissions(
         objectToDeleted.workspaces,
         wrapperOptions.request,
         [WorkspacePermissionMode.LibraryWrite, WorkspacePermissionMode.Management]
       );
-      return await client.delete(type, id, options);
+      return await wrapperOptions.client.delete(type, id, options);
     };
 
     const updateWithWorkspacePermissionControl = async <T = unknown>(
@@ -211,7 +194,7 @@ export class WorkspaceSavedObjectsClientWrapper {
           WorkspacePermissionMode.Management,
         ]);
       }
-      return await client.update(type, id, attributes, options);
+      return await wrapperOptions.client.update(type, id, attributes, options);
     };
 
     const bulkUpdateWithWorkspacePermissionControl = async <T = unknown>(
@@ -233,7 +216,7 @@ export class WorkspaceSavedObjectsClientWrapper {
         throw generateWorkspacePermissionError();
       }
 
-      return await client.bulkUpdate(objects, options);
+      return await wrapperOptions.client.bulkUpdate(objects, options);
     };
 
     const bulkCreateWithWorkspacePermissionControl = async <T = unknown>(
@@ -246,7 +229,7 @@ export class WorkspaceSavedObjectsClientWrapper {
           WorkspacePermissionMode.Management,
         ]);
       }
-      return await client.bulkCreate(objects, options);
+      return await wrapperOptions.client.bulkCreate(objects, options);
     };
 
     const createWithWorkspacePermissionControl = async <T = unknown>(
@@ -261,7 +244,7 @@ export class WorkspaceSavedObjectsClientWrapper {
           [WorkspacePermissionMode.LibraryWrite, WorkspacePermissionMode.Management]
         );
       }
-      return await client.create(type, attributes, options);
+      return await wrapperOptions.client.create(type, attributes, options);
     };
 
     const getWithWorkspacePermissionControl = async <T = unknown>(
@@ -269,7 +252,7 @@ export class WorkspaceSavedObjectsClientWrapper {
       id: string,
       options: SavedObjectsBaseOptions = {}
     ): Promise<SavedObject<T>> => {
-      const objectToGet = await client.get<T>(type, id, options);
+      const objectToGet = await wrapperOptions.client.get<T>(type, id, options);
       await this.validateAtLeastOnePermittedWorkspaces(
         objectToGet.workspaces,
         wrapperOptions.request,
@@ -286,7 +269,7 @@ export class WorkspaceSavedObjectsClientWrapper {
       objects: SavedObjectsBulkGetObject[] = [],
       options: SavedObjectsBaseOptions = {}
     ): Promise<SavedObjectsBulkResponse<T>> => {
-      const objectToBulkGet = await client.bulkGet<T>(objects, options);
+      const objectToBulkGet = await wrapperOptions.client.bulkGet<T>(objects, options);
       for (const object of objectToBulkGet.saved_objects) {
         await this.validateAtLeastOnePermittedWorkspaces(
           object.workspaces,
@@ -395,7 +378,7 @@ export class WorkspaceSavedObjectsClientWrapper {
         }
       }
 
-      return await client.find<T>(options);
+      return await wrapperOptions.client.find<T>(options);
     };
 
     const addToWorkspacesWithPermissionControl = async (
@@ -422,24 +405,24 @@ export class WorkspaceSavedObjectsClientWrapper {
         throw generateSavedObjectsPermissionError();
       }
 
-      return await client.addToWorkspaces(objects, targetWorkspaces, options);
+      return await wrapperOptions.client.addToWorkspaces(objects, targetWorkspaces, options);
     };
 
     const isDashboardAdmin = this.isDashboardAdmin(wrapperOptions.request);
 
     if (isDashboardAdmin) {
-      return client;
+      return wrapperOptions.client;
     }
 
     return {
-      ...client,
+      ...wrapperOptions.client,
       get: getWithWorkspacePermissionControl,
-      checkConflicts: client.checkConflicts,
+      checkConflicts: wrapperOptions.client.checkConflicts,
       find: findWithWorkspacePermissionControl,
       bulkGet: bulkGetWithWorkspacePermissionControl,
-      errors: client.errors,
-      addToNamespaces: client.addToNamespaces,
-      deleteFromNamespaces: client.deleteFromNamespaces,
+      errors: wrapperOptions.client.errors,
+      addToNamespaces: wrapperOptions.client.addToNamespaces,
+      deleteFromNamespaces: wrapperOptions.client.deleteFromNamespaces,
       create: createWithWorkspacePermissionControl,
       bulkCreate: bulkCreateWithWorkspacePermissionControl,
       delete: deleteWithWorkspacePermissionControl,