[Vega] Expensive queries are causing unnecessary load and delays on E…

…lasticsearch (#99023) (#110448) * [Vega] Expensive queries are causing unnecessary load and delays on Elasticsearch Part of: #93770 * Update get_usage_collector.test.ts * add namespaces: ['*'] Co-authored-by: Kibana Machine <[email protected]> Co-authored-by: Kibana Machine <[email protected]>
elastic · Aug 30, 2021 · 949e123 · 949e123
1 parent 4bbb377
commit 949e123
Show file tree

Hide file tree

Showing 4 changed files with 112 additions and 129 deletions.
diff --git a/src/plugins/vis_type_vega/server/usage_collector/get_usage_collector.test.ts b/src/plugins/vis_type_vega/server/usage_collector/get_usage_collector.test.ts
@@ -7,67 +7,63 @@
  */
 
 import { getStats } from './get_usage_collector';
-import { HomeServerPluginSetup } from '../../../home/server';
-import { createCollectorFetchContextMock } from 'src/plugins/usage_collection/server/mocks';
+import { createCollectorFetchContextMock } from '../../../usage_collection/server/mocks';
+import type { HomeServerPluginSetup } from '../../../home/server';
+import type { SavedObjectsClientContract } from '../../../../core/server';
 
 const mockedSavedObjects = [
   // vega-lite lib spec
   {
-    _id: 'visualization:vega-1',
-    _source: {
-      type: 'visualization',
-      visualization: {
-        visState: JSON.stringify({
-          type: 'vega',
-          params: {
-            spec: '{"$schema": "https://vega.github.io/schema/vega-lite/v5.json" }',
-          },
-        }),
-      },
+    attributes: {
+      visState: JSON.stringify({
+        type: 'vega',
+        params: {
+          spec: '{"$schema": "https://vega.github.io/schema/vega-lite/v5.json" }',
+        },
+      }),
     },
   },
   // vega lib spec
   {
-    _id: 'visualization:vega-2',
-    _source: {
-      type: 'visualization',
-      visualization: {
-        visState: JSON.stringify({
-          type: 'vega',
-          params: {
-            spec: '{"$schema": "https://vega.github.io/schema/vega/v5.json" }',
-          },
-        }),
-      },
+    attributes: {
+      visState: JSON.stringify({
+        type: 'vega',
+        params: {
+          spec: '{"$schema": "https://vega.github.io/schema/vega/v5.json" }',
+        },
+      }),
     },
   },
   // map layout
   {
-    _id: 'visualization:vega-3',
-    _source: {
-      type: 'visualization',
-      visualization: {
-        visState: JSON.stringify({
-          type: 'vega',
-          params: {
-            spec:
-              '{"$schema": "https://vega.github.io/schema/vega/v3.json" \n "config": { "kibana" : { "type": "map" }} }',
-          },
-        }),
-      },
+    attributes: {
+      visState: JSON.stringify({
+        type: 'vega',
+        params: {
+          spec:
+            '{"$schema": "https://vega.github.io/schema/vega/v3.json" \n "config": { "kibana" : { "type": "map" }} }',
+        },
+      }),
     },
   },
 ];
 
-const getMockCollectorFetchContext = (hits?: unknown[]) => {
+const getMockCollectorFetchContext = (savedObjects?: unknown[]) => {
   const fetchParamsMock = createCollectorFetchContextMock();
 
-  fetchParamsMock.esClient.search = jest.fn().mockResolvedValue({ body: { hits: { hits } } });
+  fetchParamsMock.soClient = ({
+    createPointInTimeFinder: jest.fn().mockResolvedValue({
+      close: jest.fn(),
+      find: function* asyncGenerator() {
+        yield { saved_objects: savedObjects };
+      },
+    }),
+  } as unknown) as SavedObjectsClientContract;
+
   return fetchParamsMock;
 };
 
 describe('Vega visualization usage collector', () => {
-  const mockIndex = 'mock_index';
   const mockDeps = {
     home: ({
       sampleData: {
@@ -94,13 +90,13 @@ describe('Vega visualization usage collector', () => {
   };
 
   test('Returns undefined when no results found (undefined)', async () => {
-    const result = await getStats(getMockCollectorFetchContext().esClient, mockIndex, mockDeps);
+    const result = await getStats(getMockCollectorFetchContext().soClient, mockDeps);
 
     expect(result).toBeUndefined();
   });
 
   test('Returns undefined when no results found (0 results)', async () => {
-    const result = await getStats(getMockCollectorFetchContext([]).esClient, mockIndex, mockDeps);
+    const result = await getStats(getMockCollectorFetchContext([]).soClient, mockDeps);
 
     expect(result).toBeUndefined();
   });
@@ -115,38 +111,34 @@ describe('Vega visualization usage collector', () => {
         },
       },
     ]);
-    const result = await getStats(mockCollectorFetchContext.esClient, mockIndex, mockDeps);
+    const result = await getStats(mockCollectorFetchContext.soClient, mockDeps);
 
     expect(result).toBeUndefined();
   });
 
   test('Should ingnore sample data visualizations', async () => {
     const mockCollectorFetchContext = getMockCollectorFetchContext([
       {
-        _id: 'visualization:sampledata-123',
-        _source: {
-          type: 'visualization',
-          visualization: {
-            visState: JSON.stringify({
-              type: 'vega',
-              title: 'sample vega visualization',
-              params: {
-                spec: '{"$schema": "https://vega.github.io/schema/vega/v5.json" }',
-              },
-            }),
-          },
+        attributes: {
+          visState: JSON.stringify({
+            type: 'vega',
+            title: 'sample vega visualization',
+            params: {
+              spec: '{"$schema": "https://vega.github.io/schema/vega/v5.json" }',
+            },
+          }),
         },
       },
     ]);
 
-    const result = await getStats(mockCollectorFetchContext.esClient, mockIndex, mockDeps);
+    const result = await getStats(mockCollectorFetchContext.soClient, mockDeps);
 
     expect(result).toBeUndefined();
   });
 
   test('Summarizes visualizations response data', async () => {
     const mockCollectorFetchContext = getMockCollectorFetchContext(mockedSavedObjects);
-    const result = await getStats(mockCollectorFetchContext.esClient, mockIndex, mockDeps);
+    const result = await getStats(mockCollectorFetchContext.soClient, mockDeps);
 
     expect(result).toMatchObject({
       vega_lib_specs_total: 2,

diff --git a/src/plugins/vis_type_vega/server/usage_collector/get_usage_collector.ts b/src/plugins/vis_type_vega/server/usage_collector/get_usage_collector.ts
@@ -7,14 +7,20 @@
  */
 
 import { parse } from 'hjson';
-import type { ElasticsearchClient } from 'src/core/server';
 
-import { VegaSavedObjectAttributes, VisTypeVegaPluginSetupDependencies } from '../types';
+import type { SavedObjectsClientContract, SavedObjectsFindResult } from '../../../../core/server';
+import type { SavedVisState } from '../../../visualizations/common';
+import type { VegaSavedObjectAttributes, VisTypeVegaPluginSetupDependencies } from '../types';
 
 type UsageCollectorDependencies = Pick<VisTypeVegaPluginSetupDependencies, 'home'>;
-
 type VegaType = 'vega' | 'vega-lite';
 
+export interface VegaUsage {
+  vega_lib_specs_total: number;
+  vega_lite_lib_specs_total: number;
+  vega_use_map_total: number;
+}
+
 function isVegaType(attributes: any): attributes is VegaSavedObjectAttributes {
   return attributes && attributes.type === 'vega' && attributes.params?.spec;
 }
@@ -45,15 +51,8 @@ const getDefaultVegaVisualizations = (home: UsageCollectorDependencies['home'])
   return titles;
 };
 
-export interface VegaUsage {
-  vega_lib_specs_total: number;
-  vega_lite_lib_specs_total: number;
-  vega_use_map_total: number;
-}
-
 export const getStats = async (
-  esClient: ElasticsearchClient,
-  index: string,
+  soClient: SavedObjectsClientContract,
   { home }: UsageCollectorDependencies
 ): Promise<VegaUsage | undefined> => {
   let shouldPublishTelemetry = false;
@@ -64,58 +63,54 @@ export const getStats = async (
     vega_use_map_total: 0,
   };
 
-  const searchParams = {
-    size: 10000,
-    index,
-    ignoreUnavailable: true,
-    filterPath: ['hits.hits._id', 'hits.hits._source.visualization'],
-    body: {
-      query: {
-        bool: {
-          filter: { term: { type: 'visualization' } },
-        },
-      },
-    },
-  };
-
-  const { body: esResponse } = await esClient.search<{ visualization: { visState: string } }>(
-    searchParams
-  );
-  const size = esResponse?.hits?.hits?.length ?? 0;
-
-  if (!size) {
-    return;
-  }
-
   // we want to exclude the Vega Sample Data visualizations from the stats
   // in order to have more accurate results
   const excludedFromStatsVisualizations = getDefaultVegaVisualizations(home);
-  for (const hit of esResponse.hits.hits) {
-    const visualization = hit._source?.visualization;
-    const visState = JSON.parse(visualization?.visState ?? '{}');
 
-    if (isVegaType(visState) && !excludedFromStatsVisualizations.includes(visState.title)) {
-      try {
-        const spec = parse(visState.params.spec, { legacyRoot: false });
+  const finder = await soClient.createPointInTimeFinder({
+    type: 'visualization',
+    perPage: 1000,
+    namespaces: ['*'],
+  });
 
-        if (spec) {
-          shouldPublishTelemetry = true;
+  const doTelemetry = ({ params }: SavedVisState) => {
+    try {
+      const spec = parse(params.spec, { legacyRoot: false });
 
-          if (checkVegaSchemaType(spec.$schema, 'vega')) {
-            vegaUsage.vega_lib_specs_total++;
-          }
-          if (checkVegaSchemaType(spec.$schema, 'vega-lite')) {
-            vegaUsage.vega_lite_lib_specs_total++;
-          }
-          if (spec.config?.kibana?.type === 'map') {
-            vegaUsage.vega_use_map_total++;
-          }
+      if (spec) {
+        shouldPublishTelemetry = true;
+
+        if (checkVegaSchemaType(spec.$schema, 'vega')) {
+          vegaUsage.vega_lib_specs_total++;
+        }
+        if (checkVegaSchemaType(spec.$schema, 'vega-lite')) {
+          vegaUsage.vega_lite_lib_specs_total++;
+        }
+        if (spec.config?.kibana?.type === 'map') {
+          vegaUsage.vega_use_map_total++;
         }
-      } catch (e) {
-        // Let it go, the data is invalid and we'll don't need to handle it
       }
+    } catch (e) {
+      // Let it go, the data is invalid and we'll don't need to handle it
     }
+  };
+
+  for await (const response of finder.find()) {
+    (response.saved_objects || []).forEach(({ attributes }: SavedObjectsFindResult<any>) => {
+      if (attributes?.visState) {
+        try {
+          const visState: SavedVisState = JSON.parse(attributes.visState);
+
+          if (isVegaType(visState) && !excludedFromStatsVisualizations.includes(visState.title)) {
+            doTelemetry(visState);
+          }
+        } catch {
+          // nothing to be here, "so" not valid
+        }
+      }
+    });
   }
+  await finder.close();
 
   return shouldPublishTelemetry ? vegaUsage : undefined;
 };
diff --git a/src/plugins/vis_type_vega/server/usage_collector/register_vega_collector.test.ts b/src/plugins/vis_type_vega/server/usage_collector/register_vega_collector.test.ts
@@ -6,27 +6,28 @@
  * Side Public License, v 1.
  */
 
-import { of } from 'rxjs';
+import {
+  createUsageCollectionSetupMock,
+  createCollectorFetchContextMock,
+} from '../../../usage_collection/server/mocks';
 import { mockStats, mockGetStats } from './get_usage_collector.mock';
-import { createUsageCollectionSetupMock } from 'src/plugins/usage_collection/server/mocks';
-import { createCollectorFetchContextMock } from 'src/plugins/usage_collection/server/mocks';
-import { HomeServerPluginSetup } from '../../../home/server';
 import { registerVegaUsageCollector } from './register_vega_collector';
-import { ConfigObservable } from '../types';
+
+import type { HomeServerPluginSetup } from '../../../home/server';
+import type { ConfigObservable } from '../types';
 
 describe('registerVegaUsageCollector', () => {
-  const mockIndex = 'mock_index';
   const mockDeps = { home: ({} as unknown) as HomeServerPluginSetup };
-  const mockConfig = of({ kibana: { index: mockIndex } }) as ConfigObservable;
+  const mockConfig = {} as ConfigObservable;
 
-  it('makes a usage collector and registers it`', () => {
+  test('makes a usage collector and registers it`', () => {
     const mockCollectorSet = createUsageCollectionSetupMock();
     registerVegaUsageCollector(mockCollectorSet, mockConfig, mockDeps);
     expect(mockCollectorSet.makeUsageCollector).toBeCalledTimes(1);
     expect(mockCollectorSet.registerCollector).toBeCalledTimes(1);
   });
 
-  it('makeUsageCollector configs fit the shape', () => {
+  test('makeUsageCollector configs fit the shape', () => {
     const mockCollectorSet = createUsageCollectionSetupMock();
     registerVegaUsageCollector(mockCollectorSet, mockConfig, mockDeps);
     expect(mockCollectorSet.makeUsageCollector).toHaveBeenCalledWith({
@@ -39,21 +40,21 @@ describe('registerVegaUsageCollector', () => {
     expect(usageCollectorConfig.isReady()).toBe(true);
   });
 
-  it('makeUsageCollector config.isReady returns true', () => {
+  test('makeUsageCollector config.isReady returns true', () => {
     const mockCollectorSet = createUsageCollectionSetupMock();
     registerVegaUsageCollector(mockCollectorSet, mockConfig, mockDeps);
     const usageCollectorConfig = mockCollectorSet.makeUsageCollector.mock.calls[0][0];
     expect(usageCollectorConfig.isReady()).toBe(true);
   });
 
-  it('makeUsageCollector config.fetch calls getStats', async () => {
+  test('makeUsageCollector config.fetch calls getStats', async () => {
     const mockCollectorSet = createUsageCollectionSetupMock();
     registerVegaUsageCollector(mockCollectorSet, mockConfig, mockDeps);
     const usageCollector = mockCollectorSet.makeUsageCollector.mock.results[0].value;
     const mockedCollectorFetchContext = createCollectorFetchContextMock();
     const fetchResult = await usageCollector.fetch(mockedCollectorFetchContext);
     expect(mockGetStats).toBeCalledTimes(1);
-    expect(mockGetStats).toBeCalledWith(mockedCollectorFetchContext.esClient, mockIndex, mockDeps);
+    expect(mockGetStats).toBeCalledWith(mockedCollectorFetchContext.soClient, mockDeps);
     expect(fetchResult).toBe(mockStats);
   });
 });
diff --git a/src/plugins/vis_type_vega/server/usage_collector/register_vega_collector.ts b/src/plugins/vis_type_vega/server/usage_collector/register_vega_collector.ts
@@ -6,10 +6,9 @@
  * Side Public License, v 1.
  */
 
-import { UsageCollectionSetup } from 'src/plugins/usage_collection/server';
-import { first } from 'rxjs/operators';
 import { getStats, VegaUsage } from './get_usage_collector';
-import { ConfigObservable, VisTypeVegaPluginSetupDependencies } from '../types';
+import type { UsageCollectionSetup } from '../../../usage_collection/server';
+import type { ConfigObservable, VisTypeVegaPluginSetupDependencies } from '../types';
 
 export function registerVegaUsageCollector(
   collectorSet: UsageCollectionSetup,
@@ -24,11 +23,7 @@ export function registerVegaUsageCollector(
       vega_lite_lib_specs_total: { type: 'long' },
       vega_use_map_total: { type: 'long' },
     },
-    fetch: async ({ esClient }) => {
-      const { index } = (await config.pipe(first()).toPromise()).kibana;
-
-      return await getStats(esClient, index, dependencies);
-    },
+    fetch: async ({ soClient }) => await getStats(soClient, dependencies),
   });
 
   collectorSet.registerCollector(collector);