Optimize augment-vis saved obj searching by adding arg to saved obj client (#4595)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>

Author: ohltyler

CHANGELOG.md

@@ -37,6 +37,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Dashboard De-Angularization ([#4502](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/4502))
 - New management overview page and rename stack management to dashboard management ([#4287](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/4287))
 - [Vis Augmenter] Update base vis height in view events flyout ([#4535](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/4535))
+- Optimize `augment-vis` saved obj searching by adding arg to saved obj client ([#4595](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/4595))
+
 
 ### 🐛 Bug Fixes
 

src/core/public/saved_objects/saved_objects_client.test.ts

@@ -471,10 +471,7 @@ describe('SavedObjectsClient', () => {
                 "fields": Array [
                   "title",
                 ],
-                "has_reference": Object {
-                  "id": "1",
-                  "type": "reference",
-                },
+                "has_reference": "{\\"id\\":\\"1\\",\\"type\\":\\"reference\\"}",
                 "page": 10,
                 "per_page": 100,
                 "search": "what is the meaning of life?|life",

src/core/public/saved_objects/saved_objects_client.ts

@@ -348,7 +348,13 @@ export class SavedObjectsClient {
     };
 
     const renamedQuery = renameKeys<SavedObjectsFindOptions, any>(renameMap, options);
-    const query = pick.apply(null, [renamedQuery, ...Object.values<string>(renameMap)]);
+    const query = pick.apply(null, [renamedQuery, ...Object.values<string>(renameMap)]) as Partial<
+      Record<string, any>
+    >;
+
+    // has_reference needs post-processing since it is an object that needs to be read as
+    // a query param
+    if (query.has_reference) query.has_reference = JSON.stringify(query.has_reference);
 
     const request: ReturnType<SavedObjectsApi['find']> = this.savedObjectsFetch(path, {
       method: 'GET',

src/plugins/saved_objects/public/saved_object/saved_object_loader.ts

@@ -128,7 +128,12 @@ export class SavedObjectLoader {
    * @param fields
    * @returns {Promise}
    */
-  findAll(search: string = '', size: number = 100, fields?: string[]) {
+  findAll(
+    search: string = '',
+    size: number = 100,
+    fields?: string[],
+    hasReference?: SavedObjectsFindOptions['hasReference']
+  ) {
     return this.savedObjectsClient
       .find<Record<string, unknown>>({
         type: this.lowercaseType,
@@ -138,6 +143,7 @@ export class SavedObjectLoader {
         searchFields: ['title^3', 'description'],
         defaultSearchOperator: 'AND',
         fields,
+        hasReference,
       } as SavedObjectsFindOptions)
       .then((resp) => {
         return {

src/plugins/vis_augmenter/public/actions/saved_object_delete_action.test.ts

@@ -18,6 +18,20 @@ jest.mock('src/plugins/vis_augmenter/public/services.ts', () => {
         },
       };
     },
+    getUISettings: () => {
+      return {
+        get: (config: string) => {
+          switch (config) {
+            case 'visualization:enablePluginAugmentation':
+              return true;
+            case 'visualization:enablePluginAugmentation.maxPluginObjects':
+              return 10;
+            default:
+              throw new Error(`Accessing ${config} is not supported in the mock.`);
+          }
+        },
+      };
+    },
   };
 });
 

src/plugins/vis_augmenter/public/actions/saved_object_delete_action.ts

@@ -7,7 +7,7 @@ import { isEmpty } from 'lodash';
 import { i18n } from '@osd/i18n';
 import { EuiIconType } from '@elastic/eui/src/components/icon/icon';
 import { Action, IncompatibleActionError } from '../../../ui_actions/public';
-import { getAllAugmentVisSavedObjs } from '../utils';
+import { getAugmentVisSavedObjs } from '../utils';
 import { getSavedAugmentVisLoader } from '../services';
 import { SavedObjectDeleteContext } from '../ui_actions_bootstrap';
 
@@ -45,12 +45,19 @@ export class SavedObjectDeleteAction implements Action<SavedObjectDeleteContext>
       throw new IncompatibleActionError();
     }
 
-    const loader = getSavedAugmentVisLoader();
-    const allAugmentVisObjs = await getAllAugmentVisSavedObjs(loader);
-    const augmentVisIdsToDelete = allAugmentVisObjs
-      .filter((augmentVisObj) => augmentVisObj.visId === savedObjectId)
-      .map((augmentVisObj) => augmentVisObj.id as string);
-
-    if (!isEmpty(augmentVisIdsToDelete)) loader.delete(augmentVisIdsToDelete);
+    try {
+      const loader = getSavedAugmentVisLoader();
+      const augmentVisObjs = await getAugmentVisSavedObjs(savedObjectId, loader);
+      const augmentVisIdsToDelete = augmentVisObjs.map(
+        (augmentVisObj) => augmentVisObj.id as string
+      );
+
+      if (!isEmpty(augmentVisIdsToDelete)) loader.delete(augmentVisIdsToDelete);
+      // silently fail. this is because this is doing extra cleanup on objects unrelated
+      // to the user flow so we dont want to add confusing errors on UI.
+    } catch (e) {
+      // eslint-disable-next-line no-console
+      console.error(e);
+    }
   }
 }

src/plugins/vis_augmenter/public/saved_augment_vis/saved_augment_vis.ts

@@ -4,7 +4,7 @@
  */
 
 import { get, isEmpty } from 'lodash';
-import { IUiSettingsClient } from 'opensearch-dashboards/public';
+import { IUiSettingsClient, SavedObjectsFindOptions } from 'opensearch-dashboards/public';
 import {
   SavedObjectLoader,
   SavedObjectOpenSearchDashboardsServices,
@@ -91,9 +91,14 @@ export class SavedObjectLoaderAugmentVis extends SavedObjectLoader {
    * @param fields
    * @returns {Promise}
    */
-  findAll(search: string = '', size: number = 100, fields?: string[]) {
+  findAll(
+    search: string = '',
+    size: number = 100,
+    fields?: string[],
+    hasReference?: SavedObjectsFindOptions['hasReference']
+  ) {
     this.isAugmentationEnabled();
-    return super.findAll(search, size, fields);
+    return super.findAll(search, size, fields, hasReference);
   }
 
   find(search: string = '', size: number = 100) {

src/plugins/vis_augmenter/public/saved_augment_vis/utils/helpers.ts

@@ -33,23 +33,24 @@ export const createAugmentVisSavedObject = async (
   }
   const maxAssociatedCount = config.get(PLUGIN_AUGMENTATION_MAX_OBJECTS_SETTING);
 
-  await loader.findAll().then(async (resp) => {
-    if (resp !== undefined) {
-      const savedAugmentObjects = get(resp, 'hits', []);
-      // gets all the saved object for this visualization
-      const savedObjectsForThisVisualization = savedAugmentObjects.filter(
-        (savedObj) => get(savedObj, 'visId', '') === AugmentVis.visId
-      );
+  await loader
+    .findAll('', 100, [], {
+      type: 'visualization',
+      id: AugmentVis.visId as string,
+    })
+    .then(async (resp) => {
+      if (resp !== undefined) {
+        const savedObjectsForThisVisualization = get(resp, 'hits', []);
 
-      if (maxAssociatedCount <= savedObjectsForThisVisualization.length) {
-        throw new Error(
-          `Cannot associate the plugin resource to the visualization due to the limit of the max
+        if (maxAssociatedCount <= savedObjectsForThisVisualization.length) {
+          throw new Error(
+            `Cannot associate the plugin resource to the visualization due to the limit of the max
           amount of associated plugin resources (${maxAssociatedCount}) with
           ${savedObjectsForThisVisualization.length} associated to the visualization`
-        );
+          );
+        }
       }
-    }
-  });
+    });
 
   return await loader.get((AugmentVis as any) as Record<string, unknown>);
 };

src/plugins/vis_augmenter/public/utils/utils.test.ts

@@ -304,12 +304,6 @@ describe('utils', () => {
       pluginResource
     );
 
-    it('returns no matching saved objs with filtering', async () => {
-      const loader = createSavedAugmentVisLoader({
-        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2, obj3]),
-      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
-      expect((await getAugmentVisSavedObjs(visId3, loader)).length).toEqual(0);
-    });
     it('returns no matching saved objs when client returns empty list', async () => {
       const loader = createSavedAugmentVisLoader({
         savedObjectsClient: getMockAugmentVisSavedObjectClient([]),
@@ -349,18 +343,12 @@ describe('utils', () => {
       } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
       expect((await getAugmentVisSavedObjs(visId1, loader)).length).toEqual(1);
     });
-    it('returns multiple matching saved objs without filtering', async () => {
+    it('returns multiple matching saved objs', async () => {
       const loader = createSavedAugmentVisLoader({
         savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
       } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
       expect((await getAugmentVisSavedObjs(visId1, loader)).length).toEqual(2);
     });
-    it('returns multiple matching saved objs with filtering', async () => {
-      const loader = createSavedAugmentVisLoader({
-        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2, obj3]),
-      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
-      expect((await getAugmentVisSavedObjs(visId1, loader)).length).toEqual(2);
-    });
   });
 
   describe('buildPipelineFromAugmentVisSavedObjs', () => {

src/plugins/vis_augmenter/public/utils/utils.ts

@@ -54,7 +54,7 @@ export const isEligibleForVisLayers = (vis: Vis, uiSettingsClient?: IUiSettingsC
 
 /**
  * Using a SavedAugmentVisLoader, fetch all saved objects that are of 'augment-vis' type.
- * Filter by vis ID.
+ * Filter by vis ID by passing in a 'hasReferences' obj with the vis ID to the findAll() fn call.
  */
 export const getAugmentVisSavedObjs = async (
   visId: string | undefined,
@@ -69,18 +69,11 @@ export const getAugmentVisSavedObjs = async (
       'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
     );
   }
-  const allSavedObjects = await getAllAugmentVisSavedObjs(loader);
-  return allSavedObjects.filter((hit: ISavedAugmentVis) => hit.visId === visId);
-};
-
-/**
- * Using a SavedAugmentVisLoader, fetch all saved objects that are of 'augment-vis' type.
- */
-export const getAllAugmentVisSavedObjs = async (
-  loader: SavedAugmentVisLoader | undefined
-): Promise<ISavedAugmentVis[]> => {
   try {
-    const resp = await loader?.findAll();
+    const resp = await loader?.findAll('', 100, [], {
+      type: 'visualization',
+      id: visId as string,
+    });
     return (get(resp, 'hits', []) as any[]) as ISavedAugmentVis[];
   } catch (e) {
     return [] as ISavedAugmentVis[];