[Backport 1.x] Add option to configure available bucket agg types + jest tests (#1196)

src/plugins/data/common/search/aggs/agg_types_registry.test.ts

@@ -41,15 +41,20 @@ describe('AggTypesRegistry', () => {
   let registry: AggTypesRegistry;
   let setup: AggTypesRegistrySetup;
   let start: ReturnType<AggTypesRegistry['start']>;
+  const uiSettings: any = {
+    get: jest
+      .fn()
+      .mockImplementation((key) => key === 'visualize:disableBucketAgg' && ['significant_terms']),
+  };
 
   beforeEach(() => {
     registry = new AggTypesRegistry();
     setup = registry.setup();
-    start = registry.start();
   });
 
   it('registerBucket adds new buckets', () => {
     setup.registerBucket('terms', bucketType);
+    start = registry.start({ uiSettings });
     expect(start.getAll().buckets).toEqual([bucketType]);
   });
 
@@ -69,6 +74,7 @@ describe('AggTypesRegistry', () => {
 
   it('registerMetric adds new metrics', () => {
     setup.registerMetric('count', metricType);
+    start = registry.start({ uiSettings });
     expect(start.getAll().metrics).toEqual([metricType]);
   });
 
@@ -89,16 +95,35 @@ describe('AggTypesRegistry', () => {
   it('gets either buckets or metrics by id', () => {
     setup.registerBucket('terms', bucketType);
     setup.registerMetric('count', metricType);
+    start = registry.start({ uiSettings });
     expect(start.get('terms')).toEqual(bucketType);
     expect(start.get('count')).toEqual(metricType);
   });
 
   it('getAll returns all buckets and metrics', () => {
     setup.registerBucket('terms', bucketType);
     setup.registerMetric('count', metricType);
+    start = registry.start({ uiSettings });
     expect(start.getAll()).toEqual({
       buckets: [bucketType],
       metrics: [metricType],
     });
   });
+
+  it('getAll returns all buckets besides disabled bucket type', () => {
+    const termsBucket = () => ({ name: 'terms', type: 'buckets' } as BucketAggType<any>);
+    const histogramBucket = () => ({ name: 'histogram', type: 'buckets' } as BucketAggType<any>);
+    const significantTermsBucket = () =>
+      ({ name: 'significant_terms', type: 'buckets' } as BucketAggType<any>);
+    setup.registerBucket('terms', termsBucket);
+    setup.registerBucket('histogram', histogramBucket);
+    setup.registerBucket('significant_terms', significantTermsBucket);
+    start = registry.start({ uiSettings });
+
+    expect(start.getAll()).toEqual({
+      buckets: [termsBucket, histogramBucket],
+      metrics: [],
+    });
+    expect(start.get('significant_terms')).toEqual(undefined);
+  });
 });

src/plugins/data/common/search/aggs/agg_types_registry.ts

@@ -30,6 +30,9 @@
  * GitHub history for details.
  */
 
+import { IUiSettingsClient as IUiSettingsClientPublic } from 'src/core/public';
+// eslint-disable-next-line @osd/eslint/no-restricted-paths
+import { IUiSettingsClient as IUiSettingsClientServer } from 'src/core/server';
 import { BucketAggType } from './buckets/bucket_agg_type';
 import { MetricAggType } from './metrics/metric_agg_type';
 import { AggTypesDependencies } from './agg_types';
@@ -48,6 +51,10 @@ export interface AggTypesRegistryStart {
   getAll: () => { buckets: Array<BucketAggType<any>>; metrics: Array<MetricAggType<any>> };
 }
 
+export interface AggTypesRegistryStartDependencies {
+  uiSettings: IUiSettingsClientPublic | IUiSettingsClientServer;
+}
+
 export class AggTypesRegistry {
   private readonly bucketAggs = new Map();
   private readonly metricAggs = new Map();
@@ -81,7 +88,15 @@ export class AggTypesRegistry {
     };
   };
 
-  start = () => {
+  start = ({ uiSettings }: AggTypesRegistryStartDependencies) => {
+    const disabledBucketAgg = uiSettings.get('visualize:disableBucketAgg');
+
+    if (disabledBucketAgg !== undefined && Array.isArray(disabledBucketAgg)) {
+      for (const k of this.bucketAggs.keys()) {
+        if (disabledBucketAgg.includes(k)) this.bucketAggs.delete(k);
+      }
+    }
+
     return {
       get: (name: string) => {
         return this.bucketAggs.get(name) || this.metricAggs.get(name);

src/plugins/data/common/search/aggs/aggs_service.test.ts

@@ -57,6 +57,7 @@ describe('Aggs service', () => {
     };
     startDeps = {
       getConfig: jest.fn(),
+      uiSettings: { get: jest.fn().mockImplementation(() => []) } as any,
     };
   });
 

src/plugins/data/common/search/aggs/aggs_service.ts

@@ -31,6 +31,9 @@
  */
 
 import { ExpressionsServiceSetup } from 'src/plugins/expressions/common';
+import { IUiSettingsClient as IUiSettingsClientPublic } from 'src/core/public';
+// eslint-disable-next-line @osd/eslint/no-restricted-paths
+import { IUiSettingsClient as IUiSettingsClientServer } from 'src/core/server';
 import { UI_SETTINGS } from '../../../common';
 import { GetConfigFn } from '../../types';
 import {
@@ -63,6 +66,7 @@ export interface AggsCommonSetupDependencies {
 /** @internal */
 export interface AggsCommonStartDependencies {
   getConfig: GetConfigFn;
+  uiSettings: IUiSettingsClientPublic | IUiSettingsClientServer;
 }
 
 /**
@@ -90,8 +94,8 @@ export class AggsCommonService {
     };
   }
 
-  public start({ getConfig }: AggsCommonStartDependencies): AggsCommonStart {
-    const aggTypesStart = this.aggTypesRegistry.start();
+  public start({ getConfig, uiSettings }: AggsCommonStartDependencies): AggsCommonStart {
+    const aggTypesStart = this.aggTypesRegistry.start({ uiSettings });
 
     return {
       calculateAutoTimeExpression: getCalculateAutoTimeExpression(getConfig),

src/plugins/data/common/search/aggs/test_helpers/mock_agg_types_registry.ts

@@ -93,7 +93,9 @@ export function mockAggTypesRegistry(deps?: AggTypesDependencies): AggTypesRegis
   aggTypes.buckets.forEach(({ name, fn }) => registrySetup.registerBucket(name, fn));
   aggTypes.metrics.forEach(({ name, fn }) => registrySetup.registerMetric(name, fn));
 
-  const registryStart = registry.start();
+  const registryStart = registry.start({
+    uiSettings: { get: jest.fn().mockImplementation(() => []) } as any,
+  });
 
   // initialize each agg type and store in memory
   registryStart.getAll().buckets.forEach((type) => {

src/plugins/data/public/search/aggs/aggs_service.ts

@@ -109,6 +109,7 @@ export class AggsService {
 
   public start({ fieldFormats, uiSettings }: AggsStartDependencies): AggsStart {
     const { calculateAutoTimeExpression, types } = this.aggsCommonService.start({
+      uiSettings,
       getConfig: this.getConfig!,
     });
 

src/plugins/data/server/search/aggs/aggs_service.ts

@@ -86,7 +86,10 @@ export class AggsService {
           return uiSettingsCache[key];
         };
 
-        const { calculateAutoTimeExpression, types } = this.aggsCommonService.start({ getConfig });
+        const { calculateAutoTimeExpression, types } = this.aggsCommonService.start({
+          getConfig,
+          uiSettings: uiSettingsClient,
+        });
 
         const aggTypesDependencies: AggTypesDependencies = {
           calculateBounds: this.calculateBounds,

src/plugins/visualizations/common/constants.ts

@@ -31,3 +31,4 @@
  */
 
 export const VISUALIZE_ENABLE_LABS_SETTING = 'visualize:enableLabs';
+export const VISUALIZE_DISABLE_BUCKET_AGG_SETTING = 'visualize:disableBucketAgg';

src/plugins/visualizations/server/index.ts

@@ -33,7 +33,10 @@
 import { PluginInitializerContext } from '../../../core/server';
 import { VisualizationsPlugin } from './plugin';
 
-export { VISUALIZE_ENABLE_LABS_SETTING } from '../common/constants';
+export {
+  VISUALIZE_ENABLE_LABS_SETTING,
+  VISUALIZE_DISABLE_BUCKET_AGG_SETTING,
+} from '../common/constants';
 
 //  This exports static code and TypeScript types,
 //  as well as, OpenSearch Dashboards Platform `plugin()` initializer.

src/plugins/visualizations/server/plugin.ts

@@ -42,7 +42,10 @@ import {
   Logger,
 } from '../../../core/server';
 
-import { VISUALIZE_ENABLE_LABS_SETTING } from '../common/constants';
+import {
+  VISUALIZE_ENABLE_LABS_SETTING,
+  VISUALIZE_DISABLE_BUCKET_AGG_SETTING,
+} from '../common/constants';
 
 import { visualizationSavedObjectType } from './saved_objects';
 
@@ -77,6 +80,18 @@ export class VisualizationsPlugin
         category: ['visualization'],
         schema: schema.boolean(),
       },
+      [VISUALIZE_DISABLE_BUCKET_AGG_SETTING]: {
+        name: i18n.translate('visualizations.advancedSettings.visualizeDisableBucketAgg', {
+          defaultMessage: 'Disable visualizations bucket aggregation types',
+        }),
+        value: [],
+        description: i18n.translate('visualizations.advancedSettings.visualizeDisableBucketAgg', {
+          defaultMessage: `A comma-separated list of bucket aggregations' names. e.g. significant_terms, terms.
+            Deactivates the specified bucket aggregations from visualizations.`,
+        }),
+        category: ['visualization'],
+        schema: schema.arrayOf(schema.string()),
+      },
     });
 
     if (plugins.usageCollection) {