[Discover] Replace RxJS firstValueFrom with RxJS lastValueFrom to prevent problem with partial results

src/plugins/discover/public/hooks/use_es_doc_search.test.tsx

@@ -6,18 +6,17 @@
  * Side Public License, v 1.
  */
 
-import { renderHook } from '@testing-library/react-hooks';
+import { renderHook, act } from '@testing-library/react-hooks';
 import { buildSearchBody, useEsDocSearch } from './use_es_doc_search';
-import { Observable } from 'rxjs';
+import { Subject } from 'rxjs';
 import { DataView } from '@kbn/data-views-plugin/public';
 import { DocProps } from '../application/doc/components/doc';
 import { ElasticRequestState } from '../application/doc/types';
 import { SEARCH_FIELDS_FROM_SOURCE as mockSearchFieldsFromSource } from '../../common';
 import { KibanaContextProvider } from '@kbn/kibana-react-plugin/public';
 import React from 'react';
 
-const mockSearchResult = new Observable();
-
+const mockSearchResult = new Subject();
 const services = {
   data: {
     search: {
@@ -171,23 +170,69 @@ describe('Test of <Doc /> helper / hook', () => {
     `);
   });
 
-  test('useEsDocSearch', async () => {
+  test('useEsDocSearch loading', async () => {
+    const indexPattern = {
+      getComputedFields: () => [],
+    };
+    const props = {
+      id: '1',
+      index: 'index1',
+      indexPattern,
+    } as unknown as DocProps;
+
+    const hook = renderHook((p: DocProps) => useEsDocSearch(p), {
+      initialProps: props,
+      wrapper: ({ children }) => (
+        <KibanaContextProvider services={services}>{children}</KibanaContextProvider>
+      ),
+    });
+
+    expect(hook.result.current.slice(0, 2)).toEqual([ElasticRequestState.Loading, null]);
+  });
+
+  test('useEsDocSearch ignore partial results', async () => {
     const indexPattern = {
       getComputedFields: () => [],
     };
+
+    const record = { test: 1 };
+
     const props = {
       id: '1',
       index: 'index1',
       indexPattern,
     } as unknown as DocProps;
 
-    const { result } = renderHook((p: DocProps) => useEsDocSearch(p), {
+    const hook = renderHook((p: DocProps) => useEsDocSearch(p), {
       initialProps: props,
       wrapper: ({ children }) => (
         <KibanaContextProvider services={services}>{children}</KibanaContextProvider>
       ),
     });
 
-    expect(result.current.slice(0, 2)).toEqual([ElasticRequestState.Loading, null]);
+    await act(async () => {
+      mockSearchResult.next({
+        isPartial: true,
+        isRunning: false,
+        rawResponse: {
+          hits: {
+            hits: [],
+          },
+        },
+      });
+      mockSearchResult.next({
+        isPartial: false,
+        isRunning: false,
+        rawResponse: {
+          hits: {
+            hits: [record],
+          },
+        },
+      });
+      mockSearchResult.complete();
+      await hook.waitForNextUpdate();
+    });
+
+    expect(hook.result.current.slice(0, 2)).toEqual([ElasticRequestState.Found, record]);
   });
 });

src/plugins/discover/public/hooks/use_es_doc_search.ts

@@ -8,7 +8,7 @@
 
 import { useCallback, useEffect, useMemo, useState } from 'react';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
-import { firstValueFrom } from 'rxjs';
+import { lastValueFrom } from 'rxjs';
 import { DataView } from '@kbn/data-views-plugin/public';
 import { DocProps } from '../application/doc/components/doc';
 import { ElasticRequestState } from '../application/doc/types';
@@ -18,47 +18,6 @@ import { useDiscoverServices } from './use_discover_services';
 
 type RequestBody = Pick<estypes.SearchRequest, 'body'>;
 
-/**
- * helper function to build a query body for Elasticsearch
- * https://www.elastic.co/guide/en/elasticsearch/reference/current//query-dsl-ids-query.html
- */
-export function buildSearchBody(
-  id: string,
-  indexPattern: DataView,
-  useNewFieldsApi: boolean,
-  requestAllFields?: boolean
-): RequestBody | undefined {
-  const computedFields = indexPattern.getComputedFields();
-  const runtimeFields = computedFields.runtimeFields as estypes.MappingRuntimeFields;
-  const request: RequestBody = {
-    body: {
-      query: {
-        ids: {
-          values: [id],
-        },
-      },
-      stored_fields: computedFields.storedFields,
-      script_fields: computedFields.scriptFields,
-      version: true,
-    },
-  };
-  if (!request.body) {
-    return undefined;
-  }
-  if (useNewFieldsApi) {
-    // @ts-expect-error
-    request.body.fields = [{ field: '*', include_unmapped: 'true' }];
-    request.body.runtime_mappings = runtimeFields ? runtimeFields : {};
-    if (requestAllFields) {
-      request.body._source = true;
-    }
-  } else {
-    request.body._source = true;
-  }
-  request.body.fields = [...(request.body?.fields || []), ...(computedFields.docvalueFields || [])];
-  return request;
-}
-
 /**
  * Custom react hook for querying a single doc in ElasticSearch
  */
@@ -75,14 +34,15 @@ export function useEsDocSearch({
 
   const requestData = useCallback(async () => {
     try {
-      const { rawResponse } = await firstValueFrom(
+      const result = await lastValueFrom(
         data.search.search({
           params: {
             index,
             body: buildSearchBody(id, indexPattern, useNewFieldsApi, requestSource)?.body,
           },
         })
       );
+      const rawResponse = result.rawResponse;
 
       const hits = rawResponse.hits;
 
@@ -109,3 +69,44 @@ export function useEsDocSearch({
 
   return [status, hit, requestData];
 }
+
+/**
+ * helper function to build a query body for Elasticsearch
+ * https://www.elastic.co/guide/en/elasticsearch/reference/current//query-dsl-ids-query.html
+ */
+export function buildSearchBody(
+  id: string,
+  indexPattern: DataView,
+  useNewFieldsApi: boolean,
+  requestAllFields?: boolean
+): RequestBody | undefined {
+  const computedFields = indexPattern.getComputedFields();
+  const runtimeFields = computedFields.runtimeFields as estypes.MappingRuntimeFields;
+  const request: RequestBody = {
+    body: {
+      query: {
+        ids: {
+          values: [id],
+        },
+      },
+      stored_fields: computedFields.storedFields,
+      script_fields: computedFields.scriptFields,
+      version: true,
+    },
+  };
+  if (!request.body) {
+    return undefined;
+  }
+  if (useNewFieldsApi) {
+    // @ts-expect-error
+    request.body.fields = [{ field: '*', include_unmapped: 'true' }];
+    request.body.runtime_mappings = runtimeFields ? runtimeFields : {};
+    if (requestAllFields) {
+      request.body._source = true;
+    }
+  } else {
+    request.body._source = true;
+  }
+  request.body.fields = [...(request.body?.fields || []), ...(computedFields.docvalueFields || [])];
+  return request;
+}

x-pack/plugins/stack_alerts/public/alert_types/es_query/expression/es_query_expression.test.tsx

@@ -7,7 +7,7 @@
 
 import React from 'react';
 import 'brace';
-import { of } from 'rxjs';
+import { of, Subject } from 'rxjs';
 import { mountWithIntl, nextTick } from '@kbn/test-jest-helpers';
 import { act } from 'react-dom/test-utils';
 import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
@@ -220,6 +220,42 @@ describe('EsQueryAlertTypeExpression', () => {
     );
   });
 
+  test('should show success message if Test Query is successful (with partial result)', async () => {
+    const partial = {
+      isRunning: true,
+      isPartial: true,
+    };
+    const complete = {
+      isRunning: false,
+      isPartial: false,
+      rawResponse: {
+        hits: {
+          total: 1234,
+        },
+      },
+    };
+    const searchResponseMock$ = new Subject();
+    dataMock.search.search.mockImplementation(() => searchResponseMock$);
+    const wrapper = await setup(defaultEsQueryExpressionParams);
+    const testQueryButton = wrapper.find('EuiButton[data-test-subj="testQuery"]');
+
+    testQueryButton.simulate('click');
+    expect(dataMock.search.search).toHaveBeenCalled();
+    await act(async () => {
+      searchResponseMock$.next(partial);
+      searchResponseMock$.next(complete);
+      searchResponseMock$.complete();
+      await nextTick();
+      wrapper.update();
+    });
+
+    expect(wrapper.find('[data-test-subj="testQuerySuccess"]').exists()).toBeTruthy();
+    expect(wrapper.find('[data-test-subj="testQueryError"]').exists()).toBeFalsy();
+    expect(wrapper.find('EuiText[data-test-subj="testQuerySuccess"]').text()).toEqual(
+      `Query matched 1234 documents in the last 15s.`
+    );
+  });
+
   test('should show error message if Test Query is throws error', async () => {
     dataMock.search.search.mockImplementation(() => {
       throw new Error('What is this query');

x-pack/plugins/stack_alerts/public/alert_types/es_query/expression/es_query_expression.tsx

@@ -6,7 +6,7 @@
  */
 
 import React, { useState, Fragment, useEffect, useCallback } from 'react';
-import { firstValueFrom } from 'rxjs';
+import { lastValueFrom } from 'rxjs';
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
 
@@ -141,7 +141,7 @@ export const EsQueryExpression = ({
     const timeWindow = parseDuration(window);
     const parsedQuery = JSON.parse(esQuery);
     const now = Date.now();
-    const { rawResponse } = await firstValueFrom(
+    const { rawResponse } = await lastValueFrom(
       data.search.search({
         params: buildSortedEventsQuery({
           index,

x-pack/plugins/stack_alerts/public/alert_types/es_query/expression/search_source_expression.test.tsx

@@ -14,8 +14,8 @@ import { EsQueryAlertParams, SearchType } from '../types';
 import { SearchSourceExpression } from './search_source_expression';
 import { chartPluginMock } from '@kbn/charts-plugin/public/mocks';
 import { act } from 'react-dom/test-utils';
-import { of } from 'rxjs';
-import { IKibanaSearchResponse, ISearchSource } from '@kbn/data-plugin/common';
+import { Subject } from 'rxjs';
+import { ISearchSource } from '@kbn/data-plugin/common';
 import { IUiSettingsClient } from '@kbn/core/public';
 import { findTestSubject } from '@elastic/eui/lib/test';
 import { EuiLoadingSpinner } from '@elastic/eui';
@@ -40,6 +40,20 @@ const defaultSearchSourceExpressionParams: EsQueryAlertParams<SearchType.searchS
   searchConfiguration: {},
 };
 
+const mockSearchResult = new Subject();
+const testResultComplete = {
+  rawResponse: {
+    hits: {
+      total: 1234,
+    },
+  },
+};
+
+const testResultPartial = {
+  partial: true,
+  running: true,
+};
+
 const searchSourceMock = {
   id: 'data_source6',
   fields: {
@@ -67,13 +81,7 @@ const searchSourceMock = {
     return searchSourceMock;
   }),
   fetch$: jest.fn(() => {
-    return of<IKibanaSearchResponse>({
-      rawResponse: {
-        hits: {
-          total: 1234,
-        },
-      },
-    });
+    return mockSearchResult;
   }),
 } as unknown as ISearchSource;
 
@@ -143,6 +151,7 @@ describe('SearchSourceAlertTypeExpression', () => {
     wrapper = await wrapper.update();
     expect(findTestSubject(wrapper, 'thresholdExpression')).toBeTruthy();
   });
+
   test('should show success message if Test Query is successful', async () => {
     let wrapper = setup(defaultSearchSourceExpressionParams);
     await act(async () => {
@@ -156,6 +165,9 @@ describe('SearchSourceAlertTypeExpression', () => {
     wrapper = await wrapper.update();
 
     await act(async () => {
+      mockSearchResult.next(testResultPartial);
+      mockSearchResult.next(testResultComplete);
+      mockSearchResult.complete();
       await nextTick();
       wrapper.update();
     });

x-pack/plugins/stack_alerts/public/alert_types/es_query/expression/search_source_expression_form.tsx

@@ -7,7 +7,7 @@
 
 import React, { Fragment, useCallback, useEffect, useMemo, useReducer, useState } from 'react';
 import deepEqual from 'fast-deep-equal';
-import { firstValueFrom } from 'rxjs';
+import { lastValueFrom } from 'rxjs';
 import { Filter } from '@kbn/es-query';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { EuiSpacer, EuiTitle } from '@elastic/eui';
@@ -183,7 +183,7 @@ export const SearchSourceExpressionForm = (props: SearchSourceExpressionFormProp
       'filter',
       timeFilter ? [timeFilter, ...ruleConfiguration.filter] : ruleConfiguration.filter
     );
-    const { rawResponse } = await firstValueFrom(testSearchSource.fetch$());
+    const { rawResponse } = await lastValueFrom(testSearchSource.fetch$());
     return { nrOfDocs: totalHitsToNumber(rawResponse.hits.total), timeWindow };
   }, [searchSource, timeWindowSize, timeWindowUnit, ruleConfiguration]);