Add new advanced settings for vis augmenter

config/opensearch_dashboards.yml

@@ -239,4 +239,7 @@
 #data_source.encryption.wrappingKey: [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
 
 # Set the value of this setting to false to hide the help menu link to the OpenSearch Dashboards user survey
-# opensearchDashboards.survey.url: "https://survey.opensearch.org"
+# opensearchDashboards.survey.url: "https://survey.opensearch.org"
+
+# Set the value of this setting to false to disable plugin augmentation on Dashboard
+# vis_augmenter.pluginAugmentationEnabled: false

src/plugins/opensearch_dashboards_usage_collection/server/collectors/management/schema.ts

@@ -40,6 +40,8 @@ export const stackManagementSchema: MakeSchemaFrom<UsageStats> = {
   'visualization:regionmap:showWarnings': { type: 'boolean' },
   'visualization:dimmingOpacity': { type: 'float' },
   'visualization:tileMap:maxPrecision': { type: 'long' },
+  'visualization:enablePluginAugmentation': { type: 'boolean' },
+  'visualization:enablePluginAugmentation.maxPluginObjects': { type: 'number' },
   'securitySolution:ipReputationLinks': { type: 'text' },
   'csv:separator': { type: 'keyword' },
   'visualization:tileMap:WMSdefaults': { type: 'text' },

src/plugins/telemetry/schema/oss_plugins.json

@@ -1372,6 +1372,12 @@
         "visualization:tileMap:maxPrecision": {
           "type": "long"
         },
+        "visualization:enablePluginAugmentation": {
+          "type": "boolean"
+        },
+        "visualization:enablePluginAugmentation.maxPluginObjects": {
+          "type": "number"
+        },
         "securitySolution:ipReputationLinks": {
           "type": "text"
         },

src/plugins/vis_augmenter/common/constants.ts

@@ -0,0 +1,8 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const PLUGIN_AUGMENTATION_ENABLE_SETTING = 'visualization:enablePluginAugmentation';
+export const PLUGIN_AUGMENTATION_MAX_OBJECTS_SETTING =
+  'visualization:enablePluginAugmentation.maxPluginObjects';

src/plugins/vis_augmenter/config.ts

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { schema, TypeOf } from '@osd/config-schema';
+
+export const configSchema = schema.object({
+  pluginAugmentationEnabled: schema.boolean({ defaultValue: true }),
+});
+
+export type VisAugmenterPluginConfigType = TypeOf<typeof configSchema>;

src/plugins/vis_augmenter/public/index.ts

@@ -30,3 +30,4 @@ export * from './constants';
 export * from './vega';
 export * from './saved_augment_vis';
 export * from './test_constants';
+export { getSavedAugmentVisLoader } from './services';

src/plugins/vis_augmenter/public/plugin.ts

@@ -7,7 +7,7 @@ import { ExpressionsSetup } from '../../expressions/public';
 import { PluginInitializerContext, CoreSetup, CoreStart, Plugin } from '../../../core/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../data/public';
 import { visLayers } from './expressions';
-import { setSavedAugmentVisLoader } from './services';
+import { setSavedAugmentVisLoader, setUISettings } from './services';
 import { createSavedAugmentVisLoader, SavedAugmentVisLoader } from './saved_augment_vis';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
@@ -40,6 +40,7 @@ export class VisAugmenterPlugin
   }
 
   public start(core: CoreStart, { data }: VisAugmenterStartDeps): VisAugmenterStart {
+    setUISettings(core.uiSettings);
     const savedAugmentVisLoader = createSavedAugmentVisLoader({
       savedObjectsClient: core.savedObjects.client,
       indexPatterns: data.indexPatterns,

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

@@ -9,8 +9,18 @@ import {
   SavedObjectOpenSearchDashboardsServicesWithAugmentVis,
 } from './saved_augment_vis';
 import { generateAugmentVisSavedObject, getMockAugmentVisSavedObjectClient } from './utils';
+import { uiSettingsServiceMock } from '../../../../core/public/mocks';
+import { setUISettings } from '../services';
+import { PLUGIN_AUGMENTATION_ENABLE_SETTING } from '../../common/constants';
+
+const uiSettingsMock = uiSettingsServiceMock.createStartContract();
+setUISettings(uiSettingsMock);
 
 describe('SavedObjectLoaderAugmentVis', () => {
+  uiSettingsMock.get.mockImplementation((key: string) => {
+    return key === PLUGIN_AUGMENTATION_ENABLE_SETTING;
+  });
+
   const fn = {
     type: VisLayerTypes.PointInTimeEvents,
     name: 'test-fn',
@@ -105,4 +115,52 @@ describe('SavedObjectLoaderAugmentVis', () => {
     expect(resp.hits[0].id).toEqual('valid-obj-id-1');
     expect(resp.hits[0].error).toEqual('visReference is missing in augment-vis saved object');
   });
+
+  it('find returns exception due to setting being disabled', async () => {
+    uiSettingsMock.get.mockImplementation((key: string) => {
+      return key !== PLUGIN_AUGMENTATION_ENABLE_SETTING;
+    });
+    const loader = createSavedAugmentVisLoader(({
+      savedObjectsClient: getMockAugmentVisSavedObjectClient([]),
+    } as unknown) as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+    try {
+      await loader.find();
+    } catch (e) {
+      expect(e.message).toBe(
+        'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
+      );
+    }
+  });
+
+  it('findAll returns exception due to setting being disabled', async () => {
+    uiSettingsMock.get.mockImplementation((key: string) => {
+      return key !== PLUGIN_AUGMENTATION_ENABLE_SETTING;
+    });
+    const loader = createSavedAugmentVisLoader(({
+      savedObjectsClient: getMockAugmentVisSavedObjectClient([]),
+    } as unknown) as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+    try {
+      await loader.findAll();
+    } catch (e) {
+      expect(e.message).toBe(
+        'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
+      );
+    }
+  });
+
+  it('get returns exception due to setting being disabled', async () => {
+    uiSettingsMock.get.mockImplementation((key: string) => {
+      return key !== PLUGIN_AUGMENTATION_ENABLE_SETTING;
+    });
+    const loader = createSavedAugmentVisLoader(({
+      savedObjectsClient: getMockAugmentVisSavedObjectClient([]),
+    } as unknown) as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+    try {
+      await loader.get();
+    } catch (e) {
+      expect(e.message).toBe(
+        'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
+      );
+    }
+  });
 });

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

@@ -4,59 +4,106 @@
  */
 
 import { get, isEmpty } from 'lodash';
+import { IUiSettingsClient } from 'opensearch-dashboards/public';
 import {
   SavedObjectLoader,
   SavedObjectOpenSearchDashboardsServices,
 } from '../../../saved_objects/public';
 import { createSavedAugmentVisClass } from './_saved_augment_vis';
 import { VisLayerTypes } from '../types';
+import { getUISettings } from '../services';
+import { PLUGIN_AUGMENTATION_ENABLE_SETTING } from '../../common/constants';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface SavedObjectOpenSearchDashboardsServicesWithAugmentVis
   extends SavedObjectOpenSearchDashboardsServices {}
 export type SavedAugmentVisLoader = ReturnType<typeof createSavedAugmentVisLoader>;
+
+export class SavedObjectLoaderAugmentVis extends SavedObjectLoader {
+  private readonly config: IUiSettingsClient = getUISettings();
+
+  mapHitSource = (source: Record<string, any>, id: string) => {
+    source.id = id;
+    source.visId = get(source, 'visReference.id', '');
+
+    if (isEmpty(source.visReference)) {
+      source.error = 'visReference is missing in augment-vis saved object';
+      return source;
+    }
+    if (isEmpty(source.visLayerExpressionFn)) {
+      source.error = 'visLayerExpressionFn is missing in augment-vis saved object';
+      return source;
+    }
+    if (!(get(source, 'visLayerExpressionFn.type', '') in VisLayerTypes)) {
+      source.error = 'Unknown VisLayer expression function type';
+      return source;
+    }
+    return source;
+  };
+
+  /**
+   * Updates hit.attributes to contain an id related to the referenced visualization
+   * (visId) and returns the updated attributes object.
+   * @param hit
+   * @returns {hit.attributes} The modified hit.attributes object, with an id and url field.
+   */
+  mapSavedObjectApiHits(hit: {
+    references: any[];
+    attributes: Record<string, unknown>;
+    id: string;
+  }) {
+    // For now we are assuming only one vis reference per saved object.
+    // If we change to multiple, we will need to dynamically handle that
+    const visReference = hit.references[0];
+    return this.mapHitSource({ ...hit.attributes, visReference }, hit.id);
+  }
+
+  /**
+   * Retrieve a saved object by id or create new one.
+   * Returns a promise that completes when the object finishes
+   * initializing. Throws exception when the setting is set to false.
+   * @param opts
+   * @returns {Promise<SavedObject>}
+   */
+  get(opts?: Record<string, unknown> | string) {
+    this.isAugmentationEnabled();
+    return super.get(opts);
+  }
+
+  /**
+   * TODO: Rather than use a hardcoded limit, implement pagination. See
+   * https://github.com/elastic/kibana/issues/8044 for reference.
+   *
+   * @param search
+   * @param size
+   * @param fields
+   * @returns {Promise}
+   */
+  findAll(search: string = '', size: number = 100, fields?: string[]) {
+    this.isAugmentationEnabled();
+    return super.findAll(search, size, fields);
+  }
+
+  find(search: string = '', size: number = 100) {
+    this.isAugmentationEnabled();
+    return super.find(search, size);
+  }
+
+  private isAugmentationEnabled() {
+    const isAugmentationEnabled = this.config.get(PLUGIN_AUGMENTATION_ENABLE_SETTING, true);
+    if (!isAugmentationEnabled) {
+      throw new Error(
+        'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
+      );
+    }
+  }
+}
+
 export function createSavedAugmentVisLoader(
   services: SavedObjectOpenSearchDashboardsServicesWithAugmentVis
 ) {
   const { savedObjectsClient } = services;
 
-  class SavedObjectLoaderAugmentVis extends SavedObjectLoader {
-    mapHitSource = (source: Record<string, any>, id: string) => {
-      source.id = id;
-      source.visId = get(source, 'visReference.id', '');
-
-      if (isEmpty(source.visReference)) {
-        source.error = 'visReference is missing in augment-vis saved object';
-        return source;
-      }
-      if (isEmpty(source.visLayerExpressionFn)) {
-        source.error = 'visLayerExpressionFn is missing in augment-vis saved object';
-        return source;
-      }
-      if (!(get(source, 'visLayerExpressionFn.type', '') in VisLayerTypes)) {
-        source.error = 'Unknown VisLayer expression function type';
-        return source;
-      }
-      return source;
-    };
-
-    /**
-     * Updates hit.attributes to contain an id related to the referenced visualization
-     * (visId) and returns the updated attributes object.
-     * @param hit
-     * @returns {hit.attributes} The modified hit.attributes object, with an id and url field.
-     */
-    mapSavedObjectApiHits(hit: {
-      references: any[];
-      attributes: Record<string, unknown>;
-      id: string;
-    }) {
-      // For now we are assuming only one vis reference per saved object.
-      // If we change to multiple, we will need to dynamically handle that
-      const visReference = hit.references[0];
-      return this.mapHitSource({ ...hit.attributes, visReference }, hit.id);
-    }
-  }
   const SavedAugmentVis = createSavedAugmentVisClass(services);
-  return new SavedObjectLoaderAugmentVis(SavedAugmentVis, savedObjectsClient) as SavedObjectLoader;
+  return new SavedObjectLoaderAugmentVis(SavedAugmentVis, savedObjectsClient);
 }

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

@@ -3,14 +3,47 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { getSavedAugmentVisLoader } from '../../services';
-import { ISavedAugmentVis } from '../../types';
+import { get } from 'lodash';
+import { getSavedAugmentVisLoader, getUISettings } from '../../services';
+import { ISavedAugmentVis } from '../types';
+import {
+  PLUGIN_AUGMENTATION_ENABLE_SETTING,
+  PLUGIN_AUGMENTATION_MAX_OBJECTS_SETTING,
+} from '../../../common/constants';
 
 /**
  * Create an augment vis saved object given an object that
  * implements the ISavedAugmentVis interface
  */
 export const createAugmentVisSavedObject = async (AugmentVis: ISavedAugmentVis): Promise<any> => {
   const loader = getSavedAugmentVisLoader();
+  const config = getUISettings();
+
+  const isAugmentationEnabled = config.get(PLUGIN_AUGMENTATION_ENABLE_SETTING);
+  if (!isAugmentationEnabled) {
+    throw new Error(
+      'Visualization augmentation is disabled, please enable visualization:enablePluginAugmentation.'
+    );
+  }
+  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
+      );
+
+      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/services.ts

@@ -3,9 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { createGetterSetter } from '../../opensearch_dashboards_utils/common';
-import { SavedObjectLoader } from '../../saved_objects/public';
+import { createGetterSetter } from '../../opensearch_dashboards_utils/public';
+import { IUiSettingsClient } from '../../../core/public';
+import { SavedObjectLoaderAugmentVis } from './saved_augment_vis';
 
 export const [getSavedAugmentVisLoader, setSavedAugmentVisLoader] = createGetterSetter<
-  SavedObjectLoader
+  SavedObjectLoaderAugmentVis
 >('savedAugmentVisLoader');
+
+export const [getUISettings, setUISettings] = createGetterSetter<IUiSettingsClient>('UISettings');

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

@@ -18,10 +18,20 @@ import {
   VisLayerTypes,
   VisLayerExpressionFn,
 } from '../';
+import { PLUGIN_AUGMENTATION_ENABLE_SETTING } from '../../common/constants';
 import { AggConfigs, AggTypesRegistryStart, IndexPattern } from '../../../data/common';
 import { mockAggTypesRegistry } from '../../../data/common/search/aggs/test_helpers';
+import { uiSettingsServiceMock } from '../../../../core/public/mocks';
+import { setUISettings } from '../services';
 
 describe('utils', () => {
+  const uiSettingsMock = uiSettingsServiceMock.createStartContract();
+  setUISettings(uiSettingsMock);
+  beforeEach(() => {
+    uiSettingsMock.get.mockImplementation((key: string) => {
+      return key === PLUGIN_AUGMENTATION_ENABLE_SETTING;
+    });
+  });
   describe('isEligibleForVisLayers', () => {
     const validConfigStates = [
       {
@@ -279,6 +289,12 @@ describe('utils', () => {
       } as unknown) as Vis;
       expect(isEligibleForVisLayers(invalidVis)).toEqual(false);
     });
+    it('vis is ineligible with valid type and disabled setting', async () => {
+      uiSettingsMock.get.mockImplementation((key: string) => {
+        return key !== PLUGIN_AUGMENTATION_ENABLE_SETTING;
+      });
+      expect(isEligibleForVisLayers(validVis)).toEqual(false);
+    });
     it('vis is eligible with valid type', async () => {
       expect(isEligibleForVisLayers(validVis)).toEqual(true);
     });