[WIP] Usage collection es clients

src/plugins/telemetry/server/plugin.ts

@@ -104,11 +104,11 @@ export class TelemetryPlugin implements Plugin<TelemetryPluginSetup, TelemetryPl
     const currentKibanaVersion = this.currentKibanaVersion;
     const config$ = this.config$;
     const isDev = this.isDev;
-    registerCollection(
-      telemetryCollectionManager,
-      elasticsearch.legacy.client,
-      () => this.elasticsearchClient
-    );
+    const collectionClients = {
+      esCluster: elasticsearch.legacy.client,
+      esClientGetter: () => this.elasticsearchClient,
+    };
+    registerCollection(telemetryCollectionManager, collectionClients);
     const router = http.createRouter();
 
     registerRoutes({

src/plugins/telemetry/server/telemetry_collection/get_cluster_stats.test.ts

@@ -21,12 +21,6 @@ import { elasticsearchServiceMock } from '../../../../../src/core/server/mocks';
 import { getClusterStats } from './get_cluster_stats';
 import { TIMEOUT } from './constants';
 
-export function mockGetClusterStats(clusterStats: any) {
-  const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
-  esClient.cluster.stats.mockResolvedValue(clusterStats);
-  return esClient;
-}
-
 describe('get_cluster_stats', () => {
   it('uses the esClient to get the response from the `cluster.stats` API', async () => {
     const response = Promise.resolve({ body: { cluster_uuid: '1234' } });

src/plugins/telemetry/server/telemetry_collection/get_cluster_stats.ts

@@ -32,9 +32,10 @@ export async function getClusterStats(esClient: ElasticsearchClient) {
 
 /**
  * Get the cluster uuids from the connected cluster.
+ * TODO: add unit tests for getClusterUuids
  */
-export const getClusterUuids: ClusterDetailsGetter = async ({ esClient }) => {
-  const { body } = await esClient.cluster.stats({ timeout: TIMEOUT });
+export const getClusterUuids: ClusterDetailsGetter = async ({ scopedClients }) => {
+  const { body } = await scopedClients.esClient.cluster.stats({ timeout: TIMEOUT });
 
   return [{ clusterUuid: body.cluster_uuid }];
 };

src/plugins/telemetry/server/telemetry_collection/get_kibana.ts

@@ -19,9 +19,10 @@
 
 import { omit } from 'lodash';
 import { UsageCollectionSetup } from 'src/plugins/usage_collection/server';
-import { LegacyAPICaller } from 'kibana/server';
-import { StatsCollectionContext } from 'src/plugins/telemetry_collection_manager/server';
-import { ElasticsearchClient } from 'src/core/server';
+import {
+  ScopedCollectionClients,
+  StatsCollectionContext,
+} from 'src/plugins/telemetry_collection_manager/server';
 
 export interface KibanaUsageStats {
   kibana: {
@@ -83,9 +84,8 @@ export function handleKibanaStats(
 
 export async function getKibana(
   usageCollection: UsageCollectionSetup,
-  callWithInternalUser: LegacyAPICaller,
-  asInternalUser: ElasticsearchClient
+  scopedClients: ScopedCollectionClients
 ): Promise<KibanaUsageStats> {
-  const usage = await usageCollection.bulkFetch(callWithInternalUser, asInternalUser);
+  const usage = await usageCollection.bulkFetch(scopedClients);
   return usageCollection.toObject(usage);
 }

src/plugins/telemetry/server/telemetry_collection/get_local_license.ts

@@ -65,8 +65,8 @@ async function getLicenseFromLocalOrMaster(esClient: ElasticsearchClient) {
   return license;
 }
 
-export const getLocalLicense: LicenseGetter = async (clustersDetails, { esClient }) => {
-  const license = await getLicenseFromLocalOrMaster(esClient);
+export const getLocalLicense: LicenseGetter = async (clustersDetails, { scopedClients }) => {
+  const license = await getLicenseFromLocalOrMaster(scopedClients.esClient);
   // It should be called only with 1 cluster element in the clustersDetails array, but doing reduce just in case.
   return clustersDetails.reduce((acc, { clusterUuid }) => ({ ...acc, [clusterUuid]: license }), {});
 };

src/plugins/telemetry/server/telemetry_collection/get_local_stats.test.ts

@@ -224,12 +224,15 @@ describe('get_local_stats', () => {
 
   describe('getLocalStats', () => {
     it('returns expected object with kibana data', async () => {
-      const callCluster = jest.fn();
       const usageCollection = mockUsageCollection(kibana);
-      const esClient = mockGetLocalStats(clusterInfo, clusterStats);
+
+      const scopedClients = {
+        callCluster: jest.fn(),
+        esClient: mockGetLocalStats(clusterInfo, clusterStats),
+      };
       const response = await getLocalStats(
         [{ clusterUuid: 'abc123' }],
-        { callCluster, usageCollection, esClient, start: '', end: '' },
+        { scopedClients, usageCollection, start: '', end: '' },
         context
       );
       const result = response[0];
@@ -244,12 +247,14 @@ describe('get_local_stats', () => {
     });
 
     it('returns an empty array when no cluster uuid is provided', async () => {
-      const callCluster = jest.fn();
       const usageCollection = mockUsageCollection(kibana);
-      const esClient = mockGetLocalStats(clusterInfo, clusterStats);
+      const scopedClients = {
+        callCluster: jest.fn(),
+        esClient: mockGetLocalStats(clusterInfo, clusterStats),
+      };
       const response = await getLocalStats(
         [],
-        { callCluster, usageCollection, esClient, start: '', end: '' },
+        { scopedClients, usageCollection, start: '', end: '' },
         context
       );
       expect(response).toBeDefined();

src/plugins/telemetry/server/telemetry_collection/get_local_stats.ts

@@ -71,15 +71,15 @@ export const getLocalStats: StatsGetter<{}, TelemetryLocalStats> = async (
   config, // contains the new esClient already scoped contains usageCollection, callCluster, esClient, start, end
   context // StatsCollectionContext contains logger and version (string)
 ) => {
-  const { callCluster, usageCollection, esClient } = config;
-
+  const { scopedClients, usageCollection } = config;
+  const { esClient } = scopedClients;
   return await Promise.all(
     clustersDetails.map(async (clustersDetail) => {
       const [clusterInfo, clusterStats, nodesUsage, kibana, dataTelemetry] = await Promise.all([
         getClusterInfo(esClient), // cluster info
         getClusterStats(esClient), // cluster stats (not to be confused with cluster _state_)
         getNodesUsage(esClient), // nodes_usage info
-        getKibana(usageCollection, callCluster, esClient),
+        getKibana(usageCollection, scopedClients),
         getDataTelemetry(esClient),
       ]);
       return handleLocalStats(

src/plugins/telemetry/server/telemetry_collection/register_collection.ts

@@ -36,21 +36,20 @@
  * under the License.
  */
 
-import { ILegacyClusterClient } from 'kibana/server';
-import { TelemetryCollectionManagerPluginSetup } from 'src/plugins/telemetry_collection_manager/server';
-import { IClusterClient } from '../../../../../src/core/server';
+import {
+  CollectionClients,
+  TelemetryCollectionManagerPluginSetup,
+} from 'src/plugins/telemetry_collection_manager/server';
 import { getLocalStats } from './get_local_stats';
 import { getClusterUuids } from './get_cluster_stats';
 import { getLocalLicense } from './get_local_license';
 
 export function registerCollection(
   telemetryCollectionManager: TelemetryCollectionManagerPluginSetup,
-  esCluster: ILegacyClusterClient,
-  esClientGetter: () => IClusterClient | undefined
+  collectionClients: CollectionClients
 ) {
   telemetryCollectionManager.setCollection({
-    esCluster,
-    esClientGetter,
+    collectionClients,
     title: 'local',
     priority: 0,
     statsGetter: getLocalStats,

src/plugins/telemetry_collection_manager/server/index.ts

@@ -38,4 +38,6 @@ export {
   ClusterDetails,
   ClusterDetailsGetter,
   LicenseGetter,
+  ScopedCollectionClients,
+  CollectionClients,
 } from './types';

src/plugins/telemetry_collection_manager/server/plugin.ts

@@ -86,8 +86,7 @@ export class TelemetryCollectionManagerPlugin
     const {
       title,
       priority,
-      esCluster,
-      esClientGetter,
+      collectionClients,
       statsGetter,
       clusterDetailsGetter,
       licenseGetter,
@@ -104,10 +103,10 @@ export class TelemetryCollectionManagerPlugin
       if (!statsGetter) {
         throw Error('Stats getter method not set.');
       }
-      if (!esCluster) {
+      if (!collectionClients.esCluster) {
         throw Error('esCluster name must be set for the getCluster method.');
       }
-      if (!esClientGetter) {
+      if (!collectionClients.esClientGetter) {
         throw Error('esClientGetter method not set.');
       }
       if (!clusterDetailsGetter) {
@@ -121,9 +120,8 @@ export class TelemetryCollectionManagerPlugin
         licenseGetter,
         statsGetter,
         clusterDetailsGetter,
-        esCluster,
+        collectionClients,
         title,
-        esClientGetter,
       });
       this.usageGetterMethodPriority = priority;
     }
@@ -137,14 +135,17 @@ export class TelemetryCollectionManagerPlugin
   ): StatsCollectionConfig {
     const { start, end, request } = config;
 
-    const callCluster = config.unencrypted
-      ? collection.esCluster.asScoped(request).callAsCurrentUser
-      : collection.esCluster.callAsInternalUser;
-    // Scope the new elasticsearch Client appropriately and pass to the stats collection config
-    const esClient = config.unencrypted
-      ? collectionEsClient.asScoped(config.request).asCurrentUser
-      : collectionEsClient.asInternalUser;
-    return { callCluster, start, end, usageCollection, esClient };
+    const scopedClients = {
+      callCluster: config.unencrypted
+        ? collection.collectionClients.esCluster?.asScoped(request).callAsCurrentUser
+        : collection.collectionClients.esCluster?.callAsInternalUser,
+      // Scope the new elasticsearch Client appropriately and pass to the stats collection config
+      esClient: config.unencrypted
+        ? collectionEsClient.asScoped(config.request).asCurrentUser
+        : collectionEsClient.asInternalUser,
+    };
+
+    return { usageCollection, scopedClients, start, end };
   }
 
   private async getOptInStats(optInStatus: boolean, config: StatsGetterConfig) {
@@ -153,7 +154,7 @@ export class TelemetryCollectionManagerPlugin
     }
     for (const collection of this.collections) {
       // first fetch the client and make sure it's not undefined.
-      const collectionEsClient = collection.esClientGetter();
+      const collectionEsClient = collection.collectionClients.esClientGetter();
       if (collectionEsClient !== undefined) {
         const statsCollectionConfig = this.getStatsCollectionConfig(
           config,
@@ -208,7 +209,7 @@ export class TelemetryCollectionManagerPlugin
       return [];
     }
     for (const collection of this.collections) {
-      const collectionEsClient = collection.esClientGetter();
+      const collectionEsClient = collection.collectionClients.esClientGetter();
       if (collectionEsClient !== undefined) {
         const statsCollectionConfig = this.getStatsCollectionConfig(
           config,

src/plugins/telemetry_collection_manager/server/types.ts

@@ -69,12 +69,16 @@ export interface ClusterDetails {
   clusterUuid: string;
 }
 
+export interface ScopedCollectionClients {
+  /** @depricated use esClient instead */
+  callCluster: LegacyAPICaller;
+  esClient: ElasticsearchClient;
+}
 export interface StatsCollectionConfig {
   usageCollection: UsageCollectionSetup;
-  callCluster: LegacyAPICaller;
+  scopedClients: ScopedCollectionClients;
   start: string | number;
   end: string | number;
-  esClient: ElasticsearchClient;
 }
 
 export interface BasicStatsPayload {
@@ -137,14 +141,18 @@ export interface CollectionConfig<
 > {
   title: string;
   priority: number;
-  esCluster: ILegacyClusterClient;
-  esClientGetter: () => IClusterClient | undefined; // --> by now we know that the client getter will return the IClusterClient but we assure that through a code check
+  collectionClients: CollectionClients;
   statsGetter: StatsGetter<CustomContext, T>;
   clusterDetailsGetter: ClusterDetailsGetter<CustomContext>;
   licenseGetter: LicenseGetter<CustomContext>;
   customContext?: CustomContext;
 }
 
+export interface CollectionClients {
+  /** @deprecated use esClientGetter instead */
+  esCluster: ILegacyClusterClient;
+  esClientGetter: () => IClusterClient | undefined; // --> by now we know that the client getter will return the IClusterClient but we assure that through a code check}
+}
 export interface Collection<
   CustomContext extends Record<string, any> = {},
   T extends BasicStatsPayload = BasicStatsPayload
@@ -153,7 +161,6 @@ export interface Collection<
   statsGetter: StatsGetter<CustomContext, T>;
   licenseGetter: LicenseGetter<CustomContext>;
   clusterDetailsGetter: ClusterDetailsGetter<CustomContext>;
-  esCluster: ILegacyClusterClient;
-  esClientGetter: () => IClusterClient | undefined; // the collection could still return undefined for the es client getter.
+  collectionClients: CollectionClients;
   title: string;
 }

src/plugins/usage_collection/server/collector/collector.ts

@@ -17,7 +17,9 @@
  * under the License.
  */
 
-import { Logger, LegacyAPICaller, ElasticsearchClient } from 'kibana/server';
+import { Logger } from 'kibana/server';
+// eslint-disable-next-line @kbn/eslint/no-restricted-paths
+import { ScopedCollectionClients } from 'src/plugins/telemetry_collection_manager/server/types';
 
 export type CollectorFormatForBulkUpload<T, U> = (result: T) => { type: string; payload: U };
 
@@ -49,7 +51,7 @@ export interface CollectorOptions<T = unknown, U = T> {
   type: string;
   init?: Function;
   schema?: MakeSchemaFrom<T>;
-  fetch: (callCluster: LegacyAPICaller, esClient?: ElasticsearchClient) => Promise<T> | T;
+  fetch: (scopedClients: ScopedCollectionClients) => Promise<T> | T;
   /*
    * A hook for allowing the fetched data payload to be organized into a typed
    * data model for internal bulk upload. See defaultFormatterForBulkUpload for

src/plugins/usage_collection/server/collector/collector_set.ts

@@ -18,7 +18,9 @@
  */
 
 import { snakeCase } from 'lodash';
-import { Logger, LegacyAPICaller, ElasticsearchClient } from 'kibana/server';
+import { Logger } from 'kibana/server';
+// eslint-disable-next-line @kbn/eslint/no-restricted-paths
+import { ScopedCollectionClients } from 'src/plugins/telemetry_collection_manager/server/types';
 import { Collector, CollectorOptions } from './collector';
 import { UsageCollector } from './usage_collector';
 
@@ -121,8 +123,7 @@ export class CollectorSet {
   // the shape of the response is different when using the new ES client as is the error handling.
   // We'll handle the refactor for using the new client in a follow up PR.
   public bulkFetch = async (
-    callCluster: LegacyAPICaller,
-    esClient: ElasticsearchClient,
+    scopedClients: ScopedCollectionClients,
     collectors: Map<string, Collector<any, any>> = this.collectors
   ) => {
     const responses = await Promise.all(
@@ -131,7 +132,7 @@ export class CollectorSet {
         try {
           return {
             type: collector.type,
-            result: await collector.fetch(callCluster, esClient), // each collector must ensure they handle the response appropriately.
+            result: await collector.fetch(scopedClients), // This is a breaking change
           };
         } catch (err) {
           this.logger.warn(err);
@@ -153,9 +154,9 @@ export class CollectorSet {
     return this.makeCollectorSetFromArray(filtered);
   };
 
-  public bulkFetchUsage = async (callCluster: LegacyAPICaller, esClient: ElasticsearchClient) => {
+  public bulkFetchUsage = async (scopedClients: ScopedCollectionClients) => {
     const usageCollectors = this.getFilteredCollectorSet((c) => c instanceof UsageCollector);
-    return await this.bulkFetch(callCluster, esClient, usageCollectors.collectors);
+    return await this.bulkFetch(scopedClients, usageCollectors.collectors);
   };
 
   // convert an array of fetched stats results into key/object

x-pack/plugins/monitoring/server/plugin.ts

@@ -146,14 +146,13 @@ export class Plugin {
 
     // Initialize telemetry
     if (plugins.telemetryCollectionManager) {
-      registerMonitoringCollection(
-        plugins.telemetryCollectionManager,
-        this.cluster,
-        () => this.telemetryElasticsearchClient,
-        {
-          maxBucketSize: config.ui.max_bucket_size,
-        }
-      );
+      const collectionClients = {
+        esCluster: this.cluster,
+        esClientGetter: () => this.telemetryElasticsearchClient,
+      };
+      registerMonitoringCollection(plugins.telemetryCollectionManager, collectionClients, {
+        maxBucketSize: config.ui.max_bucket_size,
+      });
     }
 
     // Register collector objects for stats to show up in the APIs