Add resource ID filtering in fetch augment-vis obj queries (#4608)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
(cherry picked from commit c06b298)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

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

@@ -126,21 +126,24 @@ export class SavedObjectLoader {
    * @param search
    * @param size
    * @param fields
+   * @param hasReference Optional field to specify a reference
+   * @param searchFields Optional field to specify the search fields in the query
    * @returns {Promise}
    */
   findAll(
     search: string = '',
     size: number = 100,
     fields?: string[],
-    hasReference?: SavedObjectsFindOptions['hasReference']
+    hasReference?: SavedObjectsFindOptions['hasReference'],
+    searchFields: string[] = ['title^3', 'description']
   ) {
     return this.savedObjectsClient
       .find<Record<string, unknown>>({
         type: this.lowercaseType,
         search: search ? `${search}*` : undefined,
         perPage: size,
         page: 1,
-        searchFields: ['title^3', 'description'],
+        searchFields,
         defaultSearchOperator: 'AND',
         fields,
         hasReference,

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

@@ -89,16 +89,19 @@ export class SavedObjectLoaderAugmentVis extends SavedObjectLoader {
    * @param search
    * @param size
    * @param fields
+   * @param hasReference Optional field to specify a reference
+   * @param searchFields Optional field to specify the search fields in the query
    * @returns {Promise}
    */
   findAll(
     search: string = '',
     size: number = 100,
     fields?: string[],
-    hasReference?: SavedObjectsFindOptions['hasReference']
+    hasReference?: SavedObjectsFindOptions['hasReference'],
+    searchFields?: string[]
   ) {
     this.isAugmentationEnabled();
-    return super.findAll(search, size, fields, hasReference);
+    return super.findAll(search, size, fields, hasReference, searchFields);
   }
 
   find(search: string = '', size: number = 100) {

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

@@ -349,6 +349,76 @@ describe('utils', () => {
       } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
       expect((await getAugmentVisSavedObjs(visId1, loader)).length).toEqual(2);
     });
+    it('undefined plugin resource list has no effect', async () => {
+      const loader = createSavedAugmentVisLoader({
+        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
+      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+      expect((await getAugmentVisSavedObjs(visId1, loader, undefined, undefined)).length).toEqual(
+        2
+      );
+    });
+    it('empty plugin resource list has no effect', async () => {
+      const loader = createSavedAugmentVisLoader({
+        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
+      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+      expect((await getAugmentVisSavedObjs(visId1, loader, undefined, [])).length).toEqual(2);
+    });
+    it('empty / undefined plugin resource list passes correct findAll() params', async () => {
+      const loader = createSavedAugmentVisLoader({
+        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
+      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+      loader.findAll = jest.fn().mockImplementation(loader.findAll);
+      expect((await getAugmentVisSavedObjs(visId1, loader, undefined, [])).length).toEqual(2);
+      expect(loader.findAll).toHaveBeenCalledWith(
+        '',
+        100,
+        [],
+        {
+          type: 'visualization',
+          id: visId1 as string,
+        },
+        []
+      );
+    });
+    it('single plugin resource is propagated to findAll()', async () => {
+      const loader = createSavedAugmentVisLoader({
+        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
+      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+      loader.findAll = jest.fn().mockImplementation(loader.findAll);
+      expect(
+        (await getAugmentVisSavedObjs(visId1, loader, undefined, ['resource-1'])).length
+      ).toEqual(2);
+      expect(loader.findAll).toHaveBeenCalledWith(
+        'resource-1',
+        100,
+        [],
+        {
+          type: 'visualization',
+          id: visId1 as string,
+        },
+        ['pluginResource.id']
+      );
+    });
+    it('multiple plugin resources are propagated to findAll()', async () => {
+      const loader = createSavedAugmentVisLoader({
+        savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1, obj2]),
+      } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+      loader.findAll = jest.fn().mockImplementation(loader.findAll);
+      expect(
+        (await getAugmentVisSavedObjs(visId1, loader, undefined, ['resource-1', 'resource-2']))
+          .length
+      ).toEqual(2);
+      expect(loader.findAll).toHaveBeenCalledWith(
+        'resource-1|resource-2',
+        100,
+        [],
+        {
+          type: 'visualization',
+          id: visId1 as string,
+        },
+        ['pluginResource.id']
+      );
+    });
   });
 
   describe('buildPipelineFromAugmentVisSavedObjs', () => {

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

@@ -54,12 +54,14 @@ export const isEligibleForVisLayers = (vis: Vis, uiSettingsClient?: IUiSettingsC
 
 /**
  * Using a SavedAugmentVisLoader, fetch all saved objects that are of 'augment-vis' type.
- * Filter by vis ID by passing in a 'hasReferences' obj with the vis ID to the findAll() fn call.
+ * Filter by vis ID by passing in a 'hasReferences' obj with the vis ID to the findAll() fn call,
+ * and optionally by plugin resource ID list, if specified.
  */
 export const getAugmentVisSavedObjs = async (
   visId: string | undefined,
   loader: SavedAugmentVisLoader | undefined,
-  uiSettings?: IUiSettingsClient | undefined
+  uiSettings?: IUiSettingsClient | undefined,
+  pluginResourceIds?: string[] | undefined
 ): Promise<ISavedAugmentVis[] | Error> => {
   // Using optional services provided, or the built-in services from this plugin
   const config = uiSettings !== undefined ? uiSettings : getUISettings();
@@ -70,10 +72,20 @@ export const getAugmentVisSavedObjs = async (
     );
   }
   try {
-    const resp = await loader?.findAll('', 100, [], {
-      type: 'visualization',
-      id: visId as string,
-    });
+    // If there is specified plugin resource IDs, add a search string and search field
+    // into findAll() fn call
+    const pluginResourceIdsSpecified =
+      pluginResourceIds !== undefined && pluginResourceIds.length > 0;
+    const resp = await loader?.findAll(
+      pluginResourceIdsSpecified ? pluginResourceIds.join('|') : '',
+      100,
+      [],
+      {
+        type: 'visualization',
+        id: visId as string,
+      },
+      pluginResourceIdsSpecified ? ['pluginResource.id'] : []
+    );
     return (get(resp, 'hits', []) as any[]) as ISavedAugmentVis[];
   } catch (e) {
     return [] as ISavedAugmentVis[];

src/plugins/visualizations/public/embeddable/visualize_embeddable.ts

@@ -513,19 +513,14 @@ export class VisualizeEmbeddable
    * e.g., generating other vis embeddables in the view events flyout.
    */
   public async populateVisLayers(): Promise<void> {
-    const visLayers = await this.fetchVisLayers();
-    this.visLayers =
-      this.visAugmenterConfig?.visLayerResourceIds === undefined
-        ? visLayers
-        : visLayers.filter((visLayer) =>
-            this.visAugmenterConfig.visLayerResourceIds.includes(visLayer.pluginResource.id)
-          );
+    this.visLayers = await this.fetchVisLayers();
   }
 
   /**
    * Collects any VisLayers from plugin expressions functions
-   * by fetching all AugmentVisSavedObjects that match the vis
-   * saved object ID.
+   * by fetching all AugmentVisSavedObjects that meets below criteria:
+   * - includes a reference to the vis saved object id
+   * - includes any of the plugin resource IDs, if specified
    */
   fetchVisLayers = async (): Promise<VisLayers> => {
     try {
@@ -541,7 +536,9 @@ export class VisualizeEmbeddable
       const aborted = get(this.abortController, 'signal.aborted', false) as boolean;
       const augmentVisSavedObjs = await getAugmentVisSavedObjs(
         this.vis.id,
-        this.savedAugmentVisLoader
+        this.savedAugmentVisLoader,
+        undefined,
+        this.visAugmenterConfig?.visLayerResourceIds
       );
 
       if (!isEmpty(augmentVisSavedObjs) && !aborted && isEligibleForVisLayers(this.vis)) {