[Search service] Shim total hits in async search response (#61565)

* Shim total hits in async search response

* Resolve types

* Fix types

* Fix tests

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts

@@ -19,6 +19,7 @@ import {
   getTotalLoaded,
 } from '../../../../../src/plugins/data/server';
 import { IEnhancedEsSearchRequest } from '../../common';
+import { shimHitsTotal } from './shim_hits_total';
 
 export interface AsyncSearchResponse<T> {
   id: string;
@@ -56,22 +57,27 @@ async function asyncSearch(
   request: IEnhancedEsSearchRequest,
   options?: ISearchOptions
 ) {
-  const { body = undefined, index = undefined, ...params } = request.id ? {} : request.params;
+  const { timeout = undefined, restTotalHitsAsInt = undefined, ...params } = {
+    trackTotalHits: true, // Get the exact count of hits
+    ...request.params,
+  };
 
   // If we have an ID, then just poll for that ID, otherwise send the entire request body
+  const { body = undefined, index = undefined, ...queryParams } = request.id ? {} : params;
+
   const method = request.id ? 'GET' : 'POST';
   const path = encodeURI(request.id ? `_async_search/${request.id}` : `${index}/_async_search`);
 
   // Wait up to 1s for the response to return
-  const query = toSnakeCase({ waitForCompletionTimeout: '1s', ...params });
+  const query = toSnakeCase({ waitForCompletionTimeout: '1s', ...queryParams });
 
-  const { response: rawResponse, id } = (await caller(
+  const { response, id } = (await caller(
     'transport.request',
     { method, path, body, query },
     options
   )) as AsyncSearchResponse<any>;
 
-  return { id, rawResponse, ...getTotalLoaded(rawResponse._shards) };
+  return { id, rawResponse: shimHitsTotal(response), ...getTotalLoaded(response._shards) };
 }
 
 async function rollupSearch(

x-pack/plugins/data_enhanced/server/search/shim_hits_total.test.ts

@@ -0,0 +1,56 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { shimHitsTotal } from './shim_hits_total';
+
+describe('shimHitsTotal', () => {
+  test('returns the total if it is already numeric', () => {
+    const result = shimHitsTotal({
+      hits: {
+        total: 5,
+      },
+    } as any);
+    expect(result).toEqual({
+      hits: {
+        total: 5,
+      },
+    });
+  });
+
+  test('returns the total if it is inside `value`', () => {
+    const result = shimHitsTotal({
+      hits: {
+        total: {
+          value: 5,
+        },
+      },
+    } as any);
+    expect(result).toEqual({
+      hits: {
+        total: 5,
+      },
+    });
+  });
+
+  test('returns other properties from the response', () => {
+    const result = shimHitsTotal({
+      _shards: {},
+      hits: {
+        hits: [],
+        total: {
+          value: 5,
+        },
+      },
+    } as any);
+    expect(result).toEqual({
+      _shards: {},
+      hits: {
+        hits: [],
+        total: 5,
+      },
+    });
+  });
+});

x-pack/plugins/data_enhanced/server/search/shim_hits_total.ts

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { SearchResponse } from 'elasticsearch';
+
+/**
+ * Temporary workaround until https://github.com/elastic/kibana/issues/26356 is addressed.
+ * Since we are setting `track_total_hits` in the request, `hits.total` will be an object
+ * containing the `value`.
+ */
+export function shimHitsTotal(response: SearchResponse<any>) {
+  const total = (response.hits?.total as any)?.value ?? response.hits?.total;
+  const hits = { ...response.hits, total };
+  return { ...response, hits };
+}