[Lens] Migrate legacy es client and remove total hits as int (#84340)

Author: flash1293

x-pack/plugins/lens/server/routes/existing_fields.ts

@@ -5,8 +5,9 @@
  */
 
 import Boom from '@hapi/boom';
+import { errors } from '@elastic/elasticsearch';
 import { schema } from '@kbn/config-schema';
-import { ILegacyScopedClusterClient, RequestHandlerContext } from 'src/core/server';
+import { RequestHandlerContext, ElasticsearchClient } from 'src/core/server';
 import { CoreSetup, Logger } from 'src/core/server';
 import { IndexPattern, IndexPatternsService } from 'src/plugins/data/common';
 import { BASE_API_URL } from '../../common';
@@ -68,7 +69,7 @@ export async function existingFieldsRoute(setup: CoreSetup<PluginStartContract>,
         logger.info(
           `Field existence check failed: ${isBoomError(e) ? e.output.payload.message : e.message}`
         );
-        if (e.status === 404) {
+        if (e instanceof errors.ResponseError && e.statusCode === 404) {
           return res.notFound({ body: e.message });
         }
         if (isBoomError(e)) {
@@ -111,7 +112,7 @@ async function fetchFieldExistence({
     fromDate,
     toDate,
     dslQuery,
-    client: context.core.elasticsearch.legacy.client,
+    client: context.core.elasticsearch.client.asCurrentUser,
     index: indexPattern.title,
     timeFieldName: timeFieldName || indexPattern.timeFieldName,
     fields,
@@ -149,7 +150,7 @@ async function fetchIndexPatternStats({
   toDate,
   fields,
 }: {
-  client: ILegacyScopedClusterClient;
+  client: ElasticsearchClient;
   index: string;
   dslQuery: object;
   timeFieldName?: string;
@@ -179,7 +180,7 @@ async function fetchIndexPatternStats({
   };
 
   const scriptedFields = fields.filter((f) => f.isScript);
-  const result = await client.callAsCurrentUser('search', {
+  const { body: result } = await client.search({
     index,
     body: {
       size: SAMPLE_SIZE,

x-pack/plugins/lens/server/routes/field_stats.ts

@@ -5,6 +5,7 @@
  */
 
 import Boom from '@hapi/boom';
+import { errors } from '@elastic/elasticsearch';
 import DateMath from '@elastic/datemath';
 import { schema } from '@kbn/config-schema';
 import { CoreSetup } from 'src/core/server';
@@ -47,7 +48,7 @@ export async function initFieldsRoute(setup: CoreSetup<PluginStartContract>) {
       },
     },
     async (context, req, res) => {
-      const requestClient = context.core.elasticsearch.legacy.client;
+      const requestClient = context.core.elasticsearch.client.asCurrentUser;
       const { fromDate, toDate, timeFieldName, field, dslQuery } = req.body;
 
       try {
@@ -71,18 +72,18 @@ export async function initFieldsRoute(setup: CoreSetup<PluginStartContract>) {
           },
         };
 
-        const search = (aggs: unknown) =>
-          requestClient.callAsCurrentUser('search', {
+        const search = async (aggs: unknown) => {
+          const { body: result } = await requestClient.search({
             index: req.params.indexPatternTitle,
+            track_total_hits: true,
             body: {
               query,
               aggs,
             },
-            // The hits total changed in 7.0 from number to object, unless this flag is set
-            // this is a workaround for elasticsearch response types that are from 6.x
-            restTotalHitsAsInt: true,
             size: 0,
           });
+          return result;
+        };
 
         if (field.type === 'number') {
           return res.ok({
@@ -98,7 +99,7 @@ export async function initFieldsRoute(setup: CoreSetup<PluginStartContract>) {
           body: await getStringSamples(search, field),
         });
       } catch (e) {
-        if (e.status === 404) {
+        if (e instanceof errors.ResponseError && e.statusCode === 404) {
           return res.notFound();
         }
         if (e.isBoom) {
@@ -142,8 +143,7 @@ export async function getNumberHistogram(
 
   const minMaxResult = (await aggSearchWithBody(searchBody)) as ESSearchResponse<
     unknown,
-    { body: { aggs: typeof searchBody } },
-    { restTotalHitsAsInt: true }
+    { body: { aggs: typeof searchBody } }
   >;
 
   const minValue = minMaxResult.aggregations!.sample.min_value.value;
@@ -164,7 +164,7 @@ export async function getNumberHistogram(
 
   if (histogramInterval === 0) {
     return {
-      totalDocuments: minMaxResult.hits.total,
+      totalDocuments: minMaxResult.hits.total.value,
       sampledValues: minMaxResult.aggregations!.sample.sample_count.value!,
       sampledDocuments: minMaxResult.aggregations!.sample.doc_count,
       topValues: topValuesBuckets,
@@ -187,12 +187,11 @@ export async function getNumberHistogram(
   };
   const histogramResult = (await aggSearchWithBody(histogramBody)) as ESSearchResponse<
     unknown,
-    { body: { aggs: typeof histogramBody } },
-    { restTotalHitsAsInt: true }
+    { body: { aggs: typeof histogramBody } }
   >;
 
   return {
-    totalDocuments: minMaxResult.hits.total,
+    totalDocuments: minMaxResult.hits.total.value,
     sampledDocuments: minMaxResult.aggregations!.sample.doc_count,
     sampledValues: minMaxResult.aggregations!.sample.sample_count.value!,
     histogram: {
@@ -227,12 +226,11 @@ export async function getStringSamples(
   };
   const topValuesResult = (await aggSearchWithBody(topValuesBody)) as ESSearchResponse<
     unknown,
-    { body: { aggs: typeof topValuesBody } },
-    { restTotalHitsAsInt: true }
+    { body: { aggs: typeof topValuesBody } }
   >;
 
   return {
-    totalDocuments: topValuesResult.hits.total,
+    totalDocuments: topValuesResult.hits.total.value,
     sampledDocuments: topValuesResult.aggregations!.sample.doc_count,
     sampledValues: topValuesResult.aggregations!.sample.sample_count.value!,
     topValues: {
@@ -275,12 +273,11 @@ export async function getDateHistogram(
   };
   const results = (await aggSearchWithBody(histogramBody)) as ESSearchResponse<
     unknown,
-    { body: { aggs: typeof histogramBody } },
-    { restTotalHitsAsInt: true }
+    { body: { aggs: typeof histogramBody } }
   >;
 
   return {
-    totalDocuments: results.hits.total,
+    totalDocuments: results.hits.total.value,
     histogram: {
       buckets: results.aggregations!.histo.buckets.map((bucket) => ({
         count: bucket.doc_count,

x-pack/plugins/lens/server/routes/telemetry.ts

@@ -5,6 +5,7 @@
  */
 
 import Boom from '@hapi/boom';
+import { errors } from '@elastic/elasticsearch';
 import { CoreSetup } from 'src/core/server';
 import { schema } from '@kbn/config-schema';
 import { BASE_API_URL } from '../../common';
@@ -71,7 +72,7 @@ export async function initLensUsageRoute(setup: CoreSetup<PluginStartContract>)
 
         return res.ok({ body: {} });
       } catch (e) {
-        if (e.status === 404) {
+        if (e instanceof errors.ResponseError && e.statusCode === 404) {
           return res.notFound();
         }
         if (e.isBoom) {

x-pack/plugins/lens/server/usage/task.ts

@@ -4,7 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { LegacyAPICaller, CoreSetup, Logger } from 'kibana/server';
+import { CoreSetup, Logger, ElasticsearchClient } from 'kibana/server';
 import { Observable } from 'rxjs';
 import { first } from 'rxjs/operators';
 import moment from 'moment';
@@ -69,11 +69,12 @@ async function scheduleTasks(logger: Logger, taskManager: TaskManagerStartContra
 
 export async function getDailyEvents(
   kibanaIndex: string,
-  callCluster: LegacyAPICaller
+  getEsClient: () => Promise<ElasticsearchClient>
 ): Promise<{
   byDate: Record<string, Record<string, number>>;
   suggestionsByDate: Record<string, Record<string, number>>;
 }> {
+  const esClient = await getEsClient();
   const aggs = {
     daily: {
       date_histogram: {
@@ -114,15 +115,10 @@ export async function getDailyEvents(
     },
   };
 
-  const metrics: ESSearchResponse<
-    unknown,
-    {
-      body: { aggs: typeof aggs };
-    },
-    { restTotalHitsAsInt: true }
-  > = await callCluster('search', {
+  const { body: metrics } = await esClient.search<
+    ESSearchResponse<unknown, { body: { aggs: typeof aggs } }>
+  >({
     index: kibanaIndex,
-    rest_total_hits_as_int: true,
     body: {
       query: {
         bool: {
@@ -156,9 +152,9 @@ export async function getDailyEvents(
   });
 
   // Always delete old date because we don't report it
-  await callCluster('deleteByQuery', {
+  await esClient.deleteByQuery({
     index: kibanaIndex,
-    waitForCompletion: true,
+    wait_for_completion: true,
     body: {
       query: {
         bool: {
@@ -184,18 +180,18 @@ export function telemetryTaskRunner(
 ) {
   return ({ taskInstance }: RunContext) => {
     const { state } = taskInstance;
-    const callCluster = async (...args: Parameters<LegacyAPICaller>) => {
+    const getEsClient = async () => {
       const [coreStart] = await core.getStartServices();
-      return coreStart.elasticsearch.legacy.client.callAsInternalUser(...args);
+      return coreStart.elasticsearch.client.asInternalUser;
     };
 
     return {
       async run() {
         const kibanaIndex = (await config.pipe(first()).toPromise()).kibana.index;
 
         return Promise.all([
-          getDailyEvents(kibanaIndex, callCluster),
-          getVisualizationCounts(callCluster, kibanaIndex),
+          getDailyEvents(kibanaIndex, getEsClient),
+          getVisualizationCounts(getEsClient, kibanaIndex),
         ])
           .then(([lensTelemetry, lensVisualizations]) => {
             return {

x-pack/plugins/lens/server/usage/visualization_counts.ts

@@ -4,16 +4,16 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { LegacyAPICaller } from 'kibana/server';
+import { ElasticsearchClient } from 'kibana/server';
 import { LensVisualizationUsage } from './types';
 
 export async function getVisualizationCounts(
-  callCluster: LegacyAPICaller,
+  getEsClient: () => Promise<ElasticsearchClient>,
   kibanaIndex: string
 ): Promise<LensVisualizationUsage> {
-  const results = await callCluster('search', {
+  const esClient = await getEsClient();
+  const { body: results } = await esClient.search({
     index: kibanaIndex,
-    rest_total_hits_as_int: true,
     body: {
       query: {
         bool: {

x-pack/test/api_integration/apis/lens/telemetry.ts

@@ -6,8 +6,7 @@
 
 import moment from 'moment';
 import expect from '@kbn/expect';
-import { Client, SearchParams } from 'elasticsearch';
-import { LegacyAPICaller } from 'kibana/server';
+import { Client } from '@elastic/elasticsearch';
 
 import { FtrProviderContext } from '../../ftr_provider_context';
 
@@ -20,18 +19,17 @@ const COMMON_HEADERS = {
 
 export default ({ getService }: FtrProviderContext) => {
   const supertest = getService('supertest');
-  const es: Client = getService('legacyEs');
-  const callCluster: LegacyAPICaller = (((path: 'search', searchParams: SearchParams) => {
-    return es[path].call(es, searchParams);
-  }) as unknown) as LegacyAPICaller;
+  const es: Client = getService('es');
 
   async function assertExpectedSavedObjects(num: number) {
     // Make sure that new/deleted docs are available to search
     await es.indices.refresh({
       index: '.kibana',
     });
 
-    const { count } = await es.count({
+    const {
+      body: { count },
+    } = await es.count({
       index: '.kibana',
       q: 'type:lens-ui-telemetry',
     });
@@ -44,17 +42,19 @@ export default ({ getService }: FtrProviderContext) => {
       await es.deleteByQuery({
         index: '.kibana',
         q: 'type:lens-ui-telemetry',
-        waitForCompletion: true,
-        refresh: 'wait_for',
+        wait_for_completion: true,
+        refresh: true,
+        body: {},
       });
     });
 
     afterEach(async () => {
       await es.deleteByQuery({
         index: '.kibana',
         q: 'type:lens-ui-telemetry',
-        waitForCompletion: true,
-        refresh: 'wait_for',
+        wait_for_completion: true,
+        refresh: true,
+        body: {},
       });
     });
 
@@ -107,7 +107,7 @@ export default ({ getService }: FtrProviderContext) => {
         refresh: 'wait_for',
       });
 
-      const result = await getDailyEvents('.kibana', callCluster);
+      const result = await getDailyEvents('.kibana', () => Promise.resolve(es));
 
       expect(result).to.eql({
         byDate: {},
@@ -150,7 +150,7 @@ export default ({ getService }: FtrProviderContext) => {
         ],
       });
 
-      const result = await getDailyEvents('.kibana', callCluster);
+      const result = await getDailyEvents('.kibana', () => Promise.resolve(es));
 
       expect(result).to.eql({
         byDate: {
@@ -177,7 +177,7 @@ export default ({ getService }: FtrProviderContext) => {
 
       await esArchiver.loadIfNeeded('lens/basic');
 
-      const results = await getVisualizationCounts(callCluster, '.kibana');
+      const results = await getVisualizationCounts(() => Promise.resolve(es), '.kibana');
 
       expect(results).to.have.keys([
         'saved_overall',