Convert from large terms agg to multiple smaller composite

x-pack/legacy/plugins/infra/server/lib/adapters/fields/framework_fields_adapter.ts

@@ -4,22 +4,22 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { startsWith, uniq } from 'lodash';
+import { startsWith, uniq, first } from 'lodash';
 import { InfraBackendFrameworkAdapter, InfraFrameworkRequest } from '../framework';
 import { FieldsAdapter, IndexFieldDescriptor } from './adapter_types';
 import { getAllowedListForPrefix } from '../../../../common/ecs_allowed_list';
 
 interface Bucket {
-  key: string;
+  key: { dataset: string };
   doc_count: number;
 }
 
 interface DataSetResponse {
-  modules: {
-    buckets: Bucket[];
-  };
-  dataSets: {
+  datasets: {
     buckets: Bucket[];
+    after_key: {
+      dataset: string;
+    };
   };
 }
 
@@ -59,34 +59,68 @@ export class FrameworkFieldsAdapter implements FieldsAdapter {
       allowNoIndices: true,
       ignoreUnavailable: true,
       body: {
+        size: 0,
         aggs: {
-          modules: {
-            terms: {
-              field: 'event.modules',
-              size: 1000,
-            },
-          },
-          dataSets: {
-            terms: {
-              field: 'event.dataset',
-              size: 1000,
+          datasets: {
+            composite: {
+              sources: [
+                {
+                  dataset: {
+                    terms: {
+                      field: 'event.dataset',
+                    },
+                  },
+                },
+              ],
             },
           },
         },
       },
     };
+
+    const buckets = await this.getAllCompositeData(request, params);
+    const dataSets = buckets.map(bucket => bucket.key.dataset);
+    const modules = dataSets.reduce(
+      (acc, dataset) => {
+        const module = first(dataset.split(/\./));
+        return module ? uniq([...acc, module]) : acc;
+      },
+      [] as string[]
+    );
+    return { modules, dataSets };
+  }
+
+  private async getAllCompositeData(
+    request: InfraFrameworkRequest,
+    query: any,
+    previousBuckets: Bucket[] = []
+  ): Promise<Bucket[]> {
     const response = await this.framework.callWithRequest<{}, DataSetResponse>(
       request,
       'search',
-      params
+      query
     );
+
+    // if ES doesn't return an aggregations key, something went seriously wrong.
     if (!response.aggregations) {
-      return { dataSets: [], modules: [] };
+      throw new Error('Whoops!, `aggregations` key must always be returned.');
     }
-    const { modules, dataSets } = response.aggregations;
-    return {
-      modules: modules.buckets.map(bucket => bucket.key),
-      dataSets: dataSets.buckets.map(bucket => bucket.key),
-    };
+
+    // Nothing available, return the previous buckets.
+    if (response.hits.total.value === 0) {
+      return previousBuckets;
+    }
+
+    const currentBuckets = response.aggregations.datasets.buckets;
+
+    // if there are no currentBuckets then we are finished paginating through the results
+    if (currentBuckets.length === 0) {
+      return previousBuckets;
+    }
+
+    // There is possibly more data, concat previous and current buckets and call ourselves recursively.
+    const newQuery = { ...query };
+    newQuery.body.aggs.datasets.composite.after = response.aggregations.datasets.after_key;
+    return this.getAllCompositeData(request, newQuery, previousBuckets.concat(currentBuckets));
   }
 }