[Search] Fix async search to encode index pattern in path (#61374) (#61648)

* Fix async search to encode index in path

* Update docs

* Review feedback & fixing types

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

docs/development/plugins/data/server/kibana-plugin-plugins-data-server.md

@@ -23,6 +23,7 @@
 |  Function | Description |
 |  --- | --- |
 |  [getDefaultSearchParams(config)](./kibana-plugin-plugins-data-server.getdefaultsearchparams.md) |  |
+|  [getTotalLoaded({ total, failed, successful })](./kibana-plugin-plugins-data-server.gettotalloaded.md) |  |
 |  [parseInterval(interval)](./kibana-plugin-plugins-data-server.parseinterval.md) |  |
 |  [plugin(initializerContext)](./kibana-plugin-plugins-data-server.plugin.md) | Static code to be shared externally |
 |  [shouldReadFieldFromDocValues(aggregatable, esType)](./kibana-plugin-plugins-data-server.shouldreadfieldfromdocvalues.md) |  |

src/plugins/data/server/index.ts

@@ -173,6 +173,7 @@ export {
   ISearchContext,
   TSearchStrategyProvider,
   getDefaultSearchParams,
+  getTotalLoaded,
 } from './search';
 
 // Search namespace

src/plugins/data/server/search/es_search/es_search_strategy.ts

@@ -21,7 +21,7 @@ import { APICaller } from 'kibana/server';
 import { SearchResponse } from 'elasticsearch';
 import { ES_SEARCH_STRATEGY } from '../../../common/search';
 import { ISearchStrategy, TSearchStrategyProvider } from '../i_search_strategy';
-import { getDefaultSearchParams, ISearchContext } from '..';
+import { getDefaultSearchParams, getTotalLoaded, ISearchContext } from '..';
 
 export const esSearchStrategyProvider: TSearchStrategyProvider<typeof ES_SEARCH_STRATEGY> = (
   context: ISearchContext,
@@ -46,9 +46,7 @@ export const esSearchStrategyProvider: TSearchStrategyProvider<typeof ES_SEARCH_
 
       // The above query will either complete or timeout and throw an error.
       // There is no progress indication on this api.
-      const { total, failed, successful } = rawResponse._shards;
-      const loaded = failed + successful;
-      return { total, loaded, rawResponse };
+      return { rawResponse, ...getTotalLoaded(rawResponse._shards) };
     },
   };
 };

src/plugins/data/server/search/es_search/get_total_loaded.test.ts

@@ -0,0 +1,36 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { getTotalLoaded } from './get_total_loaded';
+
+describe('getTotalLoaded', () => {
+  it('returns the total/loaded, not including skipped', () => {
+    const result = getTotalLoaded({
+      successful: 10,
+      failed: 5,
+      skipped: 5,
+      total: 100,
+    });
+
+    expect(result).toEqual({
+      total: 100,
+      loaded: 15,
+    });
+  });
+});

src/plugins/data/server/search/es_search/get_total_loaded.ts

@@ -0,0 +1,30 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { ShardsResponse } from 'elasticsearch';
+
+/**
+ * Get the `total`/`loaded` for this response (see `IKibanaSearchResponse`). Note that `skipped` is
+ * not included as it is already included in `successful`.
+ * @internal
+ */
+export function getTotalLoaded({ total, failed, successful }: ShardsResponse) {
+  const loaded = failed + successful;
+  return { total, loaded };
+}

src/plugins/data/server/search/es_search/index.ts

@@ -20,3 +20,4 @@
 export { ES_SEARCH_STRATEGY, IEsSearchRequest, IEsSearchResponse } from '../../../common/search';
 export { esSearchStrategyProvider } from './es_search_strategy';
 export { getDefaultSearchParams } from './get_default_search_params';
+export { getTotalLoaded } from './get_total_loaded';

src/plugins/data/server/search/index.ts

@@ -33,4 +33,4 @@ export { TStrategyTypes } from './strategy_types';
 
 export { TSearchStrategyProvider } from './i_search_strategy';
 
-export { getDefaultSearchParams } from './es_search';
+export { getDefaultSearchParams, getTotalLoaded } from './es_search';

src/plugins/data/server/server.api.md

@@ -125,6 +125,7 @@ import { SearchResponse } from 'elasticsearch';
 import { SearchShardsParams } from 'elasticsearch';
 import { SearchTemplateParams } from 'elasticsearch';
 import { ShallowPromise } from '@kbn/utility-types';
+import { ShardsResponse } from 'elasticsearch';
 import { SnapshotCreateParams } from 'elasticsearch';
 import { SnapshotCreateRepositoryParams } from 'elasticsearch';
 import { SnapshotDeleteParams } from 'elasticsearch';
@@ -330,6 +331,12 @@ export function getDefaultSearchParams(config: SharedGlobalConfig): {
     restTotalHitsAsInt: boolean;
 };
 
+// @internal
+export function getTotalLoaded({ total, failed, successful }: ShardsResponse): {
+    total: number;
+    loaded: number;
+};
+
 // Warning: (ae-missing-release-tag) "IFieldFormatsRegistry" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
 //
 // @public (undocumented)
@@ -732,12 +739,12 @@ export type TSearchStrategyProvider<T extends TStrategyTypes> = (context: ISearc
 // src/plugins/data/server/index.ts:102:26 - (ae-forgotten-export) The symbol "TruncateFormat" needs to be exported by the entry point index.d.ts
 // src/plugins/data/server/index.ts:130:27 - (ae-forgotten-export) The symbol "isFilterable" needs to be exported by the entry point index.d.ts
 // src/plugins/data/server/index.ts:130:27 - (ae-forgotten-export) The symbol "isNestedField" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:181:1 - (ae-forgotten-export) The symbol "dateHistogramInterval" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:182:1 - (ae-forgotten-export) The symbol "InvalidEsCalendarIntervalError" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:183:1 - (ae-forgotten-export) The symbol "InvalidEsIntervalFormatError" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:184:1 - (ae-forgotten-export) The symbol "isValidEsInterval" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:185:1 - (ae-forgotten-export) The symbol "isValidInterval" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/index.ts:188:1 - (ae-forgotten-export) The symbol "toAbsoluteDates" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:182:1 - (ae-forgotten-export) The symbol "dateHistogramInterval" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:183:1 - (ae-forgotten-export) The symbol "InvalidEsCalendarIntervalError" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:184:1 - (ae-forgotten-export) The symbol "InvalidEsIntervalFormatError" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:185:1 - (ae-forgotten-export) The symbol "isValidEsInterval" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:186:1 - (ae-forgotten-export) The symbol "isValidInterval" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/index.ts:189:1 - (ae-forgotten-export) The symbol "toAbsoluteDates" needs to be exported by the entry point index.d.ts
 // src/plugins/data/server/plugin.ts:64:14 - (ae-forgotten-export) The symbol "ISearchSetup" needs to be exported by the entry point index.d.ts
 
 // (No @packageDocumentation comment for this package)

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

@@ -0,0 +1,144 @@
+/*
+ * 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 { coreMock, pluginInitializerContextConfigMock } from '../../../../../src/core/server/mocks';
+import { enhancedEsSearchStrategyProvider } from './es_search_strategy';
+
+const mockAsyncResponse = {
+  id: 'foo',
+  response: {
+    _shards: {
+      total: 10,
+      failed: 1,
+      skipped: 2,
+      successful: 7,
+    },
+  },
+};
+
+const mockRollupResponse = {
+  _shards: {
+    total: 10,
+    failed: 1,
+    skipped: 2,
+    successful: 7,
+  },
+};
+
+describe('ES search strategy', () => {
+  const mockCoreSetup = coreMock.createSetup();
+  const mockApiCaller = jest.fn();
+  const mockSearch = jest.fn();
+  const mockConfig$ = pluginInitializerContextConfigMock<any>({}).legacy.globalConfig$;
+
+  beforeEach(() => {
+    mockApiCaller.mockClear();
+    mockSearch.mockClear();
+  });
+
+  it('returns a strategy with `search`', () => {
+    const esSearch = enhancedEsSearchStrategyProvider(
+      {
+        core: mockCoreSetup,
+        config$: mockConfig$,
+      },
+      mockApiCaller,
+      mockSearch
+    );
+
+    expect(typeof esSearch.search).toBe('function');
+  });
+
+  it('makes a POST request to async search with params when no ID is provided', async () => {
+    mockApiCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+    const params = { index: 'logstash-*', body: { query: {} } };
+    const esSearch = enhancedEsSearchStrategyProvider(
+      {
+        core: mockCoreSetup,
+        config$: mockConfig$,
+      },
+      mockApiCaller,
+      mockSearch
+    );
+
+    await esSearch.search({ params });
+
+    expect(mockApiCaller).toBeCalled();
+    expect(mockApiCaller.mock.calls[0][0]).toBe('transport.request');
+    const { method, path, body } = mockApiCaller.mock.calls[0][1];
+    expect(method).toBe('POST');
+    expect(path).toBe('logstash-*/_async_search');
+    expect(body).toEqual({ query: {} });
+  });
+
+  it('makes a GET request to async search with ID when ID is provided', async () => {
+    mockApiCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+    const params = { index: 'logstash-*', body: { query: {} } };
+    const esSearch = enhancedEsSearchStrategyProvider(
+      {
+        core: mockCoreSetup,
+        config$: mockConfig$,
+      },
+      mockApiCaller,
+      mockSearch
+    );
+
+    await esSearch.search({ id: 'foo', params });
+
+    expect(mockApiCaller).toBeCalled();
+    expect(mockApiCaller.mock.calls[0][0]).toBe('transport.request');
+    const { method, path, body } = mockApiCaller.mock.calls[0][1];
+    expect(method).toBe('GET');
+    expect(path).toBe('_async_search/foo');
+    expect(body).toEqual(undefined);
+  });
+
+  it('encodes special characters in the path', async () => {
+    mockApiCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+    const params = { index: 'foo-程', body: {} };
+    const esSearch = enhancedEsSearchStrategyProvider(
+      {
+        core: mockCoreSetup,
+        config$: mockConfig$,
+      },
+      mockApiCaller,
+      mockSearch
+    );
+
+    await esSearch.search({ params });
+
+    expect(mockApiCaller).toBeCalled();
+    expect(mockApiCaller.mock.calls[0][0]).toBe('transport.request');
+    const { method, path } = mockApiCaller.mock.calls[0][1];
+    expect(method).toBe('POST');
+    expect(path).toBe('foo-%E7%A8%8B/_async_search');
+  });
+
+  it('calls the rollup API if the index is a rollup type', async () => {
+    mockApiCaller.mockResolvedValueOnce(mockRollupResponse);
+
+    const params = { index: 'foo-程', body: {} };
+    const esSearch = enhancedEsSearchStrategyProvider(
+      {
+        core: mockCoreSetup,
+        config$: mockConfig$,
+      },
+      mockApiCaller,
+      mockSearch
+    );
+
+    await esSearch.search({ indexType: 'rollup', params });
+
+    expect(mockApiCaller).toBeCalled();
+    expect(mockApiCaller.mock.calls[0][0]).toBe('transport.request');
+    const { method, path } = mockApiCaller.mock.calls[0][1];
+    expect(method).toBe('POST');
+    expect(path).toBe('foo-%E7%A8%8B/_rollup_search');
+  });
+});

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

@@ -16,6 +16,7 @@ import {
   ISearchOptions,
   ISearchCancel,
   getDefaultSearchParams,
+  getTotalLoaded,
 } from '../../../../../src/plugins/data/server';
 import { IEnhancedEsSearchRequest } from '../../common';
 
@@ -36,31 +37,21 @@ export const enhancedEsSearchStrategyProvider: TSearchStrategyProvider<typeof ES
     const defaultParams = getDefaultSearchParams(config);
     const params = { ...defaultParams, ...request.params };
 
-    const response = await (request.indexType === 'rollup'
+    return request.indexType === 'rollup'
       ? rollupSearch(caller, { ...request, params }, options)
-      : asyncSearch(caller, { ...request, params }, options));
-
-    const rawResponse =
-      request.indexType === 'rollup'
-        ? (response as SearchResponse<any>)
-        : (response as AsyncSearchResponse<any>).response;
-
-    const id = (response as AsyncSearchResponse<any>).id;
-    const { total, failed, successful } = rawResponse._shards;
-    const loaded = failed + successful;
-    return { id, total, loaded, rawResponse };
+      : asyncSearch(caller, { ...request, params }, options);
   };
 
   const cancel: ISearchCancel<typeof ES_SEARCH_STRATEGY> = async id => {
     const method = 'DELETE';
-    const path = `_async_search/${id}`;
+    const path = encodeURI(`_async_search/${id}`);
     await caller('transport.request', { method, path });
   };
 
   return { search, cancel };
 };
 
-function asyncSearch(
+async function asyncSearch(
   caller: APICaller,
   request: IEnhancedEsSearchRequest,
   options?: ISearchOptions
@@ -69,12 +60,18 @@ function asyncSearch(
 
   // If we have an ID, then just poll for that ID, otherwise send the entire request body
   const method = request.id ? 'GET' : 'POST';
-  const path = request.id ? `_async_search/${request.id}` : `${index}/_async_search`;
+  const path = encodeURI(request.id ? `_async_search/${request.id}` : `${index}/_async_search`);
 
   // Wait up to 1s for the response to return
   const query = toSnakeCase({ waitForCompletion: '1s', ...params });
 
-  return caller('transport.request', { method, path, body, query }, options);
+  const { response: rawResponse, id } = (await caller(
+    'transport.request',
+    { method, path, body, query },
+    options
+  )) as AsyncSearchResponse<any>;
+
+  return { id, rawResponse, ...getTotalLoaded(rawResponse._shards) };
 }
 
 async function rollupSearch(
@@ -84,9 +81,16 @@ async function rollupSearch(
 ) {
   const { body, index, ...params } = request.params;
   const method = 'POST';
-  const path = `${index}/_rollup_search`;
+  const path = encodeURI(`${index}/_rollup_search`);
   const query = toSnakeCase(params);
-  return caller('transport.request', { method, path, body, query }, options);
+
+  const rawResponse = await ((caller(
+    'transport.request',
+    { method, path, body, query },
+    options
+  ) as unknown) as SearchResponse<any>);
+
+  return { rawResponse, ...getTotalLoaded(rawResponse._shards) };
 }
 
 function toSnakeCase(obj: Record<string, any>) {