TopN search: factor common query functions, more type safety (#14)

* topn: use default ES client
* introduce type on aggregations
* extract and use common query filter
* test topn search: extract query and aggs objects
* test fetching stacktraces with mget
* PR feedback: rename interface and function

Signed-off-by: inge4pres <francesco.gualazzi@elastic.co>

src/plugins/profiling/server/routes/mappings.ts

@@ -6,6 +6,93 @@
  * Side Public License, v 1.
  */
 
+import { AggregationsAggregationContainer } from '@elastic/elasticsearch/lib/api/types';
+
+interface ProjectTimeQuery {
+  bool: {
+    must: Array<
+      | {
+          term: {
+            ProjectID: {
+              value: string;
+              boost: number;
+            };
+          };
+        }
+      | {
+          range: {
+            '@timestamp': {
+              gte: string;
+              lt: string;
+              format: string;
+              boost: number;
+            };
+          };
+        }
+    >;
+  };
+}
+
+export function newProjectTimeQuery(
+  projectID: string,
+  timeFrom: string,
+  timeTo: string
+): ProjectTimeQuery {
+  return {
+    bool: {
+      must: [
+        {
+          term: {
+            ProjectID: {
+              value: projectID,
+              boost: 1.0,
+            },
+          },
+        },
+        {
+          range: {
+            '@timestamp': {
+              gte: timeFrom,
+              lt: timeTo,
+              format: 'epoch_second',
+              boost: 1.0,
+            },
+          },
+        },
+      ],
+    },
+  } as ProjectTimeQuery;
+}
+
+export function autoHistogramSumCountOnGroupByField(
+  searchField: string
+): AggregationsAggregationContainer {
+  return {
+    auto_date_histogram: {
+      field: '@timestamp',
+      buckets: 100,
+    },
+    aggs: {
+      group_by: {
+        terms: {
+          field: searchField,
+          order: {
+            Count: 'desc',
+          },
+          size: 100,
+        },
+        aggs: {
+          Count: {
+            sum: {
+              field: 'Count',
+            },
+          },
+        },
+      },
+    },
+  };
+}
+
 function getExeFileName(obj) {
   if (obj.ExeFileName === undefined) {
     return '';

src/plugins/profiling/server/routes/search_flamechart.ts

@@ -10,58 +10,7 @@ import type { IRouter } from 'kibana/server';
 import type { DataRequestHandlerContext } from '../../../data/server';
 import { getRemoteRoutePaths } from '../../common';
 import { FlameGraph } from './flamegraph';
-
-interface FilterObject {
-  bool: {
-    must: Array<
-      | {
-          term: {
-            ProjectID: {
-              value: string;
-              boost: number;
-            };
-          };
-        }
-      | {
-          range: {
-            '@timestamp': {
-              gte: string;
-              lt: string;
-              format: string;
-              boost: number;
-            };
-          };
-        }
-    >;
-  };
-}
-
-function createFilterObject(projectID: string, timeFrom: string, timeTo: string): FilterObject {
-  return {
-    bool: {
-      must: [
-        {
-          term: {
-            ProjectID: {
-              value: projectID,
-              boost: 1.0,
-            },
-          },
-        },
-        {
-          range: {
-            '@timestamp': {
-              gte: timeFrom,
-              lt: timeTo,
-              format: 'epoch_second',
-              boost: 1.0,
-            },
-          },
-        },
-      ],
-    },
-  } as FilterObject;
-}
+import { newProjectTimeQuery } from './mappings';
 
 function getSampledTraceEventsIndex(
   sampleSize: number,
@@ -126,7 +75,7 @@ export function registerFlameChartSearchRoute(router: IRouter<DataRequestHandler
 
       try {
         const esClient = context.core.elasticsearch.client.asCurrentUser;
-        const filter = createFilterObject(projectID!, timeFrom!, timeTo!);
+        const filter = newProjectTimeQuery(projectID!, timeFrom!, timeTo!);
 
         // const resp = await getCountResponse(context, filter);
         const resp = await esClient.search({

src/plugins/profiling/server/routes/search_topn.test.ts

@@ -0,0 +1,116 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { topNElasticSearchQuery } from './search_topn';
+import { DataRequestHandlerContext } from '../../../data/server';
+import { kibanaResponseFactory } from '../../../../core/server';
+import { AggregationsAggregationContainer } from '@elastic/elasticsearch/lib/api/types';
+
+const anyQuery = 'any::query';
+const index = 'test';
+const testAgg = { aggs: { test: {} } };
+
+jest.mock('./mappings', () => ({
+  newProjectTimeQuery: (proj: string, from: string, to: string) => {
+    return anyQuery;
+  },
+  autoHistogramSumCountOnGroupByField: (searchField: string): AggregationsAggregationContainer => {
+    return testAgg;
+  },
+}));
+
+function mockTopNData() {
+  return {
+    core: {
+      elasticsearch: {
+        client: {
+          asCurrentUser: {
+            search: jest.fn().mockResolvedValue({
+              body: {
+                aggregations: {
+                  histogram: {
+                    buckets: [
+                      {
+                        key_as_string: '1644506880',
+                        key: 1644506880000,
+                        doc_count: 700,
+                        group_by: {
+                          buckets: [
+                            {
+                              key: 'vyHke_Kdp2c05tXV7a_Rkg==',
+                              doc_count: 10,
+                              Count: {
+                                value: 100.0,
+                              },
+                            },
+                          ],
+                        },
+                      },
+                    ],
+                  },
+                },
+              },
+            }),
+            mget: jest.fn().mockResolvedValue({
+              body: {
+                docs: [],
+              },
+            }),
+          },
+        },
+      },
+    },
+  };
+}
+
+describe('TopN data from Elasticsearch', () => {
+  const mock = mockTopNData();
+  const queryMock = mock as unknown as DataRequestHandlerContext;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('building the query', () => {
+    it('filters by projectID and aggregates timerange on histogram', async () => {
+      await topNElasticSearchQuery(
+        queryMock,
+        index,
+        '123',
+        '456',
+        '789',
+        'field',
+        kibanaResponseFactory
+      );
+      expect(mock.core.elasticsearch.client.asCurrentUser.search).toHaveBeenCalledWith({
+        index,
+        body: {
+          query: anyQuery,
+          aggs: {
+            histogram: testAgg,
+          },
+        },
+      });
+    });
+  });
+  describe('when fetching Stack Traces', () => {
+    it('should search first then mget', async () => {
+      await topNElasticSearchQuery(
+        queryMock,
+        index,
+        '123',
+        '456',
+        '789',
+        'StackTraceID',
+        kibanaResponseFactory
+      );
+      expect(mock.core.elasticsearch.client.asCurrentUser.search).toHaveBeenCalledTimes(1);
+      expect(mock.core.elasticsearch.client.asCurrentUser.mget).toHaveBeenCalledTimes(1);
+    });
+  });
+});